[RFC] Codegen

Start with docs/guide/codegen.md and the entry point hyper_parallel/codegen/manager.py, then inspect the listed tests under tests/codegen/. The RFC describes a broad code-generation design with an existing paired pull request, so any follow-up should first compare the current implementation with the stated artifact, runtime, and training-integration contracts.

Paired: GitHub #1 ↔ GitCode !1350

Codegen 设计方案

本文描述 Codegen 的终版设计，包括代码路径、产物规则、训练接入链路和运行边界。

---

0. 目标

Codegen 的目标是把训练 YAML 中声明好的并行策略和模块替换规则，在训练启动前生成成一份可读、可
diff、可调试的 generated model。训练时，gen 后端直接导入这份 generated model，并调用其中的
hyper_parallelize() 执行冻结后的并行计划。

Codegen 不维护第二份用户配置。训练 YAML 就是生成输入，Codegen 只从 TrainerConfig 中投影会影响生成产物的
字段，避免训练配置和生成配置分离后产生不一致。

---

1. 设计边界

Codegen 负责：

• 根据 YAML 解析出的 TrainerConfig 生成或复用 artifact bundle；
• 解析 HuggingFace config 和 modeling source，记录 source sha256 与 Transformers 版本；
• 用 meta-device 模型、offline mesh 和现有 ShardingPlanner 推导并冻结 sharding plan；
• 将 plan_overrides 中的 replace_module 下沉到 generated model 的 __init__；
• 将边界通信、参数分片计划、special handler、tied 参数处理写入 generated model；
• 生成 codegen_meta.json、generated modeling 文件、diff 文件和 __init__.py；
• 训练前执行 artifact 完整性检查。

Codegen 不负责：

• DeviceMesh 的真实构建；
• PP 切分；
• FSDP / HSDP 包装；
• activation checkpoint、activation swap、torch.compile；
• checkpoint 权重读写流程；
• checkpoint transform / weight mapping 的完整静态化。该部分保留为 TODO。

---

2. 目录和文件命名

Codegen 采用 YAML 同级目录规则。给定：

/home/lxy/hyper-parallel/examples/training_demo/train_codegen_qwen3_moe.yaml

Codegen artifact bundle 默认生成在：

/home/lxy/hyper-parallel/examples/training_demo/generated/

目录规则为：

<yaml 所在目录>/generated/

同一个目录下如果存在多份 codegen YAML，它们默认共用同一个 generated/ bundle 目录。后一次生成会用目标 YAML
对应的 signature 覆盖旧 bundle。需要多份 artifact 并存时，应把 YAML 放到不同目录，或在集成代码中显式传入
codegen_artifact_dir。

生成文件名优先由 HuggingFace config 的 model_type 推导：

qwen3_moe -> Qwen3_Moe -> modeling_Qwen3_Moe_gen_npu.py

bundle 示例：

examples/training_demo/
├── train_codegen_qwen3_moe.yaml
└── generated/
    ├── __init__.py
    ├── codegen_meta.json
    ├── modeling_Qwen3_Moe_gen_npu.py
    └── modeling_Qwen3_Moe_gen_npu.py.diff

---

3. 包路径和文件清单

Codegen 包位于 hyper_parallel/codegen/。

3.1 Codegen 包文件树

hyper_parallel/codegen/
├── __init__.py
├── __main__.py
├── artifact.py
├── cli.py
├── hash.py
├── loader.py
├── manager.py
├── meta.py
├── modeling_backend.py
├── runtime.py
├── astkit/
│   ├── __init__.py
│   ├── edits.py
│   └── index.py
├── check/
│   ├── __init__.py
│   ├── drift.py
│   ├── pin.py
│   └── preflight.py
├── emit/
│   ├── __init__.py
│   ├── bundle.py
│   ├── diff.py
│   ├── modeling.py
│   ├── parallel.py
│   └── replacement.py
├── plan/
│   ├── __init__.py
│   ├── derive.py
│   ├── freeze.py
│   └── offline_mesh.py
├── source/
│   ├── __init__.py
│   ├── compat.py
│   └── resolver.py
└── spec/
    ├── __init__.py
    ├── project.py
    └── types.py

