Transformers 文件
GPU
並獲得增強的文件體驗
開始使用
GPU
GPU 因其高記憶體頻寬和並行處理能力，常用於深度學習模型訓練。根據你的 GPU 和模型大小，甚至可以訓練數十億引數的模型。關鍵在於找到 GPU 記憶體利用率（資料吞吐量/訓練時間）和訓練速度之間的平衡。
本指南將向你展示 Transformers 和 PyTorch 中可用於高效訓練 GPU 模型的特性。在許多情況下，你需要結合使用這些特性來最佳化訓練。
請參考下表，以便快速找到與你的訓練場景相關的特性。
| 特性 | 訓練速度 | 記憶體使用 |
|---|---|---|
| 批次大小 | 是的 | 是的 |
| 梯度累積 | 否 | 是的 |
| 梯度檢查點 | 否 | 是的 |
| 混合精度 | 是的 | 取決於 |
| 最佳化器 | 是的 | 是的 |
| 資料預載入 | 是的 | 否 |
| torch_empty_cache_steps | 否 | 是的 |
| torch.compile | 是的 | 否 |
| 縮放點積注意力 (SDPA) | 是的 | 是的 |
Trainer
Trainer 支援許多有用的訓練特性，可以透過 TrainingArguments 進行配置。本節重點介紹一些對最佳化訓練更重要的特性。
批次大小
批次大小是高效 GPU 訓練最重要的超引數之一，因為它會影響記憶體使用和訓練速度。較大的批次大小可以利用 GPU 的並行處理能力，從而加快訓練速度。建議使用 2 的冪次方作為批次大小，例如 8、64、128、256、512 等。批次大小取決於你的 GPU 和模型的資料型別。
在 TrainingArguments 中配置 per_device_train_batch_size()。
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=256,
per_device_eval_batch_size=256,
)請參閱 NVIDIA 的 效能 指南，瞭解輸入特徵、輸出神經元數量和批次大小如何影響效能。這些都與 GPU 執行的通用矩陣乘法 (GEMM) 有關。引數越大，並行化和效率越好。
Tensor Core 要求 一節對於根據資料型別和 GPU 選擇最大化張量乘法速度的批次大小也很有用。例如，對於 fp16，建議使用 8 的倍數，除非是 A100 GPU，這種情況下請使用 64 的倍數。
最後，考慮 維度量化效應 對於較小引數的影響。當矩陣維度不能被 GPU 的執行緒塊瓦片大小整除時，會導致瓦片量化，從而使 GPU 資源利用不足。選擇正確的批次大小乘數，使矩陣可以被瓦片大小整除，可以顯著加快訓練速度。
梯度累積
梯度累積透過在多次小批次資料上累積梯度來克服記憶體限制（對於在單個 GPU 上無法容納的超大型模型非常有用），然後在更新引數之前進行一次更新。這透過儲存更少的梯度來減少記憶體使用，並能夠使用更大的*有效批次大小*進行訓練，因為通常引數是從單個批次資料中更新的。然而，由於梯度累積引入了額外的正向和反向傳播，訓練速度可能會變慢。
在 TrainingArguments 中配置 per_device_train_batch_size() 以啟用梯度累積。
from transformers import TrainingArguments
# effective batch size of 64
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
)儘量避免過多的梯度累積步驟，因為它會大大降低訓練速度。考慮以下示例，你的 GPU 上能容納的最大批次大小是 4。你應該將批次大小保持在 4 以更好地利用 GPU。
| 批次大小 | 梯度累積步數 | 有效批次大小 | |
|---|---|---|---|
| 1 | 64 | 64 | 👎 |
| 4 | 16 | 64 | 👍 |
梯度檢查點
梯度檢查點透過在反向傳播過程中僅儲存部分中間啟用並重新計算剩餘啟用來減少記憶體使用。這避免了儲存正向傳播中的*所有*中間啟用，後者可能需要大量記憶體開銷。然而，其代價是訓練速度變慢（約 20%）。
在 TrainingArguments 中配置 gradient_checkpointing() 以啟用梯度檢查點。
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
)混合精度
混合精度透過在半精度（fp16）和全精度（fp32）下執行部分計算來加速訓練速度。半精度計算能提升訓練速度，因為它不像全精度計算那樣需要大量計算。同時，保留部分全精度計算能保持準確性。
有幾種資料型別可用於混合精度訓練。
混合精度訓練的主要優點是將啟用儲存為 fp16。
在 TrainingArguments 中配置 fp16() 以啟用 fp16 資料型別的混合精度訓練。
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
fp16=True.
)fp16 不是記憶體最佳化的，因為以 fp16 計算的梯度在最佳化步驟中會轉換回 fp32。你最終可能會使用更多的 GPU 記憶體，特別是對於小批次大小，因為現在 GPU 上有模型的兩個版本（fp16 和 fp32）。
最佳化器
Transformers 預設實現 PyTorch 的 AdamW (adamw_torch) 最佳化器。但由於它儲存了過去梯度的加權平均值，因此需要額外的記憶體，其大小與模型引數數量成比例，以儲存過去梯度。這在訓練非常大的模型時可能會成為問題，在這種情況下，你應該考慮選擇不同的最佳化器。例如，如果你在 NVIDIA 或 AMD 上安裝了 Apex，那麼使用 adamw_apex_fused 最佳化器可以為所有 AdamW 最佳化器提供最快的訓練。
在 TrainingArguments 中配置 optim() 以選擇最佳化器。
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
bf16=True,
optim="adamw_bnb_8bit"
)根據你的訓練場景，有許多最佳化器可供選擇（請參閱 OptimizerNames 獲取完整支援列表）。例如，Adafactor 可以透過儲存行或列的加權平均值而不是矩陣中的每個元素來顯著減少記憶體需求，代價是收斂速度較慢。另一個例子是使用來自 bitsandbytes 的 8 位 AdamW 最佳化器 來量化最佳化器狀態。最佳化器狀態以較低精度儲存，並在最佳化器步驟中使用之前進行反量化。
請參閱 最佳化器 指南以瞭解更多專用最佳化器。
資料預載入
資料預載入可在 CPU 上提前載入並準備批處理資料，以確保 GPU 持續工作，從而減少 GPU 空閒時間和提高利用率。有兩種方法可以預載入資料，以確保 GPU 始終處於工作狀態。
- 在 CPU 上分配固定記憶體以儲存資料，並將其直接傳輸到 GPU。
- 增加 CPU 執行緒或工作執行緒的數量以更快地預載入資料。
在 TrainingArguments 中配置 dataloader_pin_memory() 和 dataloader_num_workers() 以分配固定記憶體和增加工作執行緒數量。
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
bf16=True,
optim="adamw_bnb_8bit",
dataloader_pin_memory=True,
dataloader_num_workers=4,
)PyTorch
PyTorch 提供了多種功能來減少記憶體需求並提高訓練速度。透過僅新增少量程式碼，通常可以在 Transformers 中啟用這些功能。
torch.empty_cache_steps
torch.cuda.empty_cache 函式釋放未使用的快取記憶體，這有助於避免記憶體不足 (OOM) 錯誤，但代價是訓練速度會慢約 10%。
在 TrainingArguments 中使用 torch_empty_cache_steps()，在一定數量的訓練步驟後啟用它。
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
bf16=True,
optim="adamw_bnb_8bit",
dataloader_pin_memory=True,
dataloader_num_workers=4,
torch_empty_cache_steps=4,
)torch.compile
torch.compile 將 PyTorch 程式碼編譯成最佳化的核心，顯著加快了訓練速度。此功能依賴於 TorchDynamo 透過 Frame Evaluation API 捕獲 PyTorch 圖。該圖可以進一步編譯為針對不同後端的最佳化核心。
在 TrainingArguments 中配置 torch_compile() 以啟用它，並配置 torch_compile_backend() 以選擇要使用的後端。
from transformers import TrainingArguments
args = TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=16,
gradient_checkpointing=True,
bf16=True,
optim="adamw_bnb_8bit",
dataloader_pin_memory=True,
dataloader_num_workers=4,
torch_empty_cache_steps=4,
torch_compile=True,
torch_compile_backend="inductor"
)請參考下表，以幫助您為訓練場景選擇合適的後端。
| 後端 | 描述 | 目標 |
|---|---|---|
| eager | 使用 PyTorch 執行提取的 GraphModule | 除錯 |
| aot_eager | 對 AOTAutograd 提取的正向和反向圖使用 PyTorch eager 模式 | 除錯 |
| inductor | 透過利用 Triton 核心，將 TorchInductor 與 AOTAutograd 和 CUDA Graphs 結合使用 | 訓練和推理 |
| nvfuser | 將 nvFuser 與 TorchScript 結合使用 | 訓練和推理 |
| aot_nvfuser | 將 nvFuser 與 AOTAutograd 結合使用 | 訓練和推理 |
| aot_cudagraphs | 將 CUDA Graphs 與 AOTAutograd 結合使用 | 訓練和推理 |
| ofi | 使用 TorchScripts optimize_for_inference | 推理 |
| fx2trt | 使用 Torch-TensorRT | 推理 |
| onnxrt | 將 ONNX-RT 用於 CPU 和 GPU 推理 | 推理 |
| ipex | 將 IPEX 用於 CPU 推理 | 推理 |
縮放點積注意力
torch.nn.functional.scaled_dot_product_attention (SDPA) 是縮放點積注意力機制的 PyTorch 原生實現。SDPA 比 Transformer 模型中原始的注意力機制更高效且更最佳化。它支援三種類型的縮放點積注意力：
- 對於 torch 型別為 fp16 或 bf16 的模型，將自動啟用 FlashAttention2。請務必先將模型轉換為相應的型別。
- xFormers 或記憶體高效注意力支援 torch 型別為 fp32 的模型。
- 縮放點積注意力的 C++ 實現。
SDPA 預設在 PyTorch 2.1.1+ 中啟用，但可以透過在 from_pretrained() 中設定 attn_implementation="sdpa" 來顯式啟用。
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B", device_map="auto", attn_implementation="sdpa")