动态权重加载
检查点通常以与模型在运行时预期格式不匹配的格式进行序列化。常见场景包括:
- 融合权重:检查点中存储的是独立的
gate_proj和up_proj权重,但模型为了提高效率,使用的是融合后的gate_up_proj。 - MoE 专家整合:各个专家的权重(
experts.0.weight、experts.1.weight、……)需要堆叠成一个单一的 3D 张量。 - 旧版命名:旧检查点使用不同的命名约定(例如
LayerNorm.gamma与LayerNorm.weight)。 - 复合模型:视觉-语言模型包含两个
PreTrainedModel子模块,每个子模块都有各自的检查点约定。 - 量化:权重可能以量化格式存储,需要进行反序列化。
动态权重加载通过在加载检查点张量时应用计划好的、可逆的操作来解决此问题。Transformers 通过 WeightConverter 和 WeightRenaming 来暴露此功能,它们描述了一个或多个检查点键如何映射到一个或多个模型参数,以及应在匹配的张量上运行哪些可组合的 ConversionOps。这种方法能够适应新的权重布局,支持量化混合专家(MoE),并与张量并行集成。
本指南演示如何使用 WeightConverter 来转换张量。转换映射位于 conversion_mapping.py 中;已注册的映射以 model_type 字符串(例如 "mixtral")或类名(例如 "LlavaModel")作为键。
完整加载流程
所有模型都会经过动态权重加载系统。转换映射是该系统中的一个可选步骤,仅当模型为其类名或 model_type 注册了条目时才会激活。
Checkpoint File → from_pretrained() → convert_and_load_state_dict_in_model()
↓
┌───────────────────────────────────────────────────────────┐
│ For each weight in checkpoint: │
│ 1. Match renamed/processed source key to model parameter │
│ 2. Shard the weight and send to device (async) │
│ 3. Collect tensors with the same source_pattern together │
│ (e.g. MoE experts, gate_up_proj) │
│ 4. Apply dequantization/deserialization (if pre-quant) │
│ 5. Apply conversion (if defined) │
│ 6. Apply quantization (if enabled and step 4 not used) │
│ 7. Set parameter on model │
└───────────────────────────────────────────────────────────┘| 步骤 | 当它激活时 |
|---|---|
| 动态加载 | 始终适用于所有模型 |
| 转换映射 | 仅当模型的类或 model_type 在 _MODEL_TO_CONVERSION_PATTERN 中注册时 |
| TP 分片 | 仅当使用 DistributedConfig(tp_size=N) 且模型具有 base_model_tp_plan 时 |
| 反量化/反序列化 | 仅在加载预量化检查点时 |
| 量化 | 仅在提供了量化配置且权重未预先量化的情况下 |
稠密模型(例如 Llama)
对于大多数稠密模型,检查点格式与模型格式直接对应,因此无需进行转换映射。某些模型可能仍需重命名(例如,沿用旧的命名约定)。启用时,TP 分片规则仍然适用。
Checkpoint: Model:
model.layers.0.self_attn.q_proj.weight → model.layers.0.self_attn.q_proj.weight
model.layers.0.self_attn.k_proj.weight → model.layers.0.self_attn.k_proj.weight
model.layers.0.mlp.gate_proj.weight → model.layers.0.mlp.gate_proj.weight
model.layers.0.mlp.up_proj.weight → model.layers.0.mlp.up_proj.weight
model.layers.0.mlp.down_proj.weight → model.layers.0.mlp.down_proj.weight旧版检查点可能使用较旧的命名约定,这些约定由内置的重命名规则处理,适用于所有模型:
Checkpoint: Model:
LayerNorm.gamma → LayerNorm.weight
LayerNorm.beta → LayerNorm.biasMoE 模型(例如 Mixtral)
对于 MoE 模型,检查点格式与模型格式不同。转换映射会将分离的专家权重转换为融合后的 3D 张量,且 TP 分片在转换之后应用。
Checkpoint: Model:
experts.0.w1.weight ─┐
experts.1.w1.weight │ MergeModulelist
... ├───────────────→ experts.gate_up_proj (8, hidden, 2*intermediate)
experts.0.w3.weight │ + Concatenate
experts.1.w3.weight ─┘复合模型(例如视觉-语言模型)
PreTrainedModel 可以包含其他 PreTrainedModel 子模块。每个子模型都可以拥有自己的转换映射,该映射可以注册在其类名或 model_type 上。当加载父模型时,get_model_conversion_mapping 会以深度优先的顺序遍历子模型,收集它们的映射,并通过 scope_prefix 自动将每个转换限定在其子模块的路径范围内。
Composite model: Per-submodel mappings (auto-scoped):
LlavaForConditionalGeneration
├── vision_model: SiglipVisionModel → SiglipVisionModel mapping (scope="vision_model")
└── language_model: LlamaForCausalLM → LlamaForCausalLM mapping (scope="language_model")scope_prefix 是子模块的点分路径(例如 "vision_model"、"language_model.model" 等)。带作用域的转换仅对以 f"{scope_prefix}." 开头的键生效;在模式匹配之前会剥离该前缀,并在替换之后重新附加,因此每个子模型的映射都是相对于该子模型编写的,就好像它是根节点一样。
架构
该系统围绕 core_model_loading.py 中定义的若干关键组件构建:
阶段 1 — 按键处理(遍历检查点键):
- Walk the transform list once. Every
WeightRenamingthat matches fires (e.g.block_sparse_moe→mlp), and at most one WeightConverter may claim the key (e.g.experts.*.w1.weight). - 通过
ThreadPoolExecutor异步执行 分片(TP)并发送到设备。 - 将具有相同
source_pattern的 Collect 张量聚合在一起(例如,所有 MoE 专家权重、gate + up 投影)。
阶段 2 — 逐映射处理(遍历已收集的映射):
- 反量化/反序列化(仅适用于预量化检查点)。
- 应用 ConversionOps 链:
Chunk、Concatenate、MergeModulelist、Transpose等。 - 在运行时进行量化(如果尚未预先量化)。
- 在模型上设置参数。
WeightTransform
负责处理模式匹配和张量收集的基类:
- 模式编译:源模式是完整的正则表达式,使用
re.search()进行匹配。*通配符匹配任意可索引的组件,并将所有匹配项分组以进行批量操作。 - 捕获组与反向引用:源模式中的捕获组可以在目标模式中通过
\1、\2、… 进行引用,以便在重命名过程中保留子字符串(例如层索引)。 - 作用域:
scope_prefix(由get_model_conversion_mapping为每个子模型自动设置)将转换限制在该路径下的键。匹配前会去除前缀,替换后重新附加。 - 键重命名:
rename_source_key()应用正则表达式(并处理作用域),返回(renamed_key, source_pattern),以便加载器知道哪个转换器(如果有的话)声明了该键。 - 张量集合:
add_tensor()根据source_pattern累积已解析的张量(或Future),以便在执行操作链之前,将单次转换(例如所有 MoE 专家权重)所需的所有张量进行批量处理。 - 可逆性:
reverse_transform()会交换源 ↔ 目标模式并反转每个操作,因此将同一列表反转即可驱动保存过程。 - 跟踪:
was_used()报告在加载过程中该转换是否实际匹配到了任何键;这是必需的,以便在保存时能够对称地重新应用非双射重命名(例如,PrefixChange添加的前缀可能已经存在)。
WeightRenaming
WeightRenaming 是一种专门用于纯键名重命名(不涉及张量操作)的 WeightTransform。与 WeightConverter 不同,WeightRenaming 不会 claim 该键,因此多个重命名可以自由链式执行。在加载路径中,所有重命名都在转换器之前执行;在保存路径中(倒序列表),所有倒序重命名都在倒序转换器之后执行——请参阅上述顺序规则。
# Legacy checkpoint compatibility
WeightRenaming("LayerNorm.gamma", "LayerNorm.weight")
# Module path changes
WeightRenaming(".block_sparse_moe.", ".mlp.")PrefixChange 是 WeightRenaming 的高层封装,用于处理移除或添加整个路径组件的常见情况。可选的 model_prefix 将操作范围限定在该命名空间下的键:
# "model.layers.bad_prefix.weight" → "model.layers.weight"
PrefixChange(prefix_to_remove="bad_prefix", model_prefix="model.layers")
# "layers.0.weight" → "model.layers.0.weight"
PrefixChange(prefix_to_add="model")WeightConverter
WeightConverter 继承自 WeightTransform,并包含一个作用于所收集张量的 ConversionOps 链。支持的四种基数为:
| 基数 | 源模式 | 目标模式 | 典型操作 |
|---|---|---|---|
| 一对一 | 1 | 1 | Transpose、PermuteForRope |
| 一对多 | 1 | >1 | 分块(例如解包 qkv_proj) |
| 多对一 | >1 | 1 | Concatenate、MergeModulelist(例如融合专家) |
| 多对多 | >1 | >1 | 仅支持明确支持该功能的操作(例如 ErnieFuseAndSplitTextVisionExperts) |
WeightConverter(
source_patterns=[".experts.*.w1.weight", ".experts.*.w3.weight"],
target_patterns=".experts.gate_up_proj",
operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
)WeightConverter 也是可逆的:reverse_transform() 会交换源 ↔ 目标,并将每个操作替换为其 reverse_op,因此注册的转换映射在结构上是双向的(加载时按原样使用该列表,保存时使用其反向)。
重命名和转换器的顺序
将 WeightRenaming 条目列在 WeightConverter 条目之前,并保持它们的叶子节点互不重叠:重命名操作会规范化转换器所消费的键,但绝不能指向转换器生成的叶子节点。保存路径依赖此规则,以便分两个阶段反转映射(先反转转换器,再反转重命名)。
转换列表按顺序遍历,每个检查点键遍历一次。对于每个键:
- 每个匹配的
WeightRenaming都会触发;多个重命名可以链式执行。 - 第一个匹配 claims 中键的
WeightConverter会接管该键。后续的转换器将被跳过——这确保了张量在合并/拆分步骤中被路由到单个转换器,也是两个意图重叠的转换器会导致配置错误的原因。
weight_mapping = [
WeightRenaming("^old_prefix", "encoder"), # rename runs always
WeightConverter( # converter claims the key
"attn.qkv_proj.weight",
["attn.q_proj.weight", "attn.k_proj.weight", "attn.v_proj.weight"],
operations=[Chunk(dim=0)],
),
]
# Load: "old_prefix.attn.qkv_proj.weight"
# → WeightRenaming → "encoder.attn.qkv_proj.weight"
# → WeightConverter → "encoder.attn.{q,k,v}_proj.weight"
#
# Save: list reversed, each transform inverted:
# → rev(WeightConverter) repacks QKV → "encoder.attn.qkv_proj.weight"
# → rev(WeightRenaming) fixes prefix → "old_prefix.attn.qkv_proj.weight"转换操作
WeightConverter 类包含若干操作,在调用 from_pretrained() 时执行,用于将检查点中的源张量转换为模型的目标张量。
操作是完全可逆的。保存操作会逆转转换过程并返回原始检查点,从而方便您在不同框架之间进行工作。
| 操作 | 反转 |
|---|---|
Chunk(dim) | Concatenate(dim) |
Concatenate(dim) | Chunk(dim) |
MergeModulelist(dim) | SplitModulelist(dim) |
SplitModulelist(dim) | MergeModulelist(dim) |
Transpose(d0, d1) | Transpose(d1, d0) |
PermuteForRope() | PermuteForRope() |
Conv3dToLinear(...) | LinearToConv3d(...) |
块
Chunk 操作沿指定维度将张量拆分为大小相等的若干部分。例如,当模型期望 Q、K 和 V 作为三个独立的张量,而不是一个单独的张量时。
WeightConverter(
"self_attn.qkv_proj",
["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"],
operations=[Chunk(dim=0)],
)拼接
Concatenate 操作将多个独立的张量融合为一个张量。例如,当模型期望将 Q、K 和 V 作为单个张量输入,而不是作为独立的张量时。
WeightConverter(
["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"],
"self_attn.qkv_proj",
operations=[Concatenate(dim=0)],
)MergeModulelist
MergeModulelist 将一组 2D 张量合并为单个 3D 张量。例如,你可以将 MergeModulelist 与 Concatenate 组合使用,以堆叠 MoE 中的专家并将其打包到一个张量中。
WeightConverter(
["block_sparse_moe.experts.*.w1.weight", "block_sparse_moe.experts.*.w3.weight"],
"mlp.experts.gate_up_proj",
operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
)SplitModulelist
SplitModulelist 将 3D 张量拆分回 2D 张量列表。例如,您可以将一组专家拆分回各个独立的专家。
WeightConverter(
"mlp.experts.down_proj",
"block_sparse_moe.experts.*.w2.weight",
operations=[SplitModulelist(dim=0)],
)PermuteForRope
PermuteForRope 将权重从交错格式转换为 sin/cos 格式。例如,您可以将 Chunk 与 PermuteForRope 组合使用,以拆分融合的 QKV 张量,并对 Q 和 K 应用 sin/cos RoPE 置换。
WeightConverter(
["model.layers.*.self_attn.qkv_proj.weight"],
[
"model.layers.*.self_attn.q_proj.weight",
"model.layers.*.self_attn.k_proj.weight",
"model.layers.*.self_attn.v_proj.weight",
],
operations=[Chunk(dim=0), PermuteForRope()],
)转置
Transpose 用于交换张量的维度。在将权重布局转换为不同约定时非常有用。
WeightConverter(
source_patterns="mlp.gate.weight",
target_patterns="mlp.text_moe.gate.weight",
operations=[Transpose(dim0=0, dim1=1)],
)操作链
操作可以链式调用以执行复杂的转换。操作按顺序执行,每个操作的输出会成为下一个操作的输入。
示例:Mixtral MoE 转换
WeightConverter(
source_patterns=[
".experts.*.w1.weight", # gate_proj per expert
".experts.*.w3.weight", # up_proj per expert
],
target_patterns=".experts.gate_up_proj",
operations=[
MergeModulelist(dim=0), # Stack all experts: (n_experts, in, out)
Concatenate(dim=1), # Fuse gate+up: (n_experts, in, 2*out)
],
)数据流:
Input:
".experts.*.w1.weight": [tensor_0, tensor_1, ..., tensor_7] # 8 experts
".experts.*.w3.weight": [tensor_0, tensor_1, ..., tensor_7] # 8 experts
After MergeModulelist(dim=0):
".experts.*.w1.weight": (8, 4096, 14336) # stacked gate
".experts.*.w3.weight": (8, 4096, 14336) # stacked up
After Concatenate(dim=1):
".experts.gate_up_proj": (8, 4096, 28672) # fused gate_up模式匹配规则
- 源模式是完整的正则表达式,使用
re.search()进行评估。^锚定到 作用域 键的开头(在去除scope_prefix之后)。 *是按索引匹配的通配符,用于将同一源模式下的所有匹配张量收集起来(保留检查点顺序以确保正确拼接)。- 源模式中的捕获组可以在目标模式中通过
\1、\2、… 进行引用,以保留原始键的部分内容(例如层索引)。 - 作用域转换(设置了
scope_prefix)仅匹配以f"{scope_prefix}."开头的键。
注册转换映射
转换列表会注册到一个字符串键上——该键可以是 model_type(例如 "mixtral")或类名(例如 "LlavaModel"):
from transformers.conversion_mapping import register_checkpoint_conversion_mapping
register_checkpoint_conversion_mapping(
"my_model_type",
[
WeightRenaming(".old.", ".new."),
WeightConverter(
".experts.*.w1.weight",
".experts.gate_proj",
operations=[MergeModulelist(dim=0)],
),
],
)查找规则
当 get_model_conversion_mapping 处理一个 PreTrainedModel 时,会以 DFS 顺序遍历每个子 PreTrainedModel(即 nn.Module.named_modules() 过滤出的 PreTrainedModel 实例)。每个子模型的解析分为两个阶段:首先运行自定义代码过滤器,然后查找已注册的映射。
在查找之前,自定义代码模型会被跳过。当 is_custom_code() 返回 True 时,子模型即被视为自定义代码,这涵盖了使用 trust_remote_code=True 从 Hub 加载的模型,以及 Transformers 中不存在的任何模型类。自定义代码模型不会继承转换映射,因为其类名或 model_type 可能与原生模型发生冲突。例如,一个将 model_type="mixtral" 设置为该值的自定义架构,不会应用 Mixtral 内置的专家融合转换。
显式使用
register_checkpoint_conversion_mapping注册自定义代码模型,以应用转换映射。注册操作会将类名或model_type添加到用户注册的映射集合中,从而使该模型免于被自定义代码跳过,并使其通过下方的常规查找流程。
对于每个通过自定义代码筛选器的子模型:
- 首先尝试按类名查找,然后查找
model_type。如果同一个模块同时注册了这两者,类名映射将优先生效,该模块的model_type映射将被忽略——这使得任务头(例如LlavaForConditionalGeneration)可以覆盖共享的model_type基线("llava")。 - 所选映射的
scope_prefix被设置为子模块的点分路径(根模块为"")。 - Ancestor-based deduplication 决定是保留该映射:
- 如果某个 祖先 路径已经声明了相同的标识符(类名或
model_type),则该子模块会被 跳过 —— 祖先的无作用域或更高作用域的映射已经覆盖了此子树。 - 如果只有某个 同级模块 声明了它,则该子模块会 保留 其自身的
scope_prefix。每个同级模块都会获得自己的作用域映射。
- 如果某个 祖先 路径已经声明了相同的标识符(类名或
类名和 model_type 的已见列表是分开跟踪的,其中有一个细微之处:当某个模块是通过其类名匹配时,其 model_type 不会被添加到已见列表中。这样做是为了确保其他共享相同 model_type 但没有类特定映射的模块(例如 DetrForSegmentation 下的 DetrModel)仍然可以通过 model_type 查找机制被访问到。
基于类的映射与 model_type 别名
_MODEL_TO_CONVERSION_PATTERN 中同时存在这两种风格:
_MODEL_TO_CONVERSION_PATTERN = {
# model_type aliases (lookup by config.model_type)
"minimax": "mixtral",
"qwen3_moe": "qwen2_moe",
"mistral3": "llava",
# class-name aliases (lookup by type(submodule).__name__)
"PaliGemmaModel": "LlavaModel",
"MaskFormerDetrDecoder": "DetrModel",
...
}当 model_type 相同但特定类需要不同行为时,优先使用类名作为键。
张量并行集成
动态加载系统通过src/transformers/integrations/tensor_parallel.py中定义的TensorParallelLayer层级结构与张量并行(TP)集成。
启用 TP 后,张量会在实例化期间进行分片,而不是在实例化之后。这意味着每个 rank 仅加载其所需的张量部分。
def spawn_tp_materialize(thread_pool, tensor, sharding_method, tensor_idx, device, dtype):
def _job():
return sharding_method.shard_tensor(tensor, tensor_idx=tensor_idx, device=device, dtype=dtype)
return thread_pool.submit(_job)可用的并行方式
| 风格 | 权重分片维度 | 描述 |
|---|---|---|
colwise | -2 | 按列:输出特征分片 |
rowwise | -1 | 按行:输入特征分片 |
packed_colwise | -2 | 对于融合权重(gate_up_proj) |
packed_rowwise | -1 | 对于融合权重 |
embedding_rowwise | 0 | 词表并行 |
grouped_gemm | 0 | MoE 的专家并行 |
sequence_parallel | 无 | 无权重分片 |
打包权重处理
对于像 gate_up_proj 这样的融合权重,需要特别注意以正确地进行分片:
def get_packed_weights(param, empty_param, device_mesh, rank, dim):
"""
Interleaves gate and up shards correctly.
Packed tensor: [G0 G1 G2 G3 | U0 U1 U2 U3]
With TP=2:
- Rank 0 gets: [G0 G1 | U0 U1]
- Rank 1 gets: [G2 G3 | U2 U3]
"""TP 操作存储在 WeightTransform 中,并在转换操作之后应用:
if matched_tp_pattern := tp_plan_alt.search(renamed_key):
tp_layer = ALL_PARALLEL_STYLES[model.tp_plan[matched_tp_pattern]]
mapping.distributed_operation = tp_layer(
device_mesh=device_mesh,
rank=device_mesh.get_local_rank(),
empty_param=empty_param.clone()
)量化集成
量化通过两种方式集成到加载流程中,具体取决于检查点是否已经过量化:
- 预量化检查点:量化器提供 WeightConverter 实例(通过
get_weight_conversions())来反序列化量化张量。检查点的 dtype 会被保留,以避免不必要的类型转换。 - 即时量化:量化器提供在转换操作之后应用的量化操作,在权重加载时对其进行量化。
量化器还可以通过 update_weight_conversions(...) 在 get_model_conversion_mapping 的末尾重写整个转换列表——例如,FP8 反量化器会在每个现有转换器之前添加一个 Fp8Dequantize 操作,以便在专家合并 / 拼接操作展平每个专家的结构之前应用按块的缩放。
快速且高效的模型加载
加载模型的速度更快且占用更少的内存,因为加载器知道哪些张量是操作所必需的,并会延迟调度它们的实例化。
加载器会扫描一次检查点以发现模式匹配并收集张量。它将这些张量存储为 Future 对象,并提交到线程池中进行异步加载,从而避免阻塞 GIL。一旦有线程可用,参数便会开始加载。
如果系统正在运行其他高负载进程,多线程加载可能会拖慢而非加速加载速度。在这种情况下,请设置环境变量 HF_DEACTIVATE_ASYNC_LOAD=1 以按顺序加载权重。
默认情况下,异步参数加载使用 4 个线程。这在不同加载场景和硬件之间提供了最佳平衡。该工作主要受 I/O 限制,但根据加速器硬件以及加载时所需的
dtype,如果dtype与序列化时的不同,则可能变为 CPU/GPU 受限(这需要额外的复制操作)。
异步加载与同步加载
def spawn_materialize(thread_pool, tensor, device, dtype) -> Future | Callable:
def _job():
return _materialize_copy(tensor, device, dtype)
if thread_pool is not None:
return thread_pool.submit(_job) # Async: returns Future
else:
return _job # Sync: returns Callable (deferred execution)同步加载用于以下情况:
- 已设置
HF_DEACTIVATE_ASYNC_LOAD=1环境变量。 - 已启用磁盘卸载(内存限制要求顺序加载)。
- 已启用即时量化(避免工作线程超前于量化步骤)。
物化流程
1. Checkpoint iteration (Phase 1):
- For each key, walk the transform list once
- Submit materialization job to ThreadPoolExecutor
- Job returns Future (async) or Callable (sync)
- Collect into the matching WeightConverter / WeightRenaming
2. Per-mapping processing (Phase 2, one mapping at a time):
- materialize_tensors() waits for this mapping's Futures only
- Apply conversion operations chain (self.operations)
- Apply quantization operation (if on-the-fly)
- Set parameters on model
- Delete realized tensors immediately
3. Cleanup:
- Thread pool shutdown (with cancel_futures=True for interrupts)内存效率
在转换权重时,如果所需的张量尚未加载,转换器会等待它们全部实例化。例如,MergeModulelist 操作要求 ModuleList 中的所有权重在合并前均已加载。
张量拼接需要临时副本,因此在转换过程中,MergeModulelist 和 Concatenate 等操作需要底层张量 2 倍的内存。合并完成后,只有结果张量会保留在内存中。理论上最坏情况的内存峰值为模型大小加上最大 MergeModulelist 或 Concatenate 操作所需的张量。
这种最坏情况仅在所有其他参数在耗时的转换运行之前就已加载完毕时才会发生。有两种场景会触发这种情况。
- 在进入耗时的转换之前,所有参数均已异步加载完成(线程池的速度快于转换队列)。
- 最耗时的转换是最后一个。
例如,一个 MoE 模型在每一层使用 MergeModulelist 来管理专家,其理论上的最坏情况内存峰值为模型大小加上单层中的专家参数。
这些最坏情况并不常见。实际的内存峰值通常接近模型大小。
可逆性
该系统支持保存带有逆变换的模型,从而实现往返保存/加载。保存过程分为两个阶段:首先执行反向转换器(每个张量最多匹配一个),然后按照上述排序规则执行反向重命名。
def revert_weight_conversion(model, state_dict):
"""Applies reverse conversions for saving."""
weight_conversions = getattr(model, "_weight_conversions", None)
# Reverse all transforms
reverse_weight_conversion = [
conversion.reverse_transform() for conversion in weight_conversions
]
# Apply in reverse
for first_param_name, reversed_converter in conversion_mapping.items():
realized_value = reversed_converter.convert(first_param_name, model=model)加载时使用的 transforms 列表会缓存在模型的 _weight_conversions 中(仅保留实际生效的条目,因此像 PrefixChange 这类非双射重命名会被对称地正确重新应用)。如果模型未通过 from_pretrained 实例化(因此没有 _weight_conversions),revert_weight_conversion 会回退到通过 get_model_conversion_mapping 重新计算映射,并从中移除任何 PrefixChange(我们无法判断原始检查点是否包含该前缀)。
目标模式可能包含需要为反向处理而进行处理的正则表达式元素:
def process_target_pattern(pattern: str) -> tuple[str, str | None]:
"""
- Removes `^` and `$` anchors
- Removes negative lookahead/lookbehind
- Detects capturing groups, replaces with \1
"""实际示例
Mixtral 风格的 MoE
检查点格式:
model.layers.0.block_sparse_moe.experts.0.w1.weight # gate per expert
model.layers.0.block_sparse_moe.experts.0.w2.weight # down per expert
model.layers.0.block_sparse_moe.experts.0.w3.weight # up per expert
...
model.layers.0.block_sparse_moe.experts.7.w1.weight模型格式:
model.layers.0.mlp.experts.gate_up_proj # (8, 4096, 28672)
model.layers.0.mlp.experts.down_proj # (8, 14336, 4096)转换映射(来自 conversion_mapping.py):
"mixtral": [
WeightRenaming(".block_sparse_moe.", ".mlp."),
WeightConverter(
source_patterns=[".experts.*.w1.weight", ".experts.*.w3.weight"],
target_patterns=".experts.gate_up_proj",
operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
),
WeightConverter(
source_patterns=[".experts.*.w2.weight"],
target_patterns=".experts.down_proj",
operations=[MergeModulelist(dim=0)],
),
],复合视觉语言模型(Gemma3)
Gemma3 采用了典型的“head 模型与 base 模型使用不同的前缀逻辑”配置。这里注册了两个条目(通过指向共享的 "llava" / "LlavaModel" 列表的别名实现):
# model_type: applies to Gemma3ForConditionalGeneration / Gemma3ForSequenceClassification
"gemma3": "llava" # → adds "model." in front of language_model / vision_tower / ...
# class: applies to the inner Gemma3Model only
"Gemma3Model": "LlavaModel" # → minimal rename inside the already-prefixed namespaceGemma3ForConditionalGeneration 的 DFS 遍历:
- Root — 类查找未命中(
Gemma3ForConditionalGeneration未注册),model_type 查找命中("gemma3")→ 添加"llava"前缀重写转换,且未限定作用域,将language_model、vision_tower、multi_modal_projector置于model.*之下。 - 内部
model: Gemma3Model— 类查找命中(Gemma3Model→LlavaModel)→ 仅应用LlavaModel的映射,其作用域限定为"model"。范围更广的"gemma3"(="llava")前缀重命名在此处不会被重新应用,这正是您期望的行为:内部模型所处的命名空间中,前缀本身已经是正确的。
所有 Llava 系列 VLM(PaliGemma、InternVL、Mistral3、……)都使用相同的 pair 形状。model_type 条目负责处理头部模型的前缀修改;class 条目则保持内部基础模型的映射尽可能精简。
使用针对特定头部的覆盖实现更深层嵌套(DETR)
DetrForSegmentation 展示了一个叠加在两个嵌套层级之上的特定于 head 的覆盖:
DetrForSegmentation (class registered: segmentation-only renames)
├── detr: DetrForObjectDetection (no mapping; just walked through)
│ └── model: DetrModel (class registered: shared base transforms)
│ └── backbone, encoder, ...
└── mask_head, bbox_attention (head-specific weights)涉及两个映射(每个已注册的键对应一个):
"DetrModel": [WeightRenaming("backbone.conv_encoder", "backbone"), ...] # shared base
"DetrForSegmentation": [WeightRenaming("mask_head.lay1", "mask_head.conv1.conv"), ...] # head-specificDFS 遍历:
- 根
DetrForSegmentation按类别匹配 → 分割重命名添加了 unscoped。 detr: DetrForObjectDetection未注册 → 未添加任何 transforms;DFS 继续深入该节点。detr.model: DetrModel按类匹配 → 使用scope_prefix="detr.model"添加基础转换。
这正是 DetrForObjectDetection 无需单独定义映射的原因:其子树中唯一注册的映射(DetrModel)会自动限定到正确的路径。
基于类名的别名会复用基础映射,无需额外注册:"MaskFormerDetrDecoder": "DetrModel" 使 MaskFormer 解码器能够以其自身的类名获取相同的变换。
自定义操作(ERNIE 4.5 VL MoE)
当内置操作无法满足需求时,您可以创建一个自定义的 ConversionOps 子类。例如,ERNIE 4.5 VL MoE 需要在文本和视觉模态之间拆分共享专家列表——这是任何单个内置操作都无法处理的任务。自定义的 ErnieFuseAndSplitTextVisionExperts 操作会在两个目标键之间拆分并重新堆叠专家:
"ernie4_5_vl_moe": [
WeightRenaming("vision_model", "vision_tower"),
WeightConverter(
source_patterns=["experts.*.down_proj.weight"],
target_patterns=[
"text_moe.experts.down_proj",
"vision_moe.experts.down_proj",
],
operations=[ErnieFuseAndSplitTextVisionExperts(stack_dim=0, concat_dim=1)],
),
],自定义算子必须实现 convert() 和 reverse_op 属性,以支持往返保存/加载。
模型类型别名
许多模型共享相同的转换模式:
_MODEL_TO_CONVERSION_PATTERN = {
"mixtral": "mixtral",
"minimax": "mixtral",
"qwen2_moe": "qwen2_moe",
"deepseek_v2": "qwen2_moe",
"deepseek_v3": "qwen2_moe",
"qwen3_moe": "qwen2_moe",
"olmoe": "qwen2_moe",
...
}复用动态加载的构建模块
动态权重加载不仅限于完整的模型检查点。只要能够描述检查点键与参数之间的映射关系,并确保目标模块存在,就可以使用相同的构建模块来加载任意一组权重。
从高层次来看,该合约的结构如下:
- 准备模型命名空间。 确保您想要加载的模块/参数已存在,并且其命名方式与您的映射目标一致。对于适配器,这意味着在加载之前调用
inject_adapter_in_model(...),以确保适配器模块已存在。对于自定义头部或额外模块,请先在模型上实例化它们。 - Describe how to map weights. 构建一个转换/重命名列表(例如,在类似
_build_peft_weight_mapping(...)) 的辅助函数中使用 WeightConverter 或WeightRenaming. 在此处,您定义检查点键应如何转换、拆分、合并或重命名,以匹配您的模型命名空间。主要可以执行以下三种操作:- 将操作添加到转换器列表中:这些操作将应用于所有权重,但被任意
WeightConverter收集的权重除外。这些操作通常应为WeightRenaming操作 - 将操作添加到每个转换器的操作列表中:对于
Quantization来说,我们只需在任意WeightConverter的操作列表之后添加一个量化操作。 - 将 replace / map 操作替换为您的自定义操作:
peft正是这样做的。我们将例如mixtral的Concatenate操作替换为PeftConcatenate(该操作在 PEFT 中定义)。这样,当读取适配器检查点时,需要拼接的权重会被收集起来,并按照peft的要求进行正确的格式化。
- 将操作添加到转换器列表中:这些操作将应用于所有权重,但被任意
- Load + finalize + report. 使用核心加载器执行转换并填充张量,然后完成加载并记录结果。具体流程如下:
- 使用
LoadStateDictConfig(...)+_load_pretrained_model(...)进行加载和转换。 _finalize_load_state_dict(...)用于将任何缺失或不匹配的张量从meta中移出,对其进行初始化,并绑定权重。log_state_dict_report(...)用于报告缺失/意外/不匹配的键(以及转换错误)。
- 使用
这些 API 旨在让你能够处理自定义代码和自定义权重格式,同时确保你能充分利用 transformers API 提供的高效权重加载、分片功能以及出色的使用体验!
关键文件参考
| 文件 | 目的 |
|---|---|
src/transformers/core_model_loading.py | 核心加载逻辑,WeightConverter、WeightRenaming、ConversionOps |
src/transformers/conversion_mapping.py | 内置映射和按子模型组合(get_model_conversion_mapping) |
src/transformers/integrations/tensor_parallel.py | TP 分片类与实用工具 |
src/transformers/quantizers/base.py | 量化钩子和基类 |