【RFC】FullyShard 多维并行适配设计方案

第一章 背景介绍

1.1 现状

HyperParallel 当前在并行能力上已经具备以下基础：

1. DeviceMesh + Placement 作为统一并行表达。
2. fully_shard 已支持 1D FSDP 与 2D HSDP 语义（参数切分、前后向重分布、RS/AR 梯度通信）。
3. shard_module 已支持 TP 算子级切分，并基于 DTensor 执行。

1.2 当前问题

结合现有实现，存在两个关键缺口：

1. fully_shard 与 TP/EP 之间缺乏统一的“多维切分策略聚合”机制，导致参数布局和通信语义在不同模块中分散管理。
2. TP 场景中，前向路径是 DTensor 语义，但实际算子执行常落到 local_tensor；反向梯度路径若不显式区分通信组，容易出现梯度同步错误或重复规约。

1.3 本次需求

需求1：

1. fully_shard 接入多维并行，初始化时融合 fully_shard + tp/ep，将切分策略聚合到同一个 DTensor 描述中。
2. 算子执行前，仅对 fully_shard 维做聚合（unshard），保留 TP/EP 维度信息。
3. 反向梯度通信参考 HSDP 的 GroupInfo 机制，明确 shard group 与 replicate group 分工。

需求2：

1. 解决 TP 切分时“前向 DTensor / 实际 local_tensor 运算”导致的反向同步问题。
2. 利用 fully_shard 能力，在 unsharded_dp_group 上提供 weight 同步能力（用于兼容路径）。
3. 评审 fully_shard 如何识别“纯 TP 并行”：
   1. 方案A：接口显式传参控制是否启用 fully-shard 参数切分。
   2. 方案B：通过 mesh 自动判断，若仅 TP mesh 则允许结合；若为完整 mesh 则不启用 fully-shard 参数切分。

---

第二章 方案

2.1 基本原理

2.1.1 统一拓扑语义

统一用高维 mesh 描述并行维度，推荐角色划分：

1. fsdp_shard_dim：参数分片维（ReduceScatter 维）。
2. dp_replicate_dim：数据副本维（AllReduce 维）。
3. tp_dim：张量并行维。
4. ep_dim：专家并行维。

2.1.2 单一 DTensor 布局承载多策略

参数初始化后只维护一个“逻辑完整布局”的 placement 元组，包含 FSDP/TP/EP 信息。
示例（3D）：

1. 初始聚合布局：(Shard(fsdp_axis), Shard(tp_axis), Replicate())。
2. Forward 前 unshard 后布局：(Replicate(), Shard(tp_axis), Replicate())。
3. Forward/Backward 结束后再恢复：(Shard(fsdp_axis), Shard(tp_axis), Replicate())。

核心点：仅切换 fully_shard 对应 placement，TP/EP placement 不变。

2.1.3 前后向边界

1. 前向：按 DTensor 语义组织参数与张量分布。
2. 反向：通信阶段统一操作 local_tensor，不依赖 DTensor 通信语义。
3. 通信：先 ReduceScatter(sharded_group_info)，再 AllReduce(unsharded_group_info)（按需）。

2.2 方案设计

2.2.1 GroupInfo 双组模型

对每个参数构建：

1. sharded_group_info：切分组，负责 RS。
2. unsharded_group_info：副本组，负责 AR。
3. shard_size、dp_size：用于通信路径判定和缓冲区计算。

这与现有 HSDP 路径一致，便于复用调度和测试资产。

2.2.2 反向通信决策矩阵

按参数状态选择通信策略：

1. 参数启用 fully-shard，且 shard_size > 1：

   1. RS 于 sharded_group_info。
   2. 若 dp_size > 1 且需要副本一致性，再 AR 于 unsharded_group_info。

2. 参数未启用 fully-shard，且需要数据副本同步：

   1. 直接 AR 于 unsharded_group_info。
   2. 用于纯 TP 兼容路径（需求2）。

3. 参数既不切分也无副本：

   1. 不通信，保留本地梯度。

在纯 TP 兼容路径中，参数切分维关闭（shard_size = 1），不触发 RS，仅在 unsharded_dp_group 做 AR。

2.2.3 纯 TP 识别评审方案

方案A（显式控制）：

1. fully_shard 新增参数 enable_param_shard。
2. enable_param_shard=False 时，不做参数切分，仅复用通信框架。
3. 优点：行为可控、可灰度、问题定位简单。
4. 缺点：用户需理解开关语义。

方案B（mesh 自动判断）：

1. 当 enable_param_shard=None 时，基于 mesh 角色自动判断。
2. 若传入 mesh 仅包含 TP 维（无可用 fsdp/dp 维），自动禁用参数切分。
3. 若传入完整 mesh（含 dp/fpdp 维），按默认 fully_shard 行为执行。
4. 优点：用户心智负担低，接口简洁。
5. 缺点：推断规则复杂，边界场景易歧义。

评审建议：

1. 第一阶段采用“方案A主、方案B辅”：先显式参数保证稳定。
2. 第二阶段再完善自动推断，并允许显式参数覆盖推断结果。

2.2.4 需求2的同步修复策略

在纯 TP 或“禁用参数切分”路径下，提供 unsharded_dp_group 同步：

1. 默认同步对象为 weight.grad。
2. 兼容模式可支持 weight 本体同步（仅在明确需要时启用）。
3. 对 TP-shard 参数与 TP-replicate 参数分别处理，避免错误规约。

2.2.5 shard_size / dp_size 详细定义与来源

1. 概念定义：

   1. rank_size：当前参数所在 fully_shard 数据并行域的总并行规模。
   2. shard_size：参数切分组规模，决定 RS 维度与梯度切片份数。
   3. dp_size：副本组规模，决定 AR 维度与副本一致性范围。

