Transformers 文档

训练器

Transformers

加入 Hugging Face 社区

并获得增强型文档体验

协作模型、数据集和空间

通过加速推理获得更快的示例

在文档主题之间切换

开始使用

训练器

Transformers 库中的 Trainer 是一个用于 PyTorch 模型的完整训练和评估循环。您只需要传入训练所需的组件（模型、分词器、数据集、评估函数、训练超参数等），Trainer 类将负责其他所有操作。这使得开始训练变得更加容易，无需手动编写自己的训练循环。但同时，Trainer 非常可定制，并提供大量的训练选项，以便您可以根据您的具体训练需求进行调整。

除了 Trainer 类，Transformers 还提供了一个用于序列到序列任务（如翻译或摘要）的 Seq2SeqTrainer 类。此外还有来自 TRL 库的 SFTTrainer 类，它包装了 Trainer 类，并针对使用自回归技术训练 Llama-2 和 Mistral 等语言模型进行了优化。 SFTTrainer 还支持序列打包、LoRA、量化和 DeepSpeed 等功能，以便有效地扩展到任何模型规模。

请随时查看 API 参考，了解这些其他 Trainer 类型类，以了解更多关于何时使用哪个类的信息。一般来说，Trainer 是最通用的选项，适用于各种任务。 Seq2SeqTrainer 专为序列到序列任务而设计，SFTTrainer 专为训练语言模型而设计。

在开始之前，请确保已安装 Accelerate - 一个用于跨分布式环境启用和运行 PyTorch 训练的库。

pip install accelerate

# upgrade
pip install accelerate --upgrade

本指南概述了 Trainer 类。

基本用法

Trainer 包括您在基本训练循环中会找到的所有代码

执行训练步骤以计算损失
使用 backward 方法计算梯度
根据梯度更新权重
重复此过程，直到达到预定的迭代次数

Trainer 类将所有这些代码抽象化，因此您不必每次都手动编写训练循环，或者如果您只是刚开始使用 PyTorch 和训练。您只需要提供训练所需的必要组件，例如模型和数据集，Trainer 类将处理其他所有操作。

如果您想要指定任何训练选项或超参数，可以在 TrainingArguments 类中找到它们。例如，让我们定义在 output_dir 中保存模型的位置，并在训练后使用 push_to_hub=True 将模型推送到 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,
)

将 training_args 传递给 Trainer，以及模型、数据集、用于预处理数据集的内容（根据您的数据类型，它可能是分词器、特征提取器或图像处理器）、数据整理器和用于计算您想要在训练期间跟踪的指标的函数。

最后，调用 train() 开始训练！

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

检查点

Trainer 类会将您的模型检查点保存到在 TrainingArguments 的 output_dir 参数中指定的目录。您将在一个名为 checkpoint-000 的子文件夹中找到保存的检查点，其中末尾的数字对应于训练步骤。保存检查点对于以后恢复训练很有用。

# resume from latest checkpoint
trainer.train(resume_from_checkpoint=True)

# resume from specific checkpoint saved in output directory
trainer.train(resume_from_checkpoint="your-model/checkpoint-1000")

您可以通过在 TrainingArguments 中设置 push_to_hub=True 来将您的检查点（优化器状态默认情况下不会保存）推送到 Hub，以提交和推送它们。在 hub_strategy 参数中设置了其他决定如何保存检查点的选项

hub_strategy="checkpoint" 会将最新的检查点推送到名为“last-checkpoint”的子文件夹，您可以从该文件夹恢复训练
hub_strategy="all_checkpoints" 会将所有检查点推送到在 output_dir 中定义的目录（您将在模型库中看到一个检查点对应一个文件夹）

当您从检查点恢复训练时，Trainer 会尝试使 Python、NumPy 和 PyTorch RNG 状态与保存检查点时的状态相同。但由于 PyTorch 具有各种非确定性默认设置，因此 RNG 状态不能保证相同。如果您想要启用完全确定性，请查看控制随机性来源指南，了解您可以启用什么来使您的训练完全确定性。请记住，通过使某些设置确定性，训练可能会变慢。

