Trainer 기능

Trainer 기능 (Trainer features)

아래 각 레시피는 특정 Trainer 기능을 보여줘요: 커스텀 손실 함수, 메모리 효율적 평가, 체크포인트 전략 등등이요.

출처: 문서

본문

[!TIP] 여기에 보고 싶은 기능이나 워크플로가 있으면 issue를 열어주세요.

커스텀 손실 함수 (Custom loss function)

Trainer에 compute_loss_func을 전달하면 기본 손실 함수를 대체할 수 있어요. 이 함수는 forward pass 이후에 실행되며, 출력에서 손실을 어떻게 계산할지만 정의해요. forward pass 자체를 수정하려면 subclass로 compute_loss()를 재정의하세요.

커스텀 손실 함수는 다음 시그니처를 가져야 해요:

import torch.nn.functional as F

def my_loss_fn(outputs, labels, num_items_in_batch):
    logits = outputs["logits"]
    loss = F.cross_entropy(logits, labels, reduction="sum")
    return loss / num_items_in_batch
  • outputs는 원시 모델 출력이에요 (outputs.logits의 형상은 (batch, seq_len, vocab_size)).
  • labels는 forward pass 전에 Trainer가 입력 배치에서 pop한 토큰 id예요.
  • num_items_in_batch는 전체 누적 배치에서 예측 타겟의 개수예요. Causal LM 모델에서는 라벨 시프트 때문에 시퀀스의 위치 0마다 타겟이 없으므로 시프트된 라벨(labels[..., 1:])을 세요. 자세한 내용은 Loss scaling을 참고하세요. 커스텀 손실 함수를 제공하면 Trainer는 자동 손실 정규화를 건너뛰므로, 함수가 정규화를 직접 처리해야 해요.
trainer = Trainer(
    model=model,
    args=TrainingArguments(...),
    train_dataset=train_dataset,
    compute_loss_func=my_loss_fn,
)
trainer.train()

[!NOTE] compute_loss()를 재정의하는 더 많은 예시는 subclassing guide를 참고하세요.

시작 시 평가 (Evaluating on start)

eval_on_start=True로 설정하면 첫 번째 학습 스텝 전에 전체 평가 pass를 실행해요. 학습 전 평가는 특히 긴 실행에서 평가 파이프라인의 문제를 일찍 드러내요.

eval_on_start는 유효한 eval_strategy(예: "epoch")와 평가 데이터셋이 필요해요.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        eval_on_start=True,
    ),
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
)
trainer.train()

전체 평가는 시간이 추가되므로, 첫 실행이나 compute_metrics를 수정한 후에 가장 유용해요.

메모리 효율적 평가 (Memory-efficient evals)

평가 중에 Trainer는 모든 배치에서 forward pass를 실행하고 logits를 GPU의 단일 텐서로 연결해요. 평가 데이터셋을 완전히 처리하면 Trainer가 연결된 logits를 CPU로 옮기고 compute_metrics를 호출해요. 큰 모델이나 큰 평가 세트의 경우, 학습은 한 번에 배치 하나의 활성화만 보관하므로 같은 하드웨어에서 학습이 잘 되더라도 누적된 logits가 GPU 메모리를 고갈시킬 수 있어요.

eval_accumulation_steps

누적된 예측값을 n개의 배치마다 GPU에서 CPU로 오프로드해요. 값이 낮을수록 GPU 메모리가 줄어들지만 CPU 전송이 더 잦아져요.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        eval_accumulation_steps=16,   # move predictions to CPU every 16 batches
    ),
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
)
trainer.train()

preprocess_logits_for_metrics

평가 배치마다 GPU에서 forward pass 직후, logit 축적 전에 한 번 호출돼요. 반환값이 eval_pred.predictions의 logits를 대체해요. 배치 수준에서 계산을 실행하면 배치별 텐서 크기가 줄어들고, eval_accumulation_steps가 오프로드할 더 작은 텐서를 갖게 돼요.

import evaluate
from transformers import Trainer, TrainingArguments

metric = evaluate.load("accuracy")

def preprocess_logits_for_metrics(logits, labels):
    if isinstance(logits, tuple):
        logits = logits[0]
    return logits.argmax(dim=-1)

def compute_metrics(eval_preds):
    preds, labels = eval_preds
    labels = labels[:, 1:].reshape(-1)
    preds = preds[:, :-1].reshape(-1)
    return metric.compute(predictions=preds, references=labels)

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        eval_accumulation_steps=16,
    ),
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
    preprocess_logits_for_metrics=preprocess_logits_for_metrics,
)
trainer.train()

Dataloader 성능

기본적으로 Trainer는 dataloader_num_workers=0으로 dataloader를 만들어요. 데이터는 GPU가 유휴 상태인 동안 메인 프로세스에서 로드되므로, 배치 사이에 GPU 활용도가 낮게 나타나요.

