【RFC】hyper-parallel 支持torchtitan 风格 Module/Config (总览)

目标：在 hyper-parallel/ 仓库中引入 torchtitan 的「声明式 Module 协议 + Configurable 配置树」体系，打通从命令行 → Trainer.Config → 模型 Config 树 → config.build() → model.parallelize(mesh) → 训练循环的完整链路；接口与 torchtitan 尽量一致，必要的不一致显式记录。

---

1. 整体任务

把现有"裸 nn.Module + 字符串 _tp_plan + 过程式 parallelize_<name>()"路径，扩展为 torchtitan 风格的：

ConfigManager.parse_args()           （命令行 / yaml 兼容）
   ↓
HyperTrainer.Config                  （Configurable 配置树）
   ↓ config.build()
HyperTrainer.__init__
   ├─ model_cfg.update_from_config(trainer_config=config)
   │     └─ set_<name>_sharding_config()        # 灌 ShardingConfig
   ├─ with init_empty_weights(): model = model_cfg.build()
   ├─ model.verify_module_protocol()
   ├─ model = parallelize_<name>(model, ...)
   │     ├─ apply_cp_to_forward(...)
   │     ├─ model.parallelize(tp_mesh)           # 递归 Module.parallelize
   │     ├─ apply_ac(...)
   │     └─ apply_fsdp(...)
   ├─ model.to_empty + model.init_weights
   ├─ optimizer / lr_scheduler / tokenizer / dataloader / checkpointer
   │     ←—— 每个都走 cfg.build()
   └─ trainer.train()

3 条硬约束：

1. 零破坏：旧 yaml + LLMTrainer(args) 路径必须 100% 保留至少一个版本周期。
2. 跨后端：所有新代码不得直接 import torch.nn / import torch，必须经 platform = get_platform()；Module 基类是 (platform.Module, Configurable)。
3. 接口对齐 torchtitan：能完全照抄的接口（Configurable / Module / ShardingConfig / NamedPlacement / ModelSpec / ConfigManager）完全照抄；不一致的地方在每个模块文档显式列出。

flowchart TB
    subgraph entry [入口层]
        run_train["run_train.sh / torchrun"]
        train_py["train.py"]
        CM["ConfigManager"]
    end

    subgraph config_layer [配置层]
        CR["config_registry.py\nllama3_8b() → Trainer.Config"]
        TC["Trainer.Config\n(training/parallelism/checkpoint/...)"]
        MS["ModelSpec\n(model + parallelize_fn + pipelining_fn)"]
        MC["BaseModel.Config\n(嵌套 Module.Config 树)"]
    end

    subgraph runtime [运行时]
        Trainer["Trainer"]
        PD["ParallelDims\n(DeviceMesh)"]
        MP["model.parallelize(parallel_dims)"]
        FSDP["fully_shard / apply_fsdp"]
    end

    run_train --> train_py --> CM
    CM --> CR --> TC
    TC --> MS --> MC
    CM -->|"config.build()"| Trainer
    Trainer --> PD
    Trainer -->|"parallelize_fn"| MP --> FSDP

config驱动module的完整生命周期

① 定义 Config 树（model_registry / llama3_configs）
↓
② update_from_config() — 在 Config 树上填充/修改
（seq_len → rope.max_seq_len，set_llama3_sharding_config() → 各层 sharding_config）
↓
③ meta device 上 root Config.build() — 递归建 Module 树（无真实内存）
↓
④ verify_module_protocol() — 检查 Module 树每个节点都是 Module 子类
↓
⑤ parallelize_fn → model.parallelize(parallel_dims)
— 按 Module 树上的 _sharding_config 做 DTensor 分片
↓
⑥ to_empty(GPU) + init_states() — 按 Module 树上的 _param_init 初始化权重

2. 模块划分

按"代码依赖 + 可独立交付"切成 9 个相互独立的模块：

ID	模块	类型	作用一句话	工期	依赖
Part 1	protocols/	纯新增	协议基石：Configurable / Module / ShardingConfig / MeshAxisName / BaseModel / ModelSpec v2	3 d	无
Part 2	config/	纯新增	配置系统：ConfigManager + tyro CLI + 7 个通用 *Config + 新旧 yaml 适配	3 d	P1
Part 3	components/	纯新增	训练组件 Configurable 化：optimizer / lr_scheduler / loss / tokenizer / dataloader / checkpoint / profiler / metrics	4 d	P1
Part 4	models/common/	纯新增	通用模型组件库：Linear / RoPE / GQAttention / FeedForward / MoE / Decoder + decoder_sharding 声明式助手	7 d	P1
Part 5	trainer 接缝	新增 + 加分支	HyperTrainer.Config / HyperTrainer(BaseTrainer)，BaseTrainer.__init__ 加新分支	2.5 d	P1+P2+P3
Part 6	qwen3_5_v2/	纯新增	首个模型迁移：端到端打通新路径	3.5 d	P1+P4+P5
Part 7	CLI 入口	改 + 新增	config_registry.py + 改 scripts/train_lm.py（约 5 行）	2 d	P2+P6
Part 8	其他模型 + 切换	新增 + 切换	qwen3_5_moe_v2 / qwen3_vl_moe_v2；切换默认 spec；清理 legacy	11 d	P1–P7
Part 9	local_map	可选	core/dtensor/local_map.py 包装；按需触发	3 d	P1

