HP Tensor Parallel ParallelStyle 设计文档（转测）

0. 基本信息

项目	内容
特性名称	HP Tensor Parallel 声明式 ParallelStyle（style.py 全量）
代码位置	hyper_parallel/core/tensor_parallel/style.py
配套入口	hyper_parallel/core/tensor_parallel/api.py 的 parallelize_module
开发分支	master
适用后端	PyTorch、MindSpore
已验证设备	PyTorch NPU/HCCL 实机 ST；PyTorch CPU/Gloo level0；PyTorch mock UT。MindSpore 目前仅有 MLP Colwise+Rowwise 通信调试 ST
参考实现对齐	torch.distributed.tensor.parallel.style / parallelize_module
真实业务用法	Qwen3.5 dense TP plan：hyper_parallel/models/qwen3_5/parallelize.py
当前阶段	功能已合入，按 ParallelStyle 全量转测

本期转测对象是 style.py 导出的全部 TP Style，以及把它们应用到模块树上的 parallelize_module：

ParallelStyle
├── ColwiseParallel
├── RowwiseParallel
├── SequenceParallel
├── PrepareModuleInput
├── PrepareModuleOutput
├── PrepareModuleInputOutput
└── NoParallel

---

1. 背景

Tensor Parallel（TP）把 Linear / Embedding 的权重沿特征维切开，再靠激活上的集合通信把分片结果拼回。HyperParallel 不要求用户手写 all-gather / all-reduce，而是提供一组声明式 Style：

parallelize_module(model, tp_mesh, {
    "attn.q_proj": ColwiseParallel(),
    "attn.o_proj": RowwiseParallel(),
    "norm": SequenceParallel(),
})

这套 API 对齐 PyTorch DTensor TP：

• Style 只描述「这个 submodule 怎么切、输入输出是什么 Placement」；
• 真正的参数切分走 distribute_tensor；
• 真正的 I/O 转换走 forward pre-hook / forward hook；
• 模块树遍历、glob 匹配、src_data_rank 注入由 parallelize_module 完成。

业务侧已经按这个模型拼出完整 Transformer TP plan。以 Qwen3.5 为例：

• embed_tokens：RowwiseParallel(Replicate → Shard(1))
• input_layernorm / post_attention_layernorm / model.norm：SequenceParallel(sequence_dim=1)
• self_attn / mlp 入口：PrepareModuleInput(Shard(1) → Replicate)
• q/k/v_proj、gate/up_proj：ColwiseParallel()
• q_norm / k_norm：SequenceParallel(sequence_dim=2, use_local_output=True)
• o_proj / down_proj：RowwiseParallel(Shard(-1) → Shard(1), use_local_output=False)，部分层带 reduce_dtype=float32
• lm_head：ColwiseParallel(Shard(1) → Replicate 或 Shard(-1))

因此转测不能只测单个 Linear 切分，还必须覆盖 Style 组合、I/O 布局转换、以及 1-D TP mesh 约束。

相关资料：

• TP 使用指南
• TP Styles 文档
• PyTorch torch.distributed.tensor.parallel.style

---

2. 本期目标与非目标

2.1 本期目标

1. 公开 ParallelStyle 及上述 7 个具体 Style，均可从 hyper_parallel 顶层导入。
2. ColwiseParallel / RowwiseParallel 对 Linear、Embedding 做列切 / 行切，并挂上输入输出 hook。
3. SequenceParallel 复制参数，按序列维切激活，供 LayerNorm / Dropout / RMSNorm 使用。
4. PrepareModuleInput / PrepareModuleOutput / PrepareModuleInputOutput 只做边界布局转换，不切参数。
5. NoParallel 把参数做成全复制 DTensor，保持 DTensor 语义但不切分，供 Router / 关闭 SP 的 Norm 使用。
6. parallelize_module 支持：
   • 单个 Style 作用在根模块；
   • dict[str, ParallelStyle]，key 为 FQN，支持 fnmatch glob；
   • 仅 1-D DeviceMesh；
   • src_data_rank 写入 Style 后再 apply。
7. PyTorch / MindSpore 使用同一套 Style 语义；模块类型检测走 platform.is_linear_module / platform.is_embedding_module。

2.2 本期非目标

以下内容不在本次 style.py 转测范围内，但测试时需要知道它们会通过同一个 parallelize_module / ParallelStyle.apply 入口出现：

