
如何用 monkey patching 在 Transformers 中全局替换模型组件【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers当你需要改变模型内部模块的结构或权重布局——例如为量化库兼容而把q_proj/k_proj/v_proj融合成单个qkv_proj或替换 MoE 模型的 experts 模块——直接改modeling_*.py源文件会破坏库的可升级性。Transformers 提供的 monkey patching 功能文档中标注为 experimental feature允许你不动原始模型代码把替换类注册到全局补丁映射里之后任何通过PreTrainedModel.from_pretrained或PreTrainedModel.from_config加载的模型都会在初始化时自动应用补丁。本文以替换LlamaAttention为主路径给出从注册、加载、验证到清理的完整操作路径。实现与测试分别在 src/transformers/monkey_patching.py 和 tests/test_monkey_patching.py 中完整文档见 docs/source/en/monkey_patching.md。先确认是否真的需要 monkey patching文档明确警告monkey patching 应当是最后手段只在你确实需要改变模块布局或权重结构、且无法仅靠自定义 forward 实现时使用。对于多数定制和优化需求文档建议优先考虑以下替代机制Attention interfaceExperts interfaceKernels registry文档给出的适用场景示例量化库兼容quantization library compatibility、融合层fusing layers、架构实验architectural experiments。如果你的需求只是换 attention 计算方式而不动结构应走 attention interface 而不是本文路径。最短路径定义替换类、注册、加载并验证替换类必须继承nn.ModuleAPI 会自动校验传入非法类会得到明确的错误信息。下面以文档中的 Quick start 示例为准定义一个LlamaAttention的子类注册后加载模型补丁在初始化期间自动应用无需额外操作。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最后一行输出是文档示例加载成功后model.model.layers[0].self_attn的类型应显示为你注册的替换类。注意补丁只对transformersmodeling 模块中的类生效例如LlamaAttention、LlamaMLP。管理全局补丁映射用register_patch_mapping注册的补丁是全局的注册后持续存在影响之后所有模型加载直到你清理掉。一次可以注册多个补丁如果某个类名已经注册过需要显式传overwriteTrue才会覆盖否则会报错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 overwriteTrue, )配套的管理函数from transformers.monkey_patching import unregister_patch_mapping, clear_patch_mapping, get_patch_mapping # 注销单个补丁keys 必须与注册时使用的名称或模式完全一致 unregister_patch_mapping(keys[Qwen2MoeExperts]) # 一次注销多个 unregister_patch_mapping(keys[Qwen2MoeExperts, Qwen2MoeAttention]) # 清空所有已注册补丁 clear_patch_mapping() # 查看当前已注册的补丁 current_patches get_patch_mapping() print(current_patches)在测试、notebook 或长驻应用中文档建议用 try/finally 保证清理避免补丁泄漏影响其他代码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用正则模式一次匹配多个类mapping 的 key 既可以是精确类名也可以是正则表达式。模式用re.search()匹配因此.*Attention会命中类名中任意位置包含 Attention 的类用^和$可以锚定首尾。同时注册精确名和模式时精确匹配优先LlamaAttention命中精确替换其余匹配模式.*Attention的类走模式替换。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\\dAttention$: CustomAttention, # Llama2Attention, Llama3Attention, etc. } )文档特别警告宽泛模式会静默地弄坏模型以 BERT 为例BertSelfAttention和BertCrossAttention是内层 attention 实现而外层BertAttention模块内部包含其中一个内层类。用同一个自定义 attention 把三者全替换后外层模块不再包装内层 attention 而是它本身self、output等预期子模块消失模型直接损坏。文档建议优先用窄模式如.*SelfAttention$或精确类名。手动构建模型时需要 apply_patches 上下文管理器自动应用只发生在from_pretrained/from_config路径库内部通过上下文管理器完成。如果你直接Model(config)手动构造必须自己进入apply_patches上下文退出上下文后原始类会被还原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也就是说同一个注册状态from_pretrained永远应用补丁手动构造只在上下文内应用。权重布局变化时注册权重转换一个关键限制monkey patching只替换类不替换权重。如果你的替换类改变了权重布局例如融合 q/k/v 为单个qkv_proj加载预训练权重时会因形状不匹配而失败必须另外注册权重转换映射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_namellama, mapping[ WeightConverter( source_patterns[q_proj, k_proj, v_proj], target_patterns[qkv_proj], operations[ Concatenate(dim0), ], ) ], overwriteTrue, )文档的完整示例Complete example演示了在qwen2_moe上同时重构 expertsQwen2MoeExperts→ 自定义ModuleListExperts和 attentionQwen2MoeAttention→ 融合的FusedQKVAttention并注册覆盖原转换映射的register_checkpoint_conversion_mapping(model_type_or_class_nameqwen2_moe, ...)最后用AutoModelForCausalLM.from_pretrained(Qwen/Qwen1.5-MoE-A2.7B)加载。自定义类需要保持与原模块相同的接口契约。该示例还展示了同一机制可用于 MoE 路由记录与回放RLHF 场景见 monkey_patching.md 中 Recording and replaying MoE expert routing 一节。补丁没生效或不确定是否生效时怎么查文档的 Troubleshooting 部分给出三个检查点核对类名或模式精确名必须与模型源码中完全一致大小写敏感模式必须是合法正则。可以在模型源模块里列出类名确认from transformers.models.llama import modeling_llama print(dir(modeling_llama)) # Look for the class name确认注册成功get_patch_mapping()打印当前全部映射文档示例输出形如{LlamaAttention: class CustomAttention, .*MLP: class CustomMLP}文档示例。检查加载出的模型查单个模块类型或遍历全部模块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)})如果报权重形状不匹配走上一节的register_checkpoint_conversion_mapping路径而不是反复调整补丁类。限制与注意事项该功能在文档标题中即标注为experimental feature。只有transformersmodeling 模块中的类可被补丁替换类必须是nn.Module子类否则注册时抛出带明确信息的错误。注册具有全局性且会持续存在用完务必clear_patch_mapping()。所有注册、注销、应用操作都是线程安全的可以从多线程环境调用。正则匹配基于re.search()不带锚点时会匹配类名任意位置。下一步如果当前需求不满足结构性改变这一门槛回到 Attention interface、Experts interface 或 Kernels registry涉及权重布局转换时参考 weight converter 文档 了解WeightConverter的完整能力。【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考