Transformers 文档

训练师

训练器

Trainer 是 Transformers 库中实现的 PyTorch 模型的完整训练和评估循环。你只需要传递训练所需的必要部分(模型、分词器、数据集、评估函数、训练超参数等),Trainer 类会处理其余部分。这使得你可以更快地开始训练,而无需手动编写自己的训练循环。但与此同时,Trainer 非常可定制,并提供了大量的训练选项,因此你可以根据确切的训练需求进行调整。

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


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

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

pip install accelerate

# upgrade
pip install accelerate --upgrade

本指南提供了Trainer类的概述。

基本用法

Trainer 包含你在基本训练循环中找到的所有代码:

  1. 执行训练步骤以计算损失
  2. 使用backward方法计算梯度
  3. 根据梯度更新权重
  4. 重复这个过程,直到达到预定的epochs数量

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"],
    processing_class=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

检查点

Trainer 类会将您的模型检查点保存到 TrainingArgumentsoutput_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的方法可以被子类化和重写,以支持您想要的功能,而无需从头重写整个训练循环来适应它。这些方法包括:

例如,如果你想自定义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的另一个选项是使用callbacks。回调不会改变训练循环中的任何内容。它们检查训练循环的状态,然后根据状态执行一些操作(提前停止、记录结果等)。换句话说,回调不能用于实现像自定义损失函数这样的功能,你需要子类化并重写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 {}

然后将其传递给Trainercallback参数。

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,
    callbacks=[EarlyStoppingCallback()],
)

日志记录

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

Trainer 默认设置为 logging.INFO,它会报告错误、警告和其他基本信息。在分布式环境中,Trainer 的副本设置为 logging.WARNING,它只报告错误和警告。你可以通过 TrainingArguments 中的 log_levellog_level_replica 参数来更改日志级别。

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

TrainerTrainer.__init__() 方法中为每个节点单独设置日志级别,因此如果你在创建 Trainer 对象之前使用其他 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_levellog_level_replica组合来配置每个节点上记录的内容。

single node
multi-node
my_app.py ... --log_level warning --log_level_replica error

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)

训练后禁用NEFTune以恢复原始嵌入层,以避免任何意外行为。

Liger 内核

Liger-Kernel Kernel 是由 Linkedin 开发的一组 Triton 内核,专门为 LLM 训练设计。我们已经实现了与 Hugging Face 兼容的 RMSNorm、RoPE、SwiGLU、CrossEntropy、FusedLinearCrossEntropy 等功能,未来还会有更多功能加入。它可以有效提高多 GPU 训练的吞吐量 20%,并减少内存使用量 60%。该内核可以与 flash attention、PyTorch FSDP 和 Microsoft DeepSpeed 无缝协作。

Gain +20% throughput and reduce memory usage by 60% on LLaMA 3-8B model training. Achieve longer context lengths and larger batch sizes. It’s also useful if you want to scale up your model to multi-head training or large vocabulary sizes. Unleash multi-head training (medusa) and more. See details and examples in [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的高效实现替换,因此除了设置参数值外,您不需要做任何额外的事情。

优化器

您可以使用以下内置优化器进行训练:

from transformers import TrainingArguments
training_args = TrainingArguments(..., optim="adamw_torch")

请参阅OptimizerNames以获取完整的选项列表。我们在以下部分中包含了高级示例。

你也可以通过以下方式使用任意的PyTorch优化器:

import torch

optimizer_cls = torch.optim.AdamW
optimizer_kwargs = {
    "lr": 4e-3,
    "betas": (0.9, 0.999),
    "weight_decay": 0.05,
}

from transformers import Trainer
trainer = Trainer(..., optimizer_cls_and_kwargs=(optimizer_cls, optimizer_kwargs))

GaLore

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

首先确保安装GaLore官方仓库:

pip install galore-torch

然后只需在optim中添加["galore_adamw", "galore_adafactor", "galore_adamw_8bit"]中的一个,并与optim_target_modules一起使用,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上提交问题

LOMO 优化器

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

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

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

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优化器在IMDB数据集上微调google/gemma-2b

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优化器在IMDB数据集上微调google/gemma-2b模型。TrainingArguments配置为使用GrokAdamW,并将数据集传递给Trainer进行训练。

无调度优化器

Schedule Free 优化器已在The Road Less Scheduled中引入。 Schedule-Free 学习通过结合平均和插值来替代基础优化器的动量,从而完全消除了传统调度中需要退火学习率的需求。 SFO 支持的优化器是 "schedule_free_adamw""schedule_free_sgd"。首先从 pypi 安装 schedulefree pip install schedulefree

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

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

加速和训练器

Trainer 类由 Accelerate 提供支持,这是一个用于在分布式环境中轻松训练 PyTorch 模型的库,支持诸如 FullyShardedDataParallel (FSDP)DeepSpeed 等集成。

了解更多关于FSDP分片策略、CPU卸载等内容,请参阅Trainer完全分片数据并行指南中的内容。

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

DistributedDataParallel
FSDP
DeepSpeed
DeepSpeed with Accelerate plugin
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 命令在分布式系统上启动您的训练脚本,使用 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和自定义配置的信息。

< > Update on GitHub