Monkey patching（实验性功能）

Monkey patching 允许您在不修改原始模型代码的情况下，全局替换模型组件。注册后，在通过 from_pretrained() 或 ~PreTrainedModel.from_config 加载任何模型时，补丁会自动应用。这使您能够针对特定需求（如量化兼容性）重构模型、应用优化措施，或试验不同的架构变体。

当您需要更改模块及其权重的布局和结构时，猴子补丁应作为最后的手段。对于许多定制和优化需求，请尝试使用 Attention 接口、Experts 接口或 Kernels 注册表。仅当您需要通过自定义 forward 实现无法达成的结构性更改时（例如，为了量化库兼容性、融合层或架构实验），才使用猴子补丁。

快速入门

下面是一个简单的示例，展示了如何替换模型组件：

from transformers import AutoModelForCausalLM
from transformers.models.llama.modeling_llama import LlamaAttention
from transformers.monkey_patching import register_patch_mapping


# Define your replacement class (must inherit from nn.Module)
class CustomLlamaAttention(LlamaAttention):
    def forward(self, *args, **kwargs):
        # Your custom implementation
        print("Using custom attention!")
        return super().forward(*args, **kwargs)


# Register the patch globally (only applies to transformers modeling modules)
register_patch_mapping(mapping={"LlamaAttention": CustomLlamaAttention})

# Load a model - the patch is automatically applied during initialization
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")

# All LlamaAttention layers in the model are now CustomLlamaAttention instances
print(type(model.model.layers[0].self_attn))  # <class '__main__.CustomLlamaAttention'>

工作原理

Monkey patch 通过一个两步流程生效：

注册：调用 register_patch_mapping 将映射添加到全局注册表中。
应用：补丁会在模型初始化期间自动应用：
- from_pretrained / from_config：补丁会通过内部上下文管理器自动应用。无需额外操作！
- 手动构建（例如 Model(config)）：您必须手动使用 apply_patches 上下文管理器。

补丁注册后，它们会持续生效，并影响后续所有模型加载，直到您使用 clear_patch_mapping 将其清除。

重要限制：

只有 transformers 建模模块中的类才允许被修补（例如 LlamaAttention、LlamaMLP）。
映射键可以是精确的类名，也可以是正则表达式模式（参见下文中的模式匹配）。

全局注册

使用 register_patch_mapping 全局注册映射：

from transformers.monkey_patching import register_patch_mapping

# Register a single patch
register_patch_mapping(
    mapping={"Qwen2MoeExperts": SequentialExperts}
)

# Register multiple patches at once
register_patch_mapping(
    mapping={
        "Qwen2MoeExperts": SequentialExperts,
        "Qwen2MoeAttention": CustomAttention,
    },
    # Overwrite existing patches if they exist
    overwrite=True,
)

模式匹配

您可以使用正则表达式通过单个模式匹配多个类：

from transformers.monkey_patching import register_patch_mapping

# Match all classes containing "Attention"
register_patch_mapping(
    mapping={".*Attention": CustomAttention}
)

# More examples
register_patch_mapping(
    mapping={
        ".*MoeExperts$": CustomExperts,           # Ends with "MoeExperts"
        "^Llama\\d+Attention$": CustomAttention,  # Llama2Attention, Llama3Attention, etc.
    }
)

重要：精确匹配优先于模式匹配。如果你同时注册了 "LlamaAttention" 和 ".*Attention"，那么名为 LlamaAttention 的类将使用精确匹配替换，而其他匹配的类将使用模式匹配替换。

正则表达式模式可能会悄无声息地破坏模型。 像 ".*Attention" 这样宽泛的模式会匹配所有名称中包含 “Attention” 的类——包括那些封装了你实际想要替换的注意力机制的容器类。例如，BERT 有三个与注意力相关的类：BertSelfAttention 和 BertCrossAttention（内部注意力实现）以及 BertAttention（一个包含其中一个内部类的外部模块）。如果用同一个自定义注意力层修补这三个类，会导致模型损坏，因为外层的 BertAttention 不再封装内部的类——它本身就是一个，从而消除了 self 和 output 等预期存在的子模块。建议优先使用更精确的模式（例如 ".*SelfAttention$"）或确切的类名，以避免非预期的匹配。

