Transformers 文档
训练器
并获得增强的文档体验
开始使用
训练器
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 的子文件夹中。末尾的数字是保存检查点时的训练步数。
保存检查点对于恢复训练或在遇到错误时恢复训练进度很有用。在 train() 中设置 resume_from_checkpoint 参数以从上次检查点或特定检查点恢复训练。
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 以适应您的用例，以支持您想要添加或使用的功能，而无需从头开始重写整个训练循环。下表列出了一些可以自定义的方法。
| 方法 | 描述 |
|---|---|
| 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, inputs, return_outputs=False, num_items_in_batch=None):
labels = inputs.pop("labels")
# forward pass
outputs = model(**inputs)
logits = outputs.get("logits")
# compute custom loss for 3 labels with different weights
loss_fct = nn.CrossEntropyLoss(weight=torch.tensor([1.0, 2.0, 3.0], device=model.device))
loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
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 支持分布式训练框架，如 完全分片数据并行 (FSDP) 和 DeepSpeed。
要将 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: falsecompute_environment: LOCAL_MACHINE
tp_config:
tp_size: 4
distributed_type: TP
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 4
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/ \
--overwrite_output_dir请参阅 启动您的 Accelerate 脚本 教程，以详细了解 accelerate_launch 和自定义配置。
优化
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
梯度低秩投影 (GaLore) 显着减少了训练大型语言模型 (LLM) 时的内存使用量。与 LoRA 等低秩自适应方法不同，GaLore 的主要优势之一是全参数学习，它可以产生更好的模型性能。
pip install galore-torch trl datasets
选择 GaLore 优化器（"galore_adamw"、"galore_adafactor"、"galore_adamw_8bit”) 并将其传递给 TrainingArguments 中的 optim 参数。使用 optim_target_modules 参数指定要适配的模块（可以是字符串列表、正则表达式或完整路径）。
GaLore 支持的额外参数 rank、update_proj_gap 和 scale 应传递给 TrainingArguments 中的 optim_args 参数。
下面的示例使用 SFTTrainer 启用 GaLore，该 SFTTrainer 以正则表达式为目标 attn 和 mlp 层。
训练开始可能需要一些时间（在 NVIDIA A100 上训练 2B 模型大约需要 3 分钟）。
import torch
import datasets
import trl
from transformers import TrainingArguments, AutoConfig, AutoTokenizer, AutoModelForCausalLM
train_dataset = datasets.load_dataset('imdb', split='train')
args = TrainingArguments(
output_dir="./test-galore",
max_steps=100,
per_device_train_batch_size=2,
optim="galore_adamw",
optim_target_modules=[r".*.attn.*", r".*.mlp.*"],
optim_args="rank=64, update_proj_gap=100, scale=0.10",
)
config = AutoConfig.from_pretrained("google/gemma-2b")
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
model = AutoModelForCausalLM.from_config("google/gemma-2b").to(0)
trainer = trl.SFTTrainer(
model=model,
args=args,
train_dataset=train_dataset,
dataset_text_field='text',
max_seq_length=512,
)
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 模型。有关支持模型的最新列表，请参阅 修补 列表。
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
)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 上更新