详细任务边界、设计点、开发步骤、测试用例、与 torchtitan 接口差异，均见每个模块的独立文档。

---

3. 模块依赖关系图

              ┌─────────────────────────────────────────────────┐
              │ Part 1 protocols/                                   │
              │ Configurable / Module / BaseModel /             │
              │ ShardingConfig / NamedPlacement /               │
              │ MeshAxisName / ModelSpec v2                     │
              │ （基石，无外部依赖）                            │
              └─────────────────────────────────────────────────┘
                   │              │              │              │
       ┌───────────┘              │              │              │
       ▼                          ▼              ▼              ▼
┌──────────────┐         ┌────────────────┐ ┌─────────────────┐ ┌────────────┐
│ Part 2 config/   │         │ Part 3 components/ │ │ Part 4 models/      │ │ Part 9 local_  │
│              │         │                │ │   common/       │ │   map      │
│ ConfigMgr +  │         │ optimizer /    │ │                 │ │ （可选     │
│ tyro CLI     │         │ lr_scheduler / │ │ Linear / RoPE / │ │  扩展）    │
│              │         │ loss /         │ │ GQAttention /   │ │            │
│ 7 个通用     │         │ tokenizer /    │ │ FeedForward /   │ │            │
│ *Config      │         │ dataloader /   │ │ MoE / Decoder + │ │            │
│              │         │ checkpoint /   │ │ decoder_        │ │            │
│ legacy yaml  │         │ profiler /     │ │ sharding 助手   │ │            │
│  适配        │         │ metrics        │ │                 │ │            │
└──────────────┘         └────────────────┘ └─────────────────┘ └────────────┘
       │                         │              │
       └───────────┬─────────────┴──────┬───────┘
                   ▼                    ▼
            ┌─────────────────────────────────┐
            │ Part 5 trainer 接缝                 │
            │                                 │
            │ HyperTrainer.Config             │
            │ HyperTrainer(BaseTrainer) 新分支│
            │ BaseTrainer.__init__ 共享 helper│
            └─────────────────────────────────┘
                            │
                            ▼
            ┌─────────────────────────────────┐
            │ Part 6 qwen3_5_v2/                  │
            │                                 │
            │ model.py + sharding.py +        │
            │ parallelize.py + state_dict.py  │
            │ + __init__.py (register_spec)   │
            └─────────────────────────────────┘
                            │
                            ▼
            ┌─────────────────────────────────┐
            │ Part 7 CLI 入口打通                 │
            │                                 │
            │ qwen3_5_v2/config_registry.py + │
            │ scripts/train_lm.py 改 5 行     │
            └─────────────────────────────────┘
                            │
                            ▼
            ┌─────────────────────────────────┐
            │ Part 8 扩展到 MoE / VL + 切换默认   │
            │                                 │
            │ qwen3_5_moe_v2 /                │
            │ qwen3_vl_moe_v2 / 切默认 spec   │
            └─────────────────────────────────┘

关键路径（必须串行）：P1 → P4 → P5 → P6 → P7，最快 ~17 天。

并行机会：P1 完成后，P2 / P3 / P4 / P9 可同时推进。4 名工程师并行最快 ~12 天到 P7 完成。

---

4. 各模块一句话职责

Part 1 — hyper_parallel/protocols/（协议基石）

定义 Configurable / Configurable.Config（__init_subclass__ 自动绑 owner，Config.build() 即构造），以及 Module(platform.Module, Configurable) 三件套（init_states / parallelize(mesh) / from_nn_module）、ShardingConfig / NamedPlacement / MeshAxisName、BaseModel / ModelConfigConverter、新版 ModelSpec。纯新增，零依赖。

Part 2 — hyper_parallel/config/（配置系统）

基于 P1 的 Configurable，提供：ConfigManager.parse_args() 双入口（--module 走 tyro / 否则走旧 yaml 适配）、TrainingConfig / ParallelismConfig / ActivationCheckpointConfig / CompileConfig / CommConfig / DebugConfig 等通用 dataclass、Function(Configurable)、TORCH_DTYPE_MAP、_legacy_to_trainer_config 旧字段映射。tyro 列为可选依赖。

Part 3 — hyper_parallel/components/（训练组件 Configurable 化）

把现散落在 BaseTrainer._build_*（trainer/base.py:555-700）的"过程式 build"抽成 Configurable 子类：OptimizersContainer / LRSchedulersContainer / BaseLoss + CrossEntropyLoss / BaseTokenizer + HuggingFaceTokenizer / BaseDataLoader + HuggingFaceTextDataLoader + DummyDataLoader / CheckpointManager / Profiler / MetricsProcessor。算法实现直接搬运旧 _build_* 函数体，保证 bit-exact 等价。

