上一篇（02《并行策略全景》）给了 ZeRO-1/2/3 的字节账与通信账。本篇沿同一条 bf16 参数的生命线进 DeepSpeed 0.19.2 的源码：Stage 3 里 param.data 只是一个壳，完整的 bf16 参数在用到前由 hook all-gather 出来、用完立即打碎，常驻的是每参数 \(1/N_d\) 的 ds_tensor 碎片和扁平 fp32 分区。三个框架并排的对照在 06《三个框架对比》，本篇只读 DeepSpeed。

一个 bf16 参数在 DeepSpeed Stage 3 里存在哪里、什么时候被 all-gather、什么时候被释放、它的 fp32 主副本在哪张卡上？¹

一、来龙去脉与核心抽象

1. 来龙去脉

DeepSpeed 是微软开源的训练框架，底座是 ZeRO 论文（Rajbhandari 等，ZeRO: Memory Optimizations Toward Training Trillion Parameter Models，arXiv 1910.02054，2019）：显存里最贵的是优化器状态、梯度、参数这三份常驻，那就逐级把它们切到 \(N_d\) 份。它的另一个卖点是零侵入——模型不需要用它的层写，任意一个 nn.Module 包进 engine 就能用 ZeRO，小集群训大模型是它最常见的场景。

2. 核心抽象

二、最小任务：runs/deepspeed

train-ledger 里的练手配置是 ZeRO-3、DP 8，模型取 Llama 3.2 1B 的形状（16 层、h = 2048、32 头 / 8 个 K/V 头、FFN 8192、词表 128256、s = 8192）。三个框架并排的对账与预期数字在 06 的 train-ledger 一节，这里只留 DeepSpeed 这份配置（runs/deepspeed/ds_config.json，键名以 0.19.2 runtime/constants.py 与 zero/config.py 为准）：

{
  "train_micro_batch_size_per_gpu": 1,
  "gradient_accumulation_steps": 2,
  "gradient_clipping": 1.0,
  "steps_per_print": 1,
  "bf16": { "enabled": true },
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": 3e-4, "betas": [0.9, 0.95], "weight_decay": 0.1 }
  },
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 5e7,
    "stage3_prefetch_bucket_size": 5e7,
    "stage3_param_persistence_threshold": 1e5,
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "sub_group_size": 1e9
  },
  "wall_clock_breakdown": true,
  "memory_breakdown": false,
  "flops_profiler": { "enabled": true, "profile_step": 20, "module_depth": -1, "top_modules": 3, "detailed": false }
}

train_batch_size 省略，由 1 × 2 × 8 = 16 推出。train_ds.py 的骨架：

"""train-ledger / runs/deepspeed/train_ds.py -- ZeRO-3 on any plain-PyTorch Llama implementation."""
import deepspeed, torch
from probe_memory import MemoryProbe

def main():
    deepspeed.init_distributed()
    cfg = "runs/deepspeed/ds_config.json"
    # zero.Init: parameters are partitioned (ds_tensor) as soon as each module's __init__ returns,
    # so the full 1B model is never materialised on one GPU.
    with deepspeed.zero.Init(config_dict_or_path=cfg, dtype=torch.bfloat16):
        model = build_llama_1b()            # any plain nn.Module Llama (16 layers, h=2048, 32/8 heads, ffn 8192)
    engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=cfg)
    probe = MemoryProbe("deepspeed", profile_steps=range(20, 22))
    for step, batch in enumerate(mock_batches(seq_len=8192, vocab=128256, device=engine.device)):
        with probe.step(step):
            for micro in range(engine.gradient_accumulation_steps()):
                loss = engine(**batch[micro])
                engine.backward(loss)
                engine.step()                # only the accumulation boundary takes a real optimizer step
        if step == 50:
            break
    probe.dump("runs/deepspeed/probe.json")

模型实现不限（任何一份纯 PyTorch 的 Llama 都可以），这正是 DeepSpeed 的卖点；build_llama_1b() 与 mock_batches() 不在本篇给出。

三、进程组：groups.py 的委托与自建

并行策略落到代码里的第一件事是进程组：每张卡属于哪一个 TP 组、哪一个 DP 组、哪一个 PP 组。DeepSpeed 的设计是委托优先。