模块	位置	说明
MC2ColwiseParallel / MC2RowwiseParallel	mc2_style.py	继承 Col/Row，但会替换成 MC2Linear，并要求输入/输出为序列切分
ContextParallel / AsyncContextParallel / DSA CP / LinearAttentionContextParallel	core/context_parallel/	也是 ParallelStyle，走 CP 语义，不是 TP 权重切分
TensorParallel（EP 专家权重）	core/expert_parallel/expert_parallel.py	MoE expert 的 w1/w2/w3 切分，不是 nn.Linear Col/Row
Loss Parallel	loss_parallel.py	只影响 lm_head 输出是否保持 Shard(-1)，不是独立 Style
RaggedShard	DTensor	TP Style 当前只使用均匀 Shard / Replicate / Partial
自动并行、任意模块自动选 Style	—	parallelize_plan=None 只 warning，不做 auto-parallel

也不承诺：

• 对 Conv、RNN、自定义 Linear 的自动切分；
• 2-D/N-D mesh 直接传入 parallelize_module（必须先切 1-D 子 mesh，例如 mesh["tp"]）；
• 切分维不能被 tp_size 整除时的自动 padding；
• 整网性能、显存收益。

---

3. ParallelStyle 语义

3.1 基类契约

class ParallelStyle(ABC):
    src_data_rank: Optional[int] = 0

    @abstractmethod
    def apply(self, module, device_mesh) -> Module:
        ...

约束：

• 不能直接实例化 ParallelStyle()，必须实现 apply。
• apply 原地改模块（改参数、挂 hook），也可以返回包装后的模块。
• src_data_rank 由 parallelize_module(..., src_data_rank=...) 写入。0 表示从 rank 0 scatter/broadcast 全局参数；None 表示各 rank 本地切片、不通信。
• meta tensor 不会走源 rank 通信：_src_data_rank_for_tensor() 发现 tensor.is_meta 后强制把 src_data_rank 置为 None。

3.2 切分维约定

PyTorch nn.Linear.weight / MindSpore nn.Dense.weight 的逻辑 shape 都按 [out_features, in_features] 理解：

Style	模块	weight	bias	默认输入	默认输出
ColwiseParallel	Linear	Shard(0)，切 out_features	Shard(0)	Replicate()	Shard(-1)
ColwiseParallel	Embedding	Shard(1)，切 embedding_dim	无	Replicate()	Shard(-1)
RowwiseParallel	Linear	Shard(1)，切 in_features	Replicate()	Shard(-1)	Replicate()
RowwiseParallel	Embedding	Shard(0)，切 vocab	无	Replicate()（apply 时改 desired_input_layouts）	Replicate()
SequenceParallel	Norm / Dropout 等	全部 Replicate()	Replicate()	Shard(sequence_dim)，默认 dim=1	保持 DTensor，默认 use_local_output=False
NoParallel	任意模块	全部 Replicate()	Replicate()	Replicate()	Replicate()
PrepareModule*	任意模块	不改参数	不改参数	用户指定	用户指定

可整除约束（均匀 Shard，与 DTensor 一致）：

• Colwise Linear：out_features % tp_size == 0
• Rowwise Linear：in_features % tp_size == 0
• Colwise Embedding：embedding_dim % tp_size == 0
• Rowwise Embedding：num_embeddings % tp_size == 0
• SequenceParallel：被切的序列长度 % tp_size == 0

不满足时由 DTensor/通信层报错，Style 层不额外 padding。

3.3 激活上的通信

列切 + 行切的标准 MLP：

x: Replicate
        │
        ▼
   Colwise Linear          weight Shard(0)
        │
        ▼
   y: Shard(-1)            每个 rank 持有一部分 out_features
        │
        ▼
   Rowwise Linear          weight Shard(1)，输入最后一维已切
        │
        ▼
   Partial("sum")          本地 matmul 只覆盖部分 in_features
        │
        ▼
   Replicate               all-reduce / reduce-scatter

Qwen3.5 打开 Sequence Parallel 后，行切输出不再还原成 Replicate，而是 reduce-scatter 到 Shard(1)，让后续 Norm 继续吃序列分片。

3.4 真实示例：2 卡 Colwise Linear

tp_size = 2
Linear(in=32, out=64, bias=True)
weight global: (64, 32)
bias   global: (64,)

rank	weight local	bias local	默认输出
0	(32, 32)，对应 out [0, 32)	(32,)	最后一维 32
1	(32, 32)，对应 out [32, 64)	(32,)	最后一维 32

要和单卡参考对比，需要沿最后一维 all_gather 再 cat。

Rowwise Linear(in=32, out=24)：

