# PR #7324 完整报告

- 仓库：`verl-project/verl`
- 标题：[ckpt, fsdp] feat: sharded delta weight sync for the TorchTitan engine
- 合并时间：2026-08-21 10:15
- 原文链接：http://prhub.com.cn/verl-project/verl/pull/7324

---

# 执行摘要

- 一句话：TorchTitan 引擎接入分片 delta 权重同步，最高提速 17.2×
- 推荐动作：值得精读。三个设计决策尤其值得学习：(1) 用数学分析而非特判接纳 `_StridedShard`——右到左切割只置换 block 归属而非交错 block，`compute_local_shape_and_global_offset` 已返回置换后偏移；(2) 用 slot table 而非 probe-driven 转换器解决 lockstep gather 下的命名一致性问题；(3) gather 组的「变化 Shard 维、固定其他维」构造避免 HSDP 副本浪费，并按 rank 集合缓存保证 `dist.new_group` 的建组顺序一致。若要在 verl 中为新的 trainer 引擎接入 delta 同步，本 PR 是很好的实现模板。

# 功能与动机

delta_sharded checkpoint 引擎只同步发生变化的权重元素，要求训练引擎按 rank 导出本地分片与放置元数据；TorchTitan 引擎此前只有全量导出，选择 `delta_sharded` 会直接抛 `NotImplementedError`。PR body 用基准数据说明收益：Qwen3-8B 在 32 卡 FSDP2 下 `delta_sharded` 为 3.34 s 而 nccl 全广播为 34.26 s（10.3×），Qwen3-30B-A3B MoE 在 FSDP2×EP8 下为 8.07 s vs 43.43 s（5.4×）。此外，不修复 EFSDP>1 丢专家的问题，delta 的 seed（全量导出）本身不可信；不修复 HSDP 下的 mesh bug，任何 HSDP 运行都会在 worker 初始化阶段崩溃。

# 实现拆解

### 1. 新增分片导出入口（transformer_impl.py）

- `TorchTitanEngine.get_per_tensor_param_shard()` 遍历 `self.module` 的 `state_dict()`，通过内层 `_shard()` 完成「搬到 device → 转 bf16 → `to_local()` → 展平」，产出 `(hf_name, local_flat_bf16, ShardSpec)` 三元组；`ShardSpec.from_param()` 记录 mesh 与 placements，接收端据此推导全局 flat 坐标。
- `get_per_tensor_param_delta_shard()` 借助 `verl.workers.engine.utils.hf_delta_export` 把本次分片与 `_delta_shard_snap` 快照做逐字节差分，产出最终 HF 坐标的稀疏 delta 条目；快照每次导出时刷新，任何 rank 都不持有全量模型快照。
- 默认实现上移到 `BaseEngine`（只需分片导出加 per-parameter entry builder），删除 FSDPEngine 中两处完全相同的拷贝；Megatron 保留自己的 override（差分对象不是固定快照），veomni 保留 `_hf_delta_entry`，TorchTitanEngine 自带实现以处理专家堆栈的 slot 拆分。

### 2. HF 命名对齐与专家堆栈

- `_to_hf_named_params()` 把原 `get_per_tensor_param` 里的 `sd_adapter.to_hf()` 与 `lm_head.weight` 补回逻辑抽成共享方法，保证全量与分片导出按名字精确配对——名字只出现在一边会静默把权重冻结在 seed 值。
- `_expert_stack_slots()` 利用 `sd_adapter.from_hf_map` 结构性识别 fused expert stack（一个 torchtitan key 对应带两个 `{}` 占位的 HF 模板），整体导出堆栈并在 `spec.hf_slots` 中枚举每个专家的 HF 张量，避免 `to_hf()` 只产出本地专家导致的 lockstep 命名不一致。
- `_hf_entry_row_slots()` 处理 dim-0 identity 拆分：行 `e` 就是 `full[e]`，用一次 `searchsorted` 从全局索引恢复每个专家的计数，规避 probe-driven 转换器的复杂度。

### 3. 放置推导扩展（spec.py）

- 新增 `_shard_dim()`：torch 2.13 中 `_StridedShard` 迁到 C++ 且不再继承 `Shard`，`p.is_shard()` 会漏判；该函数同时对两种写法给出答案。
- 新增 `_assert_even_strided()`：拒绝 strided 偏移公式无法表达的 uneven cut。
- 新增 `_shard_dims_group()` 取代 `_flattened_mesh_group()`：变化 Shard 维、固定其他维构造 gather 组，按 rank 集合缓存（`dist.new_group` 是集合通信，所有 rank 以相同顺序建组），避免 HSDP 下 replicate 度被计入 wire 字节与 rank 0 的 staging 缓冲。

