2. ZeRO Stage 1 优化器状态切分

任务目标

本案例使用 DeepSpeed ZeRO Stage 1 训练一个小型 Transformer 模型（基于 torchvision 的 MNIST 图像分类任务），对比开启与关闭 ZeRO Stage 1 时的 GPU 显存占用和训练吞吐量，验证优化器状态切分带来的显存收益。

环境依赖

• Python 3.8+，PyTorch 1.13+，DeepSpeed 0.9+
• 单张 NVIDIA GPU（显存 ≥ 4GB），CUDA 11.x
• 安装：pip install deepspeed torch torchvision

核心实现

以下代码定义了一个小型 Transformer 编码器（2 层，隐藏维度 128），在 MNIST 上完成分类训练。通过 deepspeed.initialize 启用 ZeRO Stage 1，并记录训练耗时与显存峰值。

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as T
from torch.utils.data import DataLoader
import deepspeed
import time
import argparse

# ---------- 小型 Transformer 模型 ----------
class SmallTransformer(nn.Module):
    def __init__(self, img_size=28, patch_size=7, dim=128, depth=2, heads=4, num_classes=10):
        super().__init__()
        self.patch_size = patch_size
        n_patches = (img_size // patch_size) ** 2
        self.patch_embed = nn.Conv2d(1, dim, kernel_size=patch_size, stride=patch_size)
        self.pos_embed = nn.Parameter(torch.randn(1, n_patches, dim) * 0.02)
        encoder_layer = nn.TransformerEncoderLayer(d_model=dim, nhead=heads, dim_feedforward=256, dropout=0.1, batch_first=True)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=depth)
        self.head = nn.Linear(dim, num_classes)

    def forward(self, x):
        x = self.patch_embed(x)                     # [B, dim, n_patches]
        x = x.flatten(2).transpose(1, 2)            # [B, n_patches, dim]
        x = x + self.pos_embed
        x = self.encoder(x)
        x = x.mean(dim=1)                           # 全局平均池化
        return self.head(x)

# ---------- 数据加载 ----------
def get_dataloaders(batch_size=64):
    transform = T.Compose([T.ToTensor(), T.Normalize((0.1307,), (0.3081,))])
    train_ds = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=2)
    return train_loader

# ---------- 训练主函数 ----------
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--local_rank', type=int, default=-1)
    parser.add_argument('--use_zero', action='store_true', help='启用 ZeRO Stage 1')
    parser.add_argument('--epochs', type=int, default=3)
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--lr', type=float, default=1e-3)
    args = parser.parse_args()

    # 模型与优化器
    model = SmallTransformer()
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
    criterion = nn.CrossEntropyLoss()

    # DeepSpeed 配置
    if args.use_zero:
        ds_config = {
            "train_batch_size": args.batch_size,
            "optimizer": {"type": "AdamW", "params": {"lr": args.lr}},
            "zero_optimization": {
                "stage": 1,
                "reduce_bucket_size": 5e7,
                "allgather_bucket_size": 5e7
            },
            "fp16": {"enabled": True},
            "zero_allow_untested_optimizer": True
        }
    else:
        ds_config = {
            "train_batch_size": args.batch_size,
            "optimizer": {"type": "AdamW", "params": {"lr": args.lr}},
            "fp16": {"enabled": True}
        }

    # 初始化 DeepSpeed 引擎
    model_engine, optimizer, _, _ = deepspeed.initialize(
        args=args,
        model=model,
        optimizer=optimizer,
        config=ds_config
    )

    train_loader = get_dataloaders(args.batch_size)
    model_engine.train()
    total_steps = 0
    start_time = time.time()

    for epoch in range(args.epochs):
        for batch_idx, (images, labels) in enumerate(train_loader):
            images, labels = images.cuda(), labels.cuda()
            outputs = model_engine(images)
            loss = criterion(outputs, labels)
            model_engine.backward(loss)
            model_engine.step()
            total_steps += 1

            if batch_idx % 100 == 0:
                print(f"Epoch {epoch+1}/{args.epochs} | Step {batch_idx} | Loss: {loss.item():.4f}")

    elapsed = time.time() - start_time
    throughput = total_steps * args.batch_size / elapsed
    mem_allocated = torch.cuda.max_memory_allocated() / 1024**3
    print(f"\n=== 训练完成 ===")
    print(f"ZeRO Stage 1: {'启用' if args.use_zero else '关闭'}")
    print(f"总步数: {total_steps} | 耗时: {elapsed:.2f}s")
    print(f"吞吐量: {throughput:.2f} samples/s")
    print(f"GPU 显存峰值: {mem_allocated:.2f} GB")

if __name__ == "__main__":
    main()

运行与对比验证

分别以关闭和启用 ZeRO Stage 1 两种方式运行，观察输出中的显存峰值与吞吐量差异。

# 1. 关闭 ZeRO（基线）
deepspeed --num_gpus=1 train_zero1.py --epochs 3

# 2. 启用 ZeRO Stage 1
deepspeed --num_gpus=1 train_zero1.py --epochs 3 --use_zero

验证判定标准：

• 两种模式均能正常完成 3 个 epoch 训练，Loss 逐步下降，无报错退出；
• 启用 ZeRO Stage 1 后，输出的 GPU 显存峰值 应低于关闭时的数值（通常降低 20%～40%）；
• 吞吐量（samples/s）允许有小幅波动，但不应出现数量级下降。

ZeRO Stage 1 显存切分示意

常见问题

• 显存未下降：确认模型参数量足够大（建议 ≥ 100M 参数），且 batch_size 已尽量调大以占用显存；小模型下 ZeRO 收益不明显。
• fp16 相关报错：若 GPU 不支持或未安装 apex，可移除 "fp16": {"enabled": true} 配置，改用 fp32 训练（显存对比趋势不变）。
• 多卡运行：本案例默认单卡，若需多卡对比，将 --num_gpus 改为 2 或 4，并确保模型参数量足够大以体现切分收益。