要取消注册补丁，请使用 unregister_patch_mapping：

from transformers.monkey_patching import unregister_patch_mapping

# Unregister a single patch (use exact name or pattern from registration)
unregister_patch_mapping(keys=["Qwen2MoeExperts"])

# Unregister multiple patches at once
unregister_patch_mapping(keys=["Qwen2MoeExperts", "Qwen2MoeAttention"])

# Unregister a pattern
unregister_patch_mapping(keys=[".*Attention"])

要清除所有已注册的补丁，请使用 clear_patch_mapping：

from transformers.monkey_patching import clear_patch_mapping

clear_patch_mapping()

要查看当前已注册的补丁，请使用 get_patch_mapping：

from transformers.monkey_patching import get_patch_mapping

current_patches = get_patch_mapping()
print(current_patches)

手动构建模型

apply_patches 上下文管理器仅在你手动构建模型（例如 Model(config)）且未使用 from_pretrained 或 from_config 时才需要：

from transformers import LlamaModel, LlamaConfig
from transformers.monkey_patching import register_patch_mapping, apply_patches

# Register patch globally
register_patch_mapping(mapping={"LlamaAttention": CustomAttention})

# For manual construction, you need the context manager
with apply_patches():
    model = LlamaModel(LlamaConfig())  # Uses CustomAttention

# Without the context manager, manual construction uses original classes
model = LlamaModel(LlamaConfig())  # Uses LlamaAttention

# But from_pretrained and from_config will always apply registered patches
model = LlamaModel.from_pretrained("meta-llama/Llama-3.2-1B")  # Uses CustomAttention

重要说明

权重处理：Monkey patching 仅替换类，而不替换权重。如果经过补丁的类具有不同的权重布局，你需要单独处理权重转换，以确保与预训练权重兼容。请参阅下方的完整示例，了解如何将 monkey patch 与权重转换映射结合使用。
全局影响：通过 register_patch_mapping 注册的补丁会全局应用于注册后加载的所有模型。完成后请务必使用 clear_patch_mapping 进行清理，尤其是在测试、笔记本或长时间运行的应用程序中。
类验证：API 会自动验证替换类是否为 nn.Module 的子类。如果传入无效的类，你将收到清晰的错误信息。
线程安全：所有补丁操作都是线程安全的。您可以安全地从多个线程中注册、注销和应用补丁。
匹配行为：使用精确类名时，必须与模型源代码中出现的原始类名完全一致（区分大小写）。使用正则表达式模式时，会通过 re.search() 对类名进行匹配。

故障排除

我的补丁未被应用

检查类名或模式：确保映射中的类名或模式正确：

# For exact names - must match exactly (case-sensitive)
register_patch_mapping(mapping={"LlamaAttention": CustomAttention})

# For patterns - use valid regex
register_patch_mapping(mapping={".*Attention": CustomAttention})

验证注册：使用 get_patch_mapping 确认你的映射已注册：

print(get_patch_mapping())
# Shows all registered mappings: {'LlamaAttention': <class 'CustomAttention'>, '.*MLP': <class 'CustomMLP'>}

检查模型来源：在模型的源代码中找到确切的类名：

from transformers.models.llama import modeling_llama
print(dir(modeling_llama))  # Look for the class name

如何判断我的补丁是否生效？

检查已加载的模型以验证补丁：

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")

# Check the type of a specific module
print(type(model.model.layers[0].self_attn))  # Should show your custom class

# Or iterate through all modules
for name, module in model.named_modules():
    if 'attention' in name.lower():
        print(f"{name}: {type(module)}")

权重形状不匹配错误

如果您的补丁类具有不同的权重形状，请注册一个权重转换：

from transformers.conversion_mapping import register_checkpoint_conversion_mapping, WeightConverter
from transformers.monkey_patching import register_patch_mapping

register_patch_mapping(
    mapping={
        "LlamaAttention": LlamaFusedAttention,
    }
)

register_checkpoint_conversion_mapping(
    model_type_or_class_name="llama",
    mapping=[
        WeightConverter(
            source_patterns=["q_proj", "k_proj", "v_proj"],
            target_patterns=["qkv_proj"],
            operations=[
                Concatenate(dim=0),
            ],
        )
    ],
    overwrite=True,
)

