Transformers 文件
訓練器
並獲得增強的文件體驗
開始使用
Trainer
Trainer 是一個用於 Transformers PyTorch 模型的完整訓練和評估迴圈。將模型、預處理器、資料集和訓練引數插入 Trainer，讓它處理其餘部分，從而更快地開始訓練。
Trainer 還由 Accelerate 提供支援，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_losss(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 = "mean" if num_items_in_batch is not None else "sum"
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。
要將 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/ \
--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) 時的記憶體使用。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 層。
訓練開始可能需要一些時間（在 NVIDIA A100 上，2B 模型約 3 分鐘）。
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 進行訓練。這會將模型中對應的層打上 Ligers 核心補丁。
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
)您還可以使用 `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 上更新