【RFC】hyper-parallel 支持torchtitan 风格 Module/Config —— Part 4 通用module组件 + sharding

Part 4 — hyper_parallel/models/common/ 通用模型组件 + decoder_sharding 助手

把 hyper_parallel/models/modules/ 现有裸 nn.Module 组件重写为 Module 协议版，并提供 torchtitan 风格的声明式 sharding 助手。旧 models/modules/ 保留供老路径继续用。

---

1. 目标

1. 提供 Module 协议（M1）兼容的通用组件库：Linear / Embedding / RMSNorm / RoPE / GQAttention / FeedForward / MoE / TransformerBlock / Decoder。
2. 提供 9 个声明式 sharding 助手，把 ShardingConfig 灌进模型 config 树。
3. 完全不出现 TP / FSDP 代码 —— 只建图。

2. 任务边界（新增文件）

新增包目录 hyper_parallel/models/common/：

新文件	旧对应（迁移参考）	内容
linear.py	—	Linear(platform.Linear, Module) + Linear.Config(in_features, out_features, bias, _param_init: Function.Config)
embedding.py	—	Embedding(platform.Embedding, Module) + Embedding.Config
rmsnorm.py	models/modules/rmsnorm.py RMSNorm / RMSNormGated	RMSNorm(Module) 通用版 w * normed；Qwen3_5RMSNorm(RMSNorm) 残差式 (1+w) * normed
rope.py	models/modules/rope.py RotaryEmbedding / MultiModalRotaryEmbedding / apply_rotary_pos_emb	RoPE(Module)：_init_self_buffers(buffer_device) 重算 inv_freq / cos / sin cache，配合 init_states
attention.py	models/modules/attention.py:39 GroupQueryAttention	BaseAttention(Module) + GQAttention(Module)：合并现有 attn_output_gate / qk_norm 能力；初版不引入 FlexAttention / VarlenAttention
feed_forward.py	models/modules/feed_forward.py SwiGLUMLP	FeedForward(Module)：w1/w2/w3 SwiGLU
moe.py	models/modules/moe.py MoE / TopKRouter / MoEExperts / SharedExpertMoE	MoE(Module) / Router(Module) / Experts(Module) / SharedExpertMoE(Module)；端口对齐旧版字段名
decoder.py	—	TransformerBlock(Module) + Decoder(BaseModel) + DecoderBlock.Config；init_states 处理 freqs_cis
param_init.py	—	kaiming_uniform_ / xavier_normal_ / zeros_ / ones_ 跨后端封装；Function.Config 可直接持有
decoder_sharding.py	—	9 个声明式 sharding 助手（详见 §4）
__init__.py	—	集中导出

3. 关键约束

• 禁止顶层 import torch.nn as nn。Module 类来自 from hyper_parallel.protocols import Module；张量算子优先用 platform。

• 每个组件必须有 Config + __init__(self, config)。示例：

  class Linear(platform.Linear, Module):
      @dataclass(kw_only=True, slots=True)
      class Config(Module.Config):
          in_features: int
          out_features: int
          bias: bool = True
          _param_init: Function.Config = field(
              default_factory=lambda: Function.Config(fn=kaiming_uniform_)
          )
  
      def __init__(self, config: "Linear.Config"):
          platform.Linear.__init__(
              self, config.in_features, config.out_features, bias=config.bias,
          )
          Module.__init__(self, config)
  

• 完全不出现 TP / FSDP 代码，只建图。

• 数值兼容：组件 forward 必须与 models/modules/ 旧版同种子 bit-exact。

4. decoder_sharding.py 9 个助手

对应 torchtitan/models/common/decoder_sharding.py，函数签名一致。

函数	输入 / 输出
dense_param_placement(*, tp) -> NamedPlacement	返回 {TP: Shard(0)} 等通用 dense 参数 placement
dense_activation_placement(*, tp, cp=Shard(1)) -> NamedPlacement	TP / CP 联合 activation placement
colwise_config() -> ShardingConfig	列切线性层（state + in + out 全填）
rowwise_config(*, output_sp: bool) -> ShardingConfig	行切线性层
norm_config(*, enable_sp: bool) -> ShardingConfig	RMSNorm 的 SP 配置
set_qkv_linear_sharding(qkv_cfg)	inplace 给 q_proj / k_proj / v_proj.sharding_config = colwise_config()
set_gqa_attention_sharding(attn_cfg, *, enable_sp)	给 attention 子树灌 sharding，含 qk_norm、o_proj、q_norm / k_norm
set_dense_ffn_sharding(ffn_cfg, *, attn_x_placement, enable_sp)	dense FFN
set_decoder_sharding_config(cfg, *, loss_parallel, enable_sp)	顶层入口，递归走所有 layer

