跳到主要内容

数据截至 (上游 commit b6c0bfe04c82)

01 · 模型文件解剖:单文件自读与 modular 生成

这一章讲什么: transformers 怎么组织 507 个架构的代码——为什么每个模型文件都刻意重复、一个标准 modeling 文件内部长什么样、以及 v5 的 modular 系统怎么用「Python 继承当 diff 语法」同时保住可读性和可维护性。读完你可以自己读懂(或新增)任何一个架构。

1. 它要解决的小问题​

一个模型库要同时满足两个互相打架的需求:

  • 读者需求: 我想搞懂 Llama 怎么实现,最好打开一个文件就能看到全部,不用在几十个抽象基类之间跳。
  • 维护者需求: 500 多个架构里 RMSNorm、RoPE、attention 都大同小异,每个文件抄一遍,修 bug 要修几百处。

transformers 的答案是先无条件满足读者,再用代码生成机制帮维护者消重复。这一章就是这个答案的两半。

2. 思路:目录即架构,文件即教科书​

组织规则一句话:src/transformers/models/ 下一个目录 = 一个架构。本 commit 共 509 个子目录,除去 auto/(Auto 体系)与 deprecated/ 两个功能目录,即 507 个架构目录。以 llama 为例,目录里只有四类文件:

文件角色
configuration_llama.pyLlamaConfig:全部超参(hidden size、层数、RoPE 参数……)
modeling_llama.py模型本体,512 行,本目录的灵魂
tokenization_llama.py该架构专属的分词器(很多架构没有,直接用通用的)
convert_llama_weights_to_hf.py把官方原始 checkpoint 转成 HF 格式的一次性脚本

注意没有 attention.py/layers.py 之类的拆分——一切都在 modeling_llama.py 一个文件里。这就是「single file policy」:文件内部的类按从零件到整机的顺序排列,从 LlamaRMSNorm 一路排到 LlamaForCausalLM,读者从头往下读就是一堂课。

3. 一个 modeling 文件的内部结构​

modeling_llama.py 的类清单(grep "^class " 的结果)就是这个架构的解剖图:

零件层   LlamaRMSNorm            (53)   归一化
         LlamaRotaryEmbedding    (73)   预计算 RoPE 的 cos/sin
         LlamaMLP                (163)  SwiGLU 前馈
         eager_attention_forward (191)  手写版注意力(教学参考实现)
         LlamaAttention          (217)  QKV 投影 + RoPE + 调注意力函数
         LlamaDecoderLayer       (284)  = Attention + MLP + 两个残差
整机层   LlamaPreTrainedModel    (328)  基类:能力旗标集中地
         LlamaModel              (347)  embed → N 层 → final norm
         LlamaForCausalLM        (421)  LlamaModel + lm_head + loss
任务头   LlamaForSequenceClassification (495) 等 3 个

(行号均相对 src/transformers/models/llama/modeling_llama.py。)

三个观察:

  1. LlamaDecoderLayer 是教科书写法。 经典 pre-norm 结构——先 norm 再 attention、残差相加、再 norm 再 MLP、再残差相加(modeling_llama.py:284-324)。没有魔术,任何一个 decoder-only 架构的 DecoderLayer 都是这个形状。
  2. 任务头是通用 mixin。 LlamaForSequenceClassification 只是一行 class LlamaForSequenceClassification(GenericForSequenceClassification, LlamaPreTrainedModel): ...(modeling_llama.py:495);真正的分类逻辑(池化、算 loss)在通用件 GenericForSequenceClassification(modeling_layers.py:119)里,同文件还有 GenericForQuestionAnswering(:193)、GenericForTokenClassification(:254)。任务头不再逐模型重写。
  3. 每个 DecoderLayer 继承 GradientCheckpointingLayer(modeling_layers.py:52)。它重写 __call__:开启梯度检查点且处于 training 时,强制把 use_cache=False、past_key_values=None 再交给 self._gradient_checkpointing_func 执行(:77-108)——因为梯度检查点会在 backward 时重放 forward,若不掐掉 cache,KV 会被写两遍。这是「踩过坑才有」的防护。

