DeepSpeed 完全指南：ZeRO 优化与大规模模型训练

引言
1. DeepSpeed 简介
2. ZeRO 优化分阶段解析
3. ZeRO-Offload 与 ZeRO-Infinity
- 3.1 ZeRO-Offload
- 3.2 ZeRO-Infinity(NVMe 卸载)
4. DeepSpeed 配置文件完全指南
5. DeepSpeed + PyTorch 集成
6. DeepSpeed 流水线并行
7. 张量并行
- 7.1 Megatron-DeepSpeed 集成
8. 激活值检查点(Gradient Checkpointing)
- 8.1 概念与权衡
9. 混合专家(MoE)与 DeepSpeed
- 9.1 MoE 架构概览
- 9.2 专家并行设置
10. DeepSpeed Inference
11. DeepSpeed Chat(RLHF)
12. 实战示例：LLM 训练
- 12.1 完整的 Llama-2 微调示例
- 12.2 性能基准测试
13. 调试与问题排查
- 13.1 常见问题与解决方案
- 13.2 性能调优技巧
结语
参考资料

引言

要训练 GPT-4、LLaMA、Falcon 这类大规模语言模型(LLM)，单张 GPU 是不可能做到的。1750 亿参数的 GPT-3 若以 fp16 存储，需要约 350GB 的 GPU 内存，若再算上优化器状态(Adam)，则会达到 1.4TB。为解决这一问题，Microsoft 开发出了 DeepSpeed 这个库。

DeepSpeed 以 ZeRO(Zero Redundancy Optimizer)优化技术为核心，让数百亿乃至数万亿参数的模型也能在普通 GPU 集群上完成训练。本指南将从 DeepSpeed 的核心概念一路讲到实战配置，全面覆盖。

1. DeepSpeed 简介

1.1 背景与动机

深度学习模型越来越大，由此产生三个瓶颈。

内存不足：模型参数、梯度、优化器状态超出 GPU 显存
算力瓶颈：单张 GPU 的 FLOPS 上限
通信瓶颈：多 GPU 环境下梯度同步的开销

传统的数据并行(Data Parallelism)让每张 GPU 都保留一份完整的模型副本，因此无法解决内存问题。模型并行(Model Parallelism)则实现复杂，且往往效率不高。

DeepSpeed 结合了这两种方案的优点，并通过全新的 ZeRO 优化实现了革命性的内存效率。

1.2 与 PyTorch 的集成

DeepSpeed 只需对现有 PyTorch 代码做最小限度的修改即可使用。核心改动就是一次 deepspeed.initialize() 调用，它会一并处理引擎创建、分布式初始化、优化器封装。

import deepspeed

# 现有的 PyTorch 代码
model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 切换到 DeepSpeed(最小改动)
model_engine, optimizer, _, _ = deepspeed.initialize(
    args=args,
    model=model,
    model_parameters=model.parameters(),
    config="ds_config.json"
)

1.3 安装

pip install deepspeed

# 包含 CUDA 内核编译的安装
DS_BUILD_OPS=1 pip install deepspeed

# 验证安装
ds_report

2. ZeRO 优化分阶段解析

ZeRO(Zero Redundancy Optimizer)是消除数据并行训练中内存冗余的核心技术。在数据并行中，N 张 GPU 都保留着完全相同的模型副本，这是极大的内存浪费。ZeRO 逐阶段地消除这种冗余。

训练中的内存占用分析

设模型参数量为 Ψ，在混合精度(fp16)训练下，每张 GPU 的内存占用如下。

组成部分	大小	说明
参数 (fp16)	2Ψ bytes	前向/反向传播
梯度 (fp16)	2Ψ bytes	反向传播的结果
主参数 (fp32)	4Ψ bytes	供优化器使用
Adam 动量 (fp32)	4Ψ bytes	一阶矩
Adam 方差 (fp32)	4Ψ bytes	二阶矩
合计	16Ψ bytes

若是 7B 参数的模型，大约需要 112GB；若是 70B 参数的模型，则需要约 1.1TB。