自定义训练器

虽然 Trainer 类旨在易于访问和使用，但它也为更有冒险精神的用户提供了很大的可定制性。许多 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() 在测试集上进行预测（如果存在标签，则具有指标）

例如，如果您想自定义 compute_loss() 方法以使用加权损失。

from torch import nn
from transformers import Trainer

class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        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() 方法才能做到这一点。

例如，如果您想要在 10 步后将提前停止回调添加到训练循环中。

from transformers import TrainerCallback

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 的 callback 参数。

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    callback=[EarlyStoppingCallback()],
)

日志

查看日志 API 参考以获取有关不同日志级别的更多信息。

默认情况下，训练器设置为 logging.INFO，它报告错误、警告和其他基本信息。在分布式环境中，训练器副本设置为 logging.WARNING，它只报告错误和警告。您可以使用 log_level 和 log_level_replica 参数在训练参数中更改日志级别。

要为每个节点配置日志级别设置，请使用 log_on_each_node 参数来确定是在每个节点上使用日志级别还是仅在主节点上使用日志级别。

训练器在 Trainer.__init__() 方法中为每个节点分别设置日志级别，因此如果您在创建训练器对象之前使用其他 Transformers 功能，您可能需要考虑更早地设置它。

例如，要设置您的主代码和模块以根据每个节点使用相同的日志级别

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

使用 log_level 和 log_level_replica 的不同组合来配置每个节点上的日志内容。

单节点

多节点

NEFTune

NEFTune 是一种可以通过在训练期间向嵌入向量添加噪声来提高性能的技术。要在训练器中启用它，请在训练参数中设置 neftune_noise_alpha 参数以控制添加的噪声量。

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(..., neftune_noise_alpha=0.1)
trainer = Trainer(..., args=training_args)

NEFTune 在训练结束后被禁用以恢复原始嵌入层，以避免任何意外行为。

GaLore

梯度低秩投影 (GaLore) 是一种内存高效的低秩训练策略，它允许全参数学习，但比常见的低秩自适应方法（如 LoRA）更节省内存。

首先确保安装 GaLore 官方仓库

pip install galore-torch

然后只需在 optim 中添加 ["galore_adamw", "galore_adafactor", "galore_adamw_8bit"] 之一，以及 optim_target_modules，它可以是与您要自适应的目标模块名称相对应的字符串列表、正则表达式或完整路径。下面是一个端到端的示例脚本（确保执行 pip install trl datasets）

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.*"]
)

model_id = "google/gemma-2b"

config = AutoConfig.from_pretrained(model_id)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_config(config).to(0)

trainer = trl.SFTTrainer(
    model=model, 
    args=args,
    train_dataset=train_dataset,
    dataset_text_field='text',
    max_seq_length=512,
)

trainer.train()

要传递 GaLore 支持的额外参数，您应该正确传递 optim_args，例如

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",
)

model_id = "google/gemma-2b"

config = AutoConfig.from_pretrained(model_id)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_config(config).to(0)

trainer = trl.SFTTrainer(
    model=model, 
    args=args,
    train_dataset=train_dataset,
    dataset_text_field='text',
    max_seq_length=512,
)

trainer.train()

您可以在原始仓库或论文中阅读有关该方法的更多信息。

目前，您只能训练被视为 GaLore 层的线性层，这些层将使用低秩分解进行训练，而其他层将以传统方式进行优化。

请注意，在开始训练之前需要一些时间（对于在 NVIDIA A100 上运行的 2B 模型，大约需要 3 分钟），但之后的训练应该会顺利进行。

您还可以通过在优化器名称后面添加 layerwise 来执行逐层优化，如下所示

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_layerwise",
    optim_target_modules=[r".*.attn.*", r".*.mlp.*"]
)

model_id = "google/gemma-2b"