deepspeed/utils/groups.py 的设计是委托优先：模块级变量 mpu 若被设置（用户通过 deepspeed.initialize(mpu=...) 传入一个 Megatron 风格的对象，需提供 get_data_parallel_group()、get_model_parallel_group() 等方法），所有 _get_data_parallel_group()、_get_model_parallel_world_size() 都转发给它；没有 mpu 时，DP 组就是 _clone_world_group()——整个世界的一个克隆。这反映了 DeepSpeed 的定位：它自己只负责 DP 系（ZeRO）的那一维，模型并行的进程组由外部框架（历史上是 Megatron-DeepSpeed）提供。

自建的组有三类。_create_expert_and_data_parallel() / _create_expert_data_and_model_parallel() 为 MoE 建 ep 与 expt-dp 组，按 ep_size 命名存进字典（_get_expert_parallel_group(group_name)）；_create_zero_param_parallel_group() 为 ZeRO++ 的 hpZ（zero_hpz_partition_size）建节点内的”二级参数分区组”，Stage 3 反向时的参数 all-gather 走它（ds_secondary_tensor）；_init_tp_mesh_device() 为 autoTP（tensor_parallel 配置）调用 dist.initialize_mesh_device()（deepspeed/comm/comm.py）建一个 (data_parallel, tensor_parallel) 的二维 DeviceMesh——这是 DeepSpeed 里少数用到 DeviceMesh 的地方。

流水线的进程组另有一套：deepspeed/runtime/pipe/topology.py 的 ProcessTopology 用一组 axes 与 dims 描述网格，PipeModelDataParallelTopology(num_pp, num_mp, num_dp) 是三维实例，PipelineParallelGrid 从中建出 PP、DP、MP 各组，PipelineEngine 用它而不是 groups.py。所以 DeepSpeed 的进程组事实上有三个来源：外部 mpu、groups.py、pipe/topology.py，DeepSpeedEngine._configure_distributed_model() 里 self.data_parallel_group、self.seq_data_parallel_group、self.expert_data_parallel_group 各自从不同来源取。

四、源码导读：按优化器状态切

1. 目录地图

deepspeed/
  __init__.py                          initialize(model, optimizer, config, mpu, ...) → (engine, optimizer, dataloader, lr_scheduler)
  runtime/
    engine.py                          DeepSpeedEngine(Module)：forward / backward / step；_configure_distributed_model / _configure_optimizer / _configure_zero_optimizer
    config.py                          DeepSpeedConfig：解析 JSON
    constants.py                       JSON 键名常量：train_batch_size、train_micro_batch_size_per_gpu、gradient_accumulation_steps、bf16 …
    zero/
      config.py                        DeepSpeedZeroConfig：stage、overlap_comm、contiguous_gradients、reduce_bucket_size、stage3_* …
      stage_1_and_2.py                 DeepSpeedZeroOptimizer（Stage 1 / 2）
      stage3.py                        DeepSpeedZeroOptimizer_Stage3
      partition_parameters.py          Init 上下文；_convert_to_deepspeed_param 给参数挂 ds_* 属性与 all_gather / partition 方法
      parameter_offload.py             DeepSpeedZeRoOffload：注册前向 / 反向 hook
      partitioned_param_coordinator.py PartitionedParameterCoordinator：fetch / release / prefetch
    pipe/
      module.py                        PipelineModule、LayerSpec、TiedLayerSpec
      engine.py                        PipelineEngine(DeepSpeedEngine)：train_batch、_exec_* 指令实现
      schedule.py                      TrainSchedule 等 PipeSchedule；PipeInstruction 子类
      topology.py                      ProcessTopology、PipelineParallelGrid
      p2p.py                           send / recv
  utils/groups.py                      进程组（第三章）
  comm/comm.py                         deepspeed.comm：对 torch.distributed 的封装，可换后端

2. DeepSpeedEngine：包裹用户模型

DeepSpeed 的入口是 deepspeed.initialize(model=..., optimizer=..., config=...)，返回一个 DeepSpeedEngine（runtime/engine.py）。它继承 nn.Module，持有用户的 self.module，用户之后调用的是 engine(batch)、engine.backward(loss)、engine.step() 三个方法，模型代码一行不改。构造函数的顺序：

