已开启
【RFC】hyper-parallel 支持torchtitan 风格 Module/Config —— Part 2 配置系统 #140
changzherui创建于  5月15日
changzherui
changzherui成员
5月15日 创建

M2 — hyper_parallel/config/ 配置系统

提供 torchtitan 风格的 ConfigManager + tyro CLI + 通用 *Config dataclass;与现存 trainer/config.py:HyperTrainerConfig 字段建立双向映射新 / 旧两条入口并存,旧 yaml 路径 0 改动。


1. 目标

  1. 提供新的 CLI 入口语法(--module qwen3_5_v2 --config qwen3_5_v2_4b --training.steps=10)。
  2. 提供 torchtitan 风格的通用配置 dataclass(TrainingConfig / ParallelismConfig / ...)。
  3. 通过 _legacy_to_trainer_config 把旧 yaml 字段映射到新结构,老用户脚本不破坏

2. 任务边界(新增文件)

新增包目录 hyper_parallel/config/

新文件 内容 对应 torchtitan
__init__.py 导出 ConfigManager / TrainingConfig / ParallelismConfig / ActivationCheckpointConfig / CompileConfig / CommConfig / DebugConfig / Function / TORCH_DTYPE_MAP;re-export Configurable
configs.py 通用训练 dataclass(见 §3 字段对照表);全部 @dataclass(kw_only=True, slots=True) torchtitan/config/configs.py
function.py Function(Generic[R], Configurable):把任意 callable 包成 Configurable,Function.Config(fn=callable)build() 返回 fn 本身 torchtitan/config/function.py
manager.py ConfigManagerparse_args(argv=None) -> Configurable.Config。① 若 argv[0] 起始为 --module → 走 tyro 路径;② 否则走旧 yaml 路径并经 _legacy_to_trainer_config 收口 torchtitan/config/manager.py
dtype_map.py TORCH_DTYPE_MAP: dict[str, dtype],跨后端通过 platform.dtype_map 取值,避免直接 import torch torchtitan/config/dtype_map.py
tyro_rules.py 注册 list[str] 逗号分隔等 tyro custom rules;首次 tyro.cli 调用前由 manager.py 自动 install torchtitan/config/tyro_rules.py
legacy_adapter.py _legacy_to_trainer_config(legacy: HyperTrainerConfig) -> HyperTrainer.Config,把旧 yaml 字段映射到新结构