3.1 基类上的「能力旗标」​

LlamaPreTrainedModel(modeling_llama.py:328-344)几乎不含逻辑,全是声明:

# 摘自 modeling_llama.py:329-344(已精简)
class LlamaPreTrainedModel(PreTrainedModel):
    config: LlamaConfig
    base_model_prefix = "model"              # 整机的属性名,加载/绑权重靠它定位
    _no_split_modules = ["LlamaDecoderLayer"]  # device_map 切分时不许拆散的最小单元
    _supports_flash_attn = True
    _supports_sdpa = True
    _can_compile_fullgraph = True            # 可整图 torch.compile
    _can_record_outputs = {"hidden_states": LlamaDecoderLayer, "attentions": LlamaAttention}

这些旗标是模型与框架之间的能力协议:from_pretrained 读它们决定能不能用某个 attention 后端、怎么按层切到多卡;输出录制读 _can_record_outputs 决定在哪类模块上挂 hook。新增一个架构,很大程度上就是把这些旗标如实填上。

4. 注意力:一个注册表换掉所有 if-else​

4.1 小问题​

同一个架构,有人要 eager(可读)、有人要 SDPA(PyTorch 内置融合算子)、有人要 FlashAttention(省显存)、有人要 flex。在每个模型里写 if attn_implementation == "flash_attention_2": ... 会是一场灾难。

4.2 思路与实现​

把注意力核心(Q·K^T → mask → softmax → ·V)抽成一个统一签名的函数,全局注册表按名字查:

config._attn_implementation = "flash_attention_2"
                │
                ▼
ALL_ATTENTION_FUNCTIONS.get_interface(name, eager_default)
   (modeling_utils.py:5161)
                │
                ▼
   返回该名字的函数 → 模型用统一签名调用

全局注册表是 ALL_ATTENTION_FUNCTIONS: AttentionInterface = AttentionInterface()(modeling_utils.py:5185),其内建映射在 AttentionInterface._global_mapping(modeling_utils.py:5156-5168):flash_attention_4/3/2 → flash_attention_forward、sdpa → sdpa_attention_forward、flex_attention → flex_attention_forward,还有 paged|* 前缀的一族供连续批处理用。

模型侧的全部「分流」就是两行(modeling_llama.py:264-266):

attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
    self.config._attn_implementation, eager_attention_forward
)

默认值是文件里的 eager_attention_forward(modeling_llama.py:191-214)——一份 20 行、无任何优化的手写注意力,既是 fallback 也是教学参考:想知道注意力在算什么,读它就够了。

妙在哪: 注册表是 GeneralInterface(utils/generic.py:1097)的实例,支持「全局映射 + 文件局部覆盖」两级——某个模型需要特制 sdpa 行为时,在自己的 modeling 文件里 new 一个 AttentionInterface 局部注册即可,不污染全局(设计意图写在 modeling_utils.py:5147-5155 的 docstring)。新架构免费获得全部后端;新后端免费覆盖全部架构。这就是 507 个文件还能同步演进的关键基础设施。

5. modular 系统:用 Python 继承当 diff 语法​

5.1 小问题​

单文件政策带来巨量重复。Qwen3 与 Llama 的差异其实只有:RMSNorm 细节、attention 里多了 q_norm/k_norm、可选 sliding window。其余几百行一模一样。怎么让「差异」成为源码,而「全量」成为产物?

5.2 思路​

作者写继承,diff 即源码;CI 做展开,全量即产物。

作者维护的是 modular_qwen3.py,里面就是普通的 Python 继承(modular_qwen3.py:48-60):

class Qwen3RMSNorm(Qwen2RMSNorm):
    pass

class Qwen3MLP(GemmaMLP):
    pass

class Qwen3Attention(LlamaAttention):
    def __init__(self, config: Qwen3Config, layer_idx: int):
        ...
        super().__init__(config, layer_idx)
        self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)

