本文档说明 Expert Tensor Parallel（ETP） 在 HyperParallel 中的设计定位，并与 TorchTitan distributed/expert_parallel.py 中的 ExpertTensorParallel 对齐说明。不要求重复粘贴实现代码，侧重背景、职责边界、接口与 mesh 契约、测试与验收。

现状说明：HyperParallel 已在 hyper_parallel/core/expert_parallel/expert_parallel.py 中提供 ExpertTensorParallel 实现；本文档可作为设计溯源、评审与演进依据，并与用户文档 docs/expert_parallel.md 互补。

---

1. 背景

1.1 问题域：MoE 的两类并行

• Expert Parallel（EP）：按 专家维 切分权重，token 经 all-to-all 路由到持有对应专家的设备；通信发生在 EP 进程组。
• Tensor Parallel（TP）：在 单个专家内部 按 隐藏维 / 输入维 做 Megatron 式 列并行 / 行并行（SwiGLU 对应 w1/w3 列切、w2 行切），通信多为 TP 组内 all-gather / reduce-scatter。

仅 EP 时，单专家算子可能过大；仅 TP 时，专家总数受单卡显存限制。ETP 同时在 专家维 与 专家内部维 切分，需要 二维 DeviceMesh，且必须澄清：token 路由通信只随 EP 维度走，不参与 TP 维度的无关 collective。

1.2 TorchTitan 参考语义

TorchTitan 中结构如下（逻辑分层）：

• ExpertParallel：GroupedExperts 权重 Shard(0)（专家维）；_token_dispatch / _token_combine 在传入的 1D EP mesh 上做计数交换、可微 all-to-all、permute/unpermute。
• ExpertTensorParallel(ExpertParallel)：
  • _partition_fn：对 w1/w2/w3 使用 [Shard(0), Shard(1)] 或 [Shard(0), Shard(2)]（与 SwiGLU 形状约定一致：EP 切专家维，TP 切 hidden / 行维）。
  • _token_dispatch / _token_combine：传入 device_mesh["ep"]，即 仅在 EP 子 mesh 上调用父类逻辑，避免用 整幅 [ep, tp] mesh 做 token collective。

TorchTitan 注释要点：「for dp2ep with TP」——在 DP 转 EP 且仍保留 TP 的场景下，专家权重二维切分、token 仍按 EP 组交换。

1.3 HyperParallel 中的位置

• ExpertParallel：与 TorchTitan 同类 all-to-all EP 语义（实现细节可依赖平台 collective）。
• TensorParallel（MoE 专用）：仅 TP、Shard(1)/Shard(2)，无 dispatch（EP 度为 1 时切专家算子）。
• ExpertTensorParallel：继承 ExpertParallel，覆写 _partition_fn 为二维切分，覆写 _token_dispatch / _token_combine 将 device_mesh["ep"] 传给父类。

---

2. 解决的问题与不当用法

需求	ETP 是否覆盖
专家数多 + 单专家隐藏维大，需同时扩 专家并行度 与 算子切分	✅
Token 路由与 TP 混合在同一 mesh 上做 all-to-all	❌ 错误；必须 仅 EP 子 mesh
仅需 EP、无 TP	使用 ExpertParallel + 1D ep mesh，勿用 ETP
仅需切专家算子、无 EP	使用 TensorParallel + 1D tp mesh

---

3. 目标与非目标

3.1 目标

• 权重布局：与 TorchTitan 一致，在 [num_experts, …] 形状约定下 w1/w3：[Shard(0), Shard(1)]，w2：[Shard(0), Shard(2)]（维度索引与 PyTorch / HyperParallel GroupedExperts 定义一致）。
• 通信边界：Dispatch / Combine 与 ExpertParallel 完全同款算法，仅 collective group = mesh["ep"].get_group()（或平台等价 API）。
• API 形态：ParallelStyle.apply(module, device_mesh)，与 ColwiseParallel 等一致；device_mesh 为 2D，且含命名维 "ep" 与 "tp"（或与实现约定的子 mesh 名字一致）。

3.2 非目标

• Router / Gate 的并行策略（如 NoParallel、TP router）：不在本文档范围；由上层 MoE parallelize 组合。
• DeepEP / HybridEP 等自定义 dispatch：TorchTitan 另有 DeepEPExpertParallel；HyperParallel 若引入，应 独立样式，不混在本 ETP 基线语义中。
• ReordererSequenceParallel（ETP=1 时 EP 借 TP/DP 维）：TorchTitan 专用编排；HyperParallel 若支持需在 单独文档 说明与 ETP 的互斥与顺序。

---

4. 接口与 Mesh 契约

4.1 类与继承

class ExpertTensorParallel(ExpertParallel):

• 无额外构造参数（与 TorchTitan 一致；若未来增加 可选覆盖子 mesh 名字，属扩展项，需单独 RFC）。

4.2 apply(module, device_mesh)

参数	约束
module	目标为 GroupedExperts（或与其实现 相同参数名与形状约定 的 MoE 专家子模块）。
device_mesh	二维，且可通过 device_mesh["ep"]、device_mesh["tp"] 切片为 1D 子 mesh；ndim == 2，维名 ("ep", "tp")（见 init_device_mesh 约定）。

