【RFC】hyper-parallel 支持torchtitan 风格 Module/Config —— Part 1 协议层

Part 1 — hyper_parallel/protocols/ 协议层

整个工程的基石。完成后整套抽象类可独立 import、独立单测，但不接入任何已有路径。纯新增、零侵入。

---

1. 目标

实现 torchtitan 的 Configurable / Module / BaseModel / ShardingConfig / MeshAxisName / ModelSpec 一整套协议，作为后续 M2–M9 共同的基础抽象层。

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

新增包目录 hyper_parallel/protocols/：

新文件	主要内容	对应 torchtitan 文件
types.py	MeshAxisName(StrEnum)：定义 DP / DP_REPLICATE / DP_SHARD / FSDP / TP / CP / PP / EP / EFSDP。字面量与 hyper 现用 mesh 维度名对齐（"fsdp" / "dp_shard" / "tp" / "cp" / "pp" / "ep" / "dp_replicate"）	torchtitan/protocols/types.py
configurable.py	Configurable 基类 + 嵌套 Configurable.Config(@dataclass(kw_only=True, slots=True))；__init_subclass__ 自动把 Config._owner = cls，Config.build(**kwargs) 即构造 cls(self, **kwargs)；提供 replace / traverse / to_dict	torchtitan/config/configurable.py
sharding.py	NamedPlacement = dict[MeshAxisName, Placement]；@dataclass class ShardingConfig（state_shardings / in_src_shardings / in_dst_shardings / out_dst_shardings / local_map）；LocalMapConfig；resolve_placements(named, mesh_axis_names) -> list[Placement]。Placement 直接复用 hyper_parallel.core.dtensor.placement_types.{Placement, Shard, Replicate, Partial}	torchtitan/protocols/sharding.py
module.py	class Module(platform.Module, Configurable)：实现 init_states / _init_self_parameters / _init_self_buffers / _init_param / _cache_pos_arg_names / parallelize / _shard_inputs / _shard_outputs / from_nn_module；定义 ModuleList / ModuleDict / Sequential 容器版	torchtitan/protocols/module.py
model.py	class BaseModel(Module) + BaseModel.Config(Module.Config)，含 verify_module_protocol / init_weights / update_from_config(trainer_config) / get_nparams_and_flops；ModelConfigConverter(Configurable) 占位	torchtitan/protocols/model.py
state_dict_adapter.py	class BaseStateDictAdapter(ABC)，方法签名兼容现有 hyper_parallel/models/spec/state_dict_adapter.py:28 的 Protocol（load_hf_state_dict / save_hf_state_dict）	torchtitan/protocols/state_dict_adapter.py
model_spec.py	新版 ModelSpec：兼容现有 models/spec/model_spec.py:23 的 5 字段，新增 model: BaseModel.Config | None；旧字段 build_model_fn 改为 Optional，注册时二选一	torchtitan/protocols/model_spec.py
__init__.py	集中导出公共符号	—

3. 核心设计点

3.1 Configurable.__init_subclass__

class Configurable:
    Config: type["Configurable.Config"]

    @dataclass(kw_only=True, slots=True)
    class Config:
        _owner: ClassVar[type["Configurable"]] = None

        def build(self, **runtime_kwargs) -> "Configurable":
            return self._owner(self, **runtime_kwargs)

    def __init_subclass__(cls, **kw):
        super().__init_subclass__(**kw)
        if "Config" in cls.__dict__:
            cls.Config._owner = cls

3.2 Module.parallelize(tp_mesh) 流程

def parallelize(self, tp_mesh):
    sc = self.sharding_config
    if sc is None:
        # 递归 children
        for child in self.children():
            if isinstance(child, Module):
                child.parallelize(tp_mesh)
        return

    if sc.local_map is not None:
        raise NotImplementedError("local_map will be added in M9")

    # 1. 分布参数 / buffer
    for path, named_p in sc.state_shardings.items():
        param = _get_attr_by_path(self, path)
        placements = resolve_placements(named_p, tp_mesh.mesh_dim_names)
        new_local = distribute_tensor(param.data, tp_mesh, placements)
        _set_param_by_path(self, path, platform.Parameter(new_local))

    # 2. 包 forward 做 in/out reshard
    self._wrap_forward_with_reshard(tp_mesh)

    # 3. 递归子模块
    for child in self.children():
        if isinstance(child, Module):
            child.parallelize(tp_mesh)

注意 distribute_tensor 来自 hyper_parallel.core.dtensor.dtensor:409，签名 (tensor, device_mesh, placements)，与 torch 版兼容。

