Trainer 功能

以下每个配方都演示了 Trainer 的一项特定功能：自定义损失函数、内存高效的评估、检查点策略等。

如果您希望在此处看到某个功能或工作流程，请提交一个 issue。

自定义损失函数

将 compute_loss_func 传递给 Trainer 以替换默认的 loss 函数。该函数在 forward pass 之后运行，仅定义如何从输出中计算 loss。若要修改 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 的形状为 (batch, seq_len, vocab_size)）。
labels 是 Trainer 在前向传播之前从输入批次中弹出的 token id。
num_items_in_batch 是整个累积批次中预测目标的总数。对于因果语言模型，它统计的是移位后的标签（labels[..., 1:]），因为标签移位会使每个序列的位置 0 没有对应的目标。有关详细信息，请参阅损失缩放。当提供了自定义损失函数时，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()

完整的评估会增加时间，因此它最适用于首次运行或在修改 compute_metrics 之后。

内存高效的评估

在评估期间，Trainer 会对每个批次执行一次前向传播，并将 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 上每个评估批次调用一次，时机为前向传播之后、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 的数据加载器。数据在主进程中加载，而 GPU 处于空闲状态，这会导致批次之间 GPU 利用率较低。

dataloader_persistent_workers 和 dataloader_prefetch_factor 均要求 dataloader_num_workers > 0。

dataloader_persistent_workers 会在 epoch 之间保持 worker 子进程存活，以避免从头重新初始化，代价是占用更多内存。
dataloader_prefetch_factor 控制每个 worker 提前准备多少个 batch。当 dataloader_prefetch_factor=2 且 num_workers=4 时，在 GPU 训练当前 batch 的同时，内存中最多会驻留 8 个 batch。

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" 以减轻这两种影响。对于每个优化器步骤，采样器会：

按长度对批次中的样本进行排序。
将排序后的批次分片到各个设备上，以平衡每个微批次的填充 token 开销。包含较长样本的微批次会分配较少的样本，而包含较短样本的微批次则会分配更多的样本。

这减少了每个微批次内的填充，并且每个设备大致在同一时间完成一个微批次，从而减少了同步时的空闲时间，并使峰值内存低于 "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 在前向传播过程中向 token 嵌入添加随机噪声。该噪声起到正则化的作用，可以提升指令微调的性能。

通过在 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_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 会还原优化器状态、调度器状态和随机数生成器状态。

恢复检查点需要检查点目录中包含优化器和调度器状态文件。如果这些文件缺失（例如，当 save_only_model=True 时），优化器将从头开始。

JIT 检查点

在使用定期检查点（save_strategy="steps" 或 "epoch"）时，您会丢失自上次保存的检查点以来到中断时刻之间的所有训练进度。在运行可抢占工作负载（例如 Kueue）的共享集群上，任务可能随时被终止，因此这段间隙可能导致数小时的计算资源被浪费。

JIT（即时）检查点机制弥补了这一差距。当训练器接收到 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)

您必须配置编排器，为检查点完成预留足够的时间。Kubernetes 默认的优雅关闭时间仅为 30 秒，对于较大的模型来说通常不够。

Kubernetes

Slurm

所需宽限期的计算方式为：最长可能的训练步耗时加上检查点保存耗时，再加上检查点开始前的 3 秒 kill_wait 延迟。例如，如果训练步最多耗时 2 分钟，保存检查点耗时 2 分钟，则至少应设置 243 秒的宽限期。

Update on GitHub

Transformers

Trainer 功能

自定义损失函数

在开始时进行评估

内存高效的评估

eval_accumulation_steps

preprocess_logits_for_metrics

Dataloader 性能

按长度对样本进行分组

批量重新平衡采样

NEFTune

日志

检查点

恢复训练

JIT 检查点