rank	weight local	bias local	默认输出
0	(24, 16)，对应 in [0, 16)	完整 (24,)，Replicate	完整 out，各 rank 相同
1	(24, 16)，对应 in [16, 32)	完整 (24,)	完整 out，各 rank 相同

反向时，Rowwise 权重梯度在 gather 后需要除以 tp_size 再和单卡参考比（现有 ST 已按这个口径断言）。

---

4. 总体设计

4.1 架构与数据流

用户 API
  ColwiseParallel / RowwiseParallel / SequenceParallel /
  PrepareModuleInput / PrepareModuleOutput / PrepareModuleInputOutput /
  NoParallel / parallelize_module
        |
        v
parallelize_module
  校验 1-D mesh
  单 Style：写入 src_data_rank，调用 style.apply
  dict plan：按 FQN / fnmatch 找到 submodule，递归 apply
        |
        v
style.apply
  ├── Col/Row/Seq/NoParallel → distribute_module
  │     ├── partition_fn：切或复制参数
  │     ├── input_fn：local tensor → DTensor，必要时 redistribute
  │     └── output_fn：redistribute，必要时 to_local()
  └── PrepareModule* → 只注册 forward pre-hook / forward hook
        |
        v
DTensor
  distribute_tensor / from_local / redistribute / to_local
        |
        v
platform
  is_linear_module / is_embedding_module
  register_forward_pre_hook
  PT: nn.Linear / nn.Embedding
  MS: nn.Dense / nn.Embedding

设计原则：

• Style 不自己发 collective，通信发生在 distribute_tensor 和 DTensor.redistribute。
• 参数切分只发生一次；distribute_module 第二次调用同一模块会 RuntimeError。
• 未在 partition_fn 里切掉的参数，会被 distribute_module 自动做成 Replicate() DTensor。这就是 SequenceParallel / NoParallel 的参数语义来源。

4.2 parallelize_module

代码位置：hyper_parallel/core/tensor_parallel/api.py。

parallelize_module(
    module,
    device_mesh=None,
    parallelize_plan=None,
    *,
    src_data_rank=0,
) -> module

行为：

输入	行为
device_mesh=None	必须处于 with mesh: 或内部 _tensor_parallel_mesh_context
device_mesh.ndim > 1	ValueError：Tensor Parallel only accepts a 1D DeviceMesh
parallelize_plan=None	warning，原样返回 module
ParallelStyle	plan.src_data_rank = src_data_rank，然后 plan.apply(module, mesh)。会原地改传入的 Style 对象
dict	每个 value 必须是 ParallelStyle；key 按 . 拆成 FQN atom，atom 用 fnmatch 匹配 named_children / MindSpore name_cells
空 path、a..b	ValueError
dict 中某 path 匹配不到	warning，跳过该 path
其他类型	TypeError

混合并行正确用法：

mesh = init_device_mesh("npu", (dp, tp), mesh_dim_names=("dp", "tp"))
parallelize_module(model, mesh["tp"], tp_plan)
fully_shard(model, mesh=mesh["dp"])

4.3 distribute_module 与参数切分

代码位置：hyper_parallel/core/dtensor/dtensor.py。

ColwiseParallel / RowwiseParallel / SequenceParallel / NoParallel 都走它：

1. 对 named_modules 调用 partition_fn（可空）。
2. 把还不是 DTensor 的 param/buffer 复制成 Replicate()。
3. 在根模块上注册 input_fn / output_fn。
4. 打 _distribute_module_applied=True，禁止二次调用。

ColwiseParallel._partition_linear_fn：所有 param（weight、bias）distribute_tensor(..., [Shard(0)], src_data_rank=...)。

RowwiseParallel._partition_linear_fn：weight → Shard(1)，其它 param（bias）→ Replicate()。

Embedding 对应 Shard(1)（列切）或 Shard(0)（行切）。

SequenceParallel 的 partition_fn 是空操作，随后 replicate 通道把 Norm 权重做成复制 DTensor。

NoParallel 直接 partition_fn=None，全部 replicate。

4.4 I/O hook

Colwise / Rowwise / SequenceParallel / NoParallel

通过 distribute_module 的 input_fn / output_fn 注册。MindSpore pre-hook 必须返回 tuple，因此这些 Style 的 input hook 都返回 (prepared_first_input, ...)。

Colwise / Rowwise 只处理第一个位置参数：

plain tensor → DTensor.from_local(mesh, input_layouts)
if input_layouts != desired_input_layouts:
    redistribute(desired_input_layouts)

