2. ZeRO Stage 2 梯度切分

任务目标

本案例使用 DeepSpeed 的 ZeRO Stage 2 在多卡（2×GPU）环境下训练一个小型 Transformer 模型，验证梯度按数据并行度切分后，单卡梯度内存占用显著低于未启用 ZeRO 的 DDP 基线。

环境与依赖

• Python ≥ 3.8，PyTorch ≥ 1.13，DeepSpeed ≥ 0.9.0，CUDA 可用且至少 2 张 GPU
• 安装：pip install deepspeed torch

实现与关键代码

以下代码定义了一个小型 Transformer 分类模型，使用 DeepSpeed 引擎启动训练。核心逻辑：通过 zero_optimization.stage = 2 启用 ZeRO Stage 2，梯度在 backward 后按 rank 切分，每个 rank 只保留 1/2 的梯度分片。代码同时打印每步的梯度内存峰值，用于对比验证。

# train_zeRO_stage2.py
import torch
import torch.nn as nn
import deepspeed
import os

# ---- 小型 Transformer 模型（仅用于演示内存变化） ----
class SmallTransformer(nn.Module):
    def __init__(self, vocab_size=1000, d_model=128, nhead=4, num_layers=2):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward=512, batch_first=True)
        self.encoder = nn.TransformerEncoder(layer, num_layers)
        self.head = nn.Linear(d_model, 10)

    def forward(self, x):
        x = self.embed(x)
        x = self.encoder(x)
        return self.head(x.mean(dim=1))

# ---- 训练配置 ----
def create_config(use_stage2):
    if use_stage2:
        return {
            "train_batch_size": 16,
            "gradient_accumulation_steps": 1,
            "zero_optimization": {
                "stage": 2,
                "reduce_bucket_size": 5e7,
                "reduce_scatter": True
            },
            "fp16": {"enabled": True},
            "optimizer": {
                "type": "AdamW",
                "params": {"lr": 1e-4}
            }
        }
    else:
        # DDP 基线（无 ZeRO）
        return {
            "train_batch_size": 16,
            "gradient_accumulation_steps": 1,
            "fp16": {"enabled": True},
            "optimizer": {
                "type": "AdamW",
                "params": {"lr": 1e-4}
            }
        }

def train(use_stage2):
    model = SmallTransformer()
    config = create_config(use_stage2)
    engine, _, _, _ = deepspeed.initialize(
        model=model,
        model_parameters=model.parameters(),
        config=config
    )
    criterion = nn.CrossEntropyLoss()
    # 模拟输入：batch=8/卡，seq_len=32
    for step in range(5):
        x = torch.randint(0, 1000, (8, 32), device=engine.device)
        y = torch.randint(0, 10, (8,), device=engine.device)
        out = engine(x)
        loss = criterion(out, y)
        engine.backward(loss)
        engine.step()
        # 打印当前 rank 的梯度内存峰值（MB）
        grad_mem = sum(p.grad.numel() * p.grad.element_size() for p in model.parameters() if p.grad is not None) / 1e6
        if engine.global_rank == 0:
            print(f"step={step} loss={loss.item():.4f} grad_mem_rank0={grad_mem:.2f}MB")
    torch.distributed.destroy_process_group()

if __name__ == "__main__":
    import sys
    use_stage2 = sys.argv[1] == "stage2" if len(sys.argv) > 1 else True
    train(use_stage2)

运行方式

# 启动 ZeRO Stage 2 训练（2卡）
deepspeed --num_gpus=2 train_zeRO_stage2.py stage2

# 启动 DDP 基线（2卡，无 ZeRO）
deepspeed --num_gpus=2 train_zeRO_stage2.py baseline

验证方法与判定标准

• 观察输出：运行 stage2 时，每步打印的 grad_mem_rank0 应约为 baseline 运行时的 1/2（因为梯度被切分到 2 个 rank）。
• 通过条件：stage2 模式下 rank0 的梯度内存峰值显著低于 baseline（例如 baseline 为 40MB 时，stage2 约为 20MB），且 loss 正常下降，无报错。
• 注意：实际数值取决于模型大小与 batch，但比例关系应接近 1/N（N 为 GPU 数）。

常见问题

• 若显存不足，可减小 d_model 或 batch 至 4。
• 确保 deepspeed 命令可用，且 torch.cuda.device_count() ≥ 2。
• ZeRO Stage 2 的梯度切分仅在 backward 后生效，因此打印需在 backward 之后、step 之前或之后均可。