1. 课程导入

当你尝试在单张 GPU 上微调一个 7B、13B 甚至 70B 参数的大模型时，最先撞上的往往不是算法问题，而是显存墙：模型权重、梯度、优化器状态三份数据叠加，再加上激活值，显存需求会迅速膨胀到单卡无法承载的量级。DeepSpeed 正是为解决这类规模化训练与微调问题而生的深度学习优化库，它把「模型状态分片」和「内存复用」作为两条核心技术主线，让有限的显存能够支撑更大的模型和更长的序列。

一、DeepSpeed 是什么：从显存瓶颈出发的系统级优化库

DeepSpeed 是微软开源的一套面向大规模分布式训练的优化框架，它并不改变模型结构本身，而是从显存占用、通信效率、计算调度三个层面重构训练流程。它最广为人知的能力是 ZeRO（Zero Redundancy Optimizer），通过把优化器状态、梯度、参数分片到不同设备上，消除数据并行中的冗余副本；同时它还提供混合精度、梯度累积、激活检查点、CPU Offload、NVMe Offload 等内存复用与卸载手段。

从使用角度看，DeepSpeed 的接入成本并不高：通常只需要一份 JSON 配置文件，加上几行初始化代码，就能把已有的 PyTorch 训练脚本改造成分片训练脚本。下面是一个最小化的配置示例：

{
  "train_batch_size": 16,
  "gradient_accumulation_steps": 4,
  "fp16": { "enabled": true },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": { "device": "cpu" },
    "allgather_partitions": true,
    "overlap_comm": true
  }
}

这段配置的含义是：开启 FP16 混合精度，使用 ZeRO 第二阶段（分片优化器状态与梯度），并把优化器状态卸载到 CPU 内存，同时让通信与计算重叠。它已经能显著降低单卡显存占用，是很多微调任务的起点。

二、模型状态分片与内存复用：两条互补的技术主线

要理解 DeepSpeed 的价值，必须先看清「模型状态」到底占了多少显存。以 Adam 优化器为例，每个参数除了自身权重外，还需要保存一阶动量、二阶动量，以及 FP32 主权重副本。若以混合精度训练，一个 7B 模型仅优化器状态就可能超过 80GB，远超单卡容量。

模型状态分片的核心思想是：不再让每张卡都保存完整副本，而是把优化器状态、梯度、参数切分到各卡，需要时再通过通信聚合。ZeRO 的三个阶段正是这一思想的递进实现：Stage 1 分片优化器状态，Stage 2 再分片梯度，Stage 3 进一步分片参数。分片越彻底，单卡显存越低，但通信开销也越高，因此需要根据集群带宽和模型规模做权衡。

内存复用则从另一个维度解决问题：它不减少数据总量，而是让同一块显存在不同时刻被不同对象复用。典型手段包括激活检查点（用计算换显存）、CPU/NVMe Offload（把暂时不用的状态搬到更便宜的存储）、以及通信与计算重叠（让空闲时间被利用起来）。两者结合，才能在有限硬件上跑起更大的模型。

模型状态分片与内存复用的协同关系 模型状态（单卡完整副本） FP16 权重 梯度 优化器状态（动量 + 主权重） ZeRO 分片策略 Stage 1：分片优化器状态 Stage 2：分片梯度 Stage 3：分片参数 内存复用手段 激活检查点 CPU / NVMe Offload 通信与计算重叠 协同效果：显存占用下降，可训练模型规模上升 分片降低单卡冗余副本 → 复用让空闲显存被重新利用 → 通信与计算重叠隐藏部分开销 最终目标：在有限 GPU 数量下，稳定训练更大参数量的模型 本课程主线：从原理到实战，逐层拆解 原理层：ZeRO 三阶段 / 内存复用机制　→　配置层：JSON 参数与调优 实战层：微调案例 / 排障复盘　→　工具层：监控、日志与性能分析

三、方法流程：从单卡脚本到分片训练的四步改造

把普通 PyTorch 训练改造成 DeepSpeed 分片训练，通常遵循以下流程：

  1. 环境准备：安装 deepspeed、torch、transformers，确认 CUDA 与 NCCL 可用。
  2. 编写配置文件：确定 ZeRO stage、精度、batch size、offload 策略。
  3. 改造训练脚本：用 deepspeed.initialize 初始化模型、优化器和数据加载器。
  4. 启动与验证：用 deepspeed 启动器运行，观察显存、吞吐和 loss 曲线。

其中第三步的代码骨架大致如下：

import deepspeed

model, optimizer, _, lr_scheduler = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config_params="ds_config.json"
)

for step, batch in enumerate(dataloader):
    outputs = model(**batch)
    loss = outputs.loss
    model.backward(loss)
    model.step()

注意这里用 model.backward 和 model.step 替代了常规的 loss.backward 与 optimizer.step，因为 DeepSpeed 需要在反向和更新阶段插入分片通信与卸载逻辑。

四、案例场景：7B 模型微调的显存对比

假设我们在一台 4 卡 A100 40GB 的机器上微调 7B 模型，使用 Adam 优化器和 FP16 混合精度。不同策略下的单卡显存占用大致如下表所示：

策略 优化器状态 梯度 参数 单卡显存估算 是否可跑
纯数据并行 完整副本 完整副本 完整副本 约 90GB+ 否
ZeRO Stage 1 分片 完整副本 完整副本 约 45GB 否
ZeRO Stage 2 分片 分片 完整副本 约 28GB 是
ZeRO Stage 2 + CPU Offload 分片并卸载 分片 完整副本 约 18GB 是
ZeRO Stage 3 分片 分片 分片 约 12GB 是

可以看到，从纯数据并行到 ZeRO Stage 3，单卡显存需求下降了一个数量级。但代价是通信量增加，Stage 3 在低带宽集群上可能反而变慢。因此实际选择要结合硬件条件。

五、常见误区与复盘要点

误区一：分片越彻底越好。Stage 3 虽然省显存，但每次前向都要 all-gather 参数，通信开销大。若集群是低速以太网，Stage 2 往往更划算。

误区二：开了 Offload 就一定更快。CPU Offload 把数据搬到主机内存，PCIe 带宽可能成为新瓶颈，训练速度可能下降。它适合显存极度紧张、对速度不敏感的场景。

误区三：忽略 batch size 与梯度累积的配合。分片后单卡 batch 变小，若不同步调整 gradient_accumulation_steps，可能影响收敛稳定性。

复盘建议：每次调整配置后，记录三个指标——单卡峰值显存、每秒迭代步数、loss 收敛曲线。只有三者同时可接受，配置才算真正调好。

六、整体课程框架

本课程共 30 章，围绕「原理—配置—实战—排障」四条线展开：

  • 原理篇：ZeRO 三阶段的分片逻辑、内存复用机制、通信与计算重叠原理。
  • 配置篇：ds_config.json 各参数含义、stage 选择、offload 策略、混合精度设置。
  • 实战篇：单卡到多卡的改造、7B/13B 模型微调案例、吞吐与显存调优。
  • 排障篇：OOM 定位、通信死锁、速度回退、日志分析与性能剖析。

学完本课程，你将能够独立判断一个训练任务该用哪种分片策略，写出可运行的 DeepSpeed 配置，并在显存与速度之间找到适合自己硬件的平衡点。

← 上一章 📖 返回目录 下一章 →