optimum-tpu 文件
在 Google Cloud 上首次進行 TPU 訓練
加入 Hugging Face 社群
並獲得增強的文件體驗
開始使用
在 Google Cloud 上首次進行 TPU 訓練
本教程將引導您使用 optimum-tpu 包在 TPU 上設定和執行模型訓練。
先決條件
開始之前，請確保您有一個正在執行的 TPU 例項（參見TPU 設定指南）
環境設定
首先，建立並激活虛擬環境
python -m venv .venv
source .venv/bin/activate安裝所需包
# Install optimum-tpu with PyTorch/XLA support
pip install optimum-tpu -f https://storage.googleapis.com/libtpu-releases/index.html
# Install additional training dependencies
pip install transformers datasets accelerate trl peft evaluate瞭解用於 TPU 訓練的 FSDP
為了加快 TPU 上的訓練速度，您可以依賴 Optimum TPU 與 FSDP（Fully Sharded Data Parallel，完全分片資料並行）的整合。在訓練大型模型時，FSDP 會自動將模型分片（拆分）到所有可用的 TPU 工作器上，提供幾個關鍵優勢：
- 記憶體效率：每個 TPU 工作器只儲存模型引數的一部分，從而降低了每個裝置的記憶體需求。
- 自動擴充套件：FSDP 處理模型分發和梯度聚合的複雜性。
- 效能最佳化：Optimum TPU 的實現專門針對 TPU 硬體進行了調整。
當您在訓練設定中使用 fsdp_v2.get_fsdp_training_args(model) 配置時，這種分片會自動發生，從而可以輕鬆訓練無法在單個 TPU 裝置上容納的更大模型。
如何設定 FSDP
啟用 FSDP 的關鍵修改只有這幾行
+from optimum.tpu import fsdp_v2
+fsdp_v2.use_fsdp_v2()
+fsdp_training_args = fsdp_v2.get_fsdp_training_args(model)然後將這些引數包含在您的訓練器配置中
trainer = SFTTrainer(
model=model,
train_dataset=dataset,
args=TrainingArguments(
...
+ dataloader_drop_last=True, # Required for FSDPv2
+ **fsdp_training_args,
),
...
)完整示例
這是一個演示使用 FSDP 進行 TPU 訓練的完整工作示例
import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
from peft import LoraConfig
from trl import SFTTrainer
from optimum.tpu import fsdp_v2
# Enable FSDPv2 for TPU
fsdp_v2.use_fsdp_v2()
# Load model and dataset
model_id = "google/gemma-2b"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)
dataset = load_dataset("tatsu-lab/alpaca", split="train[:1000]")
# Get FSDP training arguments
fsdp_training_args = fsdp_v2.get_fsdp_training_args(model)
# Create trainer with minimal configuration
trainer = SFTTrainer(
model=model,
train_dataset=dataset,
args=TrainingArguments(
output_dir="./output",
dataloader_drop_last=True, # Required for FSDPv2
**fsdp_training_args,
),
peft_config=LoraConfig(
r=8,
target_modules=["k_proj", "v_proj"],
task_type="CAUSAL_LM",
),
)
# Start training
trainer.train()將此程式碼儲存為 train.py 並執行
python train.py您現在應該會看到訓練期間損失值下降。訓練完成後，您將獲得一個微調模型。恭喜 - 您剛剛在 TPU 上訓練了您的第一個模型！🙌
後續步驟
透過探索以下內容繼續您的 TPU 訓練之旅：