训练器
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 包含你在基本训练循环中找到的所有代码:
- 执行训练步骤以计算损失
- 使用backward方法计算梯度
- 根据梯度更新权重
- 重复这个过程,直到达到预定的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 类会将您的模型检查点保存到 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() 创建一个训练数据加载器
- 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的另一个选项是使用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 {}
然后将其传递给Trainer的callback
参数。
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_level
和 log_level_replica
参数来更改日志级别。
要为每个节点配置日志级别设置,请使用log_on_each_node
参数来确定是在每个节点上使用日志级别还是仅在主节点上使用。
Trainer 在 Trainer.__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_level
和log_level_replica
组合来配置每个节点上记录的内容。
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 无缝协作。
首先确保安装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 等集成。
要使用Trainer与Accelerate,请运行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
命令在分布式系统上启动您的训练脚本,使用 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
和自定义配置的信息。