三个 pass 类意味着「这几件和父类一模一样,只换个名字」;Qwen3Attention 覆盖了 __init__ 和 forward,把真正的差异(q_norm/k_norm)写出来。

然后工具 utils/modular_model_converter.py(仓库根目录的 utils/,注意不是 src/transformers/utils/)用 libcst(具体语法树) 做依赖追踪(find_all_dependencies,:413)与类体展开,把继承「内联」回一份自包含的 modeling_qwen3.py。产物文件头第 2 行写明:

This file was automatically generated from src/transformers/models/qwen3/modular_qwen3.py.

5.3 规模与约束​

  • 本 commit 有 287 个 modular_*.py;而 llama 目录里没有 modular 文件——Llama 是「源架构」之一,是被继承的全量真身。继承可以跨多层:Qwen3 同时从 llama、qwen2、gemma 三个架构 import 父类(modular_qwen3.py:26-39)。
  • 纪律:改 modular,不改 modeling。 产物文件头注释明说 "If any change should be done, please apply the change to the modular_qwen3.py file directly. One of our CI enforces this."(modeling_qwen3.py:4-5)——CI 会重新展开并 diff,手改产物会被打回。
作者写 modular_qwen3.py ──(继承,diff 即源码)──►  reviewer 读差异
        │
        ▼ utils/modular_model_converter.py(libcst 展开)
读者读 modeling_qwen3.py ──(自包含,无跳转)──►  单文件政策不破

6. 关键细节与坑​

  • config._attn_implementation 为 None 时会 warning。 把 LlamaAttention 当普通 nn.Module 单独拿去用(没经过 from_pretrained)时,get_interface 会提示这是 standalone 用法、回退 eager(modeling_utils.py:5170-5176)。不是 bug,是提示你走了非典型路径。
  • 「读一个文件就够」对老模型成立,对 modular 产物模型要心里有数。 modeling_qwen3.py 是自包含的,但若想改它的行为,正确入口是 modular_qwen3.py;只改 modeling 会被 CI 还原。
  • _no_split_modules 填错 = device_map 把层拆散。 这个旗标(modeling_llama.py:332)告诉多卡切分「DecoderLayer 是不可拆的最小单元」;新架构忘了填,device_map="auto" 可能把一层切到两张卡上,正确性还在但性能崩掉。
  • eager 实现是教学锚点,不是性能路径。 看到 eager_attention_forward 里 torch.matmul + 显式 softmax 不要以为是库的性能水平——生产路径在 sdpa/flash 后端里,eager 的存在价值是可读。

7. 代码地图​

主题文件路径符号名
教科书式架构实现src/transformers/models/llama/modeling_llama.pyLlamaRMSNorm、LlamaRotaryEmbedding、LlamaMLP、LlamaAttention、LlamaDecoderLayer、LlamaModel、LlamaForCausalLM
手写注意力(教学参考)src/transformers/models/llama/modeling_llama.pyeager_attention_forward、repeat_kv、apply_rotary_pos_emb
能力旗标集中地src/transformers/models/llama/modeling_llama.pyLlamaPreTrainedModel(base_model_prefix、_no_split_modules、_can_record_outputs)
通用任务头src/transformers/modeling_layers.pyGenericForSequenceClassification、GenericForQuestionAnswering、GenericForTokenClassification
梯度检查点防护src/transformers/modeling_layers.pyGradientCheckpointingLayer.__call__
注意力后端注册表src/transformers/modeling_utils.pyAttentionInterface、ALL_ATTENTION_FUNCTIONS、get_interface
注册表两级机制src/transformers/utils/generic.pyGeneralInterface
modular 示例(继承式 diff)src/transformers/models/qwen3/modular_qwen3.pyQwen3Attention、Qwen3RMSNorm
modular 展开工具utils/modular_model_converter.py(仓库根)find_all_dependencies、ClassDependencyMapper