3.2 非 Codegen 接入文件

以下清单只列直接接入 Codegen 运行链路的非 Codegen 文件。

hyper_parallel/components/modules/
└── grouped_experts.py  # materialize 后需要触发 grouped experts 的 reset，保证 Codegen 下沉替换后的 packed expert 参数初始化顺序与训练期模块语义一致。

hyper_parallel/distributed/_builder/
└── precompiled_boundary.py  # 边界重分发由 generated runtime 显式展开，基础 RedistOp 只处理单个张量，避免 tuple/list 递归入口掩盖边界结构错误。

hyper_parallel/models/_transformers/
├── auto_model.py  # gen 后端从 from_pretrained() 入口准备或复用 artifact，并把 YAML 路径、模型路径和 config_overrides 传给 Codegen。
└── model_builder.py  # gen 后端必须导入 generated modeling 文件完成模型实例化，后续仍复用权重加载、FSDP、AC、compile 等训练基础设施。

hyper_parallel/trainer/
├── base.py  # Trainer 是标准训练链路入口，需要把 codegen、modeling_backend 和完整 TrainerConfig 透传给模型构建层。
└── config/
    ├── manager.py  # 配置解析阶段只记录 YAML 绝对路径，避免 read-only 解析产生 artifact 写入副作用，并让生成目录由模型构建入口统一解析。
    └── parallelism.py  # module replacement 的可序列化来源由 Codegen spec 投影负责，配置层不再给 factory 额外挂运行期私有路径属性。

tests/codegen/
├── helpers.py  # 测试 helper 需要构造与生产一致的 artifact layout，覆盖 `<yaml_dir>/generated` 与 generated modeling 文件命名。
├── test_core_contracts.py  # 覆盖 YAML 路径记录、默认目录、Qwen3-MoE 文件名、signature 投影和 rank 协同等核心契约。
├── test_module_override.py  # 校验 replace_module 下沉到 generated `__init__` 后，meta 记录、模块替换和初始化注入保持一致。
└── test_runtime_boundary_wrap.py  # 边界 wrapper 用例跟随 generated runtime 入口，避免继续依赖旧的 tuple/list RedistOp 递归行为。

3.3 文档文件

docs/guide/
└── codegen.md

---

4. 配置来源与 Signature

Codegen spec 来自 TrainerConfig 投影，不新增独立 codegen 配置文件。核心流程是：

train.yaml
  -> load_training_config()
  -> TrainerConfig
  -> project_codegen_spec(config)
  -> signature_from_spec(spec)

signature 覆盖：

• 模型来源、resolved modeling source sha256、Transformers 版本；
• accelerator.dp/tp/cp/ep/pp_size、sequence_parallel、loss_parallel；
• plan_overrides，包括其中的 replace_module、inner_wrapper、local_compute_fn 和 placement contract；
• model.config_overrides；
• artifact target，包括 artifact dir 和 output modeling name；
• Codegen 实现本身的 sha256。

signature 不覆盖：

• 训练步数、日志、优化器、数据集等运行参数；
• YAML 注释和非 Codegen 投影字段；
• checkpoint transform / weight mapping 的完整语义。该部分保留为 TODO。

load_training_config() 会把 YAML 绝对路径记录到 config._yaml_path。当用户不显式传入 artifact_dir 或
model.codegen_artifact_dir 时，Codegen 通过 _yaml_path 将产物定位到 YAML 同级 generated/。

---

5. 生成流程

生成流程由 hyper_parallel.codegen.manager.ensure_codegen_artifact() 驱动：

1. 判断 codegen 是否开启。
2. 解析 HF config，用于 source resolution 和生成文件命名。
3. 根据 explicit artifact_dir、model.codegen_artifact_dir、yaml_path 或 config._yaml_path 解析 artifact layout。
4. 计算 signature。
5. 如果已有 meta 且 signature 一致，复用现有 bundle。
6. 如果 signature 不一致或 bundle 不存在，rank0 生成新 bundle。
7. 非 rank0 等待 rank0 写入匹配 signature 的 meta。
8. 所有 rank 在 artifact barrier 后继续。

