完全分片数据并行 (FSDP) v2

概述

当在 TPU 上微调大型语言模型 (LLM) 时，跨设备的模型分片对于内存效率和提高训练性能至关重要。optimum.tpu.fsdp_v2 模块提供了使用 SPMD（单程序多数据）实现完全分片数据并行训练的实用程序，该实用程序专为 TPU 设备优化。

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