2.1 ZeRO-1：优化器状态分片

ZeRO-1 将占内存比重最大的优化器状态分散到 N 张 GPU 上。

Adam 的情形：fp32 参数 + 动量 + 方差 = 12Ψ bytes
分散到 N 张 GPU 后，每张 GPU 的优化器状态 = 12Ψ/N bytes

内存节省效果：

原始：16Ψ bytes per GPU
ZeRO-1：4Ψ + 12Ψ/N bytes per GPU(N=64 时 ≈ 4.2Ψ bytes，约节省 4 倍)

在更新过程中，每张 GPU 只更新自己负责的参数分片，随后通过 All-gather 恢复出完整参数。

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

2.2 ZeRO-2：梯度分片

ZeRO-2 在优化器状态之外，还进一步分散了梯度。

反向传播过程中，每张 GPU 只在自己负责的参数分片上累积对应的梯度。通过 Reduce-scatter 运算，每张 GPU 只收集自己负责区间的梯度。

原始：16Ψ bytes per GPU
ZeRO-2：2Ψ + (2Ψ + 12Ψ)/N bytes per GPU(N=64 时 ≈ 2.2Ψ bytes，约节省 8 倍)

{
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "overlap_comm": true,
    "contiguous_gradients": true
  }
}

2.3 ZeRO-3：参数分片

ZeRO-3 连模型参数本身也一并分散。这是最强力的方案，但实现也最复杂。

运行原理：

每张 GPU 平时只保留全部参数的 1/N
前向传播时：在每层执行前，通过 All-gather 收集出完整参数
计算完成后：立即释放该参数
反向传播同样按需(on-demand)收集参数

ZeRO-3：(2Ψ + 2Ψ + 12Ψ)/N bytes per GPU
N=64 时：16Ψ/64 ≈ 0.25Ψ bytes(约节省 64 倍！)

缺点： All-gather 通信在每一层都会发生，通信开销随之增加。不过由于与流水线执行相重叠，实际延迟并不大。

{
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  }
}

ZeRO-3 的 Python 初始化：

import deepspeed
from transformers import AutoModelForCausalLM

# 在 ZeRO-3 下，模型创建方式同样重要
with deepspeed.zero.Init(config_dict_or_path="ds_config.json"):
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

# 这样做能让参数从一开始就处于分片状态，避免内存峰值

3. ZeRO-Offload 与 ZeRO-Infinity

3.1 ZeRO-Offload

ZeRO-Offload 把优化器状态和梯度转移到 CPU 内存，大幅减轻 GPU 内存压力。

GPU：fp16 参数 + fp16 梯度(前向/反向传播)
CPU：fp32 主参数 + Adam 状态(优化器更新)

这样一来，即便只有单张 GPU，也能训练数十亿参数的模型。

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

参数卸载(ZeRO-3 + Offload)：

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    }
  }
}

3.2 ZeRO-Infinity(NVMe 卸载)

ZeRO-Infinity 进一步利用 NVMe SSD，支持事实上不受限制的模型规模。

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "nvme",
      "nvme_path": "/local_nvme",
      "pin_memory": true,
      "buffer_count": 4,
      "fast_init": false
    },
    "offload_param": {
      "device": "nvme",
      "nvme_path": "/local_nvme",
      "pin_memory": true,
      "buffer_count": 5,
      "buffer_size": 1e8
    },
    "aio": {
      "block_size": 1048576,
      "queue_depth": 8,
      "thread_count": 1,
      "single_submit": false,
      "overlap_events": true
    }
  }
}

带宽考量：

GPU-CPU 带宽：以 PCIe 4.0 x16 为准，约 32 GB/s
CPU-NVMe 带宽：以 NVMe SSD 为准，约 7 GB/s
NVMe 比 CPU 更慢，因此会拖慢训练速度，但能让原本不可能的模型规模成为可能

4. DeepSpeed 配置文件完全指南

4.1 ds_config.json 基本结构