DeepSpeedEngine.__init__
  _configure_distributed_model(model)        把 module 移到设备、转 dtype；确定 data_parallel_group / seq_data_parallel_group / expert_data_parallel_group
                                             非 ZeRO 时广播参数使各 DP rank 一致
  _configure_optimizer(client_optimizer, model_parameters)
    _configure_basic_optimizer()             按 JSON 的 optimizer 段建 Adam / AdamW / FusedAdam …
    → zero_optimization().stage > 0 ?        _configure_zero_optimizer(basic_optimizer)
        stage 1 / 2                          DeepSpeedZeroOptimizer(...)      （stage_1_and_2.py）
        stage 3                              DeepSpeedZeroOptimizer_Stage3(...)（stage3.py）
      : bf16 ?                               _configure_bf16_optimizer()
      : fp16 ?                               _configure_fp16_optimizer()
  _configure_lr_scheduler()

所以 DeepSpeed 的”ZeRO 优化器”是一个优化器包装类：它接管用户优化器的 param_groups，把里面的参数替换成分片；ZeRO 的一切——分片、hook、通信——都在这个包装类里，engine 只在 backward() 与 step() 里调用它。这与 Megatron（DDP 负责梯度通信、优化器负责参数分片与 all-gather，两个类）和 torchtitan（FSDP 负责一切通信、优化器是普通的）都不同。三种”谁持有分片、谁发起通信”的归属画在一起（黄色 = 分片与集合通信所在的类）：

%% 图：三个框架的分工：Megatron 是 DDP + DistributedOptimizer 两个类，DeepSpeed 是 ZeRO 优化器包装类，torchtitan 是 FSDP 负责一切通信
flowchart TB
    subgraph MG["Megatron Core"]
        direction TB
        MG_S["schedules.py<br/>经 4 个回调驱动 DDP"]
        MG_D["DistributedDataParallel<br/>_ParamAndGradBuffer<br/>梯度 hook / reduce-scatter<br/>执行参数 all-gather"]
        MG_O["DistributedOptimizer<br/>fp32 主参数 1/N_d 分片<br/>Adam 后发起 all-gather"]
        MG_M["模型层：core/transformer<br/>必须用它的层写"]
        MG_S --> MG_D
        MG_D <--> MG_O
        MG_D --> MG_M
        MG_O ~~~ MG_M
    end
    subgraph DS["DeepSpeed"]
        direction TB
        DS_E["DeepSpeedEngine<br/>forward / backward / step<br/>判断梯度累积边界"]
        DS_Z["ZeRO 优化器包装类<br/>stage_1_and_2 / stage3<br/>分片 + hook + 全部通信"]
        DS_U["用户 optimizer<br/>只看到 fp32 分区"]
        DS_M["任意 nn.Module，零侵入<br/>Stage 3 下参数是空壳"]
        DS_E --> DS_Z
        DS_Z --> DS_U
        DS_E --> DS_M
        DS_U ~~~ DS_M
    end
    subgraph TT["torchtitan"]
        direction TB
        TT_T["Trainer.train_step<br/>无显式集合通信"]
        TT_F["fully_shard（FSDP2）<br/>hook 内 unshard / reshard<br/>reduce-scatter，全部通信"]
        TT_O["torch.optim.AdamW<br/>普通优化器，更新 fp32 分片"]
        TT_M["Module 协议 + ShardingConfig<br/>参数是 DTensor"]
        TT_T --> TT_F
        TT_T --> TT_O
        TT_F --> TT_M
        TT_F ~~~ TT_O
        TT_O ~~~ TT_M
    end
    classDef comm fill:#fde68a,stroke:#b45309;
    classDef plain fill:#e0f2fe,stroke:#0369a1;
    class MG_D,MG_O,DS_Z,TT_F comm;
    class MG_M,DS_M,DS_U,TT_M,TT_O plain;

3. Stage 1 / 2：扁平分区

stage_1_and_2.py 的 DeepSpeedZeroOptimizer 构造时，对每个 param group：