生成期会构造 meta-device 模型，不加载权重。计划推导复用现有 ShardingPlanner：

build_meta_model()
  -> build_offline_mesh()
  -> ShardingPlanner.plan()
  -> freeze_plan()
  -> emit_bundle()
  -> write_bundle_atomic()

write_bundle_atomic() 先把所有文件写入 generated/ 的临时兄弟目录，再整体替换目标目录。这样 reader 要么看到旧的完整
bundle，要么看到新的完整 bundle，不会导入半成品。

---

6. 服务器环境要求

Codegen 以真实训练拓扑为准，手动 generate / check 也会执行并行配置预检。

例如 examples/training_demo/train_codegen_qwen3_moe.yaml 中：

accelerator:
  tp_size: 2
  cp_size: 2
  ep_size: 2

_project_dp_size() 会根据 WORLD_SIZE 和非 DP 并行规模校验配置。单进程执行 CLI 时 WORLD_SIZE=1，会被
预检拒绝。因此这类 YAML 应在有卡服务器上用匹配并行规模的分布式方式运行：

python -m torch.distributed.run --nproc_per_node=4 \
  --module hyper_parallel.codegen --verbose generate \
  --config examples/training_demo/train_codegen_qwen3_moe.yaml

最小单卡调试 YAML 可以单进程运行，但前提是它的并行配置与 WORLD_SIZE=1 匹配。

---

7. 模块替换

模块替换复用 YAML 中的 plan_overrides，其中
replace_module 表示需要在生成模型中下沉的模块替换。

典型配置：

plan_overrides:
  - match: "*.self_attn"
    module_type: transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeAttention
    replace_module:
      _target_: hyper_parallel.models.qwen3_moe.adapter.replacements.replace_qwen3_moe_flash_attention

生成期流程：

plan_overrides
  -> entries_to_module_replacements()
  -> compile_overrides_for_meta()
  -> _apply_module_replacement_actions() on meta model
  -> ShardingPlanner.plan() on replaced meta model
  -> generated __init__ calls hyper_apply_replacements()

这样 planner 看到的模型结构与训练时 generated model 的结构保持一致。若生成期下沉了 replace_module，
codegen_meta.json 中的 covered.module_overrides 会置为 true，训练时 native module replacement 会跳过，避免重复替换。

本阶段重点支持 replace_module。独立顶层 module_overrides、replace_forward、同条 weight_mapping
静态化保留为后续扩展。

---

8. Checkpoint Transform

checkpoint transform / weight mapping 保持 TODO。

设计预留：

• TransformSpec 中保留了 operator_mappings、weights_mapping、checkpoint_transforms 字段；
• project_transform_specs() 只投影 model.config_overrides；
• checkpoint 加载仍主要依赖模型构建后的 runtime mapping；
• Codegen meta 不完整记录 checkpoint.convert_weights_on_load/save 和 weight_mapping 的静态语义。

因此，本阶段只承诺 generated model、diff、sharding plan 和模块替换下沉。checkpoint 权重转换的完整设计后续补齐。

---

9. Backend 选择

backend 是三态：

class ModelingBackend(str, Enum):
    HF = "hf"
    CUSTOM = "custom"
    GEN = "gen"

解析优先级：

force_hf=True       -> hf
modeling_backend    -> 显式 hf / custom / gen
codegen=True        -> gen
custom 可用          -> custom
否则                -> hf

显式或默认进入 custom 后，如果没有找到对应 custom model class，会 warning 并回退到 HF native。

force_hf=True 是最高优先级。即使配置了 codegen: true 或 modeling_backend: gen，只要 force_hf=True，该次
模型构建就会走 HF native。实现会记录 warning，提示 generated modeling file 不会被使用。

---

10. 训练接入点

artifact 准备在模型构建入口中完成。

实际链路：