Rowwise Embedding 比较特殊：nn.Embedding.forward 即使 weight 已切，返回的仍是普通 tensor。output hook 会把它标成 Partial("sum")，再 redistribute 到 output_layouts。非 Embedding 的普通 tensor 输出直接 TypeError。

Rowwise / PrepareModuleOutput 的 reduce_dtype：当输出仍是 Partial 且目标 layout 不同时，先把 local tensor cast 到 reduce_dtype，再 redistribute。Qwen 的 o_proj 用 reduce_dtype=torch.float32 做高精度 reduce。

PrepareModuleInput

不走 distribute_module，直接 platform.register_forward_pre_hook。

• input_layouts / desired_input_layouts 可以是单个 Placement 或 tuple。
• tuple 里的 None 表示该位置参数原样透传。
• 提供 input_kwarg_layouts 时走 with_kwargs=True 的 pre-hook，同时处理 args 和 kwargs。
• use_local_output=True 表示准备完后 to_local()，模块 forward 看到的是普通 tensor。这个名字对齐 PyTorch，容易误解：它改的是输入，不是模块输出。

构造期校验：

• 给了 input_layouts 就必须给 desired_input_layouts；
• 两个 tuple 长度必须相同；
• kwarg 两个 dict 长度必须相同。

运行期：forward 实参个数必须等于 input_layouts 长度，否则 ValueError。非 tensor 且带 layout 时 AssertionError。

PrepareModuleOutput

直接 module.register_forward_hook。单输出返回单个 tensor；多输出返回 tuple。None slot 透传。reduce_dtype 仅在 Partial 输出且需要 redistribute 时生效。

PrepareModuleInputOutput

内部组合上面两个 Style。use_local_input 映射到 PrepareModuleInput(..., use_local_output=use_local_input)。

4.5 与其它模块的交互

模块	支持的交互	不支持 / 注意
DTensor distribute_tensor / redistribute	Style 的切分和 I/O 全部依赖它	不在 Style 内重复实现 collective
parallelize_module	唯一推荐入口	直接 style.apply 也可以，但不会自动写 src_data_rank（除非调用方自己设）
FSDP / HSDP fully_shard	先 TP 再 FSDP；TP 只用 mesh["tp"]	不要把 2-D 根 mesh 传给 parallelize_module
Sequence Parallel	与 Col/Row 组合：行切输出 Shard(seq)，下一层 PrepareModuleInput 再 gather	序列长度必须能被 tp_size 整除
Loss Parallel	lm_head 用 ColwiseParallel(output_layouts=Shard(-1), use_local_output=False)	Loss Parallel 本身不是 Style
MC2 Style	子类，复用 Col/Row 的切分 hook	要求序列切分输入/输出，不能当普通 Col/Row 测
Context Parallel	也是 ParallelStyle，可出现在同一个 parallelize_module plan 里	语义是序列/头维 all-to-all，不是本次 TP 权重切分
EP TensorParallel	切 MoE expert 三维权重	不要和 ColwiseParallel 混用在同一个 nn.Linear 上
DCP	保存的是切分后的 DTensor 参数	Style 不实现自己的 checkpoint 路径
RaggedShard	无	本期 TP Style 不使用非均匀切分

---

5. 对外接口

全部可从 hyper_parallel 导入。

5.1 应用入口

from hyper_parallel import (
    ColwiseParallel, RowwiseParallel, SequenceParallel,
    PrepareModuleInput, PrepareModuleOutput, PrepareModuleInputOutput,
    NoParallel, parallelize_module, init_device_mesh,
)

tp_mesh = init_device_mesh("npu", (tp_size,), mesh_dim_names=("tp",))
parallelize_module(model, tp_mesh, {
    "layers.*.attn.q_proj": ColwiseParallel(),
    "layers.*.attn.o_proj": RowwiseParallel(),
    "layers.*.norm": SequenceParallel(),
})

5.2 构造参数

ColwiseParallel

参数	默认	含义
input_layouts	Replicate()	第一个输入的标注 layout
output_layouts	Shard(-1)	输出目标 layout
use_local_output	True	输出是否 to_local()

desired_input_layouts 固定为 (Replicate(),)，用户不能改。如果输入已经是序列切分，Colwise 会 all-gather 成复制再算。

RowwiseParallel

参数	默认	含义
input_layouts	Shard(-1)	第一个输入的标注 layout
output_layouts	Replicate()	输出目标 layout；设成 Shard(seq_dim) 即 reduce-scatter
reduce_dtype	None	Partial reduce 前的浮点精度
use_local_output	True	输出是否 to_local()