{
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": 1.0,

  "fp16": {
    "enabled": "auto",
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 16,
    "hysteresis": 2,
    "min_loss_scale": 1
  },

  "bf16": {
    "enabled": "auto"
  },

  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "overlap_comm": true,
    "contiguous_gradients": true
  },

  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": "auto",
      "betas": "auto",
      "eps": 1e-8,
      "weight_decay": "auto"
    }
  },

  "scheduler": {
    "type": "WarmupLR",
    "params": {
      "warmup_min_lr": "auto",
      "warmup_max_lr": "auto",
      "warmup_num_steps": "auto"
    }
  },

  "activation_checkpointing": {
    "partition_activations": false,
    "cpu_checkpointing": false,
    "contiguous_memory_optimization": false,
    "number_checkpoints": null,
    "synchronize_checkpoint_boundary": false,
    "profile": false
  },

  "wall_clock_breakdown": false,
  "steps_per_print": 100
}

4.2 混合精度设置详解

fp16 设置：

{
  "fp16": {
    "enabled": true,
    "auto_cast": false,
    "loss_scale": 0,
    "initial_scale_power": 16,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "consecutive_hysteresis": false,
    "min_loss_scale": 1
  }
}

loss_scale: 0：启用动态损失缩放
initial_scale_power: 16：初始损失缩放 = 2^16 = 65536
loss_scale_window：缩放系数上调前所需的成功步数
hysteresis：缩放系数下调前所需的溢出计数

bf16 设置(Ampere 及以上架构 GPU)：

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

bf16 的动态范围比 fp16 更宽，因此不需要损失缩放。推荐在 A100/H100 上使用。

4.3 梯度累积与批大小

批大小的关系式：

train_batch_size = train_micro_batch_size_per_gpu × gradient_accumulation_steps × world_size

{
  "train_batch_size": 2048,
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 64
}

上述设置在 8 张 GPU 环境下：4 × 64 × 8 = 2048 的全局批大小。

使用 "auto" 时，由 Transformers Trainer 自动填充相应的值。

4.4 学习率调度器选项

{
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "last_batch_iteration": -1,
      "total_num_steps": 100000,
      "warmup_min_lr": 0,
      "warmup_max_lr": 3e-4,
      "warmup_num_steps": 2000,
      "warmup_type": "linear"
    }
  }
}

支持的调度器类型：

LRRangeTest
OneCycle
WarmupLR
WarmupDecayLR
WarmupCosineLR(自定义)

5. DeepSpeed + PyTorch 集成

5.1 基本训练循环

import torch
import deepspeed
from torch.utils.data import DataLoader

def main():
    # 初始化分布式环境
    deepspeed.init_distributed()

    # 创建模型
    model = GPT2LMHeadModel.from_pretrained("gpt2")

    # 初始化 DeepSpeed 引擎
    model_engine, optimizer, train_loader, lr_scheduler = deepspeed.initialize(
        model=model,
        training_data=train_dataset,
        config="ds_config.json"
    )

    # 训练循环
    for epoch in range(num_epochs):
        for batch in train_loader:
            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

            # 反向传播(缩放/累积由 DeepSpeed 处理)
            model_engine.backward(loss)

            # 参数更新(包含梯度裁剪)
            model_engine.step()

    print(f"Step: {model_engine.global_steps}, Loss: {loss.item():.4f}")

5.2 检查点的保存与加载

# 保存检查点
def save_checkpoint(model_engine, save_dir, tag=None):
    # 以 DeepSpeed 格式保存(包含分布式状态)
    model_engine.save_checkpoint(save_dir, tag=tag)

# 例：每 1000 步保存一次
if model_engine.global_steps % 1000 == 0:
    save_checkpoint(model_engine, "./checkpoints", tag=f"step_{model_engine.global_steps}")

# 加载检查点
_, client_sd = model_engine.load_checkpoint("./checkpoints", tag="step_1000")

# 在 ZeRO-3 环境下提取为普通 PyTorch 格式的权重
if args.zero_stage == 3:
    state_dict = model_engine._zero3_consolidated_16bit_state_dict()
    if model_engine.local_rank == 0:
        torch.save(state_dict, "model_weights.pt")

