4、DeepSpeed架构与组件：训练引擎、推理引擎、模型并行、流水线并行
DeepSpeed 并不是一个单一的训练脚本，而是一套围绕“显存效率”和“并行效率”构建的优化系统。理解它的架构，关键是看清四个核心组件如何分工：训练引擎负责把优化器、梯度、参数分片并调度通信；推理引擎负责把大模型低延迟、高吞吐地跑起来；模型并行解决单卡放不下模型的问题；流水线并行解决层间切分后设备空转的问题。四者共同构成 DeepSpeed 从训练到推理的完整能力面。
训练引擎：ZeRO 是核心，但不是全部
训练引擎是 DeepSpeed 最常被使用的入口。它通过 deepspeed.initialize() 包装模型、优化器和数据加载器，把原本由 PyTorch 管理的训练状态接管过来。其核心是 ZeRO（Zero Redundancy Optimizer），按分片粒度分为三个阶段：
- ZeRO-1：只分片优化器状态，显存节省约 4 倍，通信量最小。
- ZeRO-2：在 ZeRO-1 基础上分片梯度，显存节省约 8 倍，适合大多数中等规模训练。
- ZeRO-3：进一步分片模型参数，显存节省与 GPU 数量成正比，但通信开销显著增加。
训练引擎还集成了混合精度、梯度累积、通信重叠和 CPU/NVMe 卸载。一个典型的配置片段如下：
{
"train_batch_size": 64,
"gradient_accumulation_steps": 4,
"fp16": {
"enabled": true,
"loss_scale": 0,
"initial_scale_power": 16
},
"zero_optimization": {
"stage": 2,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"allgather_partitions": true,
"overlap_comm": true
}
}这里 stage: 2 表示启用 ZeRO-2，offload_optimizer 把优化器状态卸载到 CPU 内存，overlap_comm 让通信与计算重叠。这些参数直接决定显存占用和训练速度的平衡点。
推理引擎：从训练到服务的桥梁
推理引擎面向的是“已经训练好的模型如何高效服务”。它支持张量并行和流水线并行推理，把大模型切分到多卡上，同时通过自定义 CUDA Kernel、KV Cache 管理和动态批处理降低单次推理延迟。与训练引擎不同，推理引擎不需要优化器状态和梯度，因此显存压力主要来自模型权重和 KV Cache。
典型使用方式是先初始化推理引擎，再调用 generate()：
import deepspeed
engine = deepspeed.init_inference(
model,
mp_size=4,
dtype=torch.half,
replace_with_kernel_inject=True
)
output = engine.generate(input_ids, max_new_tokens=128)mp_size=4 表示使用 4 路张量并行，replace_with_kernel_inject=True 会替换为 DeepSpeed 优化过的 Kernel。推理引擎的价值在于：同一套模型代码，训练时用训练引擎，上线时用推理引擎，减少重复适配成本。
模型并行与流水线并行：两种切分哲学
模型并行和流水线并行解决的是同一个问题的两个侧面：单卡放不下模型。但它们的切分维度和通信模式完全不同。
| 对比维度 | 模型并行（张量并行） | 流水线并行 |
|---|---|---|
| 切分对象 | 单层内部的权重矩阵 | 模型的不同层 |
| 通信频率 | 每层前向/反向都通信 | 仅在层边界传递激活值 |
| 通信量 | 大，All-Reduce/All-Gather | 小，点对点发送 |
| 主要挑战 | 通信开销随并行度上升 | 流水线气泡导致设备空转 |
| 适用场景 | 单层参数极大（如超大 FFN） | 层数多、层间依赖强 |
| DeepSpeed 支持 | 与 ZeRO 组合使用 | 1F1B 调度、Micro-batch |
实际训练中，两者往往组合使用：先用流水线并行把模型按层切到多组设备，再在每组内部用张量并行切分单层。DeepSpeed 的 PipelineModule 和 mp_size 参数可以同时配置这两种并行。
案例与复盘：一次显存溢出的排查
某团队训练 13B 模型时，使用 8 卡 A100 40GB，配置 ZeRO-2 + FP16，仍然出现 OOM。排查后发现两个问题：一是 train_batch_size 设为 128，梯度累积步数只有 1，导致单步激活值过大；二是没有开启 offload_optimizer，优化器状态全部留在显存。
调整方案：把 train_batch_size 降到 32，gradient_accumulation_steps 提到 4，保持全局批次不变；同时开启 CPU 卸载。调整后显存峰值从 42GB 降到 31GB，训练吞吐仅下降约 8%。这个案例说明：DeepSpeed 的配置不是越激进越好，而是要在显存、通信和计算之间找到平衡点。
总结来看，DeepSpeed 的架构可以理解为“一个训练引擎 + 一个推理引擎 + 两套并行策略”的组合。训练引擎用 ZeRO 解决显存，推理引擎用 Kernel 优化和并行解决延迟，模型并行和流水线并行则从两个维度解决模型规模问题。理解它们各自解决什么、代价是什么，才能在具体场景中做出合理选型。