Trainer 기능

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

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

사용자 정의 손실 함수

compute_loss_func을 Trainer에 전달하여 기본 손실 함수를 대체할 수 있습니다. 이 함수는 forward pass 이후에 실행되며, 출력에서 손실을 계산하는 방식만 정의합니다. forward pass 자체를 수정하려면 대신 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의 shape은 (batch, seq_len, vocab_size)입니다).
labels은 forward pass 전에 Trainer가 입력 배치에서 추출한 토큰 ID입니다.
num_items_in_batch는 누적된 전체 배치에 걸친 예측 대상의 개수입니다. Causal LM 모델의 경우, 라벨 시프트로 인해 모든 시퀀스의 0번째 위치에 예측 대상이 없으므로 시프트된 라벨(labels[..., 1:])을 카운트합니다. 자세한 내용은 Loss scaling을 참조하세요. Trainer는 커스텀 loss 함수가 제공된 경우 자동 loss 정규화를 건너뛰므로, 해당 함수에서 정규화를 직접 처리해야 합니다.

trainer = Trainer(
    model=model,
    args=TrainingArguments(...),
    train_dataset=train_dataset,
    compute_loss_func=my_loss_fn,
)
trainer.train()

compute_loss()를 오버라이드하는 추가 예제는 서브클래스 가이드를 참고하세요.

시작 시 평가

eval_on_start=True를 설정하면 첫 번째 학습 단계 전에 전체 평가 패스를 실행합니다. 사전 학습 평가는 특히 긴 실행 과정에서 평가 파이프라인의 문제를 조기에 발견하는 데 도움이 됩니다.

eval_on_start에는 유효한 eval_strategy(예: "epoch")과 eval 데이터셋이 필요합니다.

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()

전체 eval은 시간이 더 소요되므로, 최초 실행 시나 compute_metrics를 수정한 후에 가장 유용합니다.

메모리 효율적인 평가

평가 중 Trainer는 모든 배치에 대해 forward pass를 실행하고, logits를 GPU에서 단일 tensor로 연결합니다. eval 데이터셋의 처리가 모두 완료되면 Trainer는 연결된 logits를 CPU로 이동시키고 compute_metrics를 호출합니다. 대규모 모델이나 eval 세트의 경우, 동일한 하드웨어에서 학습은 정상적으로 동작하더라도 누적된 logits가 GPU 메모리를 소진할 수 있습니다. 학습은 한 번에 하나의 배치에 해당하는 activations만 유지하기 때문입니다.

eval_accumulation_steps

누적된 예측 결과를 n개 배치마다 GPU에서 CPU로 offload합니다. 값이 낮을수록 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에서 eval 배치당 한 번씩, 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
)

길이에 따라 샘플 그룹화

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에는 적용되지 않습니다.

배치 재균형 샘플링

가변 길이 데이터셋의 경우, 마이크로 배치 내부 및 디바이스 간 불균형으로 인해 디바이스가 패딩에 시간을 낭비하고 그래디언트 동기화 단계에서 유휴 상태가 됩니다.

TrainingArguments에서 train_sampling_strategy="batch_rebalance"를 설정하여 두 가지 효과를 모두 줄일 수 있습니다. 각 optimizer step마다 sampler는 다음과 같이 동작합니다:

배치의 샘플을 길이 기준으로 정렬합니다.
정렬된 배치(shard)를 디바이스에 분배하여 각 마이크로 배치의 패딩 토큰 비용을 균형 있게 맞춥니다. 긴 샘플이 포함된 마이크로 배치에는 샘플 수가 적게, 짧은 샘플이 포함된 마이크로 배치에는 샘플 수가 많게 배정됩니다.

이 방식은 각 micro-batch 내부의 padding을 줄이고, 각 디바이스가 micro-batch를 거의 동시에 완료하도록 하여 동기화 시의 유휴 시간을 줄이고 "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는 여기서는 정확한 단계별 개수가 아니라 평균값입니다. 일부 마이크로 배치에는 샘플이 적고 일부에는 많지만, 단계별로 학습되는 샘플의 총 수는 일반적인 분산 샘플링 전략과 동일하게 유지됩니다.

sampler는 배치의 정렬 및 균형을 위해 모든 샘플의 길이가 필요합니다. 기본적으로 Trainer는 학습 시작 시 전체 데이터셋을 한 번 스캔하여 이러한 길이를 계산합니다. 이 스캔을 건너뛰려면, 예컨대 전처리 단계에서 길이를 데이터셋 컬럼에 미리 계산해 두고 그 컬럼 이름을 length_column_name으로 전달하세요(기본값은 "length").

NEFTune

NEFTune는 forward pass 동안 토큰 임베딩에 무작위 노이즈를 추가합니다. 이 노이즈는 정규화 역할을 하며, 지시 fine-tuning의 성능을 향상시킬 수 있습니다.

neftune_noise_alpha를 TrainingArguments에 설정하여 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_strategy, logging_steps, report_to를 사용하여 Trainer가 로그 항목을 기록하는 시점과 위치를 제어합니다.

logging_strategy="steps"은 logging_steps()마다 옵티마이저 업데이트를 기록합니다. 대신 각 epoch이 끝날 때 기록하려면 "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()

체크포인트

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()

학습 재개

학습이 중단된 경우, 진행 상황을 잃지 않고 이어서 진행하려면 train()에 resume_from_checkpoint=True를 전달하세요. 학습은 output_dir의 최신 체크포인트에서 재개됩니다.

trainer.train(resume_from_checkpoint=True)

특정 시점부터 재개하려면 체크포인트 경로를 지정합니다.

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

재개 시 Trainer는 옵티마이저 상태, 스케줄러 상태, RNG 상태를 복원합니다.

Checkpoint 재개에는 checkpoint 디렉터리에 optimizer와 scheduler 상태 파일이 필요합니다. 해당 파일이 없는 경우(예: save_only_model=True) optimizer는 처음부터 다시 시작됩니다.

JIT 체크포인트

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

JIT(Just-In-Time) checkpointing은 이러한 격차를 해소합니다. trainer가 SIGTERM 신호를 수신하면, 학습이 중단된 정확한 시점의 checkpoint를 저장하여 진행 상황의 손실을 최소화한 상태로 재개할 수 있습니다. 주기적 checkpointing과 함께 작동합니다. 주기적 저장은 크래시 및 하드웨어 장애에 대비하는 반면, JIT 저장은 선점(preemption) 및 정상 종료에 대비합니다.

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)

오케스트레이터에 체크포인트가 완료될 때까지 충분한 시간을 허용하도록 구성해야 합니다. Kubernetes의 기본 graceful shutdown 기간은 30초로, 대형 모델의 경우 일반적으로 충분하지 않습니다.

Kubernetes

Slurm

필요한 grace period는 가장 긴 학습 단계 시간과 checkpoint 저장 시간, 그리고 checkpoint 시작 전 3초의 kill_wait 지연 시간을 더한 값으로 계산합니다. 예를 들어, 학습 단계에 최대 2분, checkpoint 저장에 2분이 걸린다면, 최소 243초의 grace time을 설정해야 합니다.

Update on GitHub

Transformers

Trainer 기능

사용자 정의 손실 함수

시작 시 평가

메모리 효율적인 평가

eval_accumulation_steps

preprocess_logits_for_metrics

Dataloader 성능

길이에 따라 샘플 그룹화

배치 재균형 샘플링

NEFTune

로깅

체크포인트

학습 재개

JIT 체크포인트