Linear 的 desired_input_layouts 默认 (Shard(-1),)；Embedding 在 apply 时改成 (Replicate(),)。

SequenceParallel

参数	默认	含义
sequence_dim	1	序列维，对应 (B, S, H) 的 S；Qwen q_norm/k_norm 用 2
use_local_output	False	对齐 PyTorch，默认保持 DTensor

PrepareModuleInput

参数	默认	含义
input_layouts	None	每个位置参数的当前 layout，None slot 透传
desired_input_layouts	None	目标 layout
input_kwarg_layouts	None	keyword 参数 layout
desired_input_kwarg_layouts	None	keyword 目标 layout
use_local_output	False	准备后是否把输入转回 local tensor

PrepareModuleOutput

参数	默认	含义
output_layouts	必填	当前输出 layout
desired_output_layouts	必填	目标输出 layout
reduce_dtype	None	Partial reduce 精度
use_local_output	True	是否 to_local()

PrepareModuleInputOutput：上述输入侧参数 + use_local_input（默认 False）+ 输出侧参数。

NoParallel

参数	默认	含义
input_layout	Replicate()	第一个输入的标注
desired_input_layout	Replicate()	输入目标 layout
output_layout	Replicate()	输出目标 layout
use_local_output	True	是否 to_local()

注意：NoParallel 用的是单数 input_layout / output_layout，不是 Col/Row 的复数 *_layouts。

---

6. 当前支持矩阵

能力	PT	MS	当前状态 / 限制
ParallelStyle 抽象契约、src_data_rank	支持	支持	不能直接实例化
ColwiseParallel Linear 前向/反向	支持	接口支持	PT 有 2 卡 NPU ST；MS 缺独立精度 ST
ColwiseParallel Embedding 前向	支持	接口支持	PT 有 2 卡 NPU ST；MS 缺 ST
RowwiseParallel Linear 前向/反向	支持	接口支持	PT 有 2 卡 NPU ST；MS 缺独立精度 ST
RowwiseParallel Embedding	支持	接口支持	PT 有 UT（含 Partial 包装）；无 NPU ST
Col + Row MLP 组合	支持	部分	PT 有精度 ST；MS 仅 CommDebugMode 通信计数 ST
SequenceParallel LayerNorm / Dropout	支持	接口支持	PT 有 2 卡 + 4 卡 ST；MS 缺 ST
SequenceParallel(sequence_dim=2)	支持	接口支持	Qwen q_norm/k_norm 使用；无独立 ST
PrepareModuleInput 位置参数 / kwargs / None slot	支持	接口支持	PT 有 UT + 2/4 卡 ST
PrepareModuleOutput 单输出 / 多输出 None slot	支持	需验证 hook	PT 有 ST；MS 走 module.register_forward_hook，需确认 Cell 路径
PrepareModuleInputOutput 链路	支持	接口支持	PT 有 identity 和 MLP block ST
NoParallel 复制 Linear + SP→NoParallel redistribute	支持	接口支持	PT 有 2 卡 NPU ST，未进 Gloo level0
reduce_dtype（Rowwise / PrepareModuleOutput）	支持	接口支持	Qwen 使用；UT/ST 均未单列覆盖
src_data_rank=0/None	支持	支持	parallelize_module 有功能 ST；meta tensor 强制 None
fnmatch glob plan	支持	支持	PT UT + 2 卡 ST
1-D mesh 校验	支持	支持	2-D mesh 明确 ValueError
非 Linear/Embedding 的 Col/Row	不支持	不支持	NotImplementedError
直接传 N-D mesh	不支持	不支持	必须先 slice
Conv / 自定义 Linear	不支持	不支持	fail-closed
CUDA/NCCL	未作为验收设备	不涉及	本期按 NPU/HCCL 和 CPU/Gloo

“接口支持”表示代码走 platform 抽象，不区分 PT/MS 分支；是否在真实 MS 多卡上数值正确，需要转测补齐。

---

7. 风险与限制

7.1 只切第一个位置参数

Colwise / Rowwise / SequenceParallel / NoParallel 的 input hook 只处理 inputs[0]。多输入模块（例如带 mask 的 Attention 根模块）应使用 PrepareModuleInput，不要假设 Col/Row 会转换全部参数。

7.2 use_local_output 语义不统一

Style	默认	含义
Colwise / Rowwise / NoParallel / PrepareModuleOutput	True	模块对外返回 local tensor
SequenceParallel	False	保持 DTensor，方便后续 Style 继续 redistribute
PrepareModuleInput	False	改的是进入模块的输入

