model.py — 模型配置

文件概述

定义 HFModelConfig，是整个框架中最基础的配置类之一，负责加载 HuggingFace 模型的配置、tokenizer 和 generation_config。

核心配置类

HFModelConfig

@dataclass
class HFModelConfig:
    path: str = ""                          # 模型路径
    local_path: str = ""                    # 本地路径
    trust_remote_code: bool = False         # 是否信任远程代码
    use_remove_padding: bool = True         # 是否使用 remove padding 优化
    enable_gradient_checkpointing: bool = True  # 梯度检查点
    use_liger: bool = False                 # 是否使用 Liger 内核
    use_fused_kernels: bool = False         # 是否使用融合内核

    # LoRA 配置
    lora_rank: int = 0                      # LoRA 秩（0=不使用）
    lora_alpha: int = 16                    # LoRA alpha
    target_modules: list = None             # LoRA 目标模块

    # 自动加载的属性
    hf_config: AutoConfig = None            # HuggingFace 模型配置
    tokenizer: PreTrainedTokenizer = None   # 分词器
    generation_config: GenerationConfig = None  # 生成配置

MtpConfig

Multi-Token Prediction (MTP) 配置：

@dataclass
class MtpConfig:
    enable: bool = False
    num_speculative_tokens: int = 1
    method: str = "eagle"

关键逻辑

模型配置加载

def __post_init__(self):
    # 自动从模型路径加载 HuggingFace 配置
    self.hf_config = AutoConfig.from_pretrained(self.path)
    # 加载 tokenizer
    self.tokenizer = AutoTokenizer.from_pretrained(self.path)
    # 加载 generation_config
    self.generation_config = GenerationConfig.from_pretrained(self.path)

与其他模块的关系

• 被所有 Worker 和 Engine 使用来获取模型信息
• hf_config 被传递给 HuggingFace / Megatron / TorchTitan 用于模型构建
• tokenizer 被 Rollout 和 Reward Manager 用于文本编解码

小结

HFModelConfig 是框架与 HuggingFace 生态对接的桥梁，自动加载模型的所有必要信息。