清理补丁

完成后请务必清理补丁，以免影响其他代码：

from transformers.monkey_patching import register_patch_mapping, clear_patch_mapping

try:
    register_patch_mapping(mapping={"LlamaAttention": CustomAttention})
    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
    # ... use model ...
finally:
    clear_patch_mapping()  # Always clean up

完整示例

下面是一个综合示例，展示了如何重构混合专家模型（qwen2_moe）中的专家模块和注意力模块，以优化并兼容量化。该示例演示了以下内容：

创建保持相同接口的自定义替换类
为多个组件注册猴子补丁
处理新结构的权重转换

from typing import Unpack

import torch
import torch.nn as nn

from transformers import AutoModelForCausalLM, Concatenate, WeightConverter
from transformers.activations import ACT2FN
from transformers.cache_utils import Cache
from transformers.conversion_mapping import register_checkpoint_conversion_mapping
from transformers.integrations.sdpa_attention import sdpa_attention_forward
from transformers.models.qwen2_moe.modeling_qwen2_moe import apply_rotary_pos_emb
from transformers.monkey_patching import register_patch_mapping
from transformers.utils.generic import TransformersKwargs


class MoeMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.num_experts = config.num_experts
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.moe_intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj


# Adapted from the original Qwen2MoeExperts
class ModuleListExperts(nn.ModuleList):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.num_experts = config.num_experts
        for _ in range(self.num_experts):
            self.append(MoeMLP(config))

    def forward(
        self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor
    ) -> torch.Tensor:
        final_hidden_states = torch.zeros_like(hidden_states)
        with torch.no_grad():
            expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
            expert_mask = expert_mask.permute(2, 1, 0)

        for expert_idx in range(self.num_experts):
            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
            current_state = hidden_states[token_idx]
            current_hidden_states = self[expert_idx](current_state)
            current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
        return final_hidden_states


# Adapted from the original Qwen2MoeAttention
class FusedQKVAttention(nn.Module):
    def __init__(self, config, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
        self.scaling = self.head_dim**-0.5
        self.attention_dropout = config.attention_dropout
        self.is_causal = True

        self.qkv_proj = nn.Linear(config.hidden_size, 3 * config.num_attention_heads * self.head_dim, bias=True)
        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)

        if self.config.layer_types[layer_idx] == "sliding_attention":
            self.sliding_window = config.sliding_window

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        attention_mask: torch.Tensor | None = None,
        past_key_values: Cache | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor, torch.Tensor]:
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states, key_states, value_states = self.qkv_proj(hidden_states).chunk(3, dim=-1)

        query_states = query_states.view(hidden_shape).transpose(1, 2)
        key_states = key_states.view(hidden_shape).transpose(1, 2)
        value_states = value_states.view(hidden_shape).transpose(1, 2)

        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        if past_key_values is not None:
            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)

        attn_output, attn_weights = sdpa_attention_forward(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            **kwargs,
        )

        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights


# Registering monkey patches for the new attention and experts modules.
register_patch_mapping(
    mapping={
        "Qwen2MoeExperts": ModuleListExperts,
        "Qwen2MoeAttention": FusedQKVAttention,
    }
)

# Registering weight conversion mappings adapted for the new modules. This registration will:
# - Override the original conversion mapping for qwen2_moe which concatenated the experts into a single parameter format.
# - Concatenate the q_proj, k_proj, v_proj weights/biases into a single qkv_proj weight/bias for the new fused attention module.
register_checkpoint_conversion_mapping(
    model_type_or_class_name="qwen2_moe",
    mapping=[
        WeightConverter(
            source_patterns=["q_proj", "k_proj", "v_proj"],
            target_patterns=["qkv_proj"],
            operations=[Concatenate(dim=0)],
        ),
    ],
    overwrite=True,
)

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen1.5-MoE-A2.7B")

记录和重放 MoE 专家路由

像 RLHF 这样的 Mixture-of-Experts 训练工作流需要记录生成过程中每个 token 被路由到了哪些专家，然后在单独的训练前向传播中精确重放该路由。您可以利用现有的 monkey patching 和输出捕获机制来端到端地实现这一流程，无需对 modeling 文件进行任何修改。