Qwen 里大量 use_local_output=False，就是为了让激活以 DTensor 形式在 Style 之间传递。测试时必须区分「返回 local」和「返回 DTensor」。

7.3 Rowwise Embedding 的 Partial 包装

Embedding 前向不会自动产出 DTensor 输出。Rowwise 依赖 output hook 把普通 tensor 标成 Partial("sum")。如果模块类型检测失败（platform.is_embedding_module 为 false），会变成 TypeError 而不是静默错误结果。

7.4 distribute_module 只能调用一次

同一模块不能先 ColwiseParallel.apply 再 PrepareModuleOutput.apply 走 distribute_module。PrepareModule* 不走 distribute_module，所以可以挂在已经被 Col/Row 切过的模块外面；但两个都会调用 distribute_module 的 Style 不能叠在同一个模块上。

7.5 MindSpore hook 差异

• PrepareModuleInput 使用 platform.register_forward_pre_hook，有 with_kwargs。
• PrepareModuleOutput 直接调用 module.register_forward_hook。
• Col/Row 的 pre-hook 返回 tuple，注释写明这是 MindSpore 要求。

转测 MS 时要把 kwargs pre-hook、多输出 forward hook 作为必测项，不能只测 PT。

7.6 梯度比较口径

• Colwise 权重梯度沿 dim 0 gather 后直接比单卡。
• Rowwise 权重梯度沿 dim 1 gather 后 除以 tp_size 再比单卡（现有 NPU ST 口径）。
• SequenceParallel 的 Norm 权重是复制的，各 rank 梯度应当一致，或按实现做 all-reduce 后再比。

7.7 混合并行顺序

推荐：先 parallelize_module(..., mesh["tp"])，再 fully_shard(..., mesh["dp"])。现有 4 卡 ST test_tp_fsdp_mlp_fwd_bwd_precision_npu 覆盖 MLP 这一路径。PP / CP 与 TP 的组合不在本次 Style 单测范围，但 plan 里可以同时出现 CP Style。

---

8. 验证设计与当前结果

8.1 已有覆盖（开发侧）

UT（tests/ut/core/tensor_parallel/，CPU mock，不启分布式）：

文件	覆盖
test_style.py	抽象类、src_data_rank、PrepareModule* 构造校验与 identity 前向
test_colwise_parallel.py	默认参数、Linear/Embedding partition、非法模块、I/O hook
test_rowwise_parallel.py	同上，含 Embedding Partial 输出路径
test_sequence_parallel.py	构造、distribute_module 回调、输入类型校验
test_no_parallel.py	构造、replicate、dict plan
test_api.py	parallelize_module glob、2-D mesh 拒绝、空 plan warning、单 Style 根应用

PyTorch ST（tests/torch/tensor_parallel/）：

启动器	关键 worker	卡数
test_tp_styles_distributed.py	Col/Row 非法模块、Col Linear fwd+bwd、Row Linear fwd+bwd、MLP 组合、Col Embedding、NoParallel	2
test_tp_sequence_parallel_distributed.py	LayerNorm chunk/gather、Dropout、无 affine、4 卡 fwd+bwd	2 / 4
test_prepare_module_io_distributed.py	Input/Output/InputOutput、kwargs、None slot、与 Col/Row 组合、4 卡 MLP block	2 / 4
test_parallelize_module_distributed.py	mesh 对齐、fnmatch、src_data_rank、单 Style 根	2
test_tp_hybrid_distributed.py	TP + FSDP MLP fwd+bwd	4

精度口径：NPU float32 vs CPU 单卡参考，rtol=1.5e-4，atol=1e-5。

MindSpore ST：

• tests/mindspore/st/dtensor/_test_comm_debug_mode_mlp.py：2 卡 Colwise+Rowwise MLP，只断言 collective 次数，不比数值。

8.2 转测建议用例

下列用例按测试可直接拆 Feature / Description / Expectation。未标注“已有”的视为转测补齐重点。

A. 接口与 fail-closed