1. 把组内所有 bf16 参数拍平拼接成一个一维张量 bit16_groups_flat[i]（flatten_dense_tensors_aligned()，pad 到 dp_size 的倍数），然后把每个参数的 .data 重新指向这个扁平张量里的视图——与 Megatron 的 param_data buffer 同一个思路，只是以 param group 而不是 dtype 为单位。
2. get_data_parallel_partitions() 把扁平张量均分成 dp_size 段，得到 parallel_partitioned_bit16_groups[i]。
3. single_partition_of_fp32_groups[i] = 本 rank 那一段的 fp32 克隆（.detach().clone().float()）——fp32 主参数，\(1/N_d\)。用户优化器的 param_groups[i]['params'] 被替换成这一个 fp32 张量。

梯度路径：create_gradient_handling_hooks() 给每个参数注册梯度 hook，reduce_ready_partitions_and_remove_grads() 把就位的梯度攒进 IPG（independent partition gradient）bucket（IPGBucket，大小 reduce_bucket_size，默认 5e8 元素），满了就 reduce_ipg_grads() → average_tensor()：Stage 2（partition_grads=True）时按分区归属把 bucket 里的梯度分别 reduce 到各自的 owner（reduce_scatter=True 时用 reduce-scatter），本 rank 只保留自己分区的梯度、其余立即释放；Stage 1 时 all-reduce、所有梯度保留。overlap_comm=True 时通信在单独的 stream 上与反向重叠；contiguous_gradients=True 时本 rank 分区的梯度拷进一段连续内存以减少碎片。

step()：检查溢出（fp16 时）→ 梯度裁剪 → 对每个 group _optimizer_step(i)：用户优化器在 fp32 分区上更新 → bit16_partitions[partition_id].data.copy_(fp32_partition.data)（fp32 → bf16，只写本 rank 那一段）→ all_gather_dp_groups() 把各 rank 的 bf16 段拼成完整的 bit16_groups_flat。因为参数的 .data 是扁平张量的视图，all-gather 完成的瞬间所有参数就是更新后的值。

与 Megatron 的 ZeRO-1 相比，机制几乎一样（扁平 buffer、按字节区间分片、fp32 分片拷贝、更新后 all-gather），差别在于粒度（DeepSpeed 以 param group 为 buffer、以 bucket 为 reduce 单位；Megatron 以 dtype 为 buffer、以 bucket group 为通信单位）和重叠方式（DeepSpeed 的 all-gather 在 step() 里同步完成，不与下一 step 前向重叠；Megatron 可以）。

4. Stage 3：ds_tensor 与 hook

Stage 3 的实现分散在四个文件里，按参数的生命周期读：

分片（partition_parameters.py）。deepspeed.zero.Init 是一个上下文管理器（类 Init，继承 InsertPostInitMethodToModuleSubClasses）：在它里面构造的每个 nn.Module，其 __init__ 结束后会被 _post_init_method() 拦截，对每个参数调用 _convert_to_deepspeed_param()。这个函数给参数挂上一组 ds_* 属性——ds_status（AVAILABLE / NOT_AVAILABLE / INFLIGHT）、ds_shape、ds_numel、ds_tensor（分片，本 rank 持有的 \(1/N_d\) 一维碎片）、ds_process_group、ds_persist（小于 stage3_param_persistence_threshold 的参数在前向、反向期间保持完整，只在 step 前后短暂分片）、ds_secondary_tensor（hpZ 的节点内二级副本）——以及一组方法：all_gather()、all_gather_coalesced()、partition()、ds_summary()。随后 _partition() 把参数的完整数据切成 dp_size 段、只留本 rank 的一段进 ds_tensor，然后 free_param() 把 param.data 换成一个空张量。从这一刻起，模型里的 nn.Parameter 只是一个壳，形状对但没有数据。若不用 zero.Init，DeepSpeedZeRoOffload.__init__（parameter_offload.py）里的 _convert_to_zero_parameters() 会对已构造好的模型做同样的转换，但那要求模型先能完整放进内存。

取回与释放（parameter_offload.py + partitioned_param_coordinator.py）。DeepSpeedZeRoOffload.setup_zero_stage3_hooks() 遍历模型的每个子 module，_register_deepspeed_module() 给它注册四个 hook：_pre_forward_module_hook、_post_forward_module_hook、_pre_backward_module_hook、_post_backward_module_hook。它们分别调用 pre_sub_module_forward_function() → PartitionedParameterCoordinator.fetch_sub_module(sub_module, forward=True)（把该 module 直接持有的参数 all-gather 回来，ds_status 变 AVAILABLE，param.data 指向拼好的完整张量）、post_sub_module_forward_function() → release_sub_module()（对不再被任何活跃 module 使用、且非 ds_persist 的参数调用 __release_param() → param.partition()，完整张量释放）；反向对称。