BaseTrainer._build_model()
  -> self.config.model.build(
       codegen=self.config.codegen,
       modeling_backend=self.config.modeling_backend,
       codegen_config=self.config,
     )
  -> HyperAutoModel.from_pretrained()
  -> resolve_modeling_backend()
  -> if backend is GEN and no codegen_artifact_dir:
       _prepare_codegen_artifact()
         -> ensure_codegen_artifact()
         -> preflight_integrity_check()
  -> _init_model(..., backend=GEN, codegen_artifact_dir=...)
  -> init_generated_model()
  -> apply_model_infrastructure()
  -> parallelize_from_generated()

from_config() 不支持 modeling_backend="gen"。generated model 需要 prepared artifact，因此 gen 后端只支持
from_pretrained() 链路。

---

11. Preflight 策略评估

preflight_integrity_check() 会在调用它的每个 rank 上执行。它检查：

• meta 必需字段；
• signature 是否匹配；
• generated 文件、diff、__init__.py 和 remote siblings 的 hash；
• generated modeling 是否可 import；
• frozen param plan 是否为空；
• 有模型实例时，param plan 是否覆盖模型参数。

各 rank 都执行 preflight 的原因是：

1. 每个 rank 后续都会导入 generated model。各 rank 独立检查可以确认 artifact 在本 rank 上真实可见、可读、可 import。
2. 多机或容器场景下，rank0 看得到文件不代表其他 rank 一定看得到同一路径。all-rank preflight 可以更早暴露同步或挂载问题。
3. preflight 主要是文件和 import 检查，不涉及大规模通信；相对模型训练启动成本，重复执行的代价可接受。
4. rank0-only 需要额外做错误广播和同步协议。没有这层协议时，rank0-only 反而可能让非 rank0 在后续 import 阶段才失败。

因此，设计采用：

rank0 负责生成或复用 bundle
all ranks 执行 preflight
all ranks 导入同一份 generated model

如果 preflight 成本明显上升，可扩展为 rank0 执行重检查、其他 rank 做轻量可见性检查，但必须补齐错误广播和
barrier 语义后才能切换。

---

12. Artifact 消费

generated model 被 hyper_parallel.codegen.loader 作为 synthetic package 导入。loader 不把 generated/ 加到
sys.path，而是根据 artifact 绝对路径构造稳定模块名，避免不同 artifact 之间的 sys.modules 冲突。

训练时，apply_model_infrastructure() 通过 meta 判断 generated artifact 是否覆盖 sharding plan：

load_codegen_meta(artifact_dir)
  -> covered.sharding_plan == true
  -> verify_codegen_signature()
  -> parallelize_from_generated()
  -> generated_module.hyper_parallelize(model, mesh_context)

如果 meta 不存在，或 covered.sharding_plan 不是 true，训练会回到常规 planner/applier 路径。

---

13. CLI

CLI 入口：

python -m hyper_parallel.codegen

常用命令：

python -m torch.distributed.run --nproc_per_node=4 \
  --module hyper_parallel.codegen --verbose generate \
  --config examples/training_demo/train_codegen_qwen3_moe.yaml

python -m torch.distributed.run --nproc_per_node=4 \
  --module hyper_parallel.codegen --verbose check \
  --config examples/training_demo/train_codegen_qwen3_moe.yaml

python -m hyper_parallel.codegen --verbose clean \
  --config examples/training_demo/train_codegen_qwen3_moe.yaml

generate 和 check 会加载训练 YAML，因此需要满足该 YAML 的并行配置要求。clean 只解析 artifact layout 并删除
bundle，不执行并行配置预检，可以单进程运行。

---

14. 测试覆盖

tests/codegen/ 覆盖：

• signature 稳定性；
• artifact layout、YAML path、文件命名等核心契约；
• generated import 与 remote sibling import；
• boundary lowering；
• module override 下沉；
• runtime boundary wrapper；
• tp_grad_info 契约。

后续补齐：

• checkpoint transform / weight mapping 静态化；
• 独立顶层 module_overrides 配置入口；
• replace_forward 下沉；
• rank0-only preflight 的广播协议。如果切换到 rank0-only，需要新增专门测试。

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