Transformers 文档
训练器
并获得增强的文档体验
开始使用
Trainer
Trainer 是一个完整的训练和评估循环，用于 Transformers 的 PyTorch 模型。将模型、预处理器、数据集和训练参数传递给 Trainer，让它处理其余部分，更快地开始训练。
Trainer 还由 Accelerate 提供支持，这是一个用于处理大型模型分布式训练的库。
本指南将向您展示 Trainer 的工作原理以及如何使用回调函数自定义它以适应您的用例。
!pip install accelerate --upgrade
Trainer 包含了训练循环的所有必要组件。
- 从训练步骤计算损失
- 使用 backward 方法计算梯度
- 根据梯度更新权重
- 重复直到达到预定的 epoch 数量
每次手动编码这个训练循环都可能很不方便，或者如果您刚开始接触机器学习，可能会成为一个障碍。 Trainer 抽象了这个过程，让您可以专注于模型、数据集和训练设计选择。
使用 TrainingArguments 中的超参数和选项来配置您的训练，它支持许多功能，例如分布式训练、torch.compile、混合精度训练以及将模型保存到 Hub。
TrainingArguments 中可用的参数数量起初可能会令人生畏。如果您想使用特定的超参数或功能，请尝试直接搜索它。否则，请随意从默认值开始，并在您熟悉训练过程后逐渐自定义它们。
下面的示例演示了一个 TrainingArguments 的例子，它会在每个 epoch 结束时评估并保存模型。它还会加载训练过程中找到的最佳模型并将其推送到 Hub。
from transformers import TrainingArguments
training_args = TrainingArguments(
output_dir="your-model",
learning_rate=2e-5,
per_device_train_batch_size=16,
per_device_eval_batch_size=16,
num_train_epochs=2,
weight_decay=0.01,
eval_strategy="epoch",
save_strategy="epoch",
load_best_model_at_end=True,
push_to_hub=True,
)将您的模型、数据集、预处理器和 TrainingArguments 传递给 Trainer，然后调用 train() 来开始训练。
有关训练过程的更完整概述，请参阅 微调 指南。
from transformers import Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=dataset["train"],
eval_dataset=dataset["test"],
processing_class=tokenizer,
data_collator=data_collator,
compute_metrics=compute_metrics,
)
trainer.train()检查点
Trainer 会将检查点（默认情况下不保存优化器状态）保存到 TrainingArguments 中 output_dir 目录下的一个名为 checkpoint-000 的子文件夹中。末尾的数字是保存检查点时的训练步数。
保存检查点有助于在遇到错误时恢复训练或挽回训练进度。将 resume_from_checkpoint 参数设置为 train() 中的值，以从最后一个检查点或特定检查点恢复训练。
trainer.train(resume_from_checkpoint=True)可以通过在 TrainingArguments 中设置 push_to_hub=True 将检查点保存到 Hub。默认方法（"every_save"）会在每次保存模型时（通常是训练结束时的最终模型）将一个检查点保存到 Hub。其他用于决定如何将检查点保存到 Hub 的选项包括：
hub_strategy="end"仅在调用 save_model() 时推送一个检查点hub_strategy="checkpoint"将最新的检查点推送到一个名为 _last-checkpoint_ 的子文件夹，可以从那里恢复训练hub_strategy="all_checkpoints"将所有检查点推送到 Hub，模型存储库中的每个子文件夹包含一个检查点
Trainer 在您从检查点恢复训练时，会尝试维护相同的 Python、NumPy 和 PyTorch RNG 状态。但 PyTorch 有各种非确定性设置，无法保证 RNG 状态完全相同。要实现完全确定性，请参阅 控制随机性源 指南，了解需要调整哪些设置才能使训练完全确定（某些设置可能会导致训练速度变慢）。
日志记录
Trainer 默认设置为 logging.INFO，用于报告错误、警告和其他基本信息。使用 log_level() 来更改日志级别和日志的详细程度。
下面的示例设置了主代码和模块使用相同的日志级别。
logger = logging.getLogger(__name__)
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
datefmt="%m/%d/%Y %H:%M:%S",
handlers=[logging.StreamHandler(sys.stdout)],
)
log_level = training_args.get_process_log_level()
logger.setLevel(log_level)
datasets.utils.logging.set_verbosity(log_level)
transformers.utils.logging.set_verbosity(log_level)
trainer = Trainer(...)在分布式环境中，Trainer 副本被设置为 logging.WARNING，只报告错误和警告。使用 log_level_replica() 来更改日志级别和日志的详细程度。要配置每个节点的日志级别，请使用 log_on_each_node() 来确定是使用每个节点的特定日志级别还是只使用主节点的日志级别。
使用 log_level 和 log_level_replica 的不同组合来配置每个节点上记录的内容。
my_app.py ... --log_level warning --log_level_replica error
日志级别在
__init__()方法中为每个节点单独设置。如果您在创建 Trainer 实例之前使用了其他 Transformers 功能，请考虑更早地设置它。
自定义
通过继承或覆盖其方法来定制 Trainer 以适应您的用例，以支持您要添加或使用的功能，而无需从头重写整个训练循环。下表列出了一些可自定义的方法。
| 方法 | description |
|---|---|
| get_train_dataloader() | 创建训练 DataLoader |
| get_eval_dataloader() | 创建评估 DataLoader |
| get_test_dataloader() | 创建测试 DataLoader |
| log() | 记录有关训练过程的信息 |
| create_optimizer_and_scheduler() | 创建优化器和学习率调度器（如果未在 __init__ 中传入，也可以通过 create_optimizer() 和 create_scheduler() 分别自定义） |
| compute_loss() | 计算训练输入批次的损失 |
| training_step() | 执行训练步骤 |
| prediction_step() | 执行预测和测试步骤 |
| evaluate() | 评估模型并返回评估指标 |
| predict() | 在测试集上进行预测（如果标签可用，则包含指标） |
例如，要使用加权损失，请在 Trainer 中重写 compute_loss()。
from torch import nn
from transformers import Trainer
class CustomTrainer(Trainer):
def compute_loss(self, model: nn.Module, inputs: dict[str, Union[torch.Tensor, Any]], return_outputs: bool = False, num_items_in_batch: Optional[torch.Tensor] = None):
labels = inputs.pop("labels")
# forward pass
outputs = model(**inputs)
logits = outputs.get("logits")
# compute custom loss for 3 labels with different weights
reduction = "sum" if num_items_in_batch is not None else "mean"
loss_fct = nn.CrossEntropyLoss(weight=torch.tensor([1.0, 2.0, 3.0], device=model.device, reduction=reduction))
loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
if num_items_in_batch is not None:
loss = loss / num_items_in_batch
return (loss, outputs) if return_outputs else loss回调函数
回调函数 是自定义 Trainer 的另一种方式，但它们不会改变训练循环内部的任何内容。相反，回调函数会检查训练循环的状态，并根据状态执行某些操作（提前停止、日志记录等）。例如，您无法使用回调函数实现自定义损失函数，因为这需要重写 compute_loss()。
要使用回调函数，请创建一个继承自 TrainerCallback 并实现您所需功能的类。然后将回调函数传递给 Trainer 的 callback 参数。下面的示例实现了一个提前停止回调函数，该函数在 10 步后停止训练。
from transformers import TrainerCallback, Trainer
class EarlyStoppingCallback(TrainerCallback):
def __init__(self, num_steps=10):
self.num_steps = num_steps
def on_step_end(self, args, state, control, **kwargs):
if state.global_step >= self.num_steps:
return {"should_training_stop": True}
else:
return {}
trainer = Trainer(
model=model,
args=training_args,
train_dataset=dataset["train"],
eval_dataset=dataset["test"],
processing_class=tokenizer,
data_collator=data_collator,
compute_metrics=compute_metrics,
callbacks=[EarlyStoppingCallback()],
)Accelerate
Accelerate 是一个简化在分布式环境和不同硬件上进行训练的库。它与 Trainer 的集成意味着 Trainer 支持分布式训练框架，如 Fully Sharded Data Parallel (FSDP) 和 DeepSpeed。
在 Trainer 的 Fully Sharded Data Parallel 指南中，了解有关 FSDP 分片策略、CPU 卸载等更多信息。
要将 Accelerate 与 Trainer 一起使用，请运行 accelerate_config 命令来配置您的训练环境。此命令创建一个 config_file.yaml 文件，该文件存储您的训练环境的配置设置，并且在您启动训练脚本时使用。下面显示了一些分布式训练配置示例。
compute_environment: LOCAL_MACHINE
distributed_type: MULTI_GPU
downcast_bf16: 'no'
gpu_ids: all
machine_rank: 0 #change rank as per the node
main_process_ip: 192.168.20.1
main_process_port: 9898
main_training_function: main
mixed_precision: fp16
num_machines: 2
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false运行 accelerate_launch 以使用 config_file.yaml 中设置的配置开始训练。此文件保存在 Accelerate 缓存文件夹中，并在您运行 accelerate_launch 时自动加载。
下面的示例使用前面显示的 FSDP 配置来启动 run_glue.py 脚本。config_file.yaml 文件中的参数也可以直接在命令行中设置。
accelerate launch \
./examples/pytorch/text-classification/run_glue.py \
--model_name_or_path google-bert/bert-base-cased \
--task_name $TASK_NAME \
--do_train \
--do_eval \
--max_seq_length 128 \
--per_device_train_batch_size 16 \
--learning_rate 5e-5 \
--num_train_epochs 3 \
--output_dir /tmp/$TASK_NAME/有关
accelerate_launch和自定义配置的更多信息，请参阅 启动您的 Accelerate 脚本 教程。
优化
Trainer 支持各种优化以提高训练性能 - 减少内存并提高训练速度 - 和模型性能。
torch.compile
torch.compile 可以显著加快训练速度并减少计算开销。在 TrainingArguments 中配置您的 torch.compile 设置。设置 torch_compile 为 True，并选择后端和编译模式。
from transformers import TrainingArguments
training_args = TrainingArguments(
torch_compile=True,
torch_compile_backend="inductor",
torch_compile_mode="default",
...,
)GaLore
Gradient Low-Rank Projection (GaLore) 显著减少了训练大型语言模型 (LLM) 时的内存使用。GaLore 的主要优点之一是全参数学习，不像低秩适应方法（如 LoRA）那样，它能产生更好的模型性能。
pip install galore-torch trl
选择一个 GaLore 优化器（"galore_adamw"、"galore_adafactor"、"galore_adamw_8bit”），并将其传递给 trl.SFTConfig 中的 optim 参数。使用 optim_target_modules 参数指定要适应的模块（可以是字符串列表、正则表达式或完整路径）。
GaLore 支持的额外参数，如 rank、update_proj_gap 和 scale，应传递给 trl.SFTConfig 中的 optim_args 参数。
下面的示例使用 SFTTrainer 启用了 GaLore，该训练器以正则表达式针对 attn 和 mlp 层。
训练可能需要一些时间才能开始（~3 分钟，针对 NVIDIA A100 上的 2B 模型）。
import datasets
from trl import SFTConfig, SFTTrainer
train_dataset = datasets.load_dataset('imdb', split='train')
args = SFTConfig(
output_dir="./test-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",
gradient_checkpointing=True,
)
trainer = SFTTrainer(
model="google/gemma-2b",
args=args,
train_dataset=train_dataset,
)
trainer.train()只有被视为 GaLore 层的线性层才能使用低秩分解进行训练。模型中其他层的优化方式保持不变。
Liger
Liger Kernel 是一系列层，如 RMSNorm、RoPE、SwiGLU、CrossEntropy、FusedLinearCrossEntropy 等，它们已集成到一个 Triton 内核中用于训练 LLM。这些内核也兼容 FlashAttention、FSDP 和 DeepSpeed。因此，Liger Kernel 可以提高多 GPU 训练吞吐量并减少内存使用。这对于多头训练以及支持更大的词汇量、更大的批次大小和更长的上下文长度非常有用。
pip install liger-kernel
通过在 TrainingArguments 中设置 use_liger_kernel=True 来启用 Liger Kernel 进行训练。这将用 Liger 的内核修补模型中相应的层。
Liger Kernel 支持 Llama、Gemma、Mistral 和 Mixtral 模型。请参阅 patching 列表以获取支持模型的最新列表。
from transformers import TrainingArguments
training_args = TrainingArguments(
output_dir="your-model",
learning_rate=2e-5,
per_device_train_batch_size=16,
per_device_eval_batch_size=16,
num_train_epochs=2,
weight_decay=0.01,
eval_strategy="epoch",
save_strategy="epoch",
load_best_model_at_end=True,
push_to_hub=True,
use_liger_kernel=True
)您还可以通过 liger_kernel_config 参数配置要应用的特定内核。此字典作为关键字参数传递给 _apply_liger_kernel_to_instance 函数，从而实现对内核使用的精细控制。可用选项因模型而异，但通常包括：rope、swiglu、cross_entropy、fused_linear_cross_entropy、rms_norm 等。
from transformers import TrainingArguments
# Apply only specific kernels
training_args = TrainingArguments(
output_dir="your-model",
learning_rate=2e-5,
per_device_train_batch_size=16,
per_device_eval_batch_size=16,
num_train_epochs=2,
weight_decay=0.01,
eval_strategy="epoch",
save_strategy="epoch",
load_best_model_at_end=True,
push_to_hub=True,
use_liger_kernel=True,
liger_kernel_config={
"rope": True,
"cross_entropy": True,
"rms_norm": False, # Don't apply Liger's RMSNorm kernel
"swiglu": True,
}
)NEFTune
NEFTune 在训练期间向嵌入向量添加噪声以提高模型性能。在 Trainer 中使用 TrainingArguments 中的 neftune_noise_alpha 参数启用它，以控制添加的噪声量。
from transformers import TrainingArguments, Trainer
training_args = TrainingArguments(..., neftune_noise_alpha=0.1)
trainer = Trainer(..., args=training_args)训练后会恢复原始嵌入层，以避免任何意外行为。
在 GitHub 上更新