config = AutoConfig.from_pretrained(model_id)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_config(config).to(0)

trainer = trl.SFTTrainer(
    model=model, 
    args=args,
    train_dataset=train_dataset,
    dataset_text_field='text',
    max_seq_length=512,
)

trainer.train()

请注意，逐层优化有点实验性，不支持 DDP（分布式数据并行），因此您只能在单个 GPU 上运行训练脚本。有关更多详细信息，请参阅此相应部分。其他功能（如梯度裁剪、DeepSpeed 等）可能无法开箱即用地支持。如果您遇到此类问题，请在 GitHub 上提出问题。

Liger 内核

Liger 内核内核是 LinkedIn 开发的一组 Triton 内核，专为 LLM 训练而设计。我们已经实现了与 Hugging Face 兼容的 RMSNorm、RoPE、SwiGLU、CrossEntropy、FusedLinearCrossEntropy 等，并将陆续推出更多内核。它可以有效地将多 GPU 训练吞吐量提高 20%，并将内存使用量减少 60%。该内核与 flash attention、PyTorch FSDP 和 Microsoft DeepSpeed 开箱即用。

在 LLaMA 3-8B 模型训练中，吞吐量提高 +20%，内存使用量减少 60%。实现更长的上下文长度和更大的批次大小。如果您想将模型扩展到多头训练或大词汇量，它也很有用。释放多头训练（美杜莎）等更多功能。在 [Liger](https://github.com/linkedin/Liger-Kernel/tree/main/examples) 中查看详细信息和示例

首先确保安装 Liger 官方仓库

pip install liger-kernel

您应该传递 use_liger_kernel=True 以将 liger 内核应用于您的模型，例如

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
)

该内核支持 Llama、Gemma、Mistral 和 Mixtral 模型架构。在此处可以找到支持模型的最新列表。当 use_liger_kernel 设置为 True 时，原始模型中的相应层将用 Liger 的高效实现进行修补，因此您无需执行任何额外操作，只需设置参数值即可。

LOMO 优化器

LOMO 优化器在使用有限资源对大型语言模型进行全参数微调和 AdaLomo：具有自适应学习率的低内存优化中被引入。它们都包含一种高效的全参数微调方法。这些优化器将梯度计算和参数更新融合到一个步骤中，以减少内存使用量。LOMO 支持的优化器为 "lomo" 和 "adalomo"。首先，从 pypi 安装 LOMO pip install lomo-optim 或从源代码安装 pip install git+https://github.com/OpenLMLab/LOMO.git。

根据作者的说法，建议使用 AdaLomo，不使用 grad_norm，以获得更好的性能和更高的吞吐量。

下面是一个简单的脚本，演示如何在全精度下对 google/gemma-2b 在 IMDB 数据集上进行微调

import torch
import datasets
from transformers import TrainingArguments, AutoTokenizer, AutoModelForCausalLM
import trl

train_dataset = datasets.load_dataset('imdb', split='train')

args = TrainingArguments(
    output_dir="./test-lomo",
    max_steps=1000,
    per_device_train_batch_size=4,
    optim="adalomo",
    gradient_checkpointing=True,
    logging_strategy="steps",
    logging_steps=1,
    learning_rate=2e-6,
    save_strategy="no",
    run_name="lomo-imdb",
)

model_id = "google/gemma-2b"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, low_cpu_mem_usage=True).to(0)

trainer = trl.SFTTrainer(
    model=model, 
    args=args,
    train_dataset=train_dataset,
    dataset_text_field='text',
    max_seq_length=1024,
)

trainer.train()

GrokAdamW 优化器

GrokAdamW 优化器旨在提高训练性能和稳定性，特别是对于从 grokking 信号函数中受益的模型。要使用 GrokAdamW，首先使用 pip install grokadamw 安装优化器包。

GrokAdamW 对于需要高级优化技术才能获得更好的性能和稳定性的模型特别有用。