### 4. 边界防护与两个预存 bug 修复

- `_assert_shard_export_supported()` 在导出边界以明确信息拒绝 PP：各 stage 持有不相交切片，导出顺序在 rank 间不一致，无法满足 lockstep gather。
- `_get_data_parallel_mesh()` 改用 `parallel_dims.get_optional_mesh("loss")` 获取 1-D 扁平 loss mesh，修复 HSDP 下 `get_data_parallel_rank()/get_data_parallel_group()` 遇到 2-D mesh 崩溃的问题。
- `get_per_tensor_param()`（全量导出）改为先对整个 DTensor 专家堆栈做 `full_tensor()` 再拆分，修复 EFSDP>1 时先 `to_hf()` 拆分再在 ep 组内 all-gather 导致静默丢失 50% 专家的 bug。

### 5. 测试、兼容与文档配套

- `tests/special_distributed/test_sharded_delta_gather.py` 扩展 2D/3D mesh 用例：FSDP×TP（`_StridedShard`）、EFSDP×EP、HSDP×TP/EP、replicate 度 4 的 sharp case，并新增断言「gather 组大小等于 Shard 维大小的乘积」。
- `sglang_rollout.py` 对 `LocalSerializedTensor` 的 import 增加 SGLang 0.5.16+ 的兼容分支。
- `docs/advance/delta_weight_sync.md` 与 `docs/workers/torchtitan_workers.rst` 更新支持矩阵、设计说明与实测性能表。

关键文件：
- `verl/workers/engine/torchtitan/transformer_impl.py`（模块 训练引擎；类别 source；类型 core-logic；符号 _hf_entry_row_slots, _to_hf_named_params, _expert_stack_slots, _hf_delta_entry）: 核心实现文件：新增分片导出与 delta 差分入口，处理 HF 命名对齐与专家堆栈 slot table，同时修复 EFSDP 丢专家与 HSDP mesh 两个预存 bug
- `verl/workers/engine/spec.py`（模块 放置推导；类别 source；类型 core-logic；符号 _shard_dim, _assert_even_strided, _shard_dims_group, _flattened_mesh_group）: 共享放置推导逻辑：_shard_dim / _assert_even_strided / _shard_dims_group 使 _StridedShard 与 HSDP 多 Shard 维成为一等公民，是所有 delta 后端的共同基础设施
- `verl/workers/rollout/sglang_rollout/sglang_rollout.py`（模块 回放端；类别 source；类型 compat-adjust；符号 _update_weights_delta_flush）: delta 应用路径的 SGLang 0.5.16+ 兼容 import，保证 TorchTitan 与 SGLang 共环境运行时不受 LocalSerializedTensor 迁移影响
- `tests/special_distributed/test_sharded_delta_gather.py`（模块 测试覆盖；类别 test；类型 test-coverage；符号 _run_case, main）: 直接覆盖 placement 数学的分布式测试，扩展 TP/EP/HSDP 组合用例与 gather 组大小断言，是引擎级测试删除后的主要防线
- `docs/advance/delta_weight_sync.md`（模块 文档；类别 docs；类型 documentation）: 支持矩阵、设计说明与实测性能数据的权威文档，明确 TorchTitan 的布局支持、PP 拒绝原因与 HSDP×EP 的上游依赖
- `docs/workers/torchtitan_workers.rst`（模块 文档；类别 docs；类型 documentation）: TorchTitan 引擎能力列表同步，列出 delta_sharded 支持的并行组合

关键符号：get_per_tensor_param_shard, get_per_tensor_param_delta_shard, _to_hf_named_params, _expert_stack_slots, _hf_entry_row_slots, _assert_shard_export_supported, _shard_dim, _assert_even_strided, _shard_dims_group, derive_dtensor_placement

## 关键源码片段

### `verl/workers/engine/torchtitan/transformer_impl.py`

核心实现文件：新增分片导出与 delta 差分入口，处理 HF 命名对齐与专家堆栈 slot table，同时修复 EFSDP 丢专家与 HSDP mesh 两个预存 bug

