옵티마이저와 스케줄러

옵티마이저와 스케줄러

옵티마이저(optimizer)는 훈련 중 모델 가중치를 업데이트합니다. 스케줄러(scheduler)는 옵티마이저를 감싸고 각 훈련 스텝에서 학습률을 조정합니다. Trainer는 create_optimizer_and_scheduler()를 호출할 때 둘 다 생성합니다.

출처: 문서

본문

옵티마이저는 훈련 중 모델 가중치를 업데이트합니다. 스케줄러는 옵티마이저를 감싸고 각 훈련 스텝에서 학습률을 조정합니다. Trainer는 create_optimizer_and_scheduler()를 호출할 때 둘 다 생성합니다.

                                    ┌────────────┐         ┌──────────────┐
                                    │ Optimizer  │         │  Scheduler   │
                                    │ (adamw_torch_fused)◄─│  (linear)    │
                                    │            │         │              │
                                    │ param_groups         |              |
                                    │  └ lr       ◄────────┤              |
                                    │  └ weight_decay      │              │
                                    └──────┬─────┘         └──────────────┘
                                           │                      
  ┌──── EACH TRAINING STEP ───────────────────────────────────────────┐
  │                                        │                          │
  │   model(batch)                         │                          │
  │       │                                │                          │
  │       ▼                                │                          │
  │     loss ──► loss.backward() ──► param.grad                       │
  │                                        │                          │
  │                          ┌─────────────┘                          │
  │                          ▼                                        │
  │              optimizer.step()                                     │
  │                          │                                        │
  │                          ▼                                        │
  │                   param.data updated                              │
  │                          │                                        │
  │                          ▼                                        │
  │              lr_scheduler.step()  ──► recalculates lr             │
  │                          │            writes to optimizer         │
  │                          ▼            .param_groups['lr']         │
  │              model.zero_grad()                                    │
  │                                                                   │
  └───────────────────────────────────────────────────────────────────┘

옵티마이저와 스케줄러 동작(예: lr_scheduler_type()와 optim())은 TrainingArguments에서 구성합니다. 기본값(adamw_torch 옵티마이저와 linear 워밍업 스케줄러)은 대부분의 파인튜닝 실행에 좋은 시작점입니다.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    # Optimizer
    optim="adamw_torch",          # or "adamw_torch_fused", "adafactor", "sgd", etc.
    learning_rate=2e-5,
    weight_decay=0.01,
    adam_beta1=0.9,
    adam_beta2=0.999,
    adam_epsilon=1e-8,
    # Scheduler
    lr_scheduler_type="cosine",   # "linear", "cosine", "constant_with_warmup", etc.
    warmup_steps=500,
    lr_scheduler_kwargs={"num_cycles": 3},  # scheduler-specific extras
)

메트릭 기반 스케줄러

일부 스케줄러는 고정된 스케줄을 따르는 대신 훈련 동역학에 적응합니다.

GreedyLR은 평가 결과에서 학습률을 업데이트합니다. 메트릭이 계속 개선되면 학습률을 factor로 나누어 높이고, 메트릭이 개선되지 않으면 학습률을 factor로 곱하여 낮춥니다. 학습률이 min_lr에서 멈추고 reset_start 스텝 후에도 개선되지 않으면 GreedyLR은 초기 상태로 재설정하고 새 주기를 시작합니다.

GreedyLR은 훈련 중 평가가 필요합니다. eval_strategy를 "steps" 또는 "epoch"로 설정하세요.

args = TrainingArguments(
+   lr_scheduler_type="greedy",
+   lr_scheduler_kwargs={"patience": 10, "factor": 0.95, "min_lr": 1e-5},
+   eval_strategy="steps",
+   eval_steps=200,
    ...  # remaining args from the TrainingArguments intro config
)

[!TIP] 기본 mode="min"은 loss에 적합합니다. accuracy처럼 값이 높을수록 좋은 메트릭을 추적한다면 lr_scheduler_kwargs에서 "mode": "max"를 전달하세요.

구성 가능한 파라미터의 전체 목록은 GreedyLR 클래스를 참조하세요.

옵티마이저 통합

Transformers는 특수 훈련 시나리오를 위해 타사 옵티마이저를 통합합니다.

옵티마이저 설치 optim="value" 설명
APOLLO apollo-torch apollo_adamw 랜덤 프로젝션을 통한 메모리 효율적인 full-param; rank-1이면 충분
GrokAdamW grokadamw grokadamw 지연된 일반화(grokking)를 대상으로 함
LOMO / AdaLomo lomo-optim lomo / adalomo 저메모리 full-param 파인튜닝을 위해 그래디언트 + 업데이트 스텝을 퓨전
Schedule Free schedulefree schedule_free_adamw, schedule_free_radam, schedule_free_sgd LR 어닐링 제거; lr_scheduler_type="constant"와 함께 사용
GaLore galore-torch galore_adamw, galore_adafactor, galore_adamw_8bit 그래디언트 저랭크 프로젝션을 통한 full-parameter 학습
StableAdamW torch-optimi stable_adamw AdamW + AdaFactor 업데이트 클리핑; 그래디언트 클리핑 불필요
pip install apollo-torch

