Repository: huggingface/transformers
Pull request: Model patching API

# What does this PR do?

Integrates a library wide monkey patching API to fix #43284 and allow things like restructuring the expert weights and fusing the qkv projections. A concrete and complete example:

<details> 

```py
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_patch import register_monkey_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


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


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,
        cache_position: torch.LongTensor | 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:
            # sin and cos are specific to RoPE models; cache_position needed for the static cache
            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)

        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_monkey_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 packed the experts into a single parameter format.
# - Convert 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="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")
```

</details>

@ArthurZucker @yiliu30

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc @MekkCyber
- kernels: @MekkCyber @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->




Changed files:
filename: docs/source/en/_toctree.yml
patch: @@ -16,6 +16,8 @@
       title: Loading models
     - local: custom_models
       title: Customizing models
+    - local: monkey_patching
+      title: Monkey patching
     - local: how_to_hack_models
       title: Customizing model components
     - local: model_sharing

filename: docs/source/en/monkey_patching.md
patch: @@ -0,0 +1,435 @@
+<!--Copyright 2026 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+
+-->
+
+# Monkey patching (experimental feature)
+
+Monkey patching allows you to replace model components globally without modifying the original model code. Once registered, patches are automatically applied when loading any model with [`~PreTrainedModel.from_pretrained`] or [`~PreTrainedModel.from_config`]. This enables you to restructure models for specific requirements like quantization compatibility, apply optimizations, or experiment with architectural variants.
+
+> [!WARNING]
+> **Monkey patching should be used as a last resort** when you need to change the layout and structure of a module and or its weights. For many customization and optimization needs, try using the [Attention interface](./attention_interface), [Experts interface](./experts_interface), or [Kernels registry](./kernel_doc/overview) instead. Only use monkey patching when you need structural changes that can't be achieved through custom forward implementations alone (e.g., for quantization library compatibility, fusing layers, or architectural experiments).
+
+## Quick start
+
+Here's a simple example showing how to replace a model component:
+
+```python
+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'>
+```
+
+## How it works
+
+Monkey patches work through a two-step process:
+
+1. **Registration**: Call [`register_patch_mapping`] to add mappings to a global registry.
+
+2. **Application**: Patches are automatically applied during model initialization:
+   - **`from_pretrained` / `from_config`**: Patches are **automatically** applied via an internal context manager. No additional action needed!
+   - **Manual construction** (e.g., `Model(config)`): You must use the [`apply_patches`] context manager manually.
+
+Once patches are registered, they persist and affect all subsequent model loads until you clear them with [`clear_patch_mapping`].
+
+**Important limitations**:
+
+- Only classes in `transformers` modeling modules are allowed to be patched (e.g., `LlamaAttention`, `LlamaMLP`).
+- The mapping keys can be either exact class names or regular expression patterns (see [Pattern matching](#pattern-matching) below).
+
+## Global registration
+
+Use [`register_patch_mapping`] to register mappings globally:
+
+```python
+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,
+)
+```
+
+## Pattern matching
+
+You can use regular expressions to match multiple classes with a single pattern:
+
+```python
+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.
+    }
+)
+```
+
+**Important**: Exact matches take precedence over patterns. If you register both `"LlamaAttention"` and `".*Attention"`, classes named `LlamaAttention` will use the exact-match replacement, while other matching classes will use the pattern-match replacement.
+
+> [!WARNING]
+> **Regex patterns can silently break models.** A broad pattern like `".*Attention"` will match *every* class whose name contains "Attention" — including container classes that wrap the attention you actually want to replace. For example, BERT has three attention-related classes: `BertSelfAttention` and `BertCrossAttention` (the inner attention implementations) and `BertAttention` (an outer module that *contains* one of those inner classes). Patching all three with the same custom attention layer produces a broken model because the outer `BertAttention` no longer wraps the inner one — it *is* one, eliminating expected sub-modules like `self` and `output`. Prefer narrow patterns (e.g., `".*SelfAttention$"`) or exact class names to avoid unintended matches.
+
+To unregister patches, use [`unregister_patch_mapping`]:
+
+```python
+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"])
+```
+
+To clear all registered patches, use [`clear_patch_mapping`]:
+
+```python
+from transformers.monkey_patching import clear_patch_mapping
+
+clear_patch_mapping()
+```
+
+To view currently registered patches, use [`get_patch_mapping`]:
+
+```python
+from transformers.monkey_patching import get_patch_mapping
+
+current_patches = get_patch_mapping()
+print(current_patches)
+```
+
+## Manual model construction
+
+The [`apply_patches`] context manager is only needed when you're constructing models **manually** (e.g., `Model(config)`) without using `from_pretrained` or `from_config`:
+
+```python
+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
+```
+
+## Important notes
+
+- **Weight handling**: Monkey patching only replaces classes, not weights. If your patched class has a different weights layout, you'll need to handle [weight conversions](./weightconverter) separately to ensure compatibility with pretrained weights. See the [Complete example](#complete-example) below for how to combine monkey patches with weight conversion mappings.
+
+- **Global effect**: Patches registered with [`register_patch_mapping`] are applied globally to all models loaded after registration. Always use [`clear_patch_mapping`] to clean up when done, especially in tests, notebooks, or long-running applications.
+
+- **Class validation**: The API automatically validates that replacement classes are `nn.Module` subclasses. If you pass an invalid class, you'll get a clear error message.
+
+- **Thread safety**: All patching operations are thread-safe. You can safely register, unregister, and apply patches from multiple threads.
+
+- **Matching behavior**: When using exact class names, they must exactly match the original class names as they appear in the model's source code (case-sensitive). When using regex patterns, they are matched against class names using `re.search()`.
+
+## Troubleshooting
+
+### My patch isn't being applied
+
+**Check class name or pattern**: Ensure the class name or pattern in your mapping is correct:
+
+```python
+# 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})
+```
+
+**Verify registration**: Use [`get_patch_mapping`] to confirm your mapping is registered:
+
+```python
+print(get_patch_mapping())
+# Shows all registered mappings: {'LlamaAttention': <class 'CustomAttention'>, '.*MLP': <class 'CustomMLP'>}
+```
+
+**Check model source**: Find the exact class name in the model's source:
+
+```python
+from transformers.models.llama import modeling_llama
+print(dir(modeling_llama))  # Look for the class name
+```
+
+### How do I know if my patch is working?
+
+Inspect the loaded model to verify the patch:
+
+```python
+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)}")
+```
+
+### Weight shape mismatch errors
+
+If your patched class has different weight shapes, register a weight conversion:
+
+```python
+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="llama",
+    mapping=[
+        WeightConverter(
+            source_patterns=["q_proj", "k_proj", "v_proj"],
+            target_patterns=["qkv_proj"],
+            operations=[
+                Concatenate(dim=0),
+            ],
+        )
+    ],
+    overwrite=True,
+)
+```
+
+### Cleaning up patches
+
+Always clean up patches when you're done to avoid affecting other code:
+
+```python
+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
+```
+
+## Complete example
+
+Here's a comprehensive example showing how to restructure both the experts and attention modules in a Mixture-of-Experts model (`qwen2_moe`) for optimization and quantization compatibility. This demonstrates:
+
+1. Creating custom replacement classes that maintain the same interface
+2. Registering monkey patches for multiple components
+3. Handling weight conversions for the new structure
+
+```python
+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,
+        cache_position: torch.LongTensor | 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:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        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="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")
+```
+
+## API reference
+
+[[autodoc]] transformers.monkey_patching.register_patch_mapping
+
+[[autodoc]] transformers.monkey_patching.unregister_patch_mapping
+
+[[autodoc]] transformers.monkey_patching.clear_patch_mapping
+
+[[autodoc]] transformers.monkey_patching.get_patch_mapping
+
+[[autodoc]] transformers.monkey_patching.apply_patches

filename: src/transformers/modeling_utils.py
patch: @@ -86,6 +86,7 @@
 from .loss.loss_utils import LOSS_MAPPING
 from .modeling_flash_attention_utils import lazy_import_flash_attention, lazy_import_paged_flash_attention
 from .modeling_rope_utils import ROPE_INIT_FUNCTIONS
+from .monkey_patching import apply_patches, patch_output_recorders
 from .pytorch_utils import id_tensor_storage
 from .quantizers import HfQuantizer
 from .quantizers.auto import get_hf_quantizer
@@ -1458,7 +1459,7 @@ def _from_config(cls, config, **kwargs):
         if "experts_implementation" in kwargs:
             config._experts_implementation = kwargs.pop("experts_implementation")
 
-        init_contexts = []
+        init_contexts = [apply_patches()]
         if dtype is not None:
             init_contexts.append(local_torch_dtype(dtype, cls.__name__))
 
@@ -1480,6 +1481,7 @@ def _from_config(cls, config, **kwargs):
         # Instantiate the model
         with ContextManagers(init_contexts):
             model = cls(config, **kwargs)
+            patch_output_recorders(model)
 
         # Under ZeRO-3, parameters were partitioned into empty tensors during construction,
         # so weight init was suppressed. Re-initialize using the ZeRO-3 variant which gathers
@@ -3591,7 +3593,7 @@ def float(self, *args):
     @classmethod
     def get_init_context(cls, dtype: torch.dtype, is_quantized: bool, _is_ds_init_called: bool):
         # Need to instantiate with correct dtype
-        init_contexts = [local_torch_dtype(dtype, cls.__name__), init.no_tie_weights()]
+        init_contexts = [local_torch_dtype(dtype, cls.__name__), init.no_tie_weights(), apply_patches()]
         if is_deepspeed_zero3_enabled():
             import deepspeed
 
@@ -4064,10 +4066,11 @@ def from_pretrained(
 
         config.name_or_path = pretrained_model_name_or_path
         model_init_context = cls.get_init_context(dtype, is_quantized, _is_ds_init_called)
+
         config = copy.deepcopy(config)  # We do not want to modify the config inplace in from_pretrained.
         with ContextManagers(model_init_context):
-            # Let's make sure we don't run the init function of buffer modules
             model = cls(config, *model_args, **model_kwargs)
+            patch_output_recorders(model)
 
             if hf_quantizer is not None:  # replace module with quantized modules (does not touch weights)
                 hf_quantizer.preprocess_model(

filename: src/transformers/monkey_patching.py
patch: @@ -0,0 +1,357 @@
+# Copyright 2026 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+import sys
+import threading
+from contextlib import contextmanager
+
+from .utils import is_torch_available, logging
+from .utils.output_capturing import OutputRecorder
+
+
+if is_torch_available():
+    import torch.nn as nn
+
+logger = logging.get_logger(__name__)
+
+_monkey_patch_mapping_cache: dict[str, type[nn.Module]] = {}
+_compiled_patterns_cache: dict[str, re.Pattern] = {}
+_monkey_patch_lock = threading.Lock()
+
+
+def _compile_pattern(pattern: str) -> re.Pattern | None:
+    """
+    Compile a regex pattern and cache it. Returns None if pattern is invalid.
+
+    Args:
+        pattern: The regex pattern string to compile
+
+    Returns:
+        Compiled regex pattern or None if invalid
+    """
+    if pattern in _compiled_patterns_cache:
+        return _compiled_patterns_cache[pattern]
+
+    try:
+        compiled = re.compile(pattern)
+        _compiled_patterns_cache[pattern] = compiled
+        return compiled
+    except re.error as e:
+        logger.warning(f"Invalid regex pattern '{pattern}': {e}. Treating as non-pattern.")
+        return None
+
+
+def _find_replacement_class(class_name: str, mapping: dict[str, type[nn.Module]]) -> type[nn.Module] | None:
+    """
+    Find replacement class for a given class name, checking exact matches first, then regex patterns.
+
+    Args:
+        class_name: The class name to find a replacement for
+        mapping: Dictionary of patterns/names to replacement classes
+
+    Returns:
+        The replacement class if found, None otherwise
+    """
+    # First check for exact match (highest priority)
+    if class_name in mapping:
+        return mapping[class_name]
+
+    # Then check regex patterns
+    for pattern, replacement_class in mapping.items():
+        # Skip if already matched as exact
+        if pattern == class_name:
+            continue
+
+        # Try to compile and match as regex
+        compiled_pattern = _compile_pattern(pattern)
+        if compiled_pattern is not None and compiled_pattern.search(class_name):
+            return replacement_class
+
+    return None
+
+
+def register_patch_mapping(mapping: dict[str, type[nn.Module]], overwrite: bool = False) -> None:
+    """
+    Register patch mappings to enable automatic patching during model creation using `from_pretrained`,
+    `from_config` or within the `apply_patches` context manager.
+
+    Use this to register class replacements that will be automatically applied when loading any model.
+    This is useful for quantization library compatibility, structural optimizations, and architectural
+    experimentation. The mapping is global, can grow with multiple calls, and can be cleared entirely.
+
+    Args:
+        mapping (`Dict[str, type[nn.Module]]`):
+            Mapping from original class names (or regex patterns) to replacement classes. Supports:
+            - Exact class names: `"Qwen2MoeExperts"` → `CustomExperts`
+            - Regex patterns: `".*Attention"` matches `LlamaAttention`, `MistralAttention`, etc.,
+            or `"^Llama\\d+Attention$"` matches `Llama2Attention`, `Llama3Attention`, etc.
+
+            Exact matches take precedence over patterns. Patterns are matched using `re.search()`,
+            so they can match anywhere in the class name unless you use anchors (`^` for start, `$` for end).
+        overwrite (`bool`, *optional*, defaults to `False`):
+            Whether to overwrite existing mappings for class names that are already registered.
+
+    Example:
+        ```python
+        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")
+        ```
+
+    Note:
+        For weight conversions, use [`~transformers.register_checkpoint_conversion_mapping`] instead.
+    """
+    global _monkey_patch_mapping_cache
+    with _monkey_patch_lock:
+        for class_name, replacement_class in mapping.items():
+            # Validate that replacement_class is actually a class and is a subclass of nn.Module
+            if not isinstance(replacement_class, type):
+                raise TypeError(
+                    f"Replacement for '{class_name}' must be a class, got {type(replacement_class).__name__}"
+                )
+            if not issubclass(replacement_class, nn.Module):
+                raise TypeError(
+                    f"Replacement class for '{class_name}' must be a subclass of nn.Module, "
+                    f"got {replacement_class.__name__} which inherits from {[c.__name__ for c in replacement_class.__mro__[1:]]}"
+                )
+
+            if class_name in _monkey_patch_mapping_cache and not overwrite:
+                raise ValueError(
+                    f"Class '{class_name}' already has a patch mapping registered. Use overwrite=True to replace it."
+                )
+            _monkey_patch_mapping_cache[class_name] = replacement_class
+
+
+def unregister_patch_mapping(keys: list[str]) -> None:
+    """
+    Unregister patch mappings to disable automatic patching.
+
+    This removes specified mappings from the global registry, preventing them from being applied
+    during model loading. You must provide the exact same name or pattern that was used during registration.
+
+    Args:
+        keys (`List[str]`):
+            List of mapping keys (class names or regex patterns) to remove from the patch mapping
+            (e.g., `["Qwen2MoeExperts"]` or `[".*Attention"]`).
+
+    Example:
+        ```python
+        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")
+        ```
+    """
+    global _monkey_patch_mapping_cache
+    with _monkey_patch_lock:
+        for key in keys:
+            if key not in _monkey_patch_mapping_cache:
+                raise ValueError(
+                    f"Class or pattern '{key}' not found in monkey patch mapping cache. "
+                    f"Cannot unregister a class that is not registered."
+                )
+            del _monkey_patch_mapping_cache[key]
+
+
+def get_patch_mapping() -> dict[str, type[nn.Module]]:
+    """
+    Get all registered patch mappings.
+
+    Returns:
+        `Dict[str, type[nn.Module]]`: Dictionary mapping class names or patterns to replacement classes.
+    """
+    with _monkey_patch_lock:
+        return _monkey_patch_mapping_cache.copy()
+
+
+def clear_patch_mapping() -> None:
+    """
+    Clear all registered patch mappings.
+
+    This removes all registered mappings from the global registry.
+
+    Example:
+        ```python
+        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()
+        ```
+    """
+    global _monkey_patch_mapping_cache
+    with _monkey_patch_lock:
+        _monkey_patch_mapping_cache.clear()
+
+
+@contextmanager
+def apply_patches():
+    """
+    Context manager to apply registered monkey patches within a block of code.
+
+    This temporarily replaces original classes with their registered replacements during the execution of the block, and restores the original classes afterward.
+
+    Example:
+        ```python
+        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())
+        ```
+    """
+    mapping = get_patch_mapping()
+    if not mapping:
+        yield
+        return
+
+    original_classes = {}
+
+    # Create list to avoid dict changed during iteration
+    for module in list(sys.modules.values()):
+        if module is None or not hasattr(module, "__name__"):
+            continue
+        if not module.__name__.startswith("transformers"):
+            continue
+
+        # Iterate through all attributes in transformers modules
+        for attr_name in dir(module):
+            # Check if this attribute name matches any pattern before accessing it
+            replacement_class = _find_replacement_class(attr_name, mapping)
+            if replacement_class is None:
+                continue
+
+            try:
+                attr = getattr(module, attr_name)
+                # Check if it's a class
+                if not isinstance(attr, type):
+                    continue
+
+                original_classes[(module.__name__, attr_name)] = attr
+                setattr(module, attr_name, replacement_class)
+            except (AttributeError, TypeError, ImportError):
+                # Skip attributes that can't be accessed or modules that can't be imported
+                continue
+
+    yield
+
+    for (module_name, class_name), original_class in original_classes.items():
+        module = sys.modules[module_name]
+        setattr(module, class_name, original_class)
+
+
+# _can_record_outputs is a class attribute so patching and unpatching it in the class won't work
+# since the model instance will still reference the original class's _can_record_outputs.
+def patch_output_recorders(model: nn.Module) -> None:
+    """
+    Patch the model instance's output recorders to use the registered replacement classes.
+
+    This function updates output recorders in a model's submodules to use monkey-patched replacement
+    classes. Output recorders are used by the transformers library to track intermediate outputs during
+    forward passes (via the `_can_record_outputs` attribute). When classes are monkey-patched, these
+    recorders need to be updated to reference the new classes.
+
+    This is automatically called during model initialization when loading with `from_pretrained` or
+    `from_config`. You typically don't need to call this manually unless you're constructing models
+    in custom ways.
+
+    Note:
+        The `_can_record_outputs` attribute is a class-level attribute that maps output names to either:
+        - `OutputRecorder` instances that have a `target_class` attribute
+        - Class types directly
+
+        This function patches both cases to use the replacement classes from the monkey patch registry.
+
+    Args:
+        model (`nn.Module`):
+            The model instance whose output recorders should be patched. All submodules will be
+            traversed to find and patch their `_can_record_outputs` attributes.
+
+    Example:
+        ```python
+        from transformers import AutoModelForCausalLM
+        from transformers.monkey_patching import register_patch_mapping, patch_output_recorders
+
+        # Register a patch
+        register_patch_mapping(mapping={"Qwen2MoeExperts": CustomExperts})
+
+        # If you construct a model manually (without from_pretrained), patch recorders
+        model = Qwen2MoeModel(config)
+        patch_output_recorders(model)  # Updates output recorders to use CustomExperts
+        ```
+    """
+
+    mapping = get_patch_mapping()
+    if not mapping:
+        return
+
+    for submodule in model.modules():
+        if hasattr(submodule, "_can_record_outputs") and submodule._can_record_outputs is not None:
+            for output, recorder in submodule._can_record_outputs.items():
+                if isinstance(recorder, OutputRecorder):
+                    # Check if target class matches any registered pattern or exact name
+                    replacement_class = _find_replacement_class(recorder.target_class.__name__, mapping)
+                    if replacement_class is not None:
+                        recorder.target_class = replacement_class
+                elif isinstance(recorder, type):
+                    # Check if class type matches any registered pattern or exact name
+                    replacement_class = _find_replacement_class(recorder.__name__, mapping)
+                    if replacement_class is not None:
+                        submodule._can_record_outputs[output] = replacement_class

filename: tests/test_monkey_patching.py
patch: @@ -0,0 +1,357 @@
+# Copyright 2026 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+import threading
+import unittest
+
+from transformers import is_torch_available
+from transformers.monkey_patching import (
+    apply_patches,
+    clear_patch_mapping,
+    get_patch_mapping,
+    patch_output_recorders,
+    register_patch_mapping,
+    unregister_patch_mapping,
+)
+from transformers.testing_utils import require_torch
+from transformers.utils.output_capturing import OutputRecorder
+
+
+if is_torch_available():
+    import torch.nn as nn
+
+
+@require_torch
+class MonkeyPatchTest(unittest.TestCase):
+    def setUp(self):
+        """Clear any existing patches before each test."""
+        clear_patch_mapping()
+
+    def tearDown(self):
+        """Clean up patches after each test."""
+        clear_patch_mapping()
+
+    def test_register_patch_mapping(self):
+        """Test basic registration of monkey patches."""
+
+        class CustomModule(nn.Module):
+            pass
+
+        # Register a patch
+        register_patch_mapping(mapping={"TestModule": CustomModule})
+
+        # Verify it was registered
+        mapping = get_patch_mapping()
+        self.assertIn("TestModule", mapping)
+        self.assertEqual(mapping["TestModule"], CustomModule)
+
+    def test_register_duplicate_with_overwrite(self):
+        """Test that registering a duplicate class with overwrite=True works."""
+
+        class CustomModule1(nn.Module):
+            pass
+
+        class CustomModule2(nn.Module):
+            pass
+
+        # Register initial patch
+        register_patch_mapping(mapping={"TestModule": CustomModule1})
+
+        # Overwrite with new patch
+        register_patch_mapping(mapping={"TestModule": CustomModule2}, overwrite=True)
+
+        # Verify the new patch is registered
+        mapping = get_patch_mapping()
+        self.assertEqual(mapping["TestModule"], CustomModule2)
+
+    def test_register_non_nn_module_raises_error(self):
+        """Test that registering a non-nn.Module class raises TypeError."""
+
+        class NotAModule:
+            pass
+
+        # Try to register a non-nn.Module class
+        with self.assertRaises(TypeError) as context:
+            register_patch_mapping(mapping={"TestModule": NotAModule})
+
+        self.assertIn("must be a subclass of nn.Module", str(context.exception))
+
+    def test_unregister_patch_mapping(self):
+        """Test unregistering monkey patches."""
+
+        class CustomModule(nn.Module):
+            pass
+
+        # Register and then unregister
+        register_patch_mapping(mapping={"TestModule": CustomModule})
+        unregister_patch_mapping(["TestModule"])
+
+        # Verify it was unregistered
+        mapping = get_patch_mapping()
+        self.assertNotIn("TestModule", mapping)
+
+    def test_unregister_nonexistent_class(self):
+        """Test unregistering a class that doesn't exist raises an error."""
+        # This should raise an error
+        with self.assertRaises(ValueError) as context:
+            unregister_patch_mapping(["NonexistentModule"])
+
+        self.assertIn("not found in monkey patch mapping cache", str(context.exception))
+        self.assertIn("Cannot unregister", str(context.exception))
+
+    def test_clear_patch_mapping(self):
+        """Test clearing all monkey patches."""
+
+        class CustomModule1(nn.Module):
+            pass
+
+        class CustomModule2(nn.Module):
+            pass
+
+        # Register multiple patches
+        register_patch_mapping(mapping={"TestModule1": CustomModule1, "TestModule2": CustomModule2})
+
+        # Clear all patches
+        clear_patch_mapping()
+
+        # Verify all were cleared
+        mapping = get_patch_mapping()
+        self.assertEqual(len(mapping), 0)
+
+    def test_get_patch_mapping_returns_copy(self):
+        """Test that get_patch_mapping returns a copy, not the original."""
+
+        class CustomModule(nn.Module):
+            pass
+
+        register_patch_mapping(mapping={"TestModule": CustomModule})
+
+        # Get mapping and modify it
+        mapping = get_patch_mapping()
+        mapping["NewModule"] = CustomModule
+
+        # Verify the internal cache was not modified
+        internal_mapping = get_patch_mapping()
+        self.assertNotIn("NewModule", internal_mapping)
+
+    def test_apply_patches_context_manager(self):
+        """Test that apply_patches context manager works correctly."""
+
+        class CustomLinear(nn.Linear):
+            pass
+
+        # Create a dummy module in transformers namespace for testing
+        import types
+
+        test_module = types.ModuleType("transformers.test_module")
+        test_module.Linear = nn.Linear
+        sys.modules["transformers.test_module"] = test_module
+
+        try:
+            # Register patch
+            register_patch_mapping(mapping={"Linear": CustomLinear})
+
+            # Outside context, original class should be used
+            self.assertEqual(test_module.Linear, nn.Linear)
+
+            # Inside context, patched class should be used
+            with apply_patches():
+                self.assertEqual(test_module.Linear, CustomLinear)
+
+            # Outside context again, original class should be restored
+            self.assertEqual(test_module.Linear, nn.Linear)
+
+        finally:
+            # Clean up the test module
+            del sys.modules["transformers.test_module"]
+
+    def test_thread_safety_concurrent_access(self):
+        """Test that concurrent reads and writes are thread-safe."""
+
+        class CustomModule(nn.Module):
+            pass
+
+        results = []
+
+        def read_mapping():
+            for _ in range(100):
+                mapping = get_patch_mapping()
+                results.append(len(mapping))
+
+        def write_mapping():
+            for i in range(100):
+                mapping = {f"Module{i}": CustomModule}
+                register_patch_mapping(mapping=mapping)
+
+        # Create threads for reading and writing
+        read_thread = threading.Thread(target=read_mapping)
+        write_thread = threading.Thread(target=write_mapping)
+
+        read_thread.start()
+        write_thread.start()
+
+        read_thread.join()
+        write_thread.join()
+
+        # Test should complete without deadlocks or errors
+        self.assertEqual(len(results), 100)
+
+    def test_patch_output_recorders_with_output_recorder_instance(self):
+        """Test patching output recorders that are OutputRecorder instances."""
+
+        class OriginalModule(nn.Module):
+            pass
+
+        class ReplacementModule(nn.Module):
+            pass
+
+        class TestModel(nn.Module):
+            # Simulate _can_record_outputs with OutputRecorder
+            _can_record_outputs = {"output": OutputRecorder(OriginalModule)}
+
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(10, 10)
+
+        model = TestModel()
+
+        # Register patch
+        register_patch_mapping(mapping={"OriginalModule": ReplacementModule})
+
+        # Patch output recorders
+        patch_output_recorders(model)
+
+        # Verify the recorder's target_class was updated
+        recorder = model._can_record_outputs["output"]
+        self.assertEqual(recorder.target_class, ReplacementModule)
+
+    def test_patch_output_recorders_with_class_type(self):
+        """Test patching output recorders that are class types directly."""
+
+        class OriginalModule(nn.Module):
+            pass
+
+        class ReplacementModule(nn.Module):
+            pass
+
+        class TestModel(nn.Module):
+            # Simulate _can_record_outputs with class type directly
+            _can_record_outputs = {"output": OriginalModule}
+
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(10, 10)
+
+        model = TestModel()
+
+        # Register patch
+        register_patch_mapping(mapping={"OriginalModule": ReplacementModule})
+
+        # Patch output recorders
+        patch_output_recorders(model)
+
+        # Verify the class was updated
+        self.assertEqual(model._can_record_outputs["output"], ReplacementModule)
+
+    def test_pattern_matching_wildcard(self):
+        """Test pattern matching with .* regex (matches any characters)."""
+
+        class CustomAttention(nn.Module):
+            pass
+
+        # Register with pattern
+        register_patch_mapping(mapping={".*Attention": CustomAttention})
+
+        # Create test modules with different attention classes
+        import types
+
+        test_module = types.ModuleType("transformers.test_pattern")
+        test_module.BertAttention = nn.Module
+        test_module.GPT2Attention = nn.Module
+        test_module.T5Attention = nn.Module
+        test_module.SomethingElse = nn.Module
+        sys.modules["transformers.test_pattern"] = test_module
+
+        try:
+            with apply_patches():
+                # All *Attention classes should be patched
+                self.assertEqual(test_module.BertAttention, CustomAttention)
+                self.assertEqual(test_module.GPT2Attention, CustomAttention)
+                self.assertEqual(test_module.T5Attention, CustomAttention)
+                # Non-matching class should not be patched
+                self.assertNotEqual(test_module.SomethingElse, CustomAttention)
+        finally:
+            del sys.modules["transformers.test_pattern"]
+
+    def test_exact_match_precedence_over_pattern(self):
+        """Test that exact matches take precedence over patterns."""
+
+        class PatternReplacement(nn.Module):
+            pass
+
+        class ExactReplacement(nn.Module):
+            pass
+
+        # Register both pattern and exact match
+        register_patch_mapping(mapping={".*Attention": PatternReplacement})
+        register_patch_mapping(mapping={"BertAttention": ExactReplacement})
+
+        import types
+
+        test_module = types.ModuleType("transformers.test_precedence")
+        test_module.BertAttention = nn.Module
+        test_module.GPT2Attention = nn.Module
+        sys.modules["transformers.test_precedence"] = test_module
+
+        try:
+            with apply_patches():
+                # Exact match should take precedence
+                self.assertEqual(test_module.BertAttention, ExactReplacement)
+                # Pattern should still match others
+                self.assertEqual(test_module.GPT2Attention, PatternReplacement)
+        finally:
+            del sys.modules["transformers.test_precedence"]
+
+    def test_pattern_with_output_recorders(self):
+        """Test pattern matching works with output recorders."""
+
+        class OriginalAttention(nn.Module):
+            pass
+
+        class ReplacementAttention(nn.Module):
+            pass
+
+        class TestModel(nn.Module):
+            _can_record_outputs = {"output": OutputRecorder(OriginalAttention)}
+
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(10, 10)
+
+        model = TestModel()
+
+        # Register with pattern
+        register_patch_mapping(mapping={".*Attention": ReplacementAttention})
+
+        # Patch output recorders
+        patch_output_recorders(model)
+
+        # Verify the recorder's target_class was updated via pattern matching
+        recorder = model._can_record_outputs["output"]
+        self.assertEqual(recorder.target_class, ReplacementAttention)
+
+
+if __name__ == "__main__":
+    unittest.main()