该模式包含三个部分：

一个可重放的 路由器子类，可选择地从实例属性中读取强制指定的专家索引。
一个上下文管理器，用于在前向传播之前为所有路由器设置这些属性，并在之后将其清除。
模型输出捕获注册表中的一个条目，使得 output_<name>=True 能够通过标准的 @capture_outputs 路径暴露索引。

from contextlib import contextmanager

import torch
import torch.nn.functional as F

from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM
from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeTopKRouter
from transformers.monkey_patching import apply_patches, register_patch_mapping
from transformers.utils.output_capturing import _CAN_RECORD_REGISTRY, OutputRecorder


class ReplayableQwen3MoeTopKRouter(Qwen3MoeTopKRouter):
    _forced_indices: torch.Tensor | None = None

    def forward(self, hidden_states):
        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
        router_logits = F.linear(hidden_states, self.weight)
        router_logits = F.softmax(router_logits, dtype=torch.float, dim=-1)

        if self._forced_indices is not None:
            router_indices = self._forced_indices.to(router_logits.device).long()
            # Megatron-style replay: preserve expert path, recompute current scores
            router_top_value = router_logits.gather(-1, router_indices)
        else:
            router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)

        if self.norm_topk_prob:
            router_top_value = router_top_value / router_top_value.sum(dim=-1, keepdim=True)
        return router_logits, router_top_value.to(router_logits.dtype), router_indices


@contextmanager
def replay_moe_routing(model, selected_experts_per_layer):
    routers = [m for m in model.modules() if isinstance(m, ReplayableQwen3MoeTopKRouter)]
    if len(routers) != len(selected_experts_per_layer):
        raise ValueError(f"Got {len(routers)} routers but {len(selected_experts_per_layer)} tensors")
    for r, t in zip(routers, selected_experts_per_layer):
        r._forced_indices = t
    try:
        yield
    finally:
        for r in routers:
            r._forced_indices = None


# Swap the router class and construct the model
register_patch_mapping({"Qwen3MoeTopKRouter": ReplayableQwen3MoeTopKRouter})
with apply_patches():
    model = Qwen3MoeForCausalLM(Qwen3MoeConfig(...)).eval()

# Expose `output_selected_experts=True` on the base model by adding an OutputRecorder
# at runtime. Index 2 of the router's tuple output is the expert indices.
inner = model.model
existing = _CAN_RECORD_REGISTRY.get(str(inner.__class__), {}) or {}
_CAN_RECORD_REGISTRY[str(inner.__class__)] = {
    **existing,
    "selected_experts": OutputRecorder(ReplayableQwen3MoeTopKRouter, index=2),
}

# Record
captured = inner(input_ids=input_ids, output_selected_experts=True)
selected_experts = captured.selected_experts  # tuple of (num_tokens, top_k) LongTensors

# Replay — same expert path regardless of current router weights
with replay_moe_routing(inner, list(selected_experts)):
    outputs = inner(input_ids=input_ids)

Replay 保留精确的专家索引，并使用当前的路由器权重重新计算路由分数，因此梯度会通过实时参数进行传播，同时专家选择保持不变。这是 Megatron 风格的 MoE 训练所使用的最小回放契约。

与 vLLM 的互操作性

vLLM 的 enable_return_routed_experts=True 会将 CompletionOutput.routed_experts 填充为一个形状为 (seq_len, num_layers, top_k) 的 np.int32 数组。使用单个表达式即可将其转换为该模式所需的逐层列表：

selected = [
    torch.from_numpy(routed_experts[:, layer, :].copy()).long()
    for layer in range(routed_experts.shape[1])
]
with replay_moe_routing(model, selected):
    loss = model(input_ids=input_ids, labels=labels).loss

同样的方法也适用于其他 MoE 系列——继承该系列的 *TopKRouter，匹配原始返回约定（通常为 (router_logits, router_scores, router_indices)），然后注册该补丁。具体签名请参见各模型的 router 类。

API 参考

transformers.monkey_patching.register_patch_mapping

< source >

( mapping: dictoverwrite: bool = False )

参数