메모리 효율적인 LLM 최적화를 위한 근사 그래디언트 스케일링 (APOLLO)은 사전 훈련 및 파인튜닝 중 full-parameter 학습을 위한 메모리 효율적인 옵티마이저입니다. SVD 대신 저렴한 랜덤 프로젝션을 사용하여 SGD 수준의 메모리 비용으로 AdamW 성능을 맞춥니다. 극단적인 메모리 절약을 위해서는 rank-1 변형인 APOLLO-Mini를 사용하세요.

optim_target_modules 파라미터로 훈련할 레이어를 지정합니다.

args = TrainingArguments(
+   optim="apollo_adamw",
+   optim_target_modules=[r".*.attn.*", r".*.mlp.*"],
    ...  # remaining args from the TrainingArguments intro config
)

optim_args를 통해 추가 하이퍼파라미터를 전달합니다.

[!TIP] scale을 n/r로 설정하세요. 여기서 n은 원래 공간 차원, r은 저랭크 공간 차원입니다. scale을 기본값으로 유지하면서 학습률을 조정하면 비슷한 효과를 얻습니다.

파라미터 설명 APOLLO APOLLO-Mini
rank 그래디언트 스케일링을 위한 보조 부분공간의 rank 256 1
scale_type 스케일링 계수가 적용되는 방식 channel (채널별 스케일링) tensor (텐서별 스케일링)
scale 훈련을 안정화하기 위해 그래디언트 업데이트 조정 1.0 128
update_proj_gap 프로젝션 행렬 업데이트 전 스텝 수 200 200
proj 프로젝션 유형 random random

rank-1 구성으로 APOLLO-Mini를 활성화합니다.

args = TrainingArguments(
    optim="apollo_adamw",
    optim_target_modules=[r".*.attn.*", r".*.mlp.*"],
    optim_args="proj=random,rank=1,scale=128.0,scale_type=tensor,update_proj_gap=200",
    ...  # remaining args from the TrainingArguments intro config
)
pip install grokadamw

GrokAdamW는 모델이 느리게 변하는 그래디언트로 인해 지연된 일반화를 보이는 grokking을 대상으로 합니다.

args = TrainingArguments(
+   optim="grokadamw",
    ...  # remaining args from the TrainingArguments intro config
)
pip install lomo-optim

저메모리 최적화 (LOMO)는 저메모리 full-parameter 파인튜닝을 위한 두 가지 옵티마이저, LOMO와 AdaLomo를 포함합니다. 둘 다 그래디언트 계산과 파라미터 업데이트를 하나의 스텝으로 퓨전합니다. AdaLomo는 Adam과 유사하게 적응형 파라미터별 학습률을 추가합니다.

[!TIP] AdaLomo는 grad_norm 없이 가장 잘 작동하며 성능과 처리량을 향상시킵니다.

args = TrainingArguments(
+   optim="adalomo",
    learning_rate=2e-6,
    ...  # remaining args from the TrainingArguments intro config
)
pip install schedulefree

Schedule Free Optimizer (SFO)는 모멘텀을 평균화와 보간의 결합으로 대체하여 학습률 어닐링의 필요성을 완전히 제거합니다.

SFO는 RAdam(schedule_free_radam), AdamW(schedule_free_adamw), SGD(schedule_free_sgd) 옵티마이저를 지원합니다. RAdam 스케줄러는 warmup_steps가 필요하지 않습니다.

SFO를 lr_scheduler_type="constant"와 함께 사용하세요. 다른 스케줄러 유형도 작동하지만 SFO의 의도된 동작에 영향을 줍니다.

args = TrainingArguments(
+   optim="schedule_free_radam",
+   lr_scheduler_type="constant",
    learning_rate=2e-6,
    ...  # remaining args from the TrainingArguments intro config
)
pip install torch-optimi

StableAdamW는 AdaFactor의 업데이트 클리핑을 AdamW로 옮겨 그래디언트 클리핑의 필요성을 제거합니다. 그 외에는 AdamW의 대체재로 사용할 수 있습니다 (drop-in replacement).

[!TIP] 큰 배치 크기로 훈련하거나 여전히 loss 스파이크가 관찰된다면 beta_2를 0.95에서 0.99 사이로 설정해 보세요.

args = TrainingArguments(
+   optim="stable_adamw",
    learning_rate=2e-6,
    ...  # remaining args from the TrainingArguments intro config
)
pip install galore-torch trl

Gradient Low-Rank Projection (GaLore)은 LLM 훈련을 위한 메모리를 줄입니다. LoRA 같은 저랭크 적응 방식과 달리 GaLore는 full-parameter 학습을 보존합니다.

