engine.py — 训练引擎配置

文件概述

定义统一训练引擎的配置类，供 engine_workers.py 使用。

核心配置类

EngineConfig (基类)

@dataclass
class EngineConfig:
    model_type: str = "language_model"    # language_model / value_model
    backend: str = "fsdp"                 # fsdp / megatron / torchtitan / veomni
    device: str = "cuda"                  # cuda / npu
    forward_only: bool = False            # 是否仅前向（Ref 模型用）
    param_offload: bool = False           # 参数卸载到 CPU
    optimizer_offload: bool = False       # 优化器卸载到 CPU

FSDPEngineConfig

@dataclass
class FSDPEngineConfig(EngineConfig):
    fsdp_size: int = -1                       # FSDP 分组大小
    ulysses_sequence_parallel_size: int = 1   # Ulysses SP 大小
    use_torch_compile: bool = False           # 是否用 torch.compile
    mixed_precision: str = "bf16"             # 混合精度
    entropy_from_logits_with_chunking: bool = False  # 分块计算 entropy

McoreEngineConfig

@dataclass
class McoreEngineConfig(EngineConfig):
    tensor_model_parallel_size: int = 1
    pipeline_model_parallel_size: int = 1
    expert_model_parallel_size: int = 1
    context_parallel_size: int = 1
    use_distributed_optimizer: bool = True
    use_fused_kernels: bool = False
    router_replay: RouterReplayConfig = ...  # 路由重放配置

TorchtitanEngineConfig

@dataclass
class TorchtitanEngineConfig(EngineConfig):
    data_parallel_shard_size: int = -1
    data_parallel_replicate_size: int = 1
    tensor_parallel_size: int = 1
    pipeline_parallel_size: int = 1
    expert_parallel_size: int = 1

VeOmniEngineConfig

TrainingWorkerConfig

组合配置，将 engine + model + optimizer + checkpoint 打包：

@dataclass
class TrainingWorkerConfig:
    engine: EngineConfig
    model: HFModelConfig
    optimizer: OptimizerConfig
    checkpoint: CheckpointConfig

与其他模块的关系

• 被 engine_workers.py 的 TrainingWorker 使用
• backend 字段决定使用哪个 Engine 实现
• 通过 EngineRegistry 匹配到具体的 Engine 类

小结

Engine 配置是统一引擎架构的关键，通过 backend + model_type + device 三元组唯一确定要使用的引擎。