mapping (Dict[str, type[nn.Module]]) — 从原始类名（或正则表达式模式）到替换类的映射。支持：
- 精确的类名："Qwen2MoeExperts" → CustomExperts
- 正则表达式模式：".*Attention" 会匹配 LlamaAttention、MistralAttention 等，或者 "^Llama\d+Attention$" 会匹配 Llama2Attention、Llama3Attention 等。
精确匹配优先于模式匹配。模式使用 re.search() 进行匹配，因此它们可以匹配类名中的任意位置，除非使用锚点（^ 表示开头，$ 表示结尾）。
overwrite (bool，可选，默认为 False) — 是否覆盖已注册的类名的现有映射。

注册补丁映射，以便在使用 from_pretrained、from_config 或在 apply_patches 上下文管理器中创建模型时自动应用补丁。

使用此功能可注册类替换，这些替换将在加载任何模型时自动应用。这对于量化库兼容性、结构优化以及架构实验非常有用。该映射是全局的，可通过多次调用不断扩展，也可以完全清除。

示例：

from transformers import AutoModelForCausalLM
from transformers.monkey_patching import register_patch_mapping

# Define custom expert implementation
class SequentialExperts(nn.Module):
    ...

# Register exact class name
register_patch_mapping(
    mapping={"Qwen2MoeExperts": SequentialExperts}
)

# Register with regex pattern to match multiple classes
register_patch_mapping(
    mapping={".*Attention": CustomAttention}  # Matches LlamaAttention, MistralAttention, etc.
)

# Match specific model versions
register_patch_mapping(
    mapping={"^Llama\d+Attention$": CustomLlamaAttention}  # Matches Llama2Attention, Llama3Attention
)

# The patch will be automatically applied during loading
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")

注意：对于权重转换，请改用 ~transformers.register_checkpoint_conversion_mapping。

transformers.monkey_patching.unregister_patch_mapping

< source >

( keys: list )

参数

keys (List[str]) — 要从补丁映射中移除的映射键（类名或正则表达式）列表（例如 ["Qwen2MoeExperts"] 或 [".*Attention"]）。

注销补丁映射以禁用自动补丁。

这将指定的映射从全局注册表中移除，从而防止它们在模型加载期间被应用。您必须提供与注册时使用的完全相同的名称或模式。

示例：

from transformers import AutoModelForCausalLM
from transformers.monkey_patching import register_patch_mapping, unregister_patch_mapping

# Register a patch
register_patch_mapping(
    mapping={"Qwen2MoeExperts": CustomExperts}
)

# Unregister the patch
unregister_patch_mapping(["Qwen2MoeExperts"])

# The patch will no longer be applied during loading
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen1.5-MoE-A2.7B")

transformers.monkey_patching.clear_patch_mapping

< source >

( )

清除所有已注册的补丁映射。

这将移除全局注册表中的所有已注册映射。

示例：

from transformers.monkey_patching import register_patch_mapping, clear_patch_mapping

# Register some patches
register_patch_mapping(
    mapping={"Qwen2MoeExperts": CustomExperts}
)

# Clear all patches
clear_patch_mapping()

transformers.monkey_patching.get_patch_mapping

< source >

( ) → Dict[str, type[nn.Module]]

返回值

Dict[str, type[nn.Module]]

将类名或模式映射到替换类的字典。

获取所有已注册的补丁映射。

transformers.monkey_patching.apply_patches

< source >

( )

用于在代码块内应用已注册的猴子补丁的上下文管理器。

在代码块执行期间，这会临时将原始类替换为其已注册的替代类，并在执行结束后恢复原始类。

示例：

from transformers import Qwen2MoeModel, Qwen2MoeConfig
from transformers.monkey_patching import register_patch_mapping, apply_patches

# Register a patch
register_patch_mapping(
    mapping={"Qwen2MoeExperts": CustomExperts}
)

# Apply patches within the context
with apply_patches():
    # The model will use CustomExperts instead of Qwen2MoeExperts
    model = Qwen2MoeModel(Qwen2MoeConfig())

# Outside the context, original classes are restored
# The model will use Qwen2MoeExperts again
model = Qwen2MoeModel(Qwen2MoeConfig())

Update on GitHub