Part 5 — trainer/ 接缝层

在不动 LLMTrainer / VLTrainer 的前提下，给 BaseTrainer 加一个"新 Config 入口"分支，以及 HyperTrainer.Config 顶层 Configurable 树。新 / 旧两条路径共享同一份算法实现（通过抽公共 helper）。

1. 目标

把 M2 + M3 + M1 拼成完整的"配置 → 训练器"链路：

HyperTrainer.Config（Configurable 树）
   .build()
      └─ HyperTrainer.__init__(config)
            ├─ init_distributed + ParallelDims
            ├─ tokenizer / dataloader / model / parallelize / optimizer / lr / ckpt / metrics / profiler
            └─ trainer.train()

2. 任务边界

3. HyperTrainer.Config 结构

from dataclasses import dataclass, field
from typing import Annotated

import tyro

from hyper_parallel.protocols import Configurable, ModelSpec
from hyper_parallel.config import (
    TrainingConfig, ParallelismConfig, ActivationCheckpointConfig,
    CompileConfig, CommConfig, DebugConfig,
)
from hyper_parallel.components.optimizer import OptimizersContainer
from hyper_parallel.components.lr_scheduler import LRSchedulersContainer
from hyper_parallel.components.loss import CrossEntropyLoss, BaseLoss
from hyper_parallel.components.tokenizer import HuggingFaceTokenizer, BaseTokenizer
from hyper_parallel.components.dataloader import BaseDataLoader
from hyper_parallel.components.checkpoint import CheckpointManager
from hyper_parallel.components.profiler import Profiler
from hyper_parallel.components.metrics import MetricsProcessor


class HyperTrainer(Configurable):

    @dataclass(kw_only=True, slots=True)
    class Config(Configurable.Config):
        # ModelSpec 持 callable / dataclass，tyro 无法解析，对 CLI 隐藏
        model_spec: Annotated[ModelSpec | None, tyro.conf.Suppress] = None

        hf_assets_path: str = "outputs/hf_assets"
        dump_folder: str = "outputs"

        tokenizer:             BaseTokenizer.Config         = field(default_factory=HuggingFaceTokenizer.Config)
        dataloader:            BaseDataLoader.Config        = field(default_factory=BaseDataLoader.Config)
        optimizer:             OptimizersContainer.Config   = field(default_factory=OptimizersContainer.Config)
        lr_scheduler:          LRSchedulersContainer.Config = field(default_factory=LRSchedulersContainer.Config)
        loss:                  BaseLoss.Config              = field(default_factory=CrossEntropyLoss.Config)
        checkpoint:            CheckpointManager.Config     = field(default_factory=CheckpointManager.Config)
        metrics:               MetricsProcessor.Config      = field(default_factory=MetricsProcessor.Config)
        profiler:              Profiler.Config              = field(default_factory=Profiler.Config)

        training:              TrainingConfig               = field(default_factory=TrainingConfig)
        parallelism:           ParallelismConfig            = field(default_factory=ParallelismConfig)
        activation_checkpoint: ActivationCheckpointConfig   = field(default_factory=ActivationCheckpointConfig)
        compile:               CompileConfig                = field(default_factory=CompileConfig)
        comm:                  CommConfig                   = field(default_factory=CommConfig)
        debug:                 DebugConfig                  = field(default_factory=DebugConfig)

4. HyperTrainer.__init__ 11 步

对应 torchtitan Trainer.__init__：

class HyperTrainer(BaseTrainer):
    def __init__(self, config: "HyperTrainer.Config"):
        self.config = config

        # 1. init_process_group + ParallelDims（复用 helper）
        self._helper_setup_distributed(config)

        # 2. tokenizer
        self.tokenizer = config.tokenizer.build(path=config.tokenizer.path)

        # 3. dataloader
        self.dataloader = config.dataloader.build(
            dp_world_size=self.parallel_dims.dp_world_size,
            dp_rank=self._dp_rank(),
            tokenizer=self.tokenizer,
            seq_len=config.training.seq_len,
        )

        # 4. model config 接入运行时
        model_config = config.model_spec.model
        model_config.update_from_config(trainer_config=config)
        # ←—— update_from_config 内部调 set_<name>_sharding_config()

        # 5. 构造模型（meta device）
        with init_empty_weights():
            model = model_config.build()

        # 6. 协议自检
        model.verify_module_protocol()

        # 7. 并行化（CP / TP / AC / FSDP）
        model = config.model_spec.parallelize_fn(
            model, self.parallel_dims,
            training=config.training,
            parallelism=config.parallelism,
            activation_checkpoint=config.activation_checkpoint,
            compile=config.compile,
            dump_folder=config.dump_folder,
        )

        # 8. 物化 + 权重初始化
        model.to_empty(device=platform.device_type())
        model.init_weights(buffer_device=platform.device_type())
        self.model = model

        # 9. optimizer / lr_scheduler
        self.optimizer = config.optimizer.build(model_parts=[model])
        self.lr_scheduler = config.lr_scheduler.build(
            optimizers=self.optimizer,
            training_steps=config.training.steps,
        )

        # 10. checkpoint
        sd_adapter = None
        if config.model_spec.state_dict_adapter is not None:
            sd_adapter = config.model_spec.state_dict_adapter(
                model_config, config.hf_assets_path,
            )
        self.checkpointer = config.checkpoint.build(
            model_parts=[model],
            optimizers=self.optimizer,
            lr_schedulers=self.lr_scheduler,
            dataloader=self.dataloader,
            sd_adapter=sd_adapter,
        )

        # 11. metrics / profiler
        self.metrics = config.metrics.build(trainer=self)
        self.profiler = config.profiler.build()

5. BaseTrainer 拆分要点

目标：算法零分叉，新 / 旧两条 __init__ 都调用同一份 helper。

旧 LLMTrainer(args) 路径保持完全不变（llm_trainer.py:47-66 不动）。

6. 与 torchtitan 接口差异说明

7. 开发步骤

1. Step 1（0.5 d）：把 BaseTrainer.__init__ 中可复用的初始化逻辑抽 helper（_helper_setup_distributed / _helper_post_parallelize / _helper_load_weights）。旧调用点保持调 helper 包装方法，行为 0 变化。
2. Step 2（0.5 d）：写 trainer_config.py 的 HyperTrainer.Config，引用 M2 / M3 的所有 Config。
3. Step 3（1 d）：写 hyper_trainer.py 的 HyperTrainer.__init__ 11 步。
4. Step 4（0.5 d）：升级 protocols/model_spec.py 到 v2 字段集（M1 已基本就位，M5 微调和测试）。

8. 验证标准

新建 tests/torch/ut/trainer/：

通过门槛：

• 3 个测试全绿。
• 旧 LLMTrainer(args) 路径 0 行为改动。
• scripts/train_lm.py 不动（M7 才动）。

9. 工期 & 依赖