ID	Feature	Description	Expectation
A1	抽象基类	ParallelStyle()	TypeError，信息含 abstract
A2	1-D mesh	把 2-D mesh 传给 parallelize_module	ValueError，提示使用 device_mesh["tp"]
A3	plan 类型	plan 为 list / 非法 dict value	TypeError
A4	空 FQN	"" 或 "a..b"	ValueError
A5	无匹配 path	dict key 在模型中不存在	warning，模块不被改
A6	Col/Row 非法模块	LayerNorm / ReLU / 自定义 Cell	NotImplementedError，信息含 Linear and Embedding
A7	Prepare 长度	input_layouts 与 desired_* 长度不同	构造期 AssertionError
A8	Prepare 前向 arity	两输入模块配单元素 layout	运行期 ValueError，含 same length
A9	SequenceParallel 输入类型	传入 Python list / int	ValueError，含 tensor or DTensor
A10	Rowwise 非 Embedding 普通 tensor 输出	Linear 路径产出非 DTensor	TypeError
A11	二次 distribute_module	对同一 Linear 连续两次 Colwise apply	RuntimeError
A12	空 plan	parallelize_plan=None	warning，返回原 module

A6、A2、A11 必须在 PT 和 MS 上都测。fail-closed 的验收表现就是立刻抛上述异常，不能静默按未切分模块继续算。

B. 参数切分几何

ID	Feature	断言
B1	Colwise Linear	weight.placements == (Shard(0),)，local shape [out/tp, in]；bias 同样 Shard(0)
B2	Rowwise Linear	weight Shard(1)，local [out, in/tp]；bias Replicate()，local 完整 out
B3	Colwise Embedding	weight Shard(1)，local [vocab, embed/tp]
B4	Rowwise Embedding	weight Shard(0)，local [vocab/tp, embed]
B5	SequenceParallel LayerNorm	weight/bias 为 Replicate() DTensor，local 完整 hidden
B6	NoParallel Linear	weight/bias 均为 Replicate()
B7	src_data_rank=0	仅 rank 0 持有正确全局权重，其它 rank 传入零 tensor，切分后各 rank local shard 与参考一致
B8	src_data_rank=None	各 rank 必须持有相同全局权重；结果与 B7 一致但不走 scatter
B9	meta 参数	is_meta 权重切分不发起源 rank 通信，不 hang

B4、B7、B8、B9 是当前缺口。

C. 数值精度（多卡 vs 单卡参考）

参考实现：同一组 float32 权重在 CPU 单进程上跑 F.linear / F.embedding / LayerNorm。

ID	Feature	卡数	说明
C1	Colwise Linear fwd+bwd	2	输出 all-gather 后比参考；wgrad cat dim0
C2	Rowwise Linear fwd+bwd	2	输出直接比；wgrad cat dim1 后 / tp_size
C3	Col+Row MLP	2	linear1 Colwise + relu + linear2 Rowwise
C4	Colwise Embedding fwd	2	输出 cat 最后一维
C5	Rowwise Embedding fwd	2	补测；输出应为完整 embedding，vocab 维 reduce
C6	SequenceParallel LayerNorm	2/4	本地 chunk 比参考切片；gather 后比完整参考；反向
C7	SequenceParallel Dropout(p=0)	2	与 identity 参考一致
C8	NoParallel Linear	2	各 rank 输出与单卡一致
C9	SP → NoParallel	2	序列切分激活 all-gather 后进复制模块
C10	PrepareInput → Colwise Linear	2/4	Shard(0)→Replicate 再列切
C11	Rowwise → PrepareOutput	2	行切后再改输出 layout
C12	PrepareInputOutput MLP block	4	输入 SP layout → 模块内 replicate → 输出再 shard
C13	RowwiseParallel(reduce_dtype=fp32)	2	补测；bf16/fp16 模块，reduce 用 fp32，对比高精度参考
C14	SequenceParallel(sequence_dim=2)	2	补测；输入 (B, H, S) 或 Qwen RMSNorm 形状
C15	TP+FSDP	4	先 TP 再 FSDP，fwd+bwd 对齐单卡

C1–C4、C6–C12、C15 在 PT 上已有；转测重点是 MS 复现 C1–C6、C8、C13、C14，以及 PT 补 C5/C13/C14。

建议 dtype：主路径 float32；C13 覆盖 fp16/bf16 + reduce_dtype=float32。

D. 组合与业务 plan

ID	Feature	Description
D1	glob	"layers.*.mlp.gate_proj": ColwiseParallel() 命中所有层
D2	单 Style 根模块	parallelize_module(linear, mesh, ColwiseParallel())
D3	Qwen 最小 block	PrepareModuleInput(Shard(1)→Replicate) + Colwise q/k/v + Rowwise o_proj(output=Shard(1), use_local_output=False) + SequenceParallel norm
D4	lm_head + Loss Parallel 布局	ColwiseParallel(input_layouts=Shard(1), output_layouts=Shard(-1), use_local_output=False)，输出保持 DTensor 且最后一维切分
D5	关闭 SP 的 Norm	NoParallel() 替代 SequenceParallel()，输入 replicate