禁止：对 1D mesh 调用 ETP（应在 parallelize_module 或 apply 内 显式校验 并给出可读错误）。

4.3 权重分片语义（GroupedExperts）

与 docs/expert_parallel.md 表格一致：

参数	Shape（逻辑）	Placement
w1	[E, H, D]	[Shard(0), Shard(1)]
w3	[E, H, D]	[Shard(0), Shard(1)]
w2	[E, D, H]	[Shard(0), Shard(2)]

其中维 0 为 专家维（EP），1/2 为 TP 切分维（列 / 行与 Megatron SwiGLU 一致）。

4.4 Token 路径语义

• Dispatch / Combine：与 ExpertParallel 相同步骤（计数 all_to_all、可微 token all_to_all、permute / unpermute）；唯一区别是 device_mesh 实参 为 device_mesh["ep"]。
• 同一 TP rank 组内，各 rank 持有 相同 token 子集与否 由上层 MoE 输入是否已 SP/TP 切分 决定；ETP 样式本身不强制输入 layout，但 须与 Router、Attention 的 parallel plan 一致，否则需在集成文档中列出 前置条件（例如 全局 token 表 vs 已按 TP 切分）。

---

5. 与 TorchTitan 的差异与对齐点

项目	TorchTitan	HyperParallel（设计意图）
基类与 hook	distribute_module + _apply	apply → distribute_module（与现有 EP 一致）
子 mesh	device_mesh["ep"]	同名
w1/w2/w3 二维 Shard	[Shard(0), Shard(1)] / [Shard(0), Shard(2)]	一致
纯 TP 专家（无 EP）	TensorParallel 类	TensorParallel（MoE）

---

6. 实现要点检查清单（落地代码时）

• parallelize_module / TP 校验：现有 api.py 可能 限制仅 1D mesh；对 ETP 需在 MoE 专用路径或 apply 入口 使用 2D mesh，避免被 TP 1D 校验误伤（若尚未放开，属 已知集成项）。
• GroupedExperts.forward：权重为 DTensor 时 to_local() 后再走 grouped matmul（已有平台约定）。
• 双重 distribute_module：避免对同一 module 重复包装（沿用全局 _distribute_module_applied 约定）。
• collective 后端：NPU/CUDA 对 all_to_all_single_autograd 的可微行为需与 ExpertParallel 单测策略一致。

---

7. 测试设计

7.1 已有单测映射（维护建议）

仓库中 tests/ut/core/expert_parallel/test_expert_parallel.py 已包含 C4：ExpertTensorParallel：

• _partition_fn：二维 Shard placement 断言；
• _token_dispatch / _token_combine：对 device_mesh["ep"] 委托的 spy/mock。

设计文档要求：分区与委托行为变更时必须同步更新上述用例。

7.2 建议增补（若尚未覆盖）

ID	内容
ETP-01	2×2 ep_tp mesh 上，apply 后 w1.placements 为 (Shard(0), Shard(1))（或平台等价表示）。
ETP-02	Mock device_mesh["ep"]，断言 dispatch/combine 未对 tp 维调用 collective。
ETP-03	分布式集成（可选）：小 MoE + ETP，前向/反向与 单卡参考或 EP-only+TP-only 分解数值对齐（容差按 dtype）。

7.3 文档

• docs/expert_parallel.md：已含 ETP 小节；本设计文档批准后，可补充 「与 TorchTitan 对齐」 一句及 mesh 命名硬性要求。
• 跨项目：auto_parallel/fast-tuner/docs/torchtitan.md 如有 MoE 策略说明，可加交叉引用。

---

8. 风险与开放问题

风险	缓解
parallelize_module 仅接受 1D TP mesh	MoE 使用 ExpertTensorParallel().apply(..., mesh_2d) 直连，或扩展 API 显式区分 TP 与 EP⊗TP。
维名非 "ep"/"tp"	首期 文档强制命名；远期可支持 ep_mesh_dim / tp_mesh_dim 字符串参数（RFC）。
专家数 / EP / TP 整除条件	与 ExpertParallel 相同：num_experts % ep_degree == 0，且隐藏维可被 tp_degree 整除（具体由 GroupedExperts 与 TP shard 维推导）。

---

9. 验收标准

• ExpertTensorParallel 行为与本文 §4.3–4.4 一致，且与 TorchTitan 权重 shard 索引 一致。
• Dispatch/Combine 仅绑定 EP 子 mesh（单测 mock 或日志断言）。
• C4 单测通过；可选 ETP-03 集成通过。
• 用户文档 mesh 维度命名 与本设计 一致。

---

10. 参考文献

• TorchTitan：torchtitan/distributed/expert_parallel.py — ExpertTensorParallel、ExpertParallel、TensorParallel（GroupedExperts）。
• HyperParallel：hyper_parallel/core/expert_parallel/expert_parallel.py、docs/expert_parallel.md。
• PyTorch：DeviceMesh 子 mesh 索引、distribute_tensor 多维 Shard。