5.3 Transformers Trainer 集成

将 DeepSpeed 与 HuggingFace Transformers 集成的最简便方式：

from transformers import TrainingArguments, Trainer, AutoModelForCausalLM

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    fp16=True,
    deepspeed="ds_config.json",  # DeepSpeed 配置文件路径
    logging_steps=10,
    save_steps=1000,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

trainer.train()

命令行执行：

deepspeed --num_gpus=8 train.py \
    --deepspeed ds_config.json \
    --model_name_or_path meta-llama/Llama-2-7b-hf \
    --dataset_name openwebtext \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --num_train_epochs 3 \
    --fp16

5.4 Accelerate 集成

from accelerate import Accelerator
from accelerate.utils import DeepSpeedPlugin

deepspeed_plugin = DeepSpeedPlugin(
    zero_stage=2,
    gradient_accumulation_steps=8,
    gradient_clipping=1.0,
)

accelerator = Accelerator(
    mixed_precision="fp16",
    deepspeed_plugin=deepspeed_plugin,
)

model, optimizer, train_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader
)

for batch in train_dataloader:
    outputs = model(**batch)
    loss = outputs.loss
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

6. DeepSpeed 流水线并行

6.1 流水线并行的概念

流水线并行把模型的各层依序分布到多张 GPU 上。GPU 0 负责靠前的层，GPU 1 负责中间的层，GPU N 负责最后的层。

1F1B 调度(One Forward, One Backward)： 将流水线气泡降到最低的调度方式。每张 GPU 在完成一个微批次的前向传播后，紧接着立即执行反向传播。

6.2 使用 PipelineModule

from deepspeed.pipe import PipelineModule, LayerSpec

class GPT2PipelineModel(PipelineModule):
    def __init__(self, num_layers, hidden_size, num_heads, vocab_size):
        layers = [
            LayerSpec(EmbeddingLayer, vocab_size, hidden_size),
        ]
        for i in range(num_layers):
            layers.append(LayerSpec(TransformerBlock, hidden_size, num_heads))
        layers.append(LayerSpec(LMHead, hidden_size, vocab_size))

        super().__init__(
            layers=layers,
            loss_fn=cross_entropy_loss,
            num_stages=4,  # 流水线阶段数(= GPU 数量)
            topology=PipeDataParallelTopology(num_pp=4, num_dp=2)
        )

# 在 ds_config 中添加流水线设置
pipeline_config = {
    "train_batch_size": 64,
    "train_micro_batch_size_per_gpu": 2,
    "gradient_accumulation_steps": 8,
    "pipeline": {
        "seed_layers": true,
        "activation_checkpoint_interval": 1
    }
}

6.3 流水线 + ZeRO 组合

model_engine, _, _, _ = deepspeed.initialize(
    model=pipeline_model,
    config={
        "train_micro_batch_size_per_gpu": 2,
        "gradient_accumulation_steps": 8,
        "zero_optimization": {"stage": 1},  # 与 ZeRO-1 组合
        "fp16": {"enabled": True}
    }
)

# 流水线训练
for step in range(num_steps):
    loss = model_engine.train_batch()
    if model_engine.global_steps % 100 == 0:
        print(f"Step {model_engine.global_steps}, Loss: {loss.item():.4f}")

7. 张量并行

7.1 Megatron-DeepSpeed 集成

张量并行把单个层分布到多张 GPU 上，把矩阵乘法沿列或行方向进行切分。

# Megatron-DeepSpeed 设置
megatron_config = {
  "tensor_model_parallel_size": 4,
  "pipeline_model_parallel_size": 2,
  "data_parallel_size": 4,  # world_size / (tp_size × pp_size)
}

在 ds_config.json 中设置张量并行：

{
  "tensor_parallel": {
    "tp_size": 4,
    "mpu": null,
    "tp_grain_size": 64
  }
}

执行命令：