PartitionedParameterCoordinator 里有两个提高效率的机制。trace：第一个 step 记录 module 的执行顺序（record_module()），之后按这个顺序预取——fetch_sub_module() 在取回当前 module 的参数时，顺带发起后面几个 module 的 all-gather（总量由 stage3_prefetch_bucket_size 控制），stage3_max_live_parameters 限制同时活着的完整参数总数、stage3_max_reuse_distance 决定”很快又要用”的参数不释放。coalesced all-gather：all_gather_coalesced()（partition_parameters.py 里的 _all_gather_coalesced()）把一个 module 的多个参数拼成一次 all_gather_into_tensor 调用，返回 AllGatherCoalescedHandle，wait() 时再拆回各参数。

四个文件在一层的前向、反向里各自出场的次序如下（ds_status 与 param.data 的身份随之变化四次）：

%% 图：ZeRO Stage 3 一层的前向与反向：pre-hook 让 ParamCoordinator all-gather 参数，post-hook 释放回分片，ds_status 变化四次
sequenceDiagram
    participant FW as 层 k forward / backward
    participant HK as DeepSpeedZeRoOffload
    participant PC as ParamCoordinator
    participant P as 层 k 的 weight
    participant Z3 as Stage3 优化器
    FW->>HK: _pre_forward_module_hook
    HK->>PC: fetch_sub_module(层 k, forward=True)
    PC->>P: all_gather_coalesced：N_d 个碎片拼成完整
    Note over P: NOT_AVAILABLE -> INFLIGHT -> AVAILABLE，param.data 指向完整张量
    PC-)PC: 按 trace 顺带预取层 k+1.. 的参数
    FW->>FW: GEMM 读完整 weight
    FW->>HK: _post_forward_module_hook
    HK->>PC: release_sub_module(层 k)
    PC->>P: partition()（非 ds_persist 且无活跃 module 使用）
    Note over P: AVAILABLE -> NOT_AVAILABLE，param.data 换回空张量
    FW->>HK: _pre_backward_module_hook
    HK->>PC: fetch_sub_module(层 k, forward=False)
    PC->>P: 再 all-gather 一次
    FW->>Z3: 梯度 hook reduce_partition_and_remove_grads
    Z3->>Z3: 进 IPG bucket，满则 reduce-scatter，本 rank 段写入 fp32 分区 .grad
    FW->>HK: _post_backward_module_hook
    HK->>PC: release_sub_module(层 k)
    PC->>P: partition()，param.grad = None

梯度（stage3.py）。create_reduce_and_remove_grad_hooks() 给每个参数注册 reduce_partition_and_remove_grads hook；梯度就位后进 IPG bucket（IPGBucketZ3），满了 __reduce_and_partition_ipg_grads() → __avg_scatter_grads()（reduce-scatter）→ partition_grads() 把本 rank 分区的梯度写进 fp32_partitioned_groups_flat[sub_group].grad 或 bf16 的梯度分区（取决于 gradient_accumulation_dtype），然后 param.grad = None。反向结束时该参数本身也已被 _post_backward_module_hook 释放。

fp32 主参数与 step（stage3.py）。参数按 sub_group_size（默认 1e9 元素）分成若干 sub-group，每个 sub-group 有一个扁平的 bf16 分区 fp16_partitioned_groups_flat[i]（各参数 ds_tensor 的拼接）和一个扁平的 fp32 分区 fp32_partitioned_groups_flat[i]（_create_fp32_partitions()，offload_optimizer 时放 CPU pinned memory）。step() 对每个 sub-group：_prepare_fp32_grad_for_sub_group() → _optimizer_step(i)（用户优化器只看到这个 fp32 分区）→ _reassign_or_swap_out_partitioned_parameters() 把 fp32 分区拷回 bf16 分区（fp16_partitioned_groups_flat[i].data.copy_(fp32_partitioned_groups_flat[i].data)）。_post_step() 只 all-gather persistent_parameters——其余参数不在 step 里 all-gather，等下一 step 前向的 hook 按需取回。这就是第二篇里”ZeRO-3 省掉优化器后那次 all-gather”的代码位置。

