PyTorch FSDP：高效分布式训练数十亿参数大模型

随着大语言模型（LLM）的参数量迈入数十亿甚至万亿的时代，单张 GPU 已经远无法承载模型的训练。PyTorch 的 Fully Sharded Data Parallel (FSDP) 应运而生，成为训练超大规模模型的行业标准。本文将通过一个逐步演进的实例，深度解析 FSDP 的核心功能与最佳实践，带你掌握开启大模型训练的钥匙。

---

导语

在分布式训练的早期，我们有 DataParallel (DP) 和 DistributedDataParallel (DDP)。它们的核心思想都是数据并行：将模型完整地复制到每一张 GPU 上，然后将数据分发到不同 GPU 上进行计算，最后同步梯度。

(图片来源: PyTorch 官方文档)

这种模式的瓶颈显而易见：模型必须能装入单张 GPU。当模型大到一张卡装不下时，DDP 就无能为力了。FSDP (Fully Sharded Data Parallel) 借鉴了 ZeRO (Zero Redundancy Optimizer) 的思想，从根本上解决了这个问题。它不再复制模型，而是将模型参数 (Parameters)、梯度 (Gradients) 和优化器状态 (Optimizer States) 全部分片 (Shard) 到所有 GPU 上。

简而言之：

• DDP: 每张卡都有一个完整的模型副本。
• FSDP: 每张卡只有模型的一部分。

这种设计极大地降低了单张 GPU 的显存峰值，使得训练数倍于单卡显存的大模型成为可能。

---

一、FSDP 核心原理与环境设置

FSDP 的工作流程可以概括为：

1. 存储：在训练开始时，每个 GPU 只存储模型的一部分参数、梯度和优化器状态。
2. All-gather: 在前向或反向计算需要某一层时，所有 GPU 通过 all-gather 操作，临时集齐该层的完整参数。
3. 计算: 使用完整的层参数进行计算。
4. 释放: 计算完毕后，立即释放临时集齐的完整参数，只保留自己负责的分片，从而释放显存。
5. Reduce-scatter: 在反向传播计算完梯度后，通过 reduce-scatter 操作将梯度均值化并分发到对应的 GPU 上，每个 GPU 只保留自己那部分参数的梯度。
6. 更新: 优化器使用分片后的梯度，更新其分片后的参数。

环境设置
FSDP 依赖于 torch.distributed 进行通信。你需要一个支持 NCCL 的 PyTorch 构建（大部分官方发行版都支持），并通过 torchrun 或 deepspeed 等启动器来运行脚本。

# a.py 是你的训练脚本
# --nproc_per_node=N 表示使用 N 张 GPU
torchrun --standalone --nproc_per_node=2 a.py

---

二、FSDP 实战：构建分布式训练脚本

我们将通过一个逐步增强的例子来学习 FSDP。我们将构建一个简单的 Transformer 模型，并一步步为其添加 FSDP 的核心功能。

案例一：基础 FSDP 包装与训练循环

这是最基础的 FSDP 脚本，展示了如何初始化进程组、包装模型和启动一个训练步骤。

# a.py

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.utils.data import DataLoader, Dataset
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

# --- 1. 初始化分布式环境 ---
def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

# --- 2. 定义模型和数据 ---
class SimpleTransformer(nn.Module):
    def __init__(self, vocab_size=1000, d_model=256, nhead=4, num_layers=3):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.output_layer = nn.Linear(d_model, vocab_size)
    
    def forward(self, x):
        x = self.embedding(x)
        x = self.transformer_encoder(x)
        return self.output_layer(x)

class RandomDataset(Dataset):
    def __init__(self, vocab_size=1000, seq_len=128, size=100):
        self.vocab_size = vocab_size
        self.seq_len = seq_len
        self.size = size

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        return torch.randint(0, self.vocab_size, (self.seq_len,))

# --- 3. 训练主函数 ---
def train(rank, world_size):
    setup(rank, world_size)
    
    # 设置当前设备
    torch.cuda.set_device(rank)
    device = torch.device(f"cuda:{rank}")

    # 创建模型实例并移动到设备
    model = SimpleTransformer().to(device)
    
    # 🌟 核心：将模型用 FSDP 包装起来 🌟
    # 这是最基础的包装方式
    fsdp_model = FSDP(model)

    # 创建优化器 (必须在模型包装之后创建！)
    optimizer = optim.AdamW(fsdp_model.parameters(), lr=0.001)

    # 创建数据加载器
    dataset = RandomDataset()
    loader = DataLoader(dataset, batch_size=4)
    loss_fn = nn.CrossEntropyLoss()
    
    # --- 训练循环 ---
    fsdp_model.train()
    for batch_idx, data in enumerate(loader):
        data = data.to(device)
        targets = torch.roll(data, -1, 1) # 简单地将序列向左滚动一位作为目标

        optimizer.zero_grad()
        output = fsdp_model(data)
        
        # FSDP 输出在 CPU，需要移回 GPU 计算 Loss
        # (这是一个简化场景，实际场景中 loss 计算可能更复杂)
        # 注意: output.view(-1, output.shape[-1]) 将 (batch, seq_len, vocab_size) 变为 (batch*seq_len, vocab_size)
        loss = loss_fn(output.view(-1, output.shape[-1]), targets.view(-1))
        
        loss.backward()
        optimizer.step()
        
        if batch_idx % 10 == 0:
            print(f"Rank {rank}, Batch {batch_idx}, Loss: {loss.item()}")
            
    cleanup()

