Trainerの機能

以下の各レシピは、Trainerの特定の機能（カスタム損失関数、メモリ効率の良い評価、チェックポイント戦略など）をデモンス

ここで見てみたい機能やワークフローがある場合は、issueを作成してください。

カスタム損失関数

compute_loss_func を Trainer に渡すことで、デフォルトの損失関数を置き換えることができます。この関数はフォワードパスの後に実行され、出力から損失をどのように計算するかを定義するだけです。フォワードパス自体を変更する場合は、代わりに 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 は、Trainer がフォワードパスの前に入力バッチから取り出したトークン 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()

compute_loss() をオーバーライドするその他の例については、サブクラス化ガイドを参照してください。

開始時の評価

eval_on_start=True を設定すると、最初の学習ステップの前に完全な評価パスが実行されます。事前学習時の評価により、評価パイプラインの問題を早期に発見できるため、長時間の実行時特に有効です。

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

完全なevalは実行時間を増やすため、compute_metricsを変更した直後や、初回の実行時に最も役立ちます。

メモリ効率の良い評価

評価中、Trainer はすべてのバッチに対してフォワードパスを実行し、logits を GPU 上の単一のテンソルに連結します。評価データセットの処理が完了すると、Trainer は連結された logits を CPU に移動し、compute_metrics を呼び出します。大規模なモデルや評価セットの場合、同じハードウェアでトレーニングが正常に動作していても、トレーニングでは一度に 1 バッチ分のアクティベーションのみを保持するため、蓄積された 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上で、eval バッチごとに1回、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" を設定すると、両方の影響を軽減できます。各オプティマイザステップにおいて、サンプラーは以下を行います:

バッチのサンプルを長さ順にソートします。
ソート済みのバッチをデバイス間で分割し、各マイクロバッチにおけるパディングトークンのコストを均一にします。長いサンプルを含むマイクロバッチにはサンプル数を少なく、短いサンプルを含むマイクロバッチにはサンプル数を多く割り当てます。

これにより、各マイクロバッチ内のパディングが削減され、各デバイスがほぼ同時にマイクロバッチの処理を完了するため、同期時のアイドル時間が短縮され、"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 はここでは平均値であり、ステップごとの正確なサンプル数ではありません。マイクロバッチごとにサンプル数が少なくなったり多くなったりしますが、1ステップあたりに学習されるサンプルの総数は、通常の分散サンプリング戦略の場合と同じです。

サンプラーは、バッチのソートとバランス調整のために、各サンプルの長さが必要です。デフォルトでは、Trainer は学習開始時にデータセット全体を一度スキャンして、これらの長さを計算します。このスキャンをスキップするには、長さを事前に計算してデータセットの列に格納し（例えば前処理の段階で）、その列名を length_column_name として渡します（デフォルトは "length"）。

NEFTune

NEFTuneは、フォワードパス中にトークン埋め込みにランダムなノイズを追加します。このノイズは正則化として機能し、指示ファインチューニングの性能向上に寄与します。

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" を使用してください。
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の状態を復元します。

チェックポイントからの再開には、チェックポイントディレクトリ内にオプティマイザーとスケジューラーの状態ファイルが必要です。これらのファイルが存在しない場合（save_only_model=Trueの場合など）、オプティマイザーは最初からやり直されます。

JIT チェックポイント

定期的なチェックポイント保存（save_strategy="steps" または "epoch"）を使用すると、最後に保存されたチェックポイントと中断の間の学習進捗が失われます。Kueue などのプリエンプティブなワークロードが利用可能な共有クラスタでは、ジョブがいつでも終了する可能性があるため、そのギャップは数時間の計算リソースの無駄を意味し得ます。

JIT（Just-In-Time）チェックポイント機能はこのギャップを埋めます。トレーナーが SIGTERM シグナルを受け取ると、トレーニングが中断された正確な時点でのチェックポイントを保存するため、進捗の損失を最小限に抑えて再開できます。この機能は、定期的なチェックポイント保存と併用して動作します。定期的な保存はクラッシュやハードウェア障害から保護し、JIT による保存はプリエンプションや正常なシャットダウンから保護します。

enable_jit_checkpoint=True を設定して有効にします。設定先は TrainingArguments です。

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 は、最長で考えられる training step の時間と checkpoint の保存時間、および checkpoint が開始される前の 3 秒の kill_wait 遅延の合計として計算します。例えば、training step に最大 2 分、checkpoint の保存に 2 分かかる場合、少なくとも 243 秒の grace time を設定してください。

Update on GitHub

Transformers

Trainerの機能

カスタム損失関数

開始時の評価

メモリ効率の良い評価

eval_accumulation_steps

preprocess_logits_for_metrics

Dataloaderの性能

長さごとにサンプルをグループ化

バッチ再バランスサンプリング

NEFTune

ログ出力

チェックポイント

トレーニングの再開

JIT チェックポイント