dataloader_persistent_workers와 dataloader_prefetch_factor 둘 다 dataloader_num_workers > 0이 필요해요.

  • dataloader_persistent_workers는 에폭 사이에 워커 하위 프로세스를 유지해서 처음부터 다시 초기화하는 것을 피하지만, 메모리는 더 많이 사용해요.
  • dataloader_prefetch_factor는 각 워커가 몇 개의 배치를 미리 준비할지 제어해요. dataloader_prefetch_factor=2와 num_workers=4를 쓰면, GPU가 현재 배치에서 학습하는 동안 최대 8개 배치가 메모리에 상주해요.
from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="out",
    dataloader_num_workers=4,            # spawn 4 worker subprocesses
    dataloader_persistent_workers=True,  # keep them alive between epochs
    dataloader_prefetch_factor=2,        # each worker preloads 2 batches ahead
)

길이별 샘플 그룹화 (Group samples by length)

train_sampling_strategy="group_by_length"를 사용하면 길이가 비슷한 예시를 배치로 묶어 패딩을 줄여요. 사전 계산된 길이를 제공하지 않으면 Trainer는 각 데이터셋 항목의 첫 번째 모델 입력에서 길이를 추론해요. 이는 프로세서 기반 다중 모달 데이터셋이 BatchFeature 객체를 반환할 때도 작동하는데, 그것들이 매핑처럼 동작하는 특징 컨테이너이기 때문이에요.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="qwen3-vl-finetuned",
    train_sampling_strategy="group_by_length",
)

Dataset에 이미 사전 계산된 길이 열이 있으면 Trainer는 그 열을 대신 사용해요. 기본 열 이름은 length예요. 데이터셋이 다른 이름을 쓰면 length_column_name을 설정하세요. 이 전략은 길이를 아는 데이터셋이 필요하며 IterableDataset에서는 무시돼요.

배치 리밸런스 샘플링 (Batch rebalance sampling)

가변 길이 데이터셋에서 마이크로 배치 내부와 장치들 사이의 불균형은 장치들이 패딩에 시간을 낭비하고 그래디언트 동기화 단계에서 유휴 상태가 되게 해요.

TrainingArguments에서 train_sampling_strategy="batch_rebalance"를 설정하면 두 효과를 모두 줄일 수 있어요. 각 옵티마이저 스텝에서 샘플러는:

  1. 배치의 샘플을 길이별로 정렬해요.
  2. 정렬된 배치를 장치들에 걸쳐 샤딩해서 각 마이크로 배치의 패딩 토큰 비용이 균형을 이루게 해요. 긴 샘플의 마이크로 배치는 샘플이 더 적고, 짧은 샘플의 마이크로 배치는 샘플이 더 많아요.

이렇게 하면 각 마이크로 배치 내의 패딩이 줄어들고, 각 장치가 거의 같은 시점에 마이크로 배치를 끝내 동기화 시 유휴 시간이 줄어들며, "group_by_length"보다 최고 메모리가 낮아져요. 이 전략은 현재 데이터 병렬 학습에서만 지원돼요 (텐서 병렬화는 아직 지원되지 않음).

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        train_sampling_strategy="batch_rebalance",  # balance padding cost across devices
        per_device_train_batch_size=8,              # average samples per micro-batch
        length_column_name="length",                # optional: dataset column with precomputed lengths
    ),
    train_dataset=train_dataset,
)
trainer.train()

여기서 per_device_train_batch_size는 정확한 스텝별 수치가 아니라 평균이에요: 일부 마이크로 배치는 샘플이 더 적고 일부는 더 많지만, 스텝당 학습되는 총 샘플 수는 일반 분산 샘플링 전략과 동일하게 유지돼요.

샘플러는 모든 샘플의 길이를 알아야 배치를 정렬하고 균형을 맞출 수 있어요. 기본적으로 Trainer는 학습 시작 시 전체 데이터셋을 한 번 스캔해서 길이를 계산해요. 이 스캔을 건너뛰려면 전처리 중에 길이를 데이터셋 열로 미리 계산하고, 그 이름을 length_column_name(기본값 "length")으로 전달해요.

NEFTune

NEFTune은 forward pass 동안 토큰 임베딩에 무작위 노이즈를 추가해요. 노이즈는 정규화 역할을 하며 명령어 파인튜닝 성능을 향상시킬 수 있어요.

TrainingArguments에서 neftune_noise_alpha를 설정하면 NEFTune을 활성화해요. 일반적인 alpha 값은 5에서 15 사이예요.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        num_train_epochs=3,
        neftune_noise_alpha=5,
    ),
    train_dataset=train_dataset,
)
trainer.train()

NEFTune은 학습에만 영향을 주며, 학습 후에는 원래 임베딩 계층이 복원돼요.

로깅 (Logging)

logging_strategy, logging_steps, report_to로 Trainer가 로그 항목을 언제 어디에 쓰는지 제어해요.

  • logging_strategy="steps"는 logging_steps()마다 옵티마이저 갱신을 로깅해요. 대신 각 에폭 끝에 로깅하려면 "epoch"를 사용해요.
  • report_to는 로그를 Trackio 같은 실험 추적기로 스트리밍해요.