if __name__ == '__main__':
    world_size = torch.cuda.device_count()
    rank = int(os.environ.get("RANK", "0"))
    # 注意：这里我们简化了 `mp.spawn` 的使用，直接从 torchrun 获取 rank 和 world_size
    # 在真实的多节点场景中，需要正确处理
    train(rank, world_size)

运行：torchrun --standalone --nproc_per_node=2 a.py
结果分析：这个基础版本已经可以工作了！FSDP 会自动处理模型的参数分片和通信。但是，它的性能和显存效率还不是最优的，因为整个模型被视为一个巨大的 FSDP 单元。

---

案例二：自动包装策略 (auto_wrap_policy)

为了获得更好的性能，我们需要对模型内部的子模块（例如 Transformer 的每一层）应用 FSDP 包装，形成嵌套的 FSDP 实例。这使得通信和计算可以更有效地重叠。auto_wrap_policy 就是用来实现这一点的。

# ... (在 train 函数中修改) ...
from torch.distributed.fsdp.wrap import (
    transformer_auto_wrap_policy,
    size_based_auto_wrap_policy,
)
from functools import partial

# ... model = SimpleTransformer().to(device) ...

# 🌟 核心：定义自动包装策略 🌟
# `transformer_auto_wrap_policy` 会自动寻找并包装指定的 Transformer Block 类
my_auto_wrap_policy = partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={
        nn.TransformerEncoderLayer, # 告诉 FSDP 要包装这个类型的模块
    },
)

# 使用包装策略
fsdp_model = FSDP(
    model,
    auto_wrap_policy=my_auto_wrap_policy,
    device_id=torch.cuda.current_device() # 指定 device_id 很重要
)
# ... (后续代码不变) ...

结果分析：通过这个策略，FSDP 不再将整个 SimpleTransformer 视为一个单元，而是会递归地进入模型，找到每一个 TransformerEncoderLayer 并将其单独包装成一个 FSDP 实例。这极大地提高了并行度和计算/通信重叠的效率，是 FSDP 的关键性能优化点。

---

案例三：混合精度训练 (MixedPrecision)

使用 bfloat16 (BF16) 或 float16 (FP16) 可以将模型显存占用减半，并利用 Tensor Core 加速计算。FSDP 提供了 MixedPrecision策略来轻松实现这一点。

# ... (在 train 函数中修改) ...
from torch.distributed.fsdp.api import MixedPrecision

# 🌟 核心：定义混合精度策略 🌟
# 对于 A100/H100 等现代 GPU，bfloat16 通常是最佳选择
# 对于 V100 等旧款 GPU，可以使用 torch.float16
mixed_precision_policy = MixedPrecision(
    param_dtype=torch.bfloat16,   # 模型参数的存储类型
    reduce_dtype=torch.bfloat16,  # 梯度同步时的计算类型
    buffer_dtype=torch.bfloat16,  # Buffer (如 a running mean) 的存储类型
)

fsdp_model = FSDP(
    model,
    auto_wrap_policy=my_auto_wrap_policy,
    mixed_precision=mixed_precision_policy, # 应用混合精度
    device_id=torch.cuda.current_device()
)
# ... (后续代码不变) ...

结果分析：添加混合精度后，模型的参数、梯度和计算都将使用 bfloat16，显存占用会显著下降，训练速度也会加快。FSDP 会自动处理类型转换和梯度缩放（对于 FP16）。

---

案例四：CPU 卸载 (cpu_offload)

当模型巨大，即使使用了 FSDP 和混合精度，显存依然不足时，可以启用 CPU Offloading。它会将当前不参与计算的参数分片从 GPU 显存移动到 CPU 内存。

# ... (在 train 函数中修改) ...
from torch.distributed.fsdp.api import CPUOffload

fsdp_model = FSDP(
    model,
    auto_wrap_policy=my_auto_wrap_policy,
    mixed_precision=mixed_precision_policy,
    cpu_offload=CPUOffload(offload_params=True), # 🌟 核心：启用 CPU Offload 🌟
    device_id=torch.cuda.current_device()
)
# ... (后续代码不变) ...

结果分析：这是一种用通信换空间的策略。它能让你训练远超 GPU 显存总和的模型，但代价是大量的 PCIe 带宽消耗，因为数据需要在 CPU 和 GPU 之间来回传输。这会显著降低训练速度，因此只在显存极度受限时使用。

---

案例五：理解分片策略 (sharding_strategy)