```python
def get_per_tensor_param_shard(self, **kwargs):
    """Yield this rank's local shard ``(hf_name, local_flat_bf16, ShardSpec)`` instead of the full tensor."""
    self._assert_shard_export_supported()
    raw = {}
    for module in self.module:
        raw.update(module.state_dict())

    # 专家堆栈必须整体导出并携带 slot table：to_hf() 只会保留本地拥有的专家，
    # 而 lockstep gather 要求所有 rank 对“当前参数是哪个”给出同一答案
    stacks = {}
    for name, param in raw.items():
        slots = self._expert_stack_slots(name, param)
        if slots is not None:
            stacks[name] = slots
    params = self._to_hf_named_params({k: v for k, v in raw.items() if k not in stacks})
    device = get_device_id()  # 开启 offload 策略时本地分片位于 CPU

    from ..spec import ShardSpec

    def _shard(param):
        # 先把参数搬到 device 再转 bf16：混合精度下 fp32 主权重保留在优化器，
        # 线路上只传输 bf16 分片，差分时才做类型对齐
        p = param.to(device, non_blocking=True)
        if p.is_floating_point():
            p = p.to(torch.bfloat16, non_blocking=True)
        local = p.to_local() if isinstance(p, DTensor) else p
        return local.reshape(-1)

    def _gen():
        # 稠密参数：按 HF 名字逐个导出，ShardSpec 携带 mesh 与 placements，
        # 供接收端 derive_dtensor_placement 推导 flat 坐标
        for name, param in params.items():
            yield name, _shard(param), ShardSpec.from_param(param)
        # 专家堆栈：以 torchtitan 原始名字为 key，spec.hf_slots 记录
        # 每个专家的 HF 张量名与形状，接收端按 slot 拆分为逐专家 delta
        for name, slots in stacks.items():
            spec = ShardSpec.from_param(raw[name])
            spec.hf_slots = slots
            yield name, _shard(raw[name]), spec

    return _gen(), None

```

### `verl/workers/engine/spec.py`

共享放置推导逻辑：_shard_dim / _assert_even_strided / _shard_dims_group 使 _StridedShard 与 HSDP 多 Shard 维成为一等公民，是所有 delta 后端的共同基础设施

```python
def _shard_dim(p) -> Optional[int]:
    """返回该 placement 切割的张量维，不切割则返回 None。

    torch 2.13 中 _StridedShard 已迁到 C++ 且不再继承 Shard，p.is_shard() 会漏判；
    _StridedShard 与 Shard 一样是单 block 切割，只是右到左的切割顺序置换 rank 与 block 的归属，而不是交错 block 内部。
    """
    if p.is_shard():
        return int(p.dim)
    return int(p.dim) if type(p).__name__ == "_StridedShard" else None


def derive_dtensor_placement(spec: ShardSpec) -> tuple[int | BlockPlacement, bool, Optional[ProcessGroup]]:
    """推导本 rank 的 (place, contributes, gather_group) 三元组。

    place 供 translate_flat_indices 做分片内 flat 坐标到全量坐标的映射；
    contributes 标记副本 rank（Replicate 维坐标非 0）以空 delta 参与 lockstep；
    gather_group 决定向哪些 rank 收集分片。
    """
    import torch.distributed as dist

    assert spec.place is None  # 显式 place 的 spec 由调用方分发，不会走到这里

    if spec.mesh is None:
        # 未分片：本地张量就是完整参数，无需 gather
        return 0, (dist.get_rank() == 0 if dist.is_initialized() else True), None

    placements = spec.placements
    # 用 _shard_dim 而非 p.is_shard()，否则 torch 2.13 下 _StridedShard 会被静默漏进
    # “无 shard 维”分支，mesh 维从 gather 组中消失
    shard_dims = [d for d, p in enumerate(placements) if _shard_dim(p) is not None]
    _assert_even_strided(spec, placements)

    coord = spec.mesh.get_coordinate()
    contributes = True
    if coord is not None:
        # 任一 Replicate 维坐标非 0 的 rank 不贡献数据，只以空 delta 保持 lockstep
        for d, p in enumerate(placements):
            if p.is_replicate() and coord[d] != 0:
                contributes = False
                break

    if not shard_dims:
        return 0, contributes, None

    local_shape, global_offset = compute_local_shape_and_global_offset(spec.full_shape, spec.mesh, list(placements))
    place = BlockPlacement(tuple(local_shape), tuple(global_offset), tuple(spec.full_shape))

    if len(shard_dims) == 1:
        # 单 Shard 维：该维子组天然只含贡献 rank，副本被跳过
        return place, contributes, spec.mesh.get_group(mesh_dim=shard_dims[0])
    # 多 Shard 维：_shard_dims_group 变化 Shard 维、固定其他维，
    # 避免 HSDP 下 replicate 度被计入 wire 字节与 rank 0 的 staging 缓冲
    return place, contributes, _shard_dims_group(spec.mesh, shard_dims)

```

# 评论区精华