2. 基本关系：

   1. dp_size = rank_size // shard_size。
   2. 当 shard_size = 1 时，表示参数不进行 fully_shard 切分，通信只可能发生在副本组。

3. 参数来源：

   1. rank_size 来自参数所在 mesh 的 fully_shard 数据并行域（按角色映射后的局部域）。
   2. shard_size 来自参数切分策略与 mesh 切分维度综合决定：
      1. 启用参数切分时，取 fsdp_shard_dim 对应组规模并结合阈值/合法性约束得到有效值。
      2. 禁用参数切分时，固定为 1。
   3. dp_size 按公式计算得到；在纯 TP 兼容路径中通常等于该同步域规模。

4. 工程语义：

   1. shard_size 控制“参数与梯度是否切片，以及切几片”。
   2. dp_size 控制“每个切片在多少副本间需要一致”。

5. 典型示例：

   1. HSDP：rank_size=8, shard_size=4, dp_size=2，反向执行 RS(4) + AR(2)。
   2. 纯 TP 兼容：rank_size=8, shard_size=1, dp_size=8，反向仅执行 AR(8)。

---

第三章 设计实现，具体实现，简单描述

3.1 涉及模块

1. hyper_parallel/core/fully_shard/hsdp_param.py
2. hyper_parallel/platform/torch/fully_shard/param.py
3. hyper_parallel/platform/torch/fully_shard/state.py
4. hyper_parallel/core/fully_shard/api.py
5. hyper_parallel/platform/torch/fully_shard/utils.py

3.2 核心实现点

1. 参数初始化阶段接入 HSDP 同款流程：

   1. _init_rank_info()
   2. _init_param_shard_size()
   3. _init_param()
   4. 计算 dp_size
   5. 创建 sharded_group_info
   6. 创建 unsharded_group_info

2. TorchHSDPParamV2 中新增 GroupInfo 属性与访问接口，通信统一经 GroupInfo 路由。

3. 前向 unshard 仅替换 fully_shard 维 placement 为 Replicate，保留 TP/EP placement。

4. 反向 post_backward 的梯度通信从“mesh类型分支”升级为“GroupInfo 分支”，减少硬编码。

5. 纯 TP 修复路径：

   1. 判定“禁用参数切分”或“纯 TP 模式”。
   2. 在 unsharded_dp_group 上执行同步操作（优先 grad，同步对象可配置）。

3.3 实现流程（简版伪代码）

fully_shard_init()
  -> parse mesh roles (fsdp/tp/ep)
  -> init param placements (single DTensor layout)
  -> init GroupInfo(sharded/unsharded)

forward_pre()
  -> unshard only on fsdp dim
  -> keep tp/ep placements

backward_post()
  -> use local grad tensor
  -> if shard_size>1: reduce_scatter(sharded_group)
  -> if dp_size>1: all_reduce(unsharded_group)
  -> pure TP mode: all_reduce in unsharded_dp_group (compat path)

---

第四章 设计接口，变化

4.1 fully_shard 接口新增项

仅新增一个参数：

enable_param_shard: Optional[bool] = None

语义：

1. True：显式启用 fully_shard 参数切分。
2. False：显式禁用 fully_shard 参数切分，仅复用统一通信框架。
3. None：按 mesh 角色自动推断（用于兼容自动识别路径）。

优先级：显式值优先于自动推断。

4.2 行为变化

1. 旧行为：默认按 FSDP/HSDP mesh 做参数切分与梯度规约。
2. 新行为：支持“仅通信不切分”的纯 TP 兼容模式。
3. 新行为：反向通信统一按 GroupInfo 驱动，减少 mesh 类型分支耦合。
4. 兼容性：不传新增参数时保持原默认行为（通过推断落回既有路径）。

4.3 纯 TP 识别与收敛策略

1. 保留显式控制与自动推断两种机制。
2. 推荐工程策略：
   1. 联调和故障定位阶段优先使用显式控制。
   2. 规模化使用阶段可启用自动推断。
3. 文档需固化边界 mesh 判定规则（例如 2D mesh 但无明确 dp 语义时的处理顺序）。

---

第五章 测试用例设计

5.1 测试目标

1. 验证多维 placement 聚合正确。
2. 验证前向仅聚合 fully_shard 维、TP/EP 信息不丢失。
3. 验证反向在 local_tensor 路径下通信正确。
4. 验证纯 TP 兼容路径可修复梯度同步问题。
5. 验证新旧接口兼容与回归稳定性。

5.2 测试矩阵

1. 1D FSDP：仅 RS，无 AR。
2. 2D HSDP(dp×fsdp)：RS + AR。
3. 2D (dp×tp)：forward 保留 TP placement，backward RS(dp_shard)+AR(dp_rep)。
4. 3D (dp×tp×ep)：unshard 仅改 fsdp 维，tp/ep placement 保持不变。
5. 纯 TP + enable_param_shard=False：unsharded_dp_group 同步生效，反向不触发 RS。
6. 纯 TP + 自动推断：自动禁用参数切分逻辑正确。
7. 完整 mesh + 自动推断：不误判为纯 TP。
8. mixed precision：reduce_dtype 下结果一致。
9. grad accumulation：requires_grad_sync=False/True 两种路径正确。
10. state_dict round-trip：参数布局、local shard 与恢复一致。

5.3 验收标准

1. 数值一致：loss、梯度、参数更新与基线对齐。
2. 通信一致：RS/AR 调用次数与分组 rank 列表符合预期。
3. 性能可接受：新增判定与分组初始化不引入明显额外开销。
4. 兼容可控：旧用法无需改动，新能力可按开关灰度启用。