deepspeed --num_nodes=4 --num_gpus=8 \
    pretrain_gpt.py \
    --tensor-model-parallel-size 4 \
    --pipeline-model-parallel-size 4 \
    --num-layers 96 \
    --hidden-size 12288 \
    --num-attention-heads 96 \
    --seq-length 2048 \
    --global-batch-size 1024 \
    --train-iters 500000 \
    --deepspeed \
    --deepspeed_config ds_config.json

8. 激活值检查点(Gradient Checkpointing)

8.1 概念与权衡

激活值检查点(activation checkpointing)在前向传播中不保存中间激活值，而是在反向传播时重新计算。

内存节省：把与层数成正比的激活值内存，降为 O(sqrt(层数))
速度损失：约多出 33% 的额外计算成本

{
  "activation_checkpointing": {
    "partition_activations": true,
    "cpu_checkpointing": true,
    "contiguous_memory_optimization": true,
    "number_checkpoints": 4,
    "synchronize_checkpoint_boundary": false,
    "profile": false
  }
}

# 在 PyTorch 层面同样可以启用
from deepspeed.runtime.activation_checkpointing import checkpointing

class TransformerLayer(nn.Module):
    def forward(self, x):
        return checkpointing.checkpoint(self._forward, x)

    def _forward(self, x):
        # 实际运算
        return self.attention(self.norm1(x)) + x

9. 混合专家(MoE)与 DeepSpeed

9.1 MoE 架构概览

Mixture of Experts 让每个输入 token 只经由一部分"专家"子网络处理，从而在保持计算量不变的同时增加参数量。

from deepspeed.moe.layer import MoE

class MoETransformerBlock(nn.Module):
    def __init__(self, hidden_size, num_experts, top_k=2):
        super().__init__()
        self.attention = MultiHeadAttention(hidden_size)
        self.moe_layer = MoE(
            hidden_size=hidden_size,
            expert=FeedForward(hidden_size, hidden_size * 4),
            num_experts=num_experts,
            ep_size=1,          # 专家并行大小
            k=top_k,            # 激活的专家数量
            capacity_factor=1.25,
            eval_capacity_factor=2.0,
            min_capacity=4,
            use_residual=False
        )

    def forward(self, x, attention_mask=None):
        x = x + self.attention(x, attention_mask)
        x, _, _ = self.moe_layer(x)
        return x

9.2 专家并行设置

{
  "zero_optimization": {
    "stage": 2
  },
  "moe_expert_parallel_size": 4,
  "moe": {
    "enabled": true,
    "ep_size": 4,
    "moe_param_group": true
  }
}

# 分离 MoE 参数组(重要！)
def create_moe_param_groups(model):
    parameters = {
        "params": [],
        "name": "parameters"
    }
    moe_parameters = {
        "params": [],
        "moe": True,
        "name": "moe_parameters"
    }
    for module_name, module in model.named_modules():
        if isinstance(module, MoE):
            moe_parameters["params"].extend(
                [p for n, p in module.named_parameters() if "expert" in n]
            )
        else:
            parameters["params"].extend(module.parameters(recurse=False))
    return [parameters, moe_parameters]

10. DeepSpeed Inference

10.1 推理优化

DeepSpeed Inference 能大幅提升已训练模型的推理速度。

import deepspeed
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 转换为 DeepSpeed 推理引擎
ds_engine = deepspeed.init_inference(
    model=model,
    mp_size=1,              # 张量并行大小
    dtype=torch.float16,    # 推理 dtype
    checkpoint=None,
    replace_with_kernel_inject=True,  # 注入优化内核
    max_tokens=2048,
    replace_method="auto"
)

model = ds_engine.module

# 执行推理
inputs = tokenizer("DeepSpeed 是 ", return_tensors="pt").to("cuda")
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=200,
        do_sample=True,
        temperature=0.7
    )
print(tokenizer.decode(outputs[0]))

10.2 张量并行推理