3.3 resolve_placements

def resolve_placements(
    named: NamedPlacement,                    # dict[MeshAxisName, Placement]
    mesh_axis_names: tuple[str, ...],         # 实际 mesh 的维度名
) -> list[Placement]:
    out = []
    for axis_name in mesh_axis_names:
        if axis_name in named:
            out.append(named[axis_name])
        else:
            out.append(Replicate())            # 未声明的轴默认 Replicate
    # 缺主轴报错
    declared = set(named.keys())
    extra = declared - set(mesh_axis_names)
    if extra:
        raise ValueError(f"NamedPlacement has axes not in mesh: {extra}")
    return out

4. 与 torchtitan 接口差异说明

#	差异点	原因
1	Module 继承 (platform.Module, Configurable) 而非 (nn.Module, Configurable)	hyper 跨后端（torch + mindspore）
2	distribute_tensor 用 hyper_parallel.core.dtensor.distribute_tensor（core/dtensor/dtensor.py:409）	hyper 自己的实现；mindspore 后端不能用 torch 版
3	MeshAxisName 字面量优先 hyper 现用名（FSDP="fsdp"），同时收 DP_SHARD="dp_shard"	与 trainer/parallel_dims.py 已建的 mesh 命名对齐
4	ShardingConfig.local_map 字段保留但 raise NotImplementedError	M9 才落，mindspore 后端可能永久 raise
5	Configurable.__init_subclass__ 的 slots=True 校验只对 Module 子树生效	hyper 存量 dataclass 没用 slots，强制开启会破坏存量代码
6	ModelSpec v2 兼容 v1：build_model_fn / model: BaseModel.Config 二选一	支持渐进迁移，老 spec 不动

5. 开发步骤

1. Step 1（0.5 d）：建包 + types.py + __init__.py。
2. Step 2（1 d）：configurable.py。
   • __init_subclass__ 自动绑 owner。
   • Config.build() 调 cls(self, **runtime_kwargs)。
   • Configurable.Config.replace(**kw) 用 dataclasses.replace 实现。
3. Step 3（0.5 d）：sharding.py。resolve_placements 是核心。
4. Step 4（1 d）：module.py 的 Module 基类 + ModuleList / ModuleDict / Sequential。
5. Step 5（0.5 d）：model.py 的 BaseModel。update_from_config 默认空实现，由子类重写。
6. Step 6（0.5 d）：state_dict_adapter.py + model_spec.py。

6. 验证标准

新建 tests/torch/ut/protocols/：

测试	断言要点
test_configurable.py	Config.build() 构造正确；replace(field=val) 不改原对象；traverse 遍历嵌套 Config 子树；to_dict 输出可 yaml 序列化
test_module.py	1-card CPU：玩具 Module 子类的 init_states()、_init_self_buffers(device)、from_nn_module(nn.Linear) 复用同一份 Config；_cache_pos_arg_names 通过 inspect.signature(self.forward) 缓存
test_module_parallelize.py	2-card：构造一个 nn.Linear 风格的 Module，配 state_shardings={"weight": {"tp": Shard(0)}}、in_dst_shardings={"input": {"tp": Replicate()}}、out_dst_shardings={"output": {"tp": Shard(-1)}}；module.parallelize(mesh) 后参数变 DTensor、forward 通过、数值与单卡一致
test_sharding_resolve.py	resolve_placements 对缺轴默认 Replicate()；对额外轴 raise ValueError；按 mesh.mesh_dim_names 顺序输出
test_module_local_map_stub.py	设 local_map 非空时调 parallelize 必须 raise NotImplementedError("local_map will be added in M9")

通过门槛：

• 5 个 UT 全绿。
• import hyper_parallel.protocols 不引入循环依赖。
• 现存 tests/torch/st/qwen3_5/ 等存量测试 0 受影响。
• hyper_parallel/__init__.py 不暴露新符号。

7. 工期 & 依赖

工期	3 天
依赖	无。可立即启动。
下游	M2 / M3 / M4 / M5 / M9 都依赖 M1

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

Start with the requested files under hyper_parallel/protocols/, especially types.py, configurable.py, and sharding.py, then inspect hyper_parallel/core/dtensor/dtensor.py:409 and the existing model-spec and state-dict adapter files. Run the tests in tests/torch/ut/protocols/ as they are added. Done means the five protocol test groups pass, hyper_parallel.protocols imports without cycles, and existing tests remain unaffected.