5. 一个 bf16 参数的一生：DeepSpeed Stage 3

以一个 Transformer 层的 nn.Linear.weight 为例（bf16.enabled=true，zero_optimization.stage=3，不 offload）：

常驻：bf16 碎片 \(2N/N_d\)、fp32 主参数 + m + v \(12N/N_d\)、梯度分区 \(2N/N_d\)（bf16 累加）或 \(4N/N_d\)——合计 \(16N/N_d\) 或 \(18N/N_d\)。临时：正在前向 / 反向的那几层的完整 bf16 参数（预取深度决定有几层）。通信：每层前向一次 all-gather、反向一次 all-gather、一次 reduce-scatter，合计 \(3N\)。

与 Megatron 的时序放在一起，最大的差别是 param.data 的身份是流动的：Megatron 的 weight.data 从头到尾是 buffer 里同一段视图；DeepSpeed Stage 3 的 weight.data 在一个 step 里被替换四次（空 → 完整 → 空 → 完整 → 空）。这带来 Stage 3 的两个著名代价：任何在 hook 之外访问参数的代码（自定义的权重初始化、param.data 的手工操作、把参数当普通张量传给别的函数）都会拿到空张量，需要用 deepspeed.zero.GatheredParameters 显式取回；以及调试困难——参数”不在”的时候栈里看不到原因。stage3_module_granularity_threshold 与 leaf_module 配置（DeepSpeedZeroLeafModuleConfig）让 hook 挂在更粗的 module 上，是对这个代价的折衷。

6. 训练循环：forward / backward / step

DeepSpeed 没有训练循环——用户写循环，engine 提供三个方法：

for step in ...:
  for micro in range(gradient_accumulation_steps):
    loss = engine(batch)                     DeepSpeedEngine.forward：计时、flops profiler、Stage 3 的 root hook 起点
    engine.backward(loss)                    loss / gradient_accumulation_steps；ZeRO 时 optimizer.scale_if_loss；loss.backward()
                                             _backward_epilogue：ZeRO 时 optimizer.backward_epilogue()（清 IPG bucket、结束 reduce）
                                                                非 ZeRO 时 allreduce_gradients()（在梯度累积边界）
    engine.step()                            is_gradient_accumulation_boundary() 才真正 _take_model_step：
                                               optimizer.step()（ZeRO 的 step，含裁剪、更新、all-gather）
                                               lr_scheduler.step(); zero_grad
                                             global_steps += 1；micro_steps += 1

gradient_accumulation_steps 由 JSON 的 train_batch_size / (train_micro_batch_size_per_gpu × dp_size) 推出，engine 用 micro_steps 计数判断边界——用户不需要自己写 no_sync。这个设计让”零侵入”成立：现有的单卡训练脚本只要把 model 换成 engine、loss.backward() 换成 engine.backward(loss)、optimizer.step() 换成 engine.step() 就能用 ZeRO-3。代价是控制权在 engine 手里：什么时候 all-gather、什么时候 reduce、什么时候释放，都由 JSON 里的阈值决定，用户能调的是数字而不是顺序。

7. pipe/ 引擎与 JSON 配置的代价

流水线在 DeepSpeed 里是另一个 engine：runtime/pipe/module.py 的 PipelineModule 要求用户把模型写成一个 LayerSpec 列表（每层的类与构造参数），_partition_layers() 按 partition_method（"uniform"、"parameters"、"type:regex"）把层分给各 stage；runtime/pipe/engine.py 的 PipelineEngine 继承 DeepSpeedEngine，train_batch() 取代 forward / backward / step 三步：它构造一个 TrainSchedule（runtime/pipe/schedule.py），_exec_schedule() 逐步执行调度产出的指令。