下面是一个简单的脚本，演示如何使用 GrokAdamW 优化器对 google/gemma-2b 在 IMDB 数据集上进行微调

import torch
import datasets
from transformers import TrainingArguments, AutoTokenizer, AutoModelForCausalLM, Trainer

# Load the IMDB dataset
train_dataset = datasets.load_dataset('imdb', split='train')

# Define the training arguments
args = TrainingArguments(
    output_dir="./test-grokadamw",
    max_steps=1000,
    per_device_train_batch_size=4,
    optim="grokadamw",
    logging_strategy="steps",
    logging_steps=1,
    learning_rate=2e-5,
    save_strategy="no",
    run_name="grokadamw-imdb",
)

# Load the model and tokenizer
model_id = "google/gemma-2b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, low_cpu_mem_usage=True).to(0)

# Initialize the Trainer
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
)

# Train the model
trainer.train()

此脚本演示如何使用 GrokAdamW 优化器对 google/gemma-2b 模型在 IMDB 数据集上进行微调。训练参数 被配置为使用 GrokAdamW，数据集被传递给 训练器 进行训练。

无调度优化器

无调度优化器在 The Road Less Scheduled 中被提出。无调度学习用平均和插值来代替基础优化器的动量，从而完全消除了对传统调度学习率退火的需要。SFO 支持的优化器为 "schedule_free_adamw" 和 "schedule_free_sgd"。首先从 pypi 安装 schedulefree：pip install schedulefree。

下面是一个简单的脚本，演示如何在全精度下对 google/gemma-2b 在 IMDB 数据集上进行微调

import torch
import datasets
from transformers import TrainingArguments, AutoTokenizer, AutoModelForCausalLM
import trl

train_dataset = datasets.load_dataset('imdb', split='train')

args = TrainingArguments(
    output_dir="./test-schedulefree",
    max_steps=1000,
    per_device_train_batch_size=4,
    optim="schedule_free_adamw",
    gradient_checkpointing=True,
    logging_strategy="steps",
    logging_steps=1,
    learning_rate=2e-6,
    save_strategy="no",
    run_name="sfo-imdb",
)

model_id = "google/gemma-2b"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, low_cpu_mem_usage=True).to(0)

trainer = trl.SFTTrainer(
    model=model, 
    args=args,
    train_dataset=train_dataset,
    dataset_text_field='text',
    max_seq_length=1024,
)

trainer.train()

加速和训练器

The Trainer 类由 Accelerate 提供支持，这是一个用于在分布式环境中轻松训练 PyTorch 模型的库，支持与完全分片数据并行 (FSDP) 和 DeepSpeed 等集成的整合。

使用 Trainer 了解更多关于 FSDP 分片策略、CPU 卸载等内容，请参考完全分片数据并行指南。

要将 Accelerate 与 Trainer 一起使用，请运行 accelerate.config 命令来为您的训练环境设置训练。此命令会创建一个 config_file.yaml 文件，并在您启动训练脚本时使用。例如，您可以设置以下一些示例配置：

分布式数据并行

FSDP

DeepSpeed

带有 Accelerate 插件的 DeepSpeed

accelerate_launch 命令是在分布式系统上使用 Accelerate 和 Trainer 启动训练脚本的推荐方法，使用 config_file.yaml 文件中指定的参数。此文件保存到 Accelerate 缓存文件夹中，并在运行 accelerate_launch 时自动加载。

例如，要使用 FSDP 配置运行 run_glue.py 训练脚本：

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

您也可以在命令行中直接指定 config_file.yaml 文件中的参数：

accelerate launch --num_processes=2 \
    --use_fsdp \
    --mixed_precision=bf16 \
    --fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP  \
    --fsdp_transformer_layer_cls_to_wrap="BertLayer" \
    --fsdp_sharding_strategy=1 \
    --fsdp_state_dict_type=FULL_STATE_DICT \
    ./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 和自定义配置。

< > 在 GitHub 上更新

←聊天模板在 Amazon SageMaker 上运行训练→