完全分片資料並行 (FSDP) v2

概述

在 TPU 上微調大型語言模型 (LLM) 時，為了提高記憶體效率和訓練效能，模型在裝置間進行分片變得至關重要。`optimum.tpu.fsdp_v2` 模組提供了使用專門為 TPU 裝置最佳化的 SPMD（單程式多資料）實現完全分片資料並行訓練的實用工具。

FSDP_v2 特性

TPU 裝置間的模型權重分片
梯度檢查點支援
常見模型架構的自動配置
與 PyTorch/XLA 的 SPMD 實現整合

基本用法

以下是如何為您的訓練啟用和配置 FSDP_v2

from optimum.tpu import fsdp_v2
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Enable FSDP_v2
fsdp_v2.use_fsdp_v2()

# Load model and tokenizer
model_id = "meta-llama/Llama-2-7b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=torch.bfloat16
)

# Get FSDP training configuration
fsdp_args = fsdp_v2.get_fsdp_training_args(model)

配置選項

`get_fsdp_training_args()` 函式返回一個包含模型特定配置的字典，例如：

{
    'fsdp': 'full_shard',
    'fsdp_config': {
        'transformer_layer_cls_to_wrap': ['LlamaDecoderLayer'],  # Model-specific
        'xla': True,
        'xla_fsdp_v2': True,
        'xla_fsdp_grad_ckpt': True
    }
}

關鍵引數

`transformer_layer_cls_to_wrap`：指定要用 FSDP 包裝哪些模型層
`xla`：啟用 XLA 最佳化
`xla_fsdp_v2`：啟用 FSDP_v2 實現
`xla_fsdp_grad_ckpt`：啟用梯度檢查點以提高記憶體效率

高階用法

自定義層包裝

您可以自定義哪些層與 FSDP 一起包裝

custom_fsdp_args = fsdp_v2.get_fsdp_training_args(
    model,
    layer_cls_to_wrap=['CustomTransformerLayer']
)

與 Transformers Trainer 整合

FSDP_v2 配置可直接與 Transformers Trainer 結合使用

from transformers import Trainer, TrainingArguments
# Or for instruction fine-tuning:
# from trl import SFTTrainer

trainer = Trainer(  # or SFTTrainer
    model=model,
    args=TrainingArguments(**fsdp_args),  # Unpack FSDP configuration
    train_dataset=dataset,
    ...
)

後續步驟

您可以檢視我們的示例筆記本，瞭解使用 optimum-tpu 進行訓練的最佳實踐
有關 PyTorch/XLA 的 FSDP 實現的更多詳細資訊，請參閱官方文件。

optimum-tpu