3. ZeRO Stage 3 参数切分

任务目标

本案例将配置DeepSpeed ZeRO Stage 3，在单机多GPU（或单GPU模拟）环境下训练模型，验证模型参数被切分到各GPU显存中，从而在有限显存下完成大模型训练。

环境依赖

• Python 3.8+，PyTorch 1.12+，DeepSpeed 0.9+，transformers 4.30+
• 安装：pip install deepspeed transformers
• 硬件：至少1块GPU（推荐4块，显存≥8GB）

核心实现

以下代码使用transformers加载（124M参数），通过DeepSpeed ZeRO Stage 3配置训练。Stage 3将模型参数、梯度和优化器状态全部切分到各GPU，每张卡仅持有1/世界大小参数分片。

import os
import torch
import deepspeed
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# 1. 初始化DeepSpeed引擎（必须在模型加载前）
deepspeed.init_distributed()

# 2. 加载模型与分词器（模型先加载到CPU，由ZeRO-3切分）
model = GPT2LMHeadModel.from_pretrained("")
tokenizer = GPT2Tokenizer.from_pretrained("")
tokenizer.pad_ = tokenizer.eos_

# 3. DeepSpeed ZeRO Stage 3配置
ds_config = {
    "train_batch_size": 8,
    "gradient_accumulation_steps": 1,
    "optimizer": {
        "type": "AdamW",
        "params": {"lr": 3e-5}
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {"device": "none"},
        "offload_param": {"device": "none"},
        "reduce_bucket_size": 5e8,
        "stage3_prefetch_bucket_size": 5e8,
        "stage3_param_persistence_threshold": 1e6
    },
    "fp16": {"enabled": True}
}

# 4. 用DeepSpeed包装模型
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config=ds_config
)

# 5. 构造训练数据（使用简单示例文本）
texts = [
    "The quick brown fox jumps over the lazy dog.",
    "DeepSpeed enables large model training with limited memory.",
    "ZeRO stage 3 partitions parameters across GPUs.",
    "Natural language processing benefits from large pretrained models."
]
inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=64)
input_ids = inputs["input_ids"].cuda()
attention_mask = inputs["attention_mask"].cuda()
labels = input_ids.clone()

# 6. 训练循环（3步验证参数分片）
model_engine.train()
for step in range(3):
    outputs = model_engine(input_ids, attention_mask=attention_mask, labels=labels)
    loss = outputs.loss
    model_engine.backward(loss)
    model_engine.step()
    if model_engine.global_steps % 1 == 0:
        print(f"Step {model_engine.global_steps}, Loss: {loss.item():.4f}")

# 7. 验证参数分片：检查每张卡上参数数量
param_count = sum(p.numel() for p in model_engine.module.parameters())
total_params = sum(p.numel() for p in GPT2LMHeadModel.from_pretrained("").parameters())
world_size = torch.distributed.get_world_size()
print(f"GPU {torch.distributed.get_rank()}: 本地参数数 {param_count}, 总参数数 {total_params}, 世界大小 {world_size}")
print(f"分片比例: {param_count / total_params:.4f} (预期约 1/{world_size})")

运行方式

# 单机4卡运行
deepspeed --num_gpus=4 train__zero3.py

# 单卡运行（模拟Stage 3，分片比例为1）
deepspeed --num_gpus=1 train__zero3.py

验证方法与判定标准

• 训练收敛：3步内loss应下降（如从约4.5降至约4.0以下），说明前向/反向传播正常。
• 参数分片验证：日志中“分片比例”应接近1/GPU数。例如4卡时约0.25，说明每卡仅持有25%参数，证明Stage 3生效。
• 显存对比：使用nvidia-smi观察每卡显存占用，4卡训练时每卡显存应显著低于单卡加载完整模型所需（约1.2GB），通常每卡约400-600MB。

常见问题

• OOM错误：若显存不足，将offload_param和offload_optimizer的device改为"cpu"，将参数/优化器状态卸载到CPU内存。
• 单卡运行：Stage 3在单卡下分片比例为1，但代码仍可运行，用于功能验证。
• 版本兼容：transformers与DeepSpeed版本不匹配时，可固定使用transformers==4.30.2和deepspeed==0.9.5。