4. ZeRO Stage 3 参数切分

任务目标

本案例使用 DeepSpeed ZeRO Stage 3 训练一个约 1.3B 参数的 风格模型，核心目标是验证：模型参数被分片存储到多个 GPU 上，而非在每个 GPU 上完整复制。通过对比 Stage 3 与 Dense（无 ZeRO）模式下的显存占用，确认参数分片生效。

环境与依赖

• Python 3.8+，PyTorch 1.13+，DeepSpeed 0.9+，CUDA 11.x，至少 2 块 GPU（推荐 4 块）
• 安装：pip install deepspeed torch

实现原理

ZeRO Stage 3 将模型参数、梯度和优化器状态全部按层切分到各 GPU。前向/反向传播时通过通信收集所需参数（all-gather），计算完成后立即释放，从而将单卡显存占用从 O(N) 降为 O(N/G)（G 为 GPU 数）。本案例通过 deepspeed.initialize 的 zero_optimization.stage=3 配置启用该机制。

完整实现代码

import torch
import torch.nn as nn
import deepspeed
import os

# ---------- 1. 定义  规模模型（简化结构） ----------
class (nn.Module):
    def __init__(self, hidden, heads):
        super().__init__()
        self.ln1 = nn.LayerNorm(hidden)
        self.attn = nn.MultiheadAttention(hidden, heads, batch_first=True)
        self.ln2 = nn.LayerNorm(hidden)
        self.mlp = nn.Sequential(
            nn.Linear(hidden, 4*hidden),
            nn.GELU(),
            nn.Linear(4*hidden, hidden)
        )
    def forward(self, x):
        x = x + self.attn(self.ln1(x), self.ln1(x), self.ln1(x))[0]
        x = x + self.mlp(self.ln2(x))
        return x

class (nn.Module):
    def __init__(self, vocab=50257, hidden=2048, layers=24, heads=16, seq_len=1024):
        super().__init__()
        self.embed = nn.Embedding(vocab, hidden)
        self.blocks = nn.ModuleList([(hidden, heads) for _ in range(layers)])
        self.ln_f = nn.LayerNorm(hidden)
        self.head = nn.Linear(hidden, vocab, bias=False)
        # 估算参数量：约 24*(2048^2*4 + 2048^2*4) + 50257*2048 ≈ 1.3B
        self.seq_len = seq_len

    def forward(self, input_ids):
        x = self.embed(input_ids)
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        return self.head(x)

# ---------- 2. 训练配置 ----------
def get_ds_config(stage):
    return {
        "train_batch_size": 8,
        "gradient_accumulation_steps": 1,
        "optimizer": {
            "type": "AdamW",
            "params": {"lr": 1e-4}
        },
        "zero_optimization": {
            "stage": stage,
            "reduce_bucket_size": 5e8,
            "stage3_prefetch_bucket_size": 5e8,
            "stage3_param_persistence_threshold": 1e6
        },
        "fp16": {"enabled": True}
    }

# ---------- 3. 主训练循环 ----------
def train(stage):
    model = ()
    # 模拟输入：随机  序列
    batch = torch.randint(0, 50257, (2, 1024)).cuda()
    labels = torch.randint(0, 50257, (2, 1024)).cuda()

    model_engine, optimizer, _, _ = deepspeed.initialize(
        model=model,
        model_parameters=model.parameters(),
        config=get_ds_config(stage)
    )

    # 统计每个 rank 上实际存储的参数张量数量
    param_count = sum(1 for p in model_engine.module.parameters() if p._ddp_if_sharded if hasattr(p, '_ddp_if_sharded'))
    # 更通用的统计方式：检查参数是否被分片（stage3 会替换为 partitioned 参数）
    sharded_params = 0
    total_params = 0
    for p in model_engine.module.parameters():
        total_params += 1
        # ZeRO-3 会将参数包装为 PartitionedParameter，可通过是否拥有 _zp_mgr 或 _ds_shard 属性判断
        if hasattr(p, '_ds_shard') or hasattr(p, '_zp_mgr'):
            sharded_params += 1

    print(f"[Stage {stage}] Rank {torch.distributed.get_rank()} - "
          f"总参数张量: {total_params}, 分片参数张量: {sharded_params}")

    # 训练一步
    for step in range(3):
        logits = model_engine(batch)
        loss = nn.CrossEntropyLoss()(logits.view(-1, 50257), labels.view(-1))
        model_engine.backward(loss)
        model_engine.step()
        if step == 0:
            print(f"[Stage {stage}] 首步 loss: {loss.item():.4f}")

    # 显存统计
    mem = torch.cuda.max_memory_allocated() / 1e9
    print(f"[Stage {stage}] Rank {torch.distributed.get_rank()} 峰值显存: {mem:.2f} GB")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--stage", type=int, default=3)
    args = parser.parse_args()
    train(args.stage)

运行方式

# 使用 DeepSpeed 启动器运行（2 卡示例）
deepspeed --num_gpus=2 train_.py --stage=3

# 对比实验：运行 Dense 模式（stage=0，即无 ZeRO）
deepspeed --num_gpus=2 train_.py --stage=0

验证方法与判定标准

1. 参数分片验证：观察日志中 分片参数张量 数量。Stage 3 运行时，该数量应等于总参数张量数（所有参数均被分片）；Stage 0 运行时，该数量应为 0。
2. 显存对比验证：对比 Stage 3 与 Stage 0 的 峰值显存 输出。Stage 3 的每卡显存应显著低于 Stage 0（例如 4 卡时约降低 3~4 倍），证明参数未在每个 GPU 上完整复制。
3. 训练可收敛性：loss 应在前几步内下降（如从 ~10 降至 ~8 以下），说明分片后的前向/反向通信正确。

常见问题

• OOM 错误：若显存不足，可降低 train_batch_size 或 seq_len，或增加 GPU 数量。
• 通信慢：Stage 3 引入 all-gather 通信开销，训练速度慢于 Dense 模式属正常现象，可通过增大 stage3_prefetch_bucket_size 缓解。
• 参数统计为 0：若 _ds_shard 属性检测失败，可改用 model_engine._config.zero_optimization.stage == 3 判断当前是否启用 Stage 3，并直接打印 torch.distributed.get_world_size() 确认多卡环境。