# 多 GPU 推理(张量并行)
ds_engine = deepspeed.init_inference(
    model=model,
    mp_size=4,              # 分散到 4 张 GPU
    dtype=torch.float16,
    replace_with_kernel_inject=True,
    injection_policy={
        LlamaDecoderLayer: ("self_attn.o_proj", "mlp.down_proj")
    }
)

10.3 ZeRO-Inference

面向单张或少数 GPU 的大规模模型推理：

# 基于 ZeRO-3 的推理(参数从 CPU/NVMe 流式加载)
ds_engine = deepspeed.init_inference(
    model=model,
    config={
        "tensor_parallel": {"tp_size": 1},
        "dtype": "fp16",
        "enable_cuda_graph": False,
        "zero": {
            "stage": 3,
            "offload_param": {
                "device": "cpu"
            }
        }
    }
)

11. DeepSpeed Chat(RLHF)

11.1 RLHF 流水线概览

DeepSpeed-Chat 提供了完整的 RLHF(Reinforcement Learning from Human Feedback)流水线。

三阶段训练：

SFT(Supervised Fine-Tuning)：有监督微调
Reward Model Training：训练奖励模型
RLHF/PPO：以 PPO 算法进行强化学习

11.2 SFT 阶段

# Step 1: SFT
deepspeed main.py \
    --data_path Dahoas/rm-static \
    --data_split 2,4,4 \
    --model_name_or_path facebook/opt-1.3b \
    --per_device_train_batch_size 8 \
    --max_seq_len 512 \
    --learning_rate 9.65e-6 \
    --weight_decay 0.1 \
    --num_train_epochs 1 \
    --gradient_accumulation_steps 1 \
    --lr_scheduler_type cosine \
    --num_warmup_steps 0 \
    --seed 1234 \
    --zero_stage 2 \
    --deepspeed \
    --output_dir /output/sft

11.3 奖励模型训练

from deepspeed_chat.reward_model import RewardModel

# 奖励模型 = SFT 模型 + 线性头
reward_model = RewardModel(
    base_model=sft_model,
    tokenizer=tokenizer,
    num_padding_at_beginning=0
)

# 用偏好对(chosen, rejected)数据训练
def compute_reward_loss(reward_model, chosen_ids, rejected_ids):
    chosen_reward = reward_model(chosen_ids)
    rejected_reward = reward_model(rejected_ids)
    # 更受偏好的回答应得到更高的奖励
    loss = -torch.mean(torch.log(torch.sigmoid(chosen_reward - rejected_reward)))
    return loss

11.4 PPO 训练

from deepspeed_chat.ppo_trainer import DeepSpeedPPOTrainer

trainer = DeepSpeedPPOTrainer(
    rlhf_engine=rlhf_engine,
    args=args
)

for epoch in range(args.num_train_epochs):
    for step, batch in enumerate(prompt_dataloader):
        out = trainer.generate_experience(batch["input_ids"])
        # 更新 Actor/Critic 模型
        actor_loss, critic_loss = trainer.train_rlhf(out)

12. 实战示例：LLM 训练

12.1 完整的 Llama-2 微调示例

import torch
import deepspeed
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForSeq2Seq
)
from datasets import load_dataset
from peft import LoraConfig, get_peft_model, TaskType

def main():
    # 加载模型与分词器
    model_name = "meta-llama/Llama-2-7b-hf"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.pad_token = tokenizer.eos_token

    # 在 ZeRO-3 语境下加载模型(提升内存效率)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map=None  # 由 DeepSpeed 直接管理
    )

    # LoRA 设置(可选：参数高效微调)
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=16,
        lora_alpha=32,
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
        lora_dropout=0.1,
        bias="none"
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # 加载数据集
    dataset = load_dataset("tatsu-lab/alpaca", split="train")

    def preprocess(example):
        prompt = f"### Instruction:\n{example['instruction']}\n\n### Response:\n{example['output']}"
        tokens = tokenizer(prompt, truncation=True, max_length=512, padding="max_length")
        tokens["labels"] = tokens["input_ids"].copy()
        return tokens

    tokenized_dataset = dataset.map(preprocess, remove_columns=dataset.column_names)

    # 带 DeepSpeed 的 TrainingArguments
    training_args = TrainingArguments(
        output_dir="./llama2-finetuned",
        num_train_epochs=3,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=8,
        learning_rate=2e-4,
        warmup_steps=100,
        logging_steps=10,
        save_steps=500,
        bf16=True,
        deepspeed="ds_zero2_config.json",
        remove_unused_columns=False,
        report_to="wandb",
        run_name="llama2-alpaca-deepspeed"
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset,
        data_collator=DataCollatorForSeq2Seq(tokenizer, padding=True)
    )

    trainer.train()
    trainer.save_model("./llama2-finetuned-final")