from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        logging_strategy="steps",
        logging_steps=50,               # write a log entry every 50 optimizer updates
        report_to="trackio",            # stream to Trackio (or "wandb", "tensorboard", …)
        run_name="model-experiment-v1", # display name in the tracker
    ),
    train_dataset=train_dataset,
)
trainer.train()

체크포인트 (Checkpointing)

Trainer는 save_steps()마다 옵티마이저 갱신 시 체크포인트를 저장하고, 모두(또는 가장 최근 ~TrainingArguments.save_total_limit개)를 보관해요.

save_strategy="best"는 메트릭에 따라 단일 최고 체크포인트만 보관해요. 추적하는 메트릭이 개선될 때만 새 체크포인트가 저장되므로 디스크 공간을 아끼고 오래된 체크포인트가 쌓이는 것을 방지해요.

from transformers import Trainer, TrainingArguments

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="out",
        eval_strategy="epoch",
        save_strategy="best",
        metric_for_best_model="perplexity",   # save when eval perplexity improves
        greater_is_better=False,              # lower perplexity is better
        load_best_model_at_end=True,          # load the best weights after training finishes
    ),
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,          # must return {"perplexity": ...}
)
trainer.train()

학습 재개 (Resume training)

학습이 중단되었고 진행 상황을 잃지 않고 재개하려면 train()에 resume_from_checkpoint=True를 전달해요. 학습은 output_dir의 최신 체크포인트에서 재개돼요.

trainer.train(resume_from_checkpoint=True)

특정 지점에서 재개하려면 체크포인트 경로를 지정하세요.

trainer.train(resume_from_checkpoint="out/checkpoint-1000")

재개할 때 Trainer는 옵티마이저 상태, 스케줄러 상태, RNG 상태를 복원해요.

체크포인트 재개는 체크포인트 디렉토리에 옵티마이저와 스케줄러 상태 파일이 필요해요. 그 파일들이 없으면(예: save_only_model=True일 때) 옵티마이저가 처음부터 다시 시작돼요.

JIT 체크포인트 (JIT checkpointing)

주기적 체크포인트(save_strategy="steps" 또는 "epoch")를 사용하면 마지막 저장된 체크포인트와 중단 사이의 모든 학습 진행을 잃어요. Kueue 같은 선점(preemptible) 워크로드가 있는 공유 클러스터에서는 작업이 언제든 종료될 수 있으므로, 그 간격은 수 시간의 낭비된 컴퓨팅을 의미할 수 있어요.

JIT(Just-In-Time) 체크포인트는 이 간격을 메워줘요. 트레이너가 SIGTERM 신호를 받으면 정확히 학습이 중단된 지점에서 체크포인트를 저장하므로, 진행 손실을 최소화하며 재개할 수 있어요. 이것은 주기적 체크포인트와 함께 작동해요. 주기적 저장은 크래시와 하드웨어 고장을 대비하고, JIT 저장은 선점과 정상 종료를 대비해요.

TrainingArguments에서 enable_jit_checkpoint=True로 설정해서 활성화해요.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="your-model",
    enable_jit_checkpoint=True,
)

SIGTERM을 받으면 Trainer는 현재 학습 스텝이 끝나기를 기다린 다음 체크포인트를 저장하고 학습을 정상적으로 중지해요. 저장이 시작될 때 센티널 파일(checkpoint-is-incomplete.txt)이 쓰여지고, 체크포인트가 완전히 쓰여지면 제거돼요. 체크포인트 디렉토리에 이 파일이 여전히 있으면 저장이 완료되기 전에 중단된 거예요. Trainer는 자동으로 확인하지 않으므로 재개 전에 직접 확인하세요.

다른 체크포인트와 같은 방식으로 JIT 체크포인트에서 재개해요.

trainer.train(resume_from_checkpoint=True)

[!WARNING] 오케스트레이터가 체크포인트가 완료될 충분한 시간을 허용하도록 구성해야 해요. 기본 Kubernetes 정상 종료 기간은 30초뿐이라 더 큰 모델에는 보통 부족해요.

Pod 또는 Job 스펙에서 terminationGracePeriodSeconds를 설정해요. 정확한 필드 위치는 트레이너에 따라 다릅니다 (Kubeflow Training Operator, Ray 등).

spec:
  template:
    spec:
      terminationGracePeriodSeconds: 300

sbatch 스크립트에서 --signal=TERM@<seconds>를 사용해 작업 시간 한도가 만료되기 전에 SIGTERM을 보내요.

#SBATCH --signal=TERM@300

필요한 유예 기간은 가장 긴 가능한 학습 스텝 시간에 체크포인트 저장 시간을 더하고, 체크포인트가 시작되기 전의 3초 kill_wait 지연을 더해 계산해요. 예를 들어 학습 스텝이 최대 2분, 체크포인트 저장이 2분 걸리면 최소 243초의 유예 시간을 설정하세요.