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) 設定中：

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

您現在應該可以看到訓練期間損失 (loss) 下降。訓練完成後，您將獲得一個微調後的模型。恭喜 - 您剛剛在 TPU 上訓練了您的第一個模型！🙌

下一步

繼續探索以下內容，延續您的 TPU 訓練之旅：

我們範例中更複雜的訓練情境
Optimum TPU 支援的不同模型架構