trl.SFTConfig에서 optim을 GaLore 옵티마이저("galore_adamw", "galore_adafactor", "galore_adamw_8bit")로 설정합니다. optim_target_modules로 대상 모듈을 지정하고 GaLore 전용 파라미터(rank, update_proj_gap, scale)를 optim_args로 지정합니다.

from trl import SFTConfig

args = SFTConfig(
    output_dir="./galore",
    max_steps=100,
    optim="galore_adamw",
    optim_target_modules=[r".*.attn.*", r".*.mlp.*"],
    optim_args="rank=64, update_proj_gap=100, scale=0.10",
)

레이어별 최적화를 위해 옵티마이저 이름에 _layerwise를 추가합니다("galore_adamw_layerwise"). GaLore가 대상으로 하는 linear 레이어만 저랭크 분해를 사용합니다. 다른 모든 레이어는 정상적으로 최적화됩니다.

from trl import SFTConfig, SFTTrainer

args = SFTConfig(
    output_dir="./galore",
    max_steps=100,
    optim="galore_adamw_layerwise",
    optim_target_modules=[r".*.attn.*", r".*.mlp.*"],
    optim_args="rank=64, update_proj_gap=100, scale=0.10",
)

레이어별 모드는 실험적입니다. 단일 GPU에서만 실행되며, DistributedDataParallel (DDP)을 지원하지 않고, 그래디언트 클리핑과 DeepSpeed가 작동하지 않을 수 있습니다.

옵티마이저와 스케줄러 커스터마이징

아직 통합되지 않은 옵티마이저를 사용하거나, 레이어별 학습률을 조정하거나, 커스텀 로직을 적용하려면 커스텀 옵티마이저와 스케줄러를 만드세요.

클래스와 kwargs 전달

~Trainer.optimizer_cls_and_kwargs는 파라미터 그룹화와 장치 배치를 Trainer에 위임하면서 커스텀 옵티마이저 클래스를 받아들입니다.

Trainer는 create_optimizer()가 실행될 때까지 옵티마이저 구축을 미루므로, 모델이 이미 올바른 장치에 있습니다.

import torch

trainer = Trainer(
    ...
    optimizer_cls_and_kwargs=(
        torch.optim.SGD,
        {"momentum": 0.9, "nesterov": True}
    ),
)

사전 빌드된 인스턴스 전달

사전 정의된 옵티마이저와 스케줄러를 ~Trainer.optimizers에 전달합니다. 사전 빌드된 인스턴스가 제공되면 Trainer는 create_optimizer()와 create_scheduler()를 건너뜁니다. 스케줄러를 전달하지 않으면 Trainer가 자동으로 하나 생성합니다.

[!WARNING] 모델을 올바른 장치에 배치한 후 옵티마이저를 구축하세요. 파라미터는 Trainer가 모델을 이동하기 전의 구성 시점에 해석됩니다. 분산 훈련에서 장치가 일치하지 않으면 조용히 잘못된 동작이 발생할 수 있습니다.

import torch
from transformers import Trainer, get_cosine_schedule_with_warmup

optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
scheduler = get_cosine_schedule_with_warmup(
    optimizer, num_warmup_steps=500, num_training_steps=10_000
)

trainer = Trainer(
    ...
    optimizers=(optimizer, scheduler),
)

사전 빌드된 인스턴스는 create_optimizer()와 create_scheduler()를 우회하므로, 자신만의 파라미터 그룹을 지정해야 합니다.

옵티마이저 및 스케줄러 메서드 재정의

전체 제어를 위해 create_optimizer()와 create_scheduler()를 서브클래스로 재정의합니다. 두 메서드 모두 train() 중 실행됩니다.

SchedulerType에 없는 OneCycleLR 같은 스케줄러를 사용하려면 create_scheduler()를 재정의하세요.

각 메서드에 대해 self에 할당하고 반환해야 합니다.

import torch
from transformers import Trainer

class MyTrainer(Trainer):

    def create_scheduler(self, num_training_steps, optimizer=None):
        optimizer = optimizer or self.optimizer
        self.lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
            optimizer,
            max_lr=0.1,
            total_steps=num_training_steps,
        )
        return self.lr_scheduler

기본 옵티마이저가 작동한다면 create_optimizer()를 재정의할 필요가 없습니다. 메서드를 super()로 확장하는 것이 완전히 교체하는 것보다 쉽습니다. 예를 들어 다른 모든 것을 유지하면서 추가 파라미터 그룹을 추가해 보세요.

class MyTrainer(Trainer):
    def create_optimizer(self, model=None):
        super().create_optimizer(model)  # builds the default two param groups
        # add extra param group
        self.optimizer.add_param_group({
            "params": self.model.classifier.parameters(),
            "lr": self.args.learning_rate * 10,
        })
        return self.optimizer

더 알아보기 (Learn more)