3. 字段对照表(旧 yaml ↔ 新 *Config

旧字段来自 hyper_parallel/trainer/config.py,行号已标注。

旧路径(hyper) 旧定义点 新路径(torchtitan 风格)
train.max_steps config.py:295 training.steps
train.global_batch_size config.py:297 training.global_batch_size
train.micro_batch_size config.py:298 training.local_batch_size
train.seed config.py:299 debug.seed
train.init_device config.py:303 training.init_device
train.comm_backend config.py:304 comm.backend
train.accelerator.dp_shard config.py:149 parallelism.data_parallel_shard_degree
train.accelerator.dp_replicate config.py:148 parallelism.data_parallel_replicate_degree
train.accelerator.tp config.py:150 parallelism.tensor_parallel_degree
train.accelerator.cp config.py:151 parallelism.context_parallel_degree
train.accelerator.pp config.py:152 parallelism.pipeline_parallel_degree
train.accelerator.ep config.py:153 parallelism.expert_parallel_degree
train.accelerator.etp config.py:154 parallelism.expert_tensor_parallel_degree
train.accelerator.reshard_after_forward config.py:156 parallelism.reshard_after_forward
train.mixed_precision.enabled / param_dtype / reduce_dtype config.py:172-175 training.mixed_precision_param / training.mixed_precision_reduce
train.gradient_checkpointing.activation_checkpoint config.py:183 activation_checkpoint.mode
train.optimizer.lr / weight_decay / eps / betas / lr_warmup_ratio / lr_decay_style / max_grad_norm config.py:195-205 optimizer.lr / wd / eps / betas + lr_scheduler.warmup_steps / decay_style / max_grad_norm
train.checkpoint.* config.py:208-214 checkpoint.*
train.profile.* config.py:242-250 profiler.*
train.debug.* config.py:271-281 debug.*
model.weights_path / tokenizer_path / freeze_modules / config_overrides config.py:67-93 model_spec.model.<field> + tokenizer.path + training.freeze_modules
data.type / train_path / max_seq_len / text_key / num_workers / shuffle config.py:99-126 dataloader.dataset / dataloader.dataset_path / training.seq_len / dataloader.text_key / dataloader.num_workers / dataloader.shuffle

转换函数 _legacy_to_trainer_config 必须双向 round-trip 等价(即 legacy → new → legacy 字段值不变,由 test_legacy_roundtrip.py 保障)。

4. 核心设计:ConfigManager.parse_args 双入口

class ConfigManager:
    def parse_args(self, argv: list[str] | None = None) -> "Configurable.Config":
        argv = argv if argv is not None else sys.argv[1:]

        # 1. 新风格:--module X --config Y [--field=value]
        if argv and argv[0].startswith("--module"):
            return self._parse_new(argv)

        # 2. 旧风格:yaml + dot-path(兼容现有 train_lm.py)
        return self._parse_legacy(argv)

    def _parse_new(self, argv):
        try:
            import tyro
        except ImportError as exc:
            raise ImportError(
                "New CLI requires `tyro`. Install with: pip install tyro"
            ) from exc
        install_tyro_rules()

        # 解析 --module / --config,从 config_registry 取 recipe
        head, remaining = _extract_module_config(argv)
        module_name, config_name = head["module"], head["config"]
        registry = importlib.import_module(
            f"hyper_parallel.models.{module_name}.config_registry"
        )
        default = getattr(registry, config_name)()    # HyperTrainer.Config 实例

        from hyper_parallel.trainer.trainer_config import HyperTrainer
        return tyro.cli(HyperTrainer.Config, args=remaining, default=default)

    def _parse_legacy(self, argv):
        from hyper_parallel.trainer.config import parse_args, HyperTrainerConfig
        from hyper_parallel.config.legacy_adapter import _legacy_to_trainer_config
        legacy = parse_args(HyperTrainerConfig)        # 现有 :547 逻辑
        return _legacy_to_trainer_config(legacy)

5. 与 torchtitan 接口差异说明

# 差异点 原因
1 dtype 字段类型用 Literal["bfloat16", "float32", "float16"] 而非 torch.dtype tyro 不能在 mindspore 环境序列化 torch.dtype;运行时 TORCH_DTYPE_MAP[s] 解析
2 ConfigManager.parse_args 同时支持新 / 旧两条入口 零破坏迁移
3 tyro 列为可选依赖(requirements.txttyro>=0.9.0; extra == "config-cli" hyper 默认安装应保持轻量;新 CLI 用到时才 import
4 local_rank 字段保留 hyper 旧语义(从 LOCAL_RANK 环境变量取),新版放 comm.local_rank torchtitan 没有这个字段;torchrun 现网都依赖环境变量
5 ModelSpec 字段在 HyperTrainer.Config 中标 tyro.conf.Suppress ModelSpec 持 callable / dataclass,tyro 不能从 CLI 解析

6. 开发步骤

  1. Step 1(1 d)configs.py —— 7 个 dataclass,按字段对照表逐项落地,写 __post_init__ 校验(如 local_batch_size > 0tensor_parallel_degree >= 1)。
  2. Step 2(0.5 d)function.py + dtype_map.py
  3. Step 3(1 d)manager.py 的双入口 + tyro_rules.py
  4. Step 4(0.5 d)legacy_adapter.py

7. 验证标准

新建 tests/torch/ut/config/

测试 断言要点
test_configs.py 默认值与 examples/yaml/train_qwen3_5_*.yaml 对齐;非法值(如 tensor_parallel_degree=-2__post_init__ raise
test_function.py Function.Config(fn=lambda x: x+1).build()(3) == 4
test_manager_legacy_yaml.py 同一 yaml 跑两次:①parse_args(HyperTrainerConfig),②ConfigManager().parse_args([yaml_path]);后者过 _legacy_to_trainer_config 后字段对齐
test_manager_new_cli.py tests/torch/ut/config/_test_dummy/config_registry.pyqwen3_5_v2_tiny() -> HyperTrainer.Config;调 ConfigManager().parse_args(["--module", "_test_dummy", "--config", "qwen3_5_v2_tiny", "--training.steps=5"]),断言覆盖生效
test_legacy_roundtrip.py legacy → new → legacy 后字段值与原对象 assertdataclass_equal
test_tyro_missing.py mock 卸 tyro,新路径必须 raise 含 "pip install tyro" 的 ImportError

通过门槛

  • 6 个 UT 全绿。
  • 旧路径 parse_args(HyperTrainerConfig)trainer/config.py:547)行为 0 改动
  • hyper_parallel.config 可独立 import;hyper_parallel/__init__.py 不导出。

8. 工期 & 依赖

工期 3 天
依赖 M1(Configurable
并行 与 M3 / M4 完全并行
下游 M5 / M7
likedislike
changzheruichangzherui成员
5月15日 修改标题为 “hyper-parallel 支持torchtitan 风格 Module/Config —— M2 配置系统”,原标题为“hyper-parallel 支持torchtitan 风格 Module/Config —— M2”
changzheruichangzherui成员
5月15日 修改标题为 “【RFC】hyper-parallel 支持torchtitan 风格 Module/Config —— M2 配置系统”,原标题为“hyper-parallel 支持torchtitan 风格 Module/Config —— M2 配置系统”
changzheruichangzherui成员
5月25日 修改标题为 “【RFC】hyper-parallel 支持torchtitan 风格 Module/Config —— Part 2 配置系统”,原标题为“【RFC】hyper-parallel 支持torchtitan 风格 Module/Config —— M2 配置系统”