外部 review 只有 wuxibin89 的 APPROVED，两条 review 评论均为作者 attack204 的自评说明。最有价值的反馈藏在 commit message 里：`ec0f35b`「drop the torchtitan delta export test」明确写道 “Review feedback. It needed 8 GPUs and a torchtitan install, so no CI job could have picked it up”——8 GPU 加 torchtitan 依赖使引擎级导出测试无法被任何 CI 任务拾取，最终删除该测试，保留直接覆盖 placement 数学的 `test_sharded_delta_gather`。作者在 `spec.py` 第 211 行的自评指出 “FSDP2 + TP is a _StridedShard, so we should modify the limitation”，即旧代码把 `_StridedShard` 当作「非单 block」拒绝，而实际上它只是置换 rank 归属的 block，由此驱动了 `_shard_dim()` 的引入与文档支持矩阵的同步更新。

- FSDP2+TP 的 _StridedShard 支持与 limitation 更新 (design): 引入 _shard_dim() 接纳 _StridedShard，并同步更新 delta_weight_sync.md 与 torchtitan_workers.rst 的支持矩阵
- 8 GPU 引擎级导出测试的可运行性 (testing): 删除 torchtitan delta 导出测试，保留直接覆盖 placement 数学的 test_sharded_delta_gather；引擎级覆盖依赖后续 CI 基建
- _hf_entry_row_slots 的专家堆栈命名转换语义 (question): 纯解释性评论，无代码变更；转换语义由 dim-0 identity 的 searchsorted 拆分解法确定

# 风险与影响

- 风险：
 1. **torch 版本敏感**：`_shard_dim()` 与 `_assert_even_strided()` 的语义依赖 torch 2.13 中 `_StridedShard` 的行为；同时 torchtitan 引擎本身需要 torch≥2.12（`DataParallelMeshDims`），而 SGLang ≤0.5.17 锁定 torch 2.11，仅在 SGLang main（torch 2.13）上收敛，环境组合较脆弱。
 2. **引擎级 e2e 测试缺失**：删除 8 GPU 导出测试后，torchtitan 引擎的分片导出路径没有 CI 覆盖，只有 placement 数学的单元测试；PR body 也注明 HSDP+EP“unit-tested; no end-to-end row”。
 3. **依赖上游修复**：HSDP×EP 需要 torchtitan 侧 MoE state dict adapter 修复（`to_hf()` 在读 `placement.dim` 前未检查类型），修复前该组合在 `dcp_load` 阶段即崩溃。
 4. **共享逻辑影响面**：`spec.py` 的 `derive_dtensor_placement()` 是 FSDP/veomni/Megatron 共享的放置推导，本次改动为向后兼容的扩展（不再拒绝 `_StridedShard`），但任何回归都会波及其他 delta 后端。
 - 影响：对 torchtitan 引擎用户，`delta_sharded` 从不可用到可用，跨节点 32 卡场景权重同步从 34 s 级降到 3 s 级，MoE 模型从 43 s 级降到 8 s 级；HSDP 用户额外获得 `_get_data_parallel_mesh()` 修复（此前任何 HSDP 运行都会在 worker 初始化崩溃）。对 nccl 全量同步用户，EFSDP>1 时全量导出静默丢失专家的预存 bug 被修复，属于正确性提升。对团队而言，17 个 commit 完整展示了从「拒绝 TP/EP」到逐步放开的设计演进、负控制验证与大量实测，文档沉淀了支持矩阵与基准数据，为后续其他引擎接入 delta 同步提供了可复用的模板。
 - 风险标记：依赖 torch 2.13 行为 , 引擎级 e2e 测试缺失 , HSDP×EP 依赖上游修复 , 共享放置推导逻辑改动

# 关联脉络

- PR #7291 [ckpt] feat: Node-local multi-sender broadcast in NCCL checkpoint engine: 同一权重同步优化方向的前序工作：NCCL 引擎多发送者中继使全量广播提速约 3 倍；本 PR 将 delta 稀疏同步推广到 torchtitan 引擎，两者共享 checkpoint 引擎与权重同步架构
- PR #7407 [megatron,veomni] feat: use torch.int16 for routed_experts: 同样深度改动 spec.py 的 ShardSpec / derive_dtensor_placement 基础设施（veomni 的 _hf_delta_entry、to_hf_chunk 探针机制），本 PR 的 slot table 方案正是对该探针机制的补充与简化
- PR #7422 [rollout] fix: preserve dummy load_format in disaggregated rollout: 分离式 rollout 的权重接收路径修复，与 delta_sharded 的 SGLang receive/apply 链路（sglang_rollout.py 的 _update_weights_delta_flush）直接相关