【RFC】HP Tensor Parallel ParallelStyle 设计文档

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 约束。

2. 本期目标与非目标

2.1 本期目标

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

对 named_modules 调用 partition_fn（可空）。
把还不是 DTensor 的 param/buffer 复制成 Replicate()。
在根模块上注册 input_fn / output_fn。
打 _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 如何验证（对应测试常问问题）

参数是否切对：读 module.weight，确认是 DTensor，看 placements 和 to_local().shape。
前向是否数值正确：把各 rank local 输出按 layout gather 成全局 tensor，和单卡参考 allclose。
通信是否发生：Rowwise 默认 Replicate 输出必须有 all-reduce / reduce-scatter；可用 CommDebugMode 或 profiler 看 collective 次数（MS 已有 MLP 通信 ST）。
fail-closed：非法模块、2-D mesh、layout 长度不匹配必须抛错，而不是按未切分模型算出一个“看起来合理”的结果。
与 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 转测通过条件。

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 335
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/335

Description