D3 是最接近 parallelize_qwen3_5_tp 的最小可测单元，建议作为 MS/PT 共同验收。

E. 平台矩阵

后端	设备	必跑
PyTorch	CPU/Gloo，2 进程	A 类 + B 类几何 + 不依赖 HCCL 的 hook
PyTorch	NPU/HCCL，2 卡	C1–C12，补 C5/C13/C14
PyTorch	NPU/HCCL，4 卡	C6 反向、C12、C15
MindSpore	NPU，2 卡	A6、B1/B2、C1–C3、C6、C8
MindSpore	NPU，2/4 卡	PrepareModule* hook、D3

8.3 如何验证（对应测试常问问题）

1. 参数是否切对：读 module.weight，确认是 DTensor，看 placements 和 to_local().shape。
2. 前向是否数值正确：把各 rank local 输出按 layout gather 成全局 tensor，和单卡参考 allclose。
3. 通信是否发生：Rowwise 默认 Replicate 输出必须有 all-reduce / reduce-scatter；可用 CommDebugMode 或 profiler 看 collective 次数（MS 已有 MLP 通信 ST）。
4. fail-closed：非法模块、2-D mesh、layout 长度不匹配必须抛错，而不是按未切分模型算出一个“看起来合理”的结果。
5. 与 FSDP 交互：TP 后参数已是 DTensor；FSDP 应作用在 DP 子 mesh。验证最终 loss/grad 与单卡同配置参考接近，且不重复 all-reduce 导致梯度放大异常。

---

9. 验收标准

9.1 功能验收

• 7 个具体 Style 均可从 hyper_parallel 导入，且都是 ParallelStyle 子类。
• Colwise / Rowwise 仅接受 Linear、Embedding；其它类型 NotImplementedError。
• Linear/Embedding 的 weight、bias Placement 与第 3.2 节表格一致。
• 默认 Colwise 前向 gather 后等于单卡 Linear/Embedding；默认 Rowwise 前向等于单卡 Linear。
• SequenceParallel 在 sequence_dim 上切激活，参数保持复制；LayerNorm 前向/反向与单卡切片一致。
• PrepareModule* 按 layout 转换，None slot 不改对应输入/输出；kwargs 路径可用。
• NoParallel 不切权重，必要时把 sharded 输入 redistribute 成 Replicate。
• parallelize_module 支持单 Style、dict、fnmatch，拒绝 N-D mesh。
• src_data_rank=0/None 都能得到正确 local shard；meta 参数不通信。

9.2 兼容性验收

• 不传 Style 时，普通 nn.Linear / nn.Dense 行为不变。
• PT / MS 公开类名、构造参数、Placement 语义一致；差异只留在 platform（模块类型、hook 注册）。
• 先 TP 再 FSDP 的 2-D mesh 切片用法可用。
• 普通 DTensor / DCP / 非 TP 模块路径不受影响。

9.3 明确报错（必须 fail-closed）

场景	异常
实例化 ParallelStyle	TypeError
Col/Row 用于非 Linear/Embedding	NotImplementedError
parallelize_module 收到 N-D mesh	ValueError
plan 不是 Style 或 dict[str, Style]	TypeError
非法 FQN	ValueError
Prepare 构造 layout 长度不一致	AssertionError
Prepare 前向参数个数不匹配	ValueError
SequenceParallel 输入不是 tensor/DTensor	ValueError
Rowwise Linear 输出不是 DTensor	TypeError
同一模块两次 distribute_module	RuntimeError
PrepareModuleInput 对非 tensor 做 layout 转换	AssertionError

9.4 精度口径

场景	建议阈值
NPU float32 vs CPU 参考	rtol=1.5e-4，atol=1e-5（与现有 ST 一致）
fp16/bf16 + reduce_dtype=fp32	以 fp32 参考为准，阈值可放宽到 rtol=1e-3，atol=1e-3，但必须优于不升精度 reduce 的误差
反向梯度	Colwise 直接比；Rowwise gather 后除 tp_size

9.5 转测完成定义

• PT：现有 UT + NPU ST 回归通过，并补齐 Rowwise Embedding、reduce_dtype、sequence_dim=2。
• MS：至少具备 2 卡 Colwise Linear、Rowwise Linear、Col+Row MLP、SequenceParallel LayerNorm、非法模块报错的数值/功能 ST。
• 上述 fail-closed 场景在 PT/MS 均有对应用例。
• 不把 MC2 / CP / EP TensorParallel 算作本次 style.py 转测通过条件。