Part 4 — hyper_parallel/models/common/（通用模型组件库）

把 hyper_parallel/models/modules/ 现有裸 nn.Module 组件重写为 Module 协议版：Linear / Embedding / RMSNorm / Qwen3_5RMSNorm / RoPE / GQAttention / FeedForward / MoE / TransformerBlock / Decoder，加 param_init.py、decoder_sharding.py（9 个声明式 sharding 助手）。旧 models/modules/ 保留供老路径继续用。

Part 5 — trainer/ 接缝层

新增 trainer_config.py（HyperTrainer.Config 顶层 Configurable 树）+ hyper_trainer.py（HyperTrainer(BaseTrainer) 新子类，11 步 __init__ 全走 cfg.build()）。BaseTrainer.__init__ 只新增一个分支，把现有 13 步 _build_* 抽 helper，旧 LLMTrainer 走旧 helper、新 HyperTrainer 走 helper 加 build —— 算法零分叉。

Part 6 — models/qwen3_5_v2/（首个模型迁移）

按 P1 / P4 重写 Qwen3_5Model(Decoder) + Qwen3_5TransformerBlock，set_qwen3_5_sharding_config(), parallelize_qwen3_5_v2()（顺序 CP → model.parallelize(tp_mesh) → AC → FSDP），Qwen3_5_v2StateDictAdapter(BaseStateDictAdapter)，register_spec("qwen3_5_v2", ...)。旧 models/qwen3_5/ 完全保留。

Part 7 — CLI 入口打通

写 models/qwen3_5_v2/config_registry.py（qwen3_5_v2_debugmodel() / _4b() / _4b_tp2_fsdp4() 等 recipe，每个返回 HyperTrainer.Config）；改 scripts/train_lm.py 入口（约 5 行）：

mgr = ConfigManager()
config = mgr.parse_args()      # 自动识别 --module / yaml
trainer = config.build()       # Configurable.Config.build() → HyperTrainer(config)
trainer.train()

Part 8 — 其他模型迁移 + 切换默认

按 Part 6 模式迁移 qwen3_5_moe_v2 / qwen3_vl_moe_v2；待全部 v2 稳定 ≥ 1 周后切换默认：register_spec("qwen3_5", ...) 切到 v2 实现，旧实现挪到 _legacy/ 加 @deprecated；hyper_parallel/__init__.py 暴露新 API；写 docs/zh/module_protocol.md。

Part 9 — local_map 扩展（可选）

新增 hyper_parallel/core/dtensor/local_map.py（torch 后端转发 torch.distributed.tensor.experimental.local_map，mindspore 后端 raise）；改 Module.parallelize 在 sharding_config.local_map is not None 时启用；新增 set_gqa_inner_attention_local_map(...) 助手。只有当某个模型必须走 q/k/v head-shard 路径时才触发。

---

5. 风险与红线

1. 任何模块 PR 必须先过 1-card UT + 8-card 数值对齐（与 main 分支同种子 loss bit-exact）才能合入。
2. hyper_parallel/__init__.py 在 Part 8.3 之前不暴露任何新 API，避免阶段未稳定就成为公共面。
3. 旧 LLMTrainer(args) 路径在 Part 8.3 切换前 0 修改，保护现网用户。
4. tyro 是新增依赖，列入可选 extras；Part 2 必须保证缺失时报清晰错误。
5. 跨后端兼容：M1 / M4 严格走 platform；Part 9 在 mindspore 后端用 capability flag 控制。

---

6. 必要的不一致（与 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	local_map 第一阶段未实现（M9 才落）	torch 实验 API，mindspore 后端可能永久 raise
5	ConfigManager 同时支持新 tyro / 旧 yaml 入口	零破坏迁移
6	dtype 字段用 Literal["bfloat16", "float32", "float16"]	tyro 不能解析 torch.dtype，且要跨后端
7	ModelSpec v2 兼容 v1：build_model_fn / model: BaseModel.Config 二选一	支持渐进迁移，老 spec 不动
8	ModelSpec 字段标 tyro.conf.Suppress	spec 持 callable / dataclass，tyro 不能从 CLI 解析
9	hyper 暂时保留 13 个 Callback 体系不动	M3 MetricsProcessor 做组件式入口，但不重写 callback；M8 再统一

---

7. issue任务清单

任务	内容
00_overview.md	本文（总览）
Part 1_protocols	协议基石
Part 2_config	配置系统
Part 3_components	训练组件
Part 4_models_common	通用模型组件 + sharding 助手
Part 5_trainer_glue	trainer 接缝层
Part 6_qwen3_5_v2	首个模型迁移
Part 7_cli_entrypoint	CLI 入口打通
Part 8_migrate_others	其他模型迁移 + 切换默认
Part 9_local_map	local_map 扩展（可选）

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

Treat this as an RFC umbrella rather than a single first contribution. Start with Part 1, protocols/, and the referenced M1_protocols.md, then inspect trainer/base.py:555-700 and the existing models/modules/ path. Done requires completing a defined module from the dependency plan with its tests, compatibility checks, and required 1-card and 8-card validation.