分布式训练实战：DeepSpeed ZeRO 与 FSDP 大模型训练指南

1. 引言

当模型参数量超过单卡显存容量时，分布式训练成为必选项。本文将深入讲解两种主流的分布式训练方案：DeepSpeed ZeRO 和 PyTorch FSDP，并通过实际代码演示如何在多卡环境下训练 7B-70B 模型。

核心挑战：

• 7B 模型 FP16 需要 14GB 存储权重
• 训练时还需要梯度（14GB）+ 优化器状态（Adam 需要 3×14GB = 42GB）
• 总计约 70GB，单张 A100 80GB 刚好放不下
• 70B 模型需要 700GB+，必须多卡分布

2. 分布式训练策略

数据并行 (Data Parallel):
  每张卡持有完整模型副本，处理不同数据批次
  → 适合小模型，通信开销大

模型并行 (Model Parallel):
  - 张量并行 (Tensor Parallel): 切分单层权重
  - 流水线并行 (Pipeline Parallel): 切分不同层
  → 适合超大模型

ZeRO (Zero Redundancy Optimizer):
  优化器状态/梯度/参数分片到不同卡
  → 通信量和显存完美平衡，主流方案

3. DeepSpeed ZeRO

3.1 ZeRO 三个阶段

ZeRO Stage 1: 分片优化器状态
  - 每张卡只保存 1/N 的优化器状态
  - 显存节省: ~4x (Adam: 12 bytes/param → 3 bytes/param)

ZeRO Stage 2: + 分片梯度
  - 每张卡只保存 1/N 的梯度
  - 显存节省: ~8x

ZeRO Stage 3: + 分片参数
  - 每张卡只保存 1/N 的模型参数
  - 显存节省: ~N 倍（理论无限扩展）
  - 代价：额外通信开销

3.2 DeepSpeed 配置

{
    "train_batch_size": 32,
    "train_micro_batch_size_per_gpu": 4,
    "gradient_accumulation_steps": 8,
    "steps_per_print": 100,

    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 5e8,
        "reduce_scatter": true,
        "reduce_bucket_size": 5e8,
        "overlap_comm": true,
        "contiguous_gradients": true
    },

    "bf16": {
        "enabled": true
    },

    "gradient_clipping": 1.0,

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 2e-5,
            "betas": [0.9, 0.999],
            "eps": 1e-8,
            "weight_decay": 0.01
        }
    },

    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_num_steps": 100,
            "total_num_steps": 10000
        }
    },

    "activation_checkpointing": {
        "partition_activations": true,
        "cpu_checkpointing": true,
        "contiguous_memory_optimization": true
    }
}

3.3 DeepSpeed 训练代码

import deepspeed
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载模型
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    torch_dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 初始化 DeepSpeed
model_engine, optimizer, _, scheduler = deepspeed.initialize(
    model=model,
    config="ds_config.json",
)

# 训练循环
for epoch in range(num_epochs):
    for batch in dataloader:
        input_ids = batch["input_ids"].to(model_engine.device)
        labels = batch["labels"].to(model_engine.device)

        outputs = model_engine(input_ids=input_ids, labels=labels)
        loss = outputs.loss

        model_engine.backward(loss)
        model_engine.step()

3.4 启动训练

deepspeed --num_gpus=4 train.py \
    --deepspeed_config ds_config.json \
    --model_name meta-llama/Llama-2-7b-hf \
    --dataset_path ./data \
    --output_dir ./output \
    --num_epochs 3

4. PyTorch FSDP

4.1 FSDP 原理

FSDP (Fully Sharded Data Parallel) 是 PyTorch 原生的分布式训练方案，原理与 ZeRO Stage 3 类似：

前向传播时：
1. All-Gather 收集完整参数
2. 执行前向计算
3. 释放非本分片参数

反向传播时：
1. All-Gather 收集完整参数
2. 计算梯度
3. Reduce-Scatter 聚合梯度到对应分片
4. 释放非本分片参数和梯度

4.2 FSDP 训练代码

import torch
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
)
from torch.distributed.fsdp.wrap import (
    transformer_auto_wrap_policy,
)
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    LlamaDecoderLayer,
)
import torch.distributed as dist

# 初始化分布式
dist.init_process_group("nccl")

# 模型加载
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
)

# FSDP 包装
auto_wrap_policy = transformer_auto_wrap_policy(
    transformer_layer_cls={LlamaDecoderLayer},
)

mixed_precision = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)

model = FSDP(
    model,
    auto_wrap_policy=auto_wrap_policy,
    mixed_precision=mixed_precision,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    device_id=torch.cuda.current_device(),
    limit_all_gathers=True,
    use_orig_params=True,  # 支持参数组
)

# 优化器
optimizer = torch.optim.AdamW(
    model.parameters(), lr=2e-5, weight_decay=0.01
)

# 训练循环
for epoch in range(num_epochs):
    model.train()
    for batch in dataloader:
        input_ids = batch["input_ids"].cuda()
        labels = batch["labels"].cuda()

        outputs = model(input_ids=input_ids, labels=labels)
        loss = outputs.loss

        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

4.3 启动 FSDP

torchrun \
    --nproc_per_node=4 \
    --master_port=29500 \
    train_fsdp.py \
    --model_name meta-llama/Llama-2-7b-hf \
    --batch_size 4 \
    --num_epochs 3

5. 混合并行策略

5.1 3D 并行

对于 70B+ 模型，通常需要 3D 并行：

3D 并行 = 数据并行 (DP) + 张量并行 (TP) + 流水线并行 (PP)

示例：8 节点 × 8 GPU = 64 GPU
- TP = 8 (节点内)
- PP = 4 (跨 4 个节点)
- DP = 2 (剩余维度)

5.2 DeepSpeed 3D 并行配置

{
    "train_batch_size": 64,
    "train_micro_batch_size_per_gpu": 1,

    "tensor_parallel": {
        "tp_size": 4
    },
    "pipeline": {
        "stages": 4,
        "partition": "best"
    },

    "zero_optimization": {
        "stage": 1
    }
}

6. 性能优化技巧

6.1 混合精度训练

# PyTorch 原生 AMP
scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
    outputs = model(input_ids=input_ids)
    loss = outputs.loss
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

6.2 梯度检查点

# 用时间换空间：重计算激活值而非存储
model.gradient_checkpointing_enable()

# 或在 FSDP 中
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper,
    CheckpointImpl,
    apply_activation_checkpointing,
)

apply_activation_checkpointing(
    model,
    checkpoint_wrapper_fn=checkpoint_wrapper,
    check_fn=lambda submodule: isinstance(submodule, LlamaDecoderLayer),
)

6.3 通信优化

# 梯度累积减少通信频率
for i, batch in enumerate(dataloader):
    loss = model(batch) / accumulation_steps
    loss.backward()

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

7. 方案选择指南

场景	推荐方案	说明
单卡 7B	原生 PyTorch	AMP + 梯度检查点
多卡 7B-13B	FSDP	PyTorch 原生，无额外依赖
多卡 30B-70B	DeepSpeed ZeRO-3	成熟稳定，CPU offload
多节点 70B+	DeepSpeed 3D 并行	TP + PP + DP
超大模型 175B+	Megatron-LM	NVIDIA 官方方案

8. 总结

分布式训练的核心在于用通信换显存：

1. ZeRO Stage 2 是最常用的方案：分片优化器和梯度，平衡性能和显存
2. FSDP 是 PyTorch 原生方案，适合不想引入额外依赖的团队
3. 梯度检查点 + 混合精度 + 梯度累积 是基础优化三件套
4. 70B+ 模型需要 3D 并行（TP + PP + DP）