Replicate / Shard / Partial 直接来自 hyper_parallel.core.dtensor.placement_types；MeshAxisName 来自 M1。

5. 与 torchtitan 接口差异说明

#	差异点	原因
1	组件名前缀避免冲突。models/modules/ 已有 GroupQueryAttention / RMSNorm / SwiGLUMLP / MoE；models/common/ 用 GQAttention / RMSNorm / FeedForward / MoE（同名但 import 路径区分）	新旧并存
2	Linear / Embedding 强制继承 platform.Linear / platform.Embedding	mindspore 后端走 Cell 等价物
3	RoPE._init_self_buffers 取代 reset_inv_freq	旧 base.py:1212 仍可调；新协议 Module 通过 init_states(buffer_device) 触发重算
4	Qwen3_5RMSNorm 单独保留（残差式 (1+w)*normed）	与 Qwen3.5 / Qwen3.5-MoE 权重存储约定耦合，存量权重不能改
5	初版不引入 FlexAttention / VarlenAttention	hyper 当前 attention 路径不依赖；按需后续补

6. 开发步骤

步	内容	工期
1	linear.py + embedding.py + param_init.py	1 d
2	rmsnorm.py + rope.py	1 d
3	attention.py（GQAttention）+ feed_forward.py	1.5 d
4	moe.py	1.5 d
5	decoder.py：TransformerBlock + Decoder(BaseModel)	1 d
6	decoder_sharding.py 9 个助手	1 d

Step 2 数值对齐示例：

# 与 models/modules/rmsnorm.py 同种子对齐：
torch.manual_seed(42)
old = RMSNorm_old(hidden=256)
new = RMSNorm.Config(hidden=256).build()
# 复制参数：
new.weight.data.copy_(old.weight.data)
x = torch.randn(2, 8, 256)
assert torch.equal(old(x), new(x))

7. 验证标准

新建 tests/torch/ut/models/common/：

测试	断言要点
test_linear_embedding.py	同种子下 Linear.Config(...).build()(x) 与 nn.Linear(...)(x) 数值一致；Embedding 同样
test_rmsnorm.py	RMSNorm.Config(...).build()(x) 与 models/modules/rmsnorm.py:RMSNorm(x) 一致；Qwen3_5RMSNorm 与 models/qwen3_5/model.py:51 一致
test_rope.py	init_states(buffer_device=cpu) 后 cos / sin cache 与旧 MultiModalRotaryEmbedding 完全一致
test_attention.py	GQAttention.Config(...).build()(hidden_states, position_ids) 与 models/modules/attention.py:39 GroupQueryAttention 同种子 bit-exact，含 attn_output_gate / qk_norm
test_feed_forward.py	同上
test_moe.py	MoE.Config(...).build() 同种子 forward 与 models/modules/moe.py:MoE 一致；router 输出 bit-exact
test_decoder.py	构造 2 层 Decoder.Config(...)，CPU init_states + forward 通过
test_decoder_sharding.py	给 Decoder.Config 调 set_decoder_sharding_config(..., loss_parallel=True, enable_sp=False)，断言每个子 module config 的 sharding_config 字段完全等于预期

通过门槛：

• 8 个 UT 全绿、数值 bit-exact。
• 存量 models/modules/、models/qwen3_5/ 0 改动。

8. 工期 & 依赖

工期	7 天
依赖	M1（Module / Configurable / ShardingConfig / MeshAxisName）
并行	与 M2 / M3 完全并行
下游	M6

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

Start by reading the M1 Module, Configurable, ShardingConfig, and MeshAxisName interfaces, then compare the existing files under hyper_parallel/models/modules/ and models/qwen3_5/. Implement the listed files under hyper_parallel/models/common/ and add the eight tests under tests/torch/ut/models/common/. Done means all tests pass bit-exactly and the existing models/modules/ and models/qwen3_5/ files remain unchanged.