if __name__ == "__main__":
    main()

ZeRO-2 配置文件(ds_zero2_config.json)：

{
  "bf16": {
    "enabled": "auto"
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": "auto",
      "betas": "auto",
      "eps": 1e-8,
      "weight_decay": "auto"
    }
  },
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "warmup_min_lr": "auto",
      "warmup_max_lr": "auto",
      "warmup_num_steps": "auto",
      "total_num_steps": "auto"
    }
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  },
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "steps_per_print": 100,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "wall_clock_breakdown": false
}

执行：

deepspeed --num_gpus=8 train_llama.py \
    --deepspeed ds_zero2_config.json

12.2 性能基准测试

在 8 张 A100 80GB 环境下的实测对比：

配置	最大模型规模	吞吐量 (tokens/sec)
基础 DDP	~7B	12,000
ZeRO-1	~14B	11,500
ZeRO-2	~35B	10,800
ZeRO-3	~175B	9,200
ZeRO-3 + CPU Offload	~350B	4,100
ZeRO-Infinity	~1T+	1,800

13. 调试与问题排查

13.1 常见问题与解决方案

OOM(内存不足)错误：

# 内存占用剖析
ds_config = {
    "memory_breakdown": True,
    "wall_clock_breakdown": True,
}

# 分步排查：
# 1. 提高 ZeRO stage(1 → 2 → 3)
# 2. 启用 CPU Offload
# 3. 启用激活值检查点
# 4. 减小批大小
# 5. 缩短序列长度

梯度溢出(fp16)：

{
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 500,
    "initial_scale_power": 12,
    "hysteresis": 2,
    "min_loss_scale": 1e-8
  }
}

ZeRO-3 参数收集错误：

# 在 ZeRO-3 下直接访问参数时
from deepspeed import zero

with zero.GatheredParameters(model.parameters()):
    # 仅在此代码块内可使用完整参数
    weight_norm = model.weight.norm()

13.2 性能调优技巧

# 1. 通信与计算重叠
zero_config = {
    "overlap_comm": True,
    "contiguous_gradients": True,
}

# 2. 优化桶大小(依据 GPU 内存与通信带宽调整)
zero_config["allgather_bucket_size"] = 5e8  # 500MB
zero_config["reduce_bucket_size"] = 5e8

# 3. ZeRO-3 预取
zero_config["stage3_prefetch_bucket_size"] = 5e7
zero_config["stage3_param_persistence_threshold"] = 1e6

# 4. 编译优化(PyTorch 2.0+)
model = torch.compile(model)

结语

DeepSpeed 是现代 LLM 训练的核心工具。借助 ZeRO 优化，原本单张 GPU 无法企及的大规模模型，如今也能以合理的成本完成训练。

核心要点：

ZeRO-1/2：分散优化器/梯度，提升数据并行的效率
ZeRO-3：连参数也一并分散，内存节省随 GPU 数量成比例增长
ZeRO-Offload：卸载到 CPU/NVMe，让小规模集群也能训练大型模型
流水线 + 张量并行：水平/垂直方向的模型分布，将通信效率最大化
DeepSpeed Inference：通过内核优化，将推理速度提升 2-5 倍

在实际项目中，ZeRO-2 + CPU Offload 的组合是不错的起点。若需要更大的模型，再转向 ZeRO-3；推理场景则可利用 DeepSpeed Inference 的 kernel injection。