完全分片数据并行 (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