TrainSchedule.steps() 是一个生成器，共 2 × (micro_batches + stages - 1) 步，每步 yield 一组 PipeInstruction——RecvActivation、ForwardPass、SendActivation、RecvGrad、BackwardPass、SendGrad、LoadMicroBatch、最后一步加 ReduceTiedGrads、ReduceGrads、OptimizerStep；_exec_schedule() 用 _INSTRUCTION_MAP 把每种指令映射到 _exec_forward_pass()、_exec_send_activations() 等方法。这是介于 Megatron 的过程式与 PyTorch pipelining 的动作表之间的第三种写法：调度是声明的（指令序列），执行是解释的。num_pipe_buffers() 返回 min(stages - stage_id, micro_batches)——就是 1F1B 在途 micro-batch 数 \(p - i\)。PP 与 ZeRO 的组合有限制：_configure_zero_optimizer() 在 PipelineModule 下强制 overlap_comm=False，Stage 3 与 PP 不兼容。

JSON 配置是 DeepSpeed 易用性的来源，也是它的代价所在。runtime/zero/config.py 的 DeepSpeedZeroConfig 有四十多个字段，其中 Stage 3 的六个阈值（stage3_prefetch_bucket_size、stage3_param_persistence_threshold、stage3_max_live_parameters、stage3_max_reuse_distance、sub_group_size、reduce_bucket_size）共同决定显存峰值与通信重叠程度，但它们之间的关系不在任何一个地方被算出来——用户要靠经验或 see_memory_usage 的日志调。而且配置里有大量隐式耦合：train_batch_size、train_micro_batch_size_per_gpu、gradient_accumulation_steps 三者给两个推第三个；bf16.enabled 与 fp16.enabled 互斥；zero_optimization.offload_optimizer 只在 Stage 1–3 有效但 Stage 1/2 与 3 的行为不同。JSON 无法表达这些约束，只能在 DeepSpeedConfig 的校验里报错。

五、前向前 all-gather：Stage 3 的代码位置

同一个”前向前把分片参数拼回来”的动作，Stage 3 的实现分在四个文件里——分片（partition_parameters.py）、hook（parameter_offload.py）、协调（partitioned_param_coordinator.py）、梯度与 step（stage3.py），本篇 第四章 已逐个读过。它与 FSDP1、FSDP2 的逐项并排在 06《三个框架对比》。

六、本文小结

一个 bf16 参数在 DeepSpeed Stage 3 里的一生：常驻 \(1/N_d\) 碎片（ds_tensor），param.data 只是形状对的空壳；前向到某个 module 时 hook 触发参数协调器 all-gather 出完整参数、用完立即释放回碎片，反向再来一次；fp32 主参数是优化器里的扁平分区，只在 step 时被 Adam 看到——合计常驻约 \(16N/N_d\)、通信 \(3N\)。param.data 身份的流动（一个 step 内换四次）是 Stage 3 一切的来源，也是它一切代价的来源。

本篇涉及的源码位置（DeepSpeed 0.19.2，另见各节目录地图）：

七、自测

1. DeepSpeed ZeRO-3 的”参数协调器”做什么？trace 预取解决了什么问题？

   答案

   在 module 的 pre-forward / pre-backward hook 里 all-gather 即将用到的参数、post hook 里释放；第一步记录 module 的执行顺序（trace），之后按顺序提前 all-gather 下一个 module 的参数——把通信藏在当前 module 的计算后面，否则每层都要停下等参数。

2. deepspeed.zero.Init 解决了什么问题？不用它会怎样？

   答案

   它在每个 nn.Module.__init__ 结束时拦截并把参数切成 ds_tensor——模型从构造起就只有分片，完整参数从不占内存。不用它的话 _convert_to_zero_parameters() 也能转换已构造好的模型，但要求完整模型先能放进内存，大模型在这一步就 OOM。

3. engine.step() 为什么不一定真的更新参数？Stage 3 的 step 里为什么不做参数 all-gather？

   答案

   is_gradient_accumulation_boundary() 为真才 _take_model_step，其余 step 只累积梯度——梯度累积边界由 JSON 推出，用户不用自己写 no_sync。Stage 3 的 step 只 all-gather ds_persist 的小参数；完整参数在下一 step 前向被 hook 按需取回，省掉一次全量 all-gather。

• 上一篇
  大规模训练工程（03）：Megatron-LM——按模型结构切的架构与源码导读
• 下一篇
  大规模训练工程（05）：torchtitan——用 PyTorch 原生原语组合并行的架构与源码导读