2、DeepSpeed基础：分布式训练、ZeRO优化器、模型并行、检查点机制

要理解 DeepSpeed 为什么会出现模型加载失败、检查点损坏或恢复后精度异常，必须先弄清它在训练时到底“动了哪些东西”。DeepSpeed 并不是一个简单的训练加速库，它同时改变了数据分布方式、参数与优化器状态的存放位置、模型切分方式，以及检查点的组织格式。这四件事任意一件配置不一致，都会在加载阶段暴露为报错或静默错误。

一、分布式训练：数据并行是默认底座

DeepSpeed 建立在 PyTorch 的分布式通信之上，最常用的模式是数据并行（DP）。每个进程持有一份完整模型副本，读取不同的数据分片，反向传播后通过 AllReduce 同步梯度。其核心流程可以概括为：前向计算 → 反向求梯度 → 梯度归约 → 优化器更新。

关键点在于：数据并行下，模型结构和参数形状在所有 rank 上必须完全一致。如果加载时某个 rank 的模型定义不同，或者词表大小、层数不一致，就会出现形状不匹配。DeepSpeed 通过 dist_init_required、local_rank 等参数接管进程组初始化，因此启动方式与普通 PyTorch DDP 不同，必须用 deepspeed 命令或正确设置环境变量。

deepspeed --num_gpus=4 train.py \
  --deepspeed ds_config.json \
  --model_name_or_path ./ckpt

二、ZeRO 优化器：把状态切开放到不同 GPU

ZeRO（Zero Redundancy Optimizer）是 DeepSpeed 的核心。它把原本每张卡都完整保存的优化器状态、梯度、模型参数按 stage 逐级切分，从而大幅降低单卡显存。理解 ZeRO 的分级，是排查检查点问题的前提。

ZeRO-3 最容易出问题：参数在训练时并不完整地存在于单卡上，保存检查点时会写成分片文件。如果加载时 stage 从 3 改成 2，或 world_size 变化，分片数量对不上，就会直接报缺少文件或形状错误。

三、模型并行：张量并行与流水线并行

当单层参数大到一张卡放不下时，需要模型并行。DeepSpeed 支持张量并行（TP）和流水线并行（PP）。张量并行把同一层的权重矩阵按列或按行切开，分到不同 GPU；流水线并行把不同层分到不同 GPU，形成 stage 流水。

模型并行会直接改变参数的命名和形状。例如一个 Linear 层在 TP 下会变成带 tp 后缀的分片权重。加载检查点时，如果并行度 tensor_parallel_size 或 pipeline_parallel_size 与保存时不同，参数名和形状都会错位。这是“模型加载失败”中非常典型的一类原因。

四、检查点机制：分片、元数据与恢复

DeepSpeed 的检查点不是单个 pytorch_model.bin，而是一个目录，通常包含：模型分片文件、优化器状态分片、训练元数据（如 latest、zero_pp_rank_* 文件）以及随机数状态。保存时由 save_checkpoint 完成，加载时由 load_checkpoint 读取。

# ds_config.json 关键片段
{
  "train_batch_size": 32,
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu" }
  },
  "checkpoint": {
    "save_interval": 1000
  }
}

加载时 DeepSpeed 会校验：分片数量是否等于 world_size、每个分片的参数名是否匹配、优化器状态是否完整。任何一项不满足，就会抛出类似 Checkpoint mismatch 或 Missing key(s) in state_dict 的错误。

五、常见误区与复盘

• 误区一：认为改了 batch size 不影响加载。实际上 batch size 会影响优化器状态中的部分统计量，跨配置恢复时可能精度异常。
• 误区二：把 ZeRO-3 的检查点直接当普通权重加载，结果参数被切分，形状对不上。
• 误区三：换卡数后直接恢复训练，未重新分片，导致分片文件数量与 world_size 不匹配。
• 误区四：忽略 latest 文件，手动指定了不完整的分片目录。

复盘时建议固定一条原则：保存与加载必须使用同一份 ds_config，且 stage、world_size、TP/PP 并行度、优化器类型保持一致。若确实需要变更，应先离线合并权重，再以新配置重新初始化训练状态。理解这四层机制，后续排查加载失败时就能快速定位是数据并行、ZeRO 分片、模型并行还是检查点元数据出了问题。