[MLOps] Optuna를 활용한 최적의 하이퍼파라미터 탐색

Optuna

Optuna

머신러닝, 딥러닝 모델의 하이퍼파라미터를 자동으로 탐색하고 최적화하는 오픈 소스 프레임워크로, 모델의 성능을 극대화할 수 있는 최적의 하이퍼파라미터 조합을 효율적으로 찾을 수 있도록 도와줍니다.

• 효율적인 하이퍼파라미터 탐색: 트리 기반의 구조화된 파라미터 최적화(Tree-structured Parzen Estimator) 알고리즘을 사용해 Grid Search나 Random Search보다 적은 시도로 더 좋은 결과를 얻을 수 있습니다. TPE는 이전 평가 결과를 바탕으로 성능이 좋을 것으로 예상되는 하이퍼파라미터 조합에 더 많은 자원을 할당합니다.
• 자동 프루닝 기능: 비효율적인 시도를 조기에 중단하는 프루닝 기능을 제공합니다. 중간 평가에서 성능이 떨어지는 시도는 학습을 일찍 종료시켜 전체 최적화 과정의 속도를 높이고 불필요한 자원 낭비를 줄입니다.
• 병렬화 및 분산 처리 지원: 여러 실험을 병렬로 실행할 수 있어 대규모 데이터나 복잡한 모델을 빠르게 실험할 수 있습니다.
• 유연한 목적 함수 정의: 사용자가 직접 목적 함수를 자유롭게 정의할 수 있어 맞춤형 최적화가 가능합니다.
• 시각화 도구 제공: 하이퍼파라미터의 탐색 과정과 중요도를 분석할 수 있는 다양한 시각화 도구를 내장하고 있습니다.

더 자세한 내용은 아래 링크에서 자세히 확인할 수 있습니다.
Optuna Tutorial

목적 함수

하이퍼파라미터를 입력받아 모델을 학습시키고 정확도, 손실 등의 결과를 반환하는 함수

목적 함수는 매 실험(trial)마다 실행되며 다음 과정을 포함합니다.

• 하이퍼파라미터 정의: trial 객체를 사용해 값 제안
• 모델 구성: 정의된 하이퍼파라미터로 모델 빌드
• 학습 및 평가: 학습 루프 또는 학습 함수 실행
• 최적화 목표 반환: validation accuracy, loss 등 최적화할 값 반환

def objective(trial):
    lr = trial.suggest_float("learning_rate", 1e-5, 1e-3, log=True)
    dropout = trial.suggest_float("dropout", 0.2, 0.5)
    
    model = build_model(dropout=dropout)
    accuracy = train_and_validate(model, lr=lr)
    return accuracy

파라미터 제안 함수

하이퍼파라미터를 탐색할 때 사용할 값의 범위나 후보군을 trial 객체를 통해 정의할 수 있으며, 각각의 함수는 다루는 데이터 타입과 목적에 따라 달라집니다.

함수	설명	예시	특징
suggest_int	정수 범위 내에서 값을 선택합니다.	trial.suggest_int(..., 1, 5)	정수형 하이퍼파라미터에 적합
suggest_float	실수 범위 내에서 값을 선택합니다.	trial.suggest_float(..., 0.1, 0.5)	연속적인 실수값 탐색에 적합
suggest_float(..., log=True)	로그 스케일로 실수 값을 선택합니다.	trial.suggest_float(..., 1e-5, 1e-2, log=True)	값의 범위가 넓을 때 효과적
suggest_categorical	정해진 후보군 중 하나를 선택합니다.	trial.suggest_categorical(..., [16, 32, 64])	이산형(범주형) 파라미터 탐색에 적합

Optuna 튜토리얼 진행하기

저는 엔비디아 GPU를 가지고 있지 않아 Colab에서 진행하였습니다.
모델은 이전에도 사용한 적 있는 Bllossom/llama-3.2-Korean-Bllossom-3B 모델입니다.

먼저, Colab에 접속해 새 노트를 생성합니다.
노트 생성 후, 런타임 유형을 GPU로 바꿔줍니다.

이제 다음과 같이 코드를 노트북에 작성합니다.

# 기존 bitsandbytes 제거 후 최신 버전 재설치 (8bit 로딩 지원 라이브러리)
!pip uninstall -y bitsandbytes
!pip install -U bitsandbytes

# 주요 학습 및 추적 관련 라이브러리 설치
!pip install -U transformers accelerate peft deepspeed datasets mpi4py mlflow gputil optuna

# 필수 라이브러리 임포트
import torch, optuna, mlflow
import optuna.visualization as vis
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer, DataCollatorForLanguageModeling
from datasets import load_dataset
from peft import get_peft_model, LoraConfig

# SQuAD 데이터셋 일부 로드
dataset = load_dataset("squad", split="train[:1000]")

# LoRA 설정: 파라미터 효율화 위한 설정
lora_config = LoraConfig(
    r=16,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],  # Llama 모델의 주요 어텐션 모듈
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

# MLflow 실험 경로 설정
mlflow.set_tracking_uri("file:/content/mlruns")
# 실험 이름 지정
mlflow.set_experiment("Optuna_Tuning")

# Optuna 목적 함수
def objective(trial):
    # 하이퍼파라미터 샘플링
    per_device_batch_size = trial.suggest_categorical("per_device_train_batch_size", [1, 2])
    gradient_accum_steps = trial.suggest_categorical("gradient_accumulation_steps", [1, 2, 4])
    learning_rate = trial.suggest_float("learning_rate", 1e-5, 2e-4, log=True)

    # 모델 및 토크나이저 로드
    model_id = "Bllossom/llama-3.2-Korean-Bllossom-3B"
    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        load_in_8bit=True,  # 8bit로 로드하여 메모리 절약
        torch_dtype=torch.float16,
        device_map="auto" # 자동으로 GPU 할당
    )
    model = get_peft_model(model, lora_config)  # LoRA 적용

    # 데이터 전처리 함수 정의
    def preprocess(examples):
        inputs = []
        for question, answer in zip(examples["question"], examples["answers"]):
            text = answer["text"][0] if answer["text"] else "정답 없음"
            inputs.append(f"질문: {question}\n답변: {text}")
        return tokenizer(inputs, padding="max_length", truncation=True, max_length=512)

    # 패딩 토큰 설정 및 전처리 적용
    tokenizer.pad_token = tokenizer.eos_token
    tokenized_dataset = dataset.map(preprocess, batched=True)
    tokenized_dataset.set_format(type="torch", columns=["input_ids", "attention_mask"])

    # 학습 인자 설정
    training_args = TrainingArguments(
        output_dir=f"./results_trial_{trial.number}",
        per_device_train_batch_size=per_device_batch_size,
        gradient_accumulation_steps=gradient_accum_steps,
        learning_rate=learning_rate,
        num_train_epochs=1,
        fp16=True,
        logging_steps=100,
        save_strategy="epoch",
        report_to="none", # 로그는 MLflow로만
    )

    # Trainer 구성
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset,
        data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False) # MLM 비활성화 (Causal LM)
    )

    # MLflow 실험 기록
    with mlflow.start_run(run_name=f"Trial_{trial.number}"):
        mlflow.log_params({
            "batch_size": per_device_batch_size,
            "accumulation": gradient_accum_steps,
            "lr": learning_rate,
        })
        train_result = trainer.train()  # 학습 실행
        loss = train_result.training_loss if hasattr(train_result, "training_loss") else train_result.metrics.get("train_loss", -1)
        mlflow.log_metric("loss", loss)
        print(f"Trial {trial.number} | Loss: {loss:.4f}")
        return loss


# Optuna 튜닝 실행
study = optuna.create_study(direction="minimize")
study.optimize(objective, n_trials=3)

# 최적 결과 출력
print("Best Trial:")
print(study.best_trial.params)

# Optuna 시각화
vis.plot_optimization_history(study).show()
vis.plot_param_importances(study).show()

!zip -r mlruns.zip /content/mlruns

코드를 실행하고 출력된 최적의 파라미터를 확인합니다.

Best Trial:
{'per_device_train_batch_size': 1, 'gradient_accumulation_steps': 1, 'learning_rate': 4.600105347650124e-05}

이 조합에서 가장 낮은 loss를 기록했다는 것을 의미합니다.

다음으로 Optuna 시각화 결과를 확인합니다.

Optimization History Plot

• X축: trial 번호 (0~2)
• Y축: objective value (loss)

Trial 0이 가장 낮은 loss를 기록해 가장 좋은 조합을 기록하였습니다.

Hyperparameter Importance

각 하이퍼파라미터가 결과에 얼마나 영향을 미쳤는지를 나타냅니다.

중요도 순위

1. per_device_train_batch_size (43%)
2. gradient_accumulation_steps (36%)
3. learning_rate (21%)

MLFlow 결과 분석

다음으로는 MLFlow로 결과를 확인합니다.
오른쪽을 확인하면 mlruns.zip 파일이 생성된 것을 볼 수 있습니다.

오른쪽 옵션 버튼을 눌러 다운로드하고 압축을 풀어줍니다.

터미널로 압축 해제한 폴더로 이동한 후, 아래 명령어로 MLflow 웹 서버를 실행합니다.

mlflow ui --backend-store-uri mlruns --host 0.0.0.0 --port 5050

localhost:5050에 접속하여 실험이 정상적으로 기록되었는지 확인합니다.
study.optimize의 n_trials이 3으로 설정되어 있으므로 실험이 3번 진행된 것을 확인할 수 있습니다.

모든 Run을 선택하고 Compare 버튼을 눌러 비교해봅니다.

Visualization에서도 시각화된 결과를 확인할 수 있습니다.

X축의 옵션을 바꿔가며 확인해봅니다.

아래를 확인해보면 실험에 대한 정보들을 자세히 볼 수 있습니다.
각 실험들이 어떤 파라미터로 실험이 진행되었는 지, 결과(loss)가 어떻게 되었는 지, 학습 시간은 얼마나 걸렸는 지 확인할 수 있습니다.

결론

Optuna를 활용하면 짧은 시간 안에 효율적으로 최적의 하이퍼파라미터를 탐색할 수 있어, 실험 시간을 절감하면서도 높은 성능을 달성할 수 있습니다.

위 예시에서는 간단한 파라미터만 탐색하였지만, 더 넓은 범위의 하이퍼파라미터를 탐색하거나 다양한 모델 구조에 적용할수록 Optuna를 더 효율적으로 사용할 수 있습니다.