FSDP 提供了多种分片策略，让你可以在显存效率和通信开销之间做权衡。

• FULL_SHARD: 默认且最节省显存的策略。将参数、梯度和优化器状态都进行分片。
• SHARD_GRAD_OP: 只分片梯度和优化器状态，模型参数在每个 GPU 上是完整的（类似于 DDP）。显存占用更高，但通信量可能更少。
• NO_SHARD: 不进行任何分片，等效于 DDP。
• HYBRID_SHARD: 在节点内进行 FULL_SHARD，在节点间进行 SHARD_GRAD_OP，是多节点训练的常用策略。

# ... (在 train 函数中修改) ...
from torch.distributed.fsdp.api import ShardingStrategy

fsdp_model = FSDP(
    model,
    auto_wrap_policy=my_auto_wrap_policy,
    mixed_precision=mixed_precision_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD, # 🌟 核心：显式指定分片策略 🌟
    device_id=torch.cuda.current_device()
)
# ... (后续代码不变) ...

结果分析：对于单节点多 GPU 训练，FULL_SHARD 通常是最佳选择。了解不同策略的存在，有助于在更复杂的场景（如多节点、异构硬件）下进行精细调优。

---

案例六：优化器状态分片

这是 FSDP 的一个隐式但至关重要的特性。在使用 AdamW 等需要为每个参数保存一阶和二阶矩的优化器时，优化器状态本身就会占用大量显存（通常是模型参数的 2 倍）。

FSDP 自动处理了这一点。当你将 fsdp_model.parameters() 传入优化器时，优化器在每个 rank 上只会看到和管理它自己负责的那部分参数的状态。你不需要做任何额外的配置。这是 FSDP 相比 DDP 节省大量显存的关键原因之一。

---

案例七：FSDP Checkpointing 的正确姿势

保存和加载 FSDP 模型的状态字典比较特殊，因为模型参数是分散在所有 GPU 上的。

# ... (在 train 函数中修改) ...
from torch.distributed.fsdp.api import StateDictType, FullStateDictConfig

# --- 保存 Checkpoint (只在 Rank 0 上执行) ---
if rank == 0:
    print("Saving checkpoint...")

# 1. 定义一个配置，告诉 FSDP 我们想要收集一个完整的、在 CPU 上的 state_dict
save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)

# 2. 使用 FSDP 的 state_dict 上下文管理器
with FSDP.state_dict_type(fsdp_model, StateDictType.FULL_STATE_DICT, save_policy):
    cpu_state = fsdp_model.state_dict()

if rank == 0:
    torch.save(cpu_state, "my_checkpoint.pt")
    print("Checkpoint saved.")


# --- 加载 Checkpoint ---
# 加载前，模型必须已经用 FSDP 包装好
# 1. 定义加载配置
load_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)

# 2. 使用 FSDP 的 state_dict 上下文管理器
with FSDP.state_dict_type(fsdp_model, StateDictType.FULL_STATE_DICT, load_policy):
    # a. 在所有 rank 上加载 state_dict 到 CPU
    if rank == 0:
        cpu_state = torch.load("my_checkpoint.pt")
    else:
        cpu_state = None
    
    # b. FSDP 会自动将 state_dict 从 Rank 0 分发并加载到所有 rank 的对应分片上
    fsdp_model.load_state_dict(cpu_state)

print(f"Rank {rank} loaded checkpoint.")

结果分析：Checkpointing 是 FSDP 中最容易出错的环节。核心思想是：在保存时，将所有分片 all-gather 到 Rank 0 的 CPU 上，然后保存；在加载时，Rank 0 读取文件，然后 FSDP 负责将完整的 state_dict 分发 scatter 到所有 GPU 的对应分片上。 必须使用 FSDP 提供的 state_dict_type 上下文管理器来确保这个过程正确执行。

---

总结与展望

PyTorch FSDP 是一个功能强大且灵活的分布式训练框架，它通过对模型参数、梯度和优化器状态的完全分片，成功地突破了单卡显存的壁垒，成为训练数十亿乃至更大规模模型的关键技术。掌握 FSDP 的自动包装、混合精度、Checkpointing 等核心实践，是每一位大模型算法工程师的必备技能。

当然，从零开始训练一个巨大的模型需要海量的计算资源和时间。在许多场景下，利用已经训练好的、强大的基础模型进行微调或直接通过 API 调用，是更具经济效益和效率的选择。

站在巨人的肩膀上：拥抱强大的 AI 服务

• 😳 0v0 AI 助手推荐：对于希望免费体验和探索各种 SOTA（State-of-the-Art）模型能力的开发者，https://0v0.pro 是一个绝佳的平台。它集成了 Llama、Qwen、gpt-4o 等多种模型，并且无次数限制，甚至每周提供 gpt-5 等旗舰模型的免费试用，非常适合进行快速原型验证和学习。

无论是通过 FSDP 亲手训练大模型，还是通过 API 利用现有的模型能力，我们都在通往更智能的未来的道路上。