From b3871906e6cc244bcd7a7a183c311db36ec21c14 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Fri, 10 Apr 2026 14:52:10 -0400 Subject: [PATCH 01/49] feat: module fusion API for kernels --- src/transformers/module_fusion.py | 168 ++++++++++++++++++++++++++++++ 1 file changed, 168 insertions(+) create mode 100644 src/transformers/module_fusion.py diff --git a/src/transformers/module_fusion.py b/src/transformers/module_fusion.py new file mode 100644 index 000000000000..d50576155a8d --- /dev/null +++ b/src/transformers/module_fusion.py @@ -0,0 +1,168 @@ +# 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 functools +import re + +from .utils import is_torch_available + + +if is_torch_available(): + import torch.nn as nn + + +class FusedModuleBase(nn.Module): + def __init__(self, modules_to_fuse: list[nn.Module], source_names: list[str]): + super().__init__() + if len(modules_to_fuse) == 0: + raise ValueError("At least one module must be provided for fusion.") + if len(modules_to_fuse) != len(source_names): + raise ValueError("Length of modules_to_fuse and source_names must match.") + + self._source_names = source_names + + for module in modules_to_fuse: + attr_name = getattr(module, "kernel_layer_name", None) + if attr_name is None: + raise ValueError(f"Module {module} does not have a 'kernel_layer_name' attribute.") + self.add_module(attr_name, module) + + self._fused_module_names = [m.kernel_layer_name for m in modules_to_fuse] + + # `kernelize` validates the kernel's forward signature against the class being replaced. + # Since the fused container sits at the position of the first module in the chain, the + # kernel's forward must match that module's signature. We patch the class-level forward + # here (via `functools.wraps`) so the signature is correct when `kernelize` inspects it. + # The body raises because this forward is always replaced by the kernel before any call. + @functools.wraps(type(modules_to_fuse[0]).forward) + def forward(self, *args, **kwargs): + raise NotImplementedError("FusedModule is a placeholder and should not be called directly.") + + self.__class__.forward = forward + + def __repr__(self): + names = ", ".join(self._fused_module_names) + return f"{self.__class__.__name__}(fused={names})" + + +@functools.cache +def make_fused_module_class(source_layer_names: tuple[str, ...], kernel_layer_name: str) -> type: + """ + Dynamically create and cache a `FusedModuleBase` subclass for a given fusion combination. + + Args: + source_layer_names (`tuple[str, ...]`): + Ordered tuple of `kernel_layer_name` values of the modules being fused + (e.g. ``("RMSNorm", "MLP")``). Used as the cache key — the same combination + always returns the same class object. + kernel_layer_name (`str`): + The name assigned to the fused class, used by `kernelize` to look up the + kernel in the mapping (e.g. ``"RMSNormMLP"``). + + Returns: + A subclass of `FusedModuleBase` with `kernel_layer_name` set as a class attribute. + """ + return type( + f"Fused_{'_'.join(source_layer_names)}", + (FusedModuleBase,), + {"kernel_layer_name": kernel_layer_name}, + ) + + +def fuse_modules( + model: nn.Module, + module_names_to_fuse: list[str], + kernel_layer_name: str, +) -> None: + """ + Fuse a sequence of submodules into a single `FusedModuleBase` subclass in-place. + + For every parent module whose immediate children match all entries in + ``module_names_to_fuse``, the function: + + - replaces the first module with a `FusedModuleBase` subclass instance that holds + all source modules as named children, + - replaces the remaining modules with `nn.Identity()` pass-throughs. + + The fused container's ``forward`` signature is patched to match the first source + module's ``forward``, satisfying the ``kernelize`` signature check. + + Args: + model (`nn.Module`): + The model to modify in-place. + module_names_to_fuse (`list[str]`): + Glob-style paths of the modules to fuse, e.g. + ``["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"]``. + Integer indices are replaced with ``*`` so the same pattern applies to + every repeated block. + kernel_layer_name (`str`): + The ``kernel_layer_name`` assigned to the fused class, used by ``kernelize`` + to look up the kernel in the mapping (e.g. ``"RMSNormMLP"``). + + Example:: + + fuse_modules( + model, + ["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"], + "RMSNormMLP", + ) + """ + pattern = re.compile(r"\d+") + for module_name, module in model.named_modules(): + generic_children = { + re.sub(pattern, "*", f"{module_name}.{n}" if module_name else n): (n, child) + for n, child in module.named_children() + } + if not all(p in generic_children for p in module_names_to_fuse): + continue + + child_names = [generic_children[p][0] for p in module_names_to_fuse] + modules_to_fuse = [generic_children[p][1] for p in module_names_to_fuse] + + source_layer_names = tuple(getattr(m, "kernel_layer_name") for m in modules_to_fuse) + FusedClass = make_fused_module_class(source_layer_names, kernel_layer_name) + fused_instance = FusedClass(modules_to_fuse, child_names) + + module.add_module(child_names[0], fused_instance) + for child_name in child_names[1:]: + module.add_module(child_name, nn.Identity()) + + +def unfuse_modules(model: nn.Module) -> None: + """ + Revert a previous `fuse_modules` call in-place, restoring the original modules. + + For each `FusedModuleBase` instance found in the model tree, the function: + + - restores the original first module at the fused container's position, + - restores the remaining original modules at their original positions + (replacing the `nn.Identity()` pass-throughs). + + Args: + model (`nn.Module`): The model to restore in-place. + + Example:: + + fuse_modules(model, ["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"], "RMSNormMLP") + # ... kernelized forward pass ... + unfuse_modules(model) # back to original + """ + for parent in model.modules(): + for name, child in list(parent.named_children()): + if not isinstance(child, FusedModuleBase): + continue + orig_modules = [getattr(child, layer_name) for layer_name in child._fused_module_names] + parent.add_module(name, orig_modules[0]) + for sibling_name, orig_module in zip(child._source_names[1:], orig_modules[1:]): + parent.add_module(sibling_name, orig_module) From 6bc9402140c508bf15cf7958b6874744494bb1df Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Fri, 10 Apr 2026 15:11:47 -0400 Subject: [PATCH 02/49] fix: improve __repr__ for fused modules --- src/transformers/module_fusion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/module_fusion.py b/src/transformers/module_fusion.py index d50576155a8d..a370b496a94c 100644 --- a/src/transformers/module_fusion.py +++ b/src/transformers/module_fusion.py @@ -53,7 +53,7 @@ def forward(self, *args, **kwargs): def __repr__(self): names = ", ".join(self._fused_module_names) - return f"{self.__class__.__name__}(fused={names})" + return f"{self.__class__.__name__}(fused=({names}))" @functools.cache From 62d4454588f15f2c8b45d4059d0f50a3c1740d4d Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Fri, 10 Apr 2026 16:15:05 -0400 Subject: [PATCH 03/49] wip: integration to KernelConfig --- src/transformers/modeling_utils.py | 4 ++ .../models/qwen3/modeling_qwen3.py | 4 ++ src/transformers/utils/kernel_config.py | 39 +++++++++++++++++++ 3 files changed, 47 insertions(+) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 27fcc3eaae1b..10d6d86f5c8d 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -3662,6 +3662,10 @@ def set_use_kernels(self, use_kernels, kernel_config: KernelConfig | None = None register_kernel_mapping_transformers() if kernel_config is not None and isinstance(kernel_config, KernelConfig): + # For n-to-1 entries (tuple keys), fuse the corresponding modules in the model + # before validation so that FusedModuleBase instances are visible to sanitize. + kernel_config.apply_fusions(self) + # This will make sure the mapping is valid, and the layers are registered in the model kernel_config.sanitize_kernel_mapping(self) diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index 91715a33cf9d..63ecfab0eff1 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -67,6 +67,7 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" +@use_kernel_forward_from_hub("MLP") class Qwen3MLP(nn.Module): def __init__(self, config): super().__init__() @@ -443,6 +444,9 @@ class Qwen3ForCausalLM(Qwen3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_gather_output"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + _kernel_fusion_patterns = { + "RMSNormMLP": ["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"], + } def __init__(self, config): super().__init__(config) diff --git a/src/transformers/utils/kernel_config.py b/src/transformers/utils/kernel_config.py index bb4f965ddbf4..fb45e085a431 100644 --- a/src/transformers/utils/kernel_config.py +++ b/src/transformers/utils/kernel_config.py @@ -116,6 +116,45 @@ def update_kernel(self, repo_id, registered_name, layer_name, device, mode, revi } } + def apply_fusions(self, model): + """ + For each n-to-1 entry (tuple key) in the kernel mapping, find the fusion patterns + registered on the model, fuse the corresponding modules in-place, then replace the + tuple key with the resolved kernel layer name so the rest of the pipeline is unchanged. + """ + from ..module_fusion import fuse_modules + + new_mapping = {} + for layer_name, kernel in self.kernel_mapping.items(): + if not isinstance(layer_name, tuple): + new_mapping[layer_name] = kernel + continue + + # Parse the target kernel layer name from the repo string (the part after ':') + repo_str = kernel if isinstance(kernel, str) else next(iter(kernel.values())) + if isinstance(repo_str, dict): + repo_str = next(iter(repo_str.values())) + kernel_layer_name = repo_str.split(":")[1] if ":" in repo_str else repo_str.split("/")[-1] + + # Only fuse if a kernel is available for the current device + current_device = infer_device(model) + has_kernel_for_device = isinstance(kernel, str) or current_device in kernel + if not has_kernel_for_device: + continue + + # Look up fusion patterns registered on the model class + fusion_patterns = getattr(model, "_kernel_fusion_patterns", {}) + if kernel_layer_name not in fusion_patterns: + raise ValueError( + f"{type(model).__name__} does not define fusion patterns for '{kernel_layer_name}'. " + f'Add `_kernel_fusion_patterns = {{"{kernel_layer_name}": [...]}}` to the model class.' + ) + + fuse_modules(model, fusion_patterns[kernel_layer_name], kernel_layer_name) + new_mapping[kernel_layer_name] = kernel + + self.kernel_mapping = new_mapping + def store_registered_layer_names(self, model): for name, module in model.named_modules(): if hasattr(module, "kernel_layer_name"): From 4082fe15e78f36160d311cc4baa8bc804f274568 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Fri, 10 Apr 2026 16:19:18 -0400 Subject: [PATCH 04/49] wip: add temporary example --- fused_qwen_example.py | 65 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 fused_qwen_example.py diff --git a/fused_qwen_example.py b/fused_qwen_example.py new file mode 100644 index 000000000000..58467cbe3f62 --- /dev/null +++ b/fused_qwen_example.py @@ -0,0 +1,65 @@ +import copy +import torch +import torch.nn as nn + +from kernels import Mode, register_kernel_mapping + +from transformers import AutoModelForCausalLM, AutoTokenizer, KernelConfig +from transformers.module_fusion import unfuse_modules + + +model_id = "michaelbenayoun/qwen3-tiny-4kv-heads-4layers-random" +tokenizer = AutoTokenizer.from_pretrained(model_id) + + +class FakeRMSNormMLP(nn.Module): + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + print("Using fake RMSNormMLP kernel") + hidden_states = self.RMSNorm(hidden_states) + hidden_states = self.MLP(hidden_states) + return hidden_states + + +class _InMemoryRepo: + """Minimal fake repository that returns a local class instead of downloading from the hub.""" + + def __init__(self, layer_cls: type): + self.layer_name = layer_cls.__name__ + self._layer_cls = layer_cls + + def load(self) -> type: + return self._layer_cls + + def __hash__(self): + return hash(self._layer_cls) + + def __eq__(self, other): + return isinstance(other, _InMemoryRepo) and self._layer_cls is other._layer_cls + + +# In production this would be a real hub repo string e.g. "kernels-community/rmsnorm-mlp:RMSNormMLP". +# For testing we pre-register a fake in-memory kernel so no hub download is needed. +register_kernel_mapping({ + "RMSNormMLP": { + "cuda": {Mode.INFERENCE: _InMemoryRepo(FakeRMSNormMLP)}, + } +}) + +kernel_config = KernelConfig({ + ("RMSNorm", "MLP"): "fake/repo:RMSNormMLP", +}) + +model = AutoModelForCausalLM.from_pretrained(model_id, use_kernels=True, kernel_config=kernel_config, device_map="cuda") +model.eval() + +input_ids = tokenizer("Hello, how are you?", return_tensors="pt").input_ids.to(model.device) + +original_model = copy.deepcopy(model) +unfuse_modules(original_model) +original_model.eval() + +with torch.no_grad(): + fused_out = model(input_ids).logits + original_out = original_model(input_ids).logits + +print("Max diff fused vs original:", (fused_out - original_out).abs().max().item()) From ac4a6994a35cc054c557909546434110014de0dc Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Mon, 13 Apr 2026 15:15:56 -0400 Subject: [PATCH 05/49] wip: pattern matching in KernelConfig and actual kernel repo --- fused_qwen_example.py | 43 +------- .../models/qwen3/modeling_qwen3.py | 4 - src/transformers/module_fusion.py | 101 ++++++++++++++++-- src/transformers/utils/kernel_config.py | 30 ++++-- 4 files changed, 116 insertions(+), 62 deletions(-) diff --git a/fused_qwen_example.py b/fused_qwen_example.py index 58467cbe3f62..33a00915811a 100644 --- a/fused_qwen_example.py +++ b/fused_qwen_example.py @@ -1,8 +1,5 @@ import copy import torch -import torch.nn as nn - -from kernels import Mode, register_kernel_mapping from transformers import AutoModelForCausalLM, AutoTokenizer, KernelConfig from transformers.module_fusion import unfuse_modules @@ -11,46 +8,14 @@ model_id = "michaelbenayoun/qwen3-tiny-4kv-heads-4layers-random" tokenizer = AutoTokenizer.from_pretrained(model_id) - -class FakeRMSNormMLP(nn.Module): - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - print("Using fake RMSNormMLP kernel") - hidden_states = self.RMSNorm(hidden_states) - hidden_states = self.MLP(hidden_states) - return hidden_states - - -class _InMemoryRepo: - """Minimal fake repository that returns a local class instead of downloading from the hub.""" - - def __init__(self, layer_cls: type): - self.layer_name = layer_cls.__name__ - self._layer_cls = layer_cls - - def load(self) -> type: - return self._layer_cls - - def __hash__(self): - return hash(self._layer_cls) - - def __eq__(self, other): - return isinstance(other, _InMemoryRepo) and self._layer_cls is other._layer_cls - - -# In production this would be a real hub repo string e.g. "kernels-community/rmsnorm-mlp:RMSNormMLP". -# For testing we pre-register a fake in-memory kernel so no hub download is needed. -register_kernel_mapping({ - "RMSNormMLP": { - "cuda": {Mode.INFERENCE: _InMemoryRepo(FakeRMSNormMLP)}, - } -}) - kernel_config = KernelConfig({ - ("RMSNorm", "MLP"): "fake/repo:RMSNormMLP", + ( + ("RMSNorm", "model.layers.*.post_attention_layernorm"), + ("MLP", "model.layers.*.mlp"), + ): "michaelbenayoun/dummy-rmsnorm-mlp:RMSNormMLP", }) model = AutoModelForCausalLM.from_pretrained(model_id, use_kernels=True, kernel_config=kernel_config, device_map="cuda") -model.eval() input_ids = tokenizer("Hello, how are you?", return_tensors="pt").input_ids.to(model.device) diff --git a/src/transformers/models/qwen3/modeling_qwen3.py b/src/transformers/models/qwen3/modeling_qwen3.py index 63ecfab0eff1..91715a33cf9d 100644 --- a/src/transformers/models/qwen3/modeling_qwen3.py +++ b/src/transformers/models/qwen3/modeling_qwen3.py @@ -67,7 +67,6 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" -@use_kernel_forward_from_hub("MLP") class Qwen3MLP(nn.Module): def __init__(self, config): super().__init__() @@ -444,9 +443,6 @@ class Qwen3ForCausalLM(Qwen3PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} _tp_plan = {"lm_head": "colwise_gather_output"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} - _kernel_fusion_patterns = { - "RMSNormMLP": ["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"], - } def __init__(self, config): super().__init__(config) diff --git a/src/transformers/module_fusion.py b/src/transformers/module_fusion.py index a370b496a94c..c0bcf76ee427 100644 --- a/src/transformers/module_fusion.py +++ b/src/transformers/module_fusion.py @@ -22,8 +22,72 @@ import torch.nn as nn +# Module-level registry: model class → {kernel_layer_name → [glob patterns]} +# Populated via `register_fusion_patterns` for models that cannot be modified directly. +_FUSION_PATTERNS_REGISTRY: dict[type, dict[str, list[str]]] = {} + + +def register_fusion_patterns( + model_class_or_instance, + patterns: dict[str, list[str]], +) -> None: + """ + Register kernel fusion patterns for a model class without modifying it directly. + + This is an alternative to setting ``_kernel_fusion_patterns`` as a class attribute, + useful when the model class is frozen or comes from an external library. + + Args: + model_class_or_instance: + The model class (or an instance of it) for which patterns are being registered. + patterns (`dict[str, list[str]]`): + Mapping from ``kernel_layer_name`` to a list of glob-style module paths, + identical in format to ``_kernel_fusion_patterns``. For example:: + + { + "RMSNormMLP": [ + "model.layers.*.post_attention_layernorm", + "model.layers.*.mlp", + ] + } + + Example:: + + from transformers.module_fusion import register_fusion_patterns + from transformers.models.qwen3 import Qwen3ForCausalLM + + register_fusion_patterns( + Qwen3ForCausalLM, + { + "RMSNormMLP": [ + "model.layers.*.post_attention_layernorm", + "model.layers.*.mlp", + ] + }, + ) + """ + if not isinstance(model_class_or_instance, type): + model_class_or_instance = type(model_class_or_instance) + _FUSION_PATTERNS_REGISTRY[model_class_or_instance] = patterns + + class FusedModuleBase(nn.Module): - def __init__(self, modules_to_fuse: list[nn.Module], source_names: list[str]): + def __init__( + self, + modules_to_fuse: list[nn.Module], + source_names: list[str], + fused_module_names: list[str] | None = None, + ): + """ + Args: + modules_to_fuse: The source modules to fuse together. + source_names: The attribute names under which each module lives in its parent + (used to restore them on ``unfuse_modules``). + fused_module_names: The names under which each source module is registered as a + child of this container (i.e. ``self.``). When ``None``, the + ``kernel_layer_name`` attribute of each source module is used. Pass this + explicitly when the source modules do not carry ``@use_kernel_forward_from_hub``. + """ super().__init__() if len(modules_to_fuse) == 0: raise ValueError("At least one module must be provided for fusion.") @@ -32,13 +96,24 @@ def __init__(self, modules_to_fuse: list[nn.Module], source_names: list[str]): self._source_names = source_names - for module in modules_to_fuse: - attr_name = getattr(module, "kernel_layer_name", None) - if attr_name is None: - raise ValueError(f"Module {module} does not have a 'kernel_layer_name' attribute.") - self.add_module(attr_name, module) - - self._fused_module_names = [m.kernel_layer_name for m in modules_to_fuse] + if fused_module_names is not None: + if len(fused_module_names) != len(modules_to_fuse): + raise ValueError("Length of fused_module_names and modules_to_fuse must match.") + for module, name in zip(modules_to_fuse, fused_module_names): + self.add_module(name, module) + self._fused_module_names = list(fused_module_names) + else: + for module in modules_to_fuse: + attr_name = getattr(module, "kernel_layer_name", None) + if attr_name is None: + raise ValueError( + f"Module {module} does not have a 'kernel_layer_name' attribute. " + f"Either decorate it with @use_kernel_forward_from_hub or provide " + f"explicit names via the inline pattern format: " + f'(("", ""), ...).' + ) + self.add_module(attr_name, module) + self._fused_module_names = [m.kernel_layer_name for m in modules_to_fuse] # `kernelize` validates the kernel's forward signature against the class being replaced. # Since the fused container sits at the position of the first module in the chain, the @@ -84,6 +159,7 @@ def fuse_modules( model: nn.Module, module_names_to_fuse: list[str], kernel_layer_name: str, + source_layer_names: list[str] | None = None, ) -> None: """ Fuse a sequence of submodules into a single `FusedModuleBase` subclass in-place. @@ -130,9 +206,12 @@ def fuse_modules( child_names = [generic_children[p][0] for p in module_names_to_fuse] modules_to_fuse = [generic_children[p][1] for p in module_names_to_fuse] - source_layer_names = tuple(getattr(m, "kernel_layer_name") for m in modules_to_fuse) - FusedClass = make_fused_module_class(source_layer_names, kernel_layer_name) - fused_instance = FusedClass(modules_to_fuse, child_names) + if source_layer_names is not None: + resolved_names = tuple(source_layer_names) + else: + resolved_names = tuple(getattr(m, "kernel_layer_name") for m in modules_to_fuse) + FusedClass = make_fused_module_class(resolved_names, kernel_layer_name) + fused_instance = FusedClass(modules_to_fuse, child_names, fused_module_names=list(resolved_names)) module.add_module(child_names[0], fused_instance) for child_name in child_names[1:]: diff --git a/src/transformers/utils/kernel_config.py b/src/transformers/utils/kernel_config.py index fb45e085a431..8f66411480c0 100644 --- a/src/transformers/utils/kernel_config.py +++ b/src/transformers/utils/kernel_config.py @@ -142,15 +142,29 @@ def apply_fusions(self, model): if not has_kernel_for_device: continue - # Look up fusion patterns registered on the model class - fusion_patterns = getattr(model, "_kernel_fusion_patterns", {}) - if kernel_layer_name not in fusion_patterns: - raise ValueError( - f"{type(model).__name__} does not define fusion patterns for '{kernel_layer_name}'. " - f'Add `_kernel_fusion_patterns = {{"{kernel_layer_name}": [...]}}` to the model class.' - ) + # Detect inline format: tuple of (kernel_layer_name, glob_pattern) pairs. + # e.g. (("RMSNorm", "model.layers.*.post_attention_layernorm"), ("MLP", "model.layers.*.mlp")) + is_inline = all(isinstance(item, tuple) and len(item) == 2 for item in layer_name) + if is_inline: + source_names = [item[0] for item in layer_name] + patterns = [item[1] for item in layer_name] + fuse_modules(model, patterns, kernel_layer_name, source_layer_names=source_names) + else: + # Legacy format: ("RMSNorm", "MLP") — look up patterns from model class or registry. + from ..module_fusion import _FUSION_PATTERNS_REGISTRY - fuse_modules(model, fusion_patterns[kernel_layer_name], kernel_layer_name) + fusion_patterns = getattr(model, "_kernel_fusion_patterns", None) or _FUSION_PATTERNS_REGISTRY.get( + type(model), {} + ) + if kernel_layer_name not in fusion_patterns: + raise ValueError( + f"{type(model).__name__} does not define fusion patterns for '{kernel_layer_name}'. " + f'Either add `_kernel_fusion_patterns = {{"{kernel_layer_name}": [...]}}` to the model class, ' + f"call `register_fusion_patterns({type(model).__name__}, ...)` before loading the model, " + f"or use the inline pattern format: " + f'(("{kernel_layer_name}", ""), ...).' + ) + fuse_modules(model, fusion_patterns[kernel_layer_name], kernel_layer_name) new_mapping[kernel_layer_name] = kernel self.kernel_mapping = new_mapping From e13111fbf358335fff766842d8374a96e9fb7ca9 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Mon, 13 Apr 2026 15:21:59 -0400 Subject: [PATCH 06/49] refactor: move relevant code to hub_kernels.py --- fused_qwen_example.py | 2 +- src/transformers/integrations/__init__.py | 10 + src/transformers/integrations/hub_kernels.py | 253 ++++++++++++++++++- src/transformers/module_fusion.py | 247 ------------------ src/transformers/utils/kernel_config.py | 5 +- 5 files changed, 260 insertions(+), 257 deletions(-) delete mode 100644 src/transformers/module_fusion.py diff --git a/fused_qwen_example.py b/fused_qwen_example.py index 33a00915811a..0b7abd98439c 100644 --- a/fused_qwen_example.py +++ b/fused_qwen_example.py @@ -2,7 +2,7 @@ import torch from transformers import AutoModelForCausalLM, AutoTokenizer, KernelConfig -from transformers.module_fusion import unfuse_modules +from transformers.integrations import unfuse_modules model_id = "michaelbenayoun/qwen3-tiny-4kv-heads-4layers-random" diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index 336db3773f76..f74836d890f7 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -67,10 +67,15 @@ ], "hqq": ["prepare_for_hqq_linear"], "hub_kernels": [ + "FusedModuleBase", "LayerRepository", + "fuse_modules", "lazy_load_kernel", + "make_fused_module_class", + "register_fusion_patterns", "register_kernel_mapping", "replace_kernel_forward_from_hub", + "unfuse_modules", "use_kernel_forward_from_hub", "use_kernel_func_from_hub", "use_kernelized_func", @@ -221,10 +226,15 @@ from .higgs import HiggsLinear, dequantize_higgs, quantize_with_higgs, replace_with_higgs_linear from .hqq import prepare_for_hqq_linear from .hub_kernels import ( + FusedModuleBase, LayerRepository, + fuse_modules, lazy_load_kernel, + make_fused_module_class, + register_fusion_patterns, register_kernel_mapping, replace_kernel_forward_from_hub, + unfuse_modules, use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func, diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 39c1448af02b..fb5f2064917b 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -11,6 +11,7 @@ # 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 functools import importlib.metadata import os import re @@ -21,10 +22,14 @@ from packaging import version as pkg_version from ..utils import ENV_VARS_TRUE_VALUES, logging -from ..utils.import_utils import is_kernels_available +from ..utils.import_utils import is_kernels_available, is_torch_available from .flash_attention import flash_attention_forward +if is_torch_available(): + import torch.nn as nn + + logger = logging.get_logger(__name__) try: @@ -483,14 +488,252 @@ def allow_all_hub_kernels(): ALLOW_ALL_KERNELS = False +# --------------------------------------------------------------------------- +# Module fusion helpers +# --------------------------------------------------------------------------- + +# Model class → {kernel_layer_name → [glob patterns]} +# Populated via `register_fusion_patterns` for models that cannot be modified directly. +_FUSION_PATTERNS_REGISTRY: dict[type, dict[str, list[str]]] = {} + + +def register_fusion_patterns( + model_class_or_instance, + patterns: dict[str, list[str]], +) -> None: + """ + Register kernel fusion patterns for a model class without modifying it directly. + + This is an alternative to setting ``_kernel_fusion_patterns`` as a class attribute, + useful when the model class is frozen or comes from an external library. + + Args: + model_class_or_instance: + The model class (or an instance of it) for which patterns are being registered. + patterns (`dict[str, list[str]]`): + Mapping from ``kernel_layer_name`` to a list of glob-style module paths, + identical in format to ``_kernel_fusion_patterns``. For example:: + + { + "RMSNormMLP": [ + "model.layers.*.post_attention_layernorm", + "model.layers.*.mlp", + ] + } + + Example:: + + from transformers.integrations import register_fusion_patterns + from transformers.models.qwen3 import Qwen3ForCausalLM + + register_fusion_patterns( + Qwen3ForCausalLM, + { + "RMSNormMLP": [ + "model.layers.*.post_attention_layernorm", + "model.layers.*.mlp", + ] + }, + ) + """ + if not isinstance(model_class_or_instance, type): + model_class_or_instance = type(model_class_or_instance) + _FUSION_PATTERNS_REGISTRY[model_class_or_instance] = patterns + + +class FusedModuleBase(nn.Module): + def __init__( + self, + modules_to_fuse: list["nn.Module"], + source_names: list[str], + fused_module_names: list[str] | None = None, + ): + """ + Args: + modules_to_fuse: The source modules to fuse together. + source_names: The attribute names under which each module lives in its parent + (used to restore them on ``unfuse_modules``). + fused_module_names: The names under which each source module is registered as a + child of this container (i.e. ``self.``). When ``None``, the + ``kernel_layer_name`` attribute of each source module is used. Pass this + explicitly when the source modules do not carry ``@use_kernel_forward_from_hub``. + """ + super().__init__() + if len(modules_to_fuse) == 0: + raise ValueError("At least one module must be provided for fusion.") + if len(modules_to_fuse) != len(source_names): + raise ValueError("Length of modules_to_fuse and source_names must match.") + + self._source_names = source_names + + if fused_module_names is not None: + if len(fused_module_names) != len(modules_to_fuse): + raise ValueError("Length of fused_module_names and modules_to_fuse must match.") + for module, name in zip(modules_to_fuse, fused_module_names): + self.add_module(name, module) + self._fused_module_names = list(fused_module_names) + else: + for module in modules_to_fuse: + attr_name = getattr(module, "kernel_layer_name", None) + if attr_name is None: + raise ValueError( + f"Module {module} does not have a 'kernel_layer_name' attribute. " + f"Either decorate it with @use_kernel_forward_from_hub or provide " + f"explicit names via the inline pattern format: " + f'(("", ""), ...).' + ) + self.add_module(attr_name, module) + self._fused_module_names = [m.kernel_layer_name for m in modules_to_fuse] + + # `kernelize` validates the kernel's forward signature against the class being replaced. + # Since the fused container sits at the position of the first module in the chain, the + # kernel's forward must match that module's signature. We patch the class-level forward + # here (via `functools.wraps`) so the signature is correct when `kernelize` inspects it. + # The body raises because this forward is always replaced by the kernel before any call. + @functools.wraps(type(modules_to_fuse[0]).forward) + def forward(self, *args, **kwargs): + raise NotImplementedError("FusedModule is a placeholder and should not be called directly.") + + self.__class__.forward = forward + + def __repr__(self): + names = ", ".join(self._fused_module_names) + return f"{self.__class__.__name__}(fused=({names}))" + + +@functools.cache +def make_fused_module_class(source_layer_names: tuple[str, ...], kernel_layer_name: str) -> type: + """ + Dynamically create and cache a `FusedModuleBase` subclass for a given fusion combination. + + Args: + source_layer_names (`tuple[str, ...]`): + Ordered tuple of `kernel_layer_name` values of the modules being fused + (e.g. ``("RMSNorm", "MLP")``). Used as the cache key — the same combination + always returns the same class object. + kernel_layer_name (`str`): + The name assigned to the fused class, used by `kernelize` to look up the + kernel in the mapping (e.g. ``"RMSNormMLP"``). + + Returns: + A subclass of `FusedModuleBase` with `kernel_layer_name` set as a class attribute. + """ + return type( + f"Fused_{'_'.join(source_layer_names)}", + (FusedModuleBase,), + {"kernel_layer_name": kernel_layer_name}, + ) + + +def fuse_modules( + model: "nn.Module", + module_names_to_fuse: list[str], + kernel_layer_name: str, + source_layer_names: list[str] | None = None, +) -> None: + """ + Fuse a sequence of submodules into a single `FusedModuleBase` subclass in-place. + + For every parent module whose immediate children match all entries in + ``module_names_to_fuse``, the function: + + - replaces the first module with a `FusedModuleBase` subclass instance that holds + all source modules as named children, + - replaces the remaining modules with `nn.Identity()` pass-throughs. + + The fused container's ``forward`` signature is patched to match the first source + module's ``forward``, satisfying the ``kernelize`` signature check. + + Args: + model (`nn.Module`): + The model to modify in-place. + module_names_to_fuse (`list[str]`): + Glob-style paths of the modules to fuse, e.g. + ``["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"]``. + Integer indices are replaced with ``*`` so the same pattern applies to + every repeated block. + kernel_layer_name (`str`): + The ``kernel_layer_name`` assigned to the fused class, used by ``kernelize`` + to look up the kernel in the mapping (e.g. ``"RMSNormMLP"``). + source_layer_names (`list[str]`, *optional*): + Explicit names for the child modules inside the fused container + (e.g. ``["RMSNorm", "MLP"]``). When ``None``, the ``kernel_layer_name`` + attribute of each source module is used. + + Example:: + + fuse_modules( + model, + ["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"], + "RMSNormMLP", + ) + """ + pattern = re.compile(r"\d+") + for module_name, module in model.named_modules(): + generic_children = { + re.sub(pattern, "*", f"{module_name}.{n}" if module_name else n): (n, child) + for n, child in module.named_children() + } + if not all(p in generic_children for p in module_names_to_fuse): + continue + + child_names = [generic_children[p][0] for p in module_names_to_fuse] + modules_to_fuse = [generic_children[p][1] for p in module_names_to_fuse] + + if source_layer_names is not None: + resolved_names = tuple(source_layer_names) + else: + resolved_names = tuple(getattr(m, "kernel_layer_name") for m in modules_to_fuse) + FusedClass = make_fused_module_class(resolved_names, kernel_layer_name) + fused_instance = FusedClass(modules_to_fuse, child_names, fused_module_names=list(resolved_names)) + + module.add_module(child_names[0], fused_instance) + for child_name in child_names[1:]: + module.add_module(child_name, nn.Identity()) + + +def unfuse_modules(model: "nn.Module") -> None: + """ + Revert a previous `fuse_modules` call in-place, restoring the original modules. + + For each `FusedModuleBase` instance found in the model tree, the function: + + - restores the original first module at the fused container's position, + - restores the remaining original modules at their original positions + (replacing the `nn.Identity()` pass-throughs). + + Args: + model (`nn.Module`): The model to restore in-place. + + Example:: + + fuse_modules(model, ["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"], "RMSNormMLP") + # ... kernelized forward pass ... + unfuse_modules(model) # back to original + """ + for parent in model.modules(): + for name, child in list(parent.named_children()): + if not isinstance(child, FusedModuleBase): + continue + orig_modules = [getattr(child, layer_name) for layer_name in child._fused_module_names] + parent.add_module(name, orig_modules[0]) + for sibling_name, orig_module in zip(child._source_names[1:], orig_modules[1:]): + parent.add_module(sibling_name, orig_module) + + __all__ = [ + "FusedModuleBase", "LayerRepository", - "use_kernel_forward_from_hub", - "use_kernel_func_from_hub", + "fuse_modules", + "get_kernel", + "lazy_load_kernel", + "make_fused_module_class", + "register_fusion_patterns", "register_kernel_mapping", "register_kernel_mapping_transformers", "replace_kernel_forward_from_hub", - "lazy_load_kernel", - "get_kernel", + "unfuse_modules", + "use_kernel_forward_from_hub", + "use_kernel_func_from_hub", "use_kernelized_func", ] # type: ignore diff --git a/src/transformers/module_fusion.py b/src/transformers/module_fusion.py deleted file mode 100644 index c0bcf76ee427..000000000000 --- a/src/transformers/module_fusion.py +++ /dev/null @@ -1,247 +0,0 @@ -# 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 functools -import re - -from .utils import is_torch_available - - -if is_torch_available(): - import torch.nn as nn - - -# Module-level registry: model class → {kernel_layer_name → [glob patterns]} -# Populated via `register_fusion_patterns` for models that cannot be modified directly. -_FUSION_PATTERNS_REGISTRY: dict[type, dict[str, list[str]]] = {} - - -def register_fusion_patterns( - model_class_or_instance, - patterns: dict[str, list[str]], -) -> None: - """ - Register kernel fusion patterns for a model class without modifying it directly. - - This is an alternative to setting ``_kernel_fusion_patterns`` as a class attribute, - useful when the model class is frozen or comes from an external library. - - Args: - model_class_or_instance: - The model class (or an instance of it) for which patterns are being registered. - patterns (`dict[str, list[str]]`): - Mapping from ``kernel_layer_name`` to a list of glob-style module paths, - identical in format to ``_kernel_fusion_patterns``. For example:: - - { - "RMSNormMLP": [ - "model.layers.*.post_attention_layernorm", - "model.layers.*.mlp", - ] - } - - Example:: - - from transformers.module_fusion import register_fusion_patterns - from transformers.models.qwen3 import Qwen3ForCausalLM - - register_fusion_patterns( - Qwen3ForCausalLM, - { - "RMSNormMLP": [ - "model.layers.*.post_attention_layernorm", - "model.layers.*.mlp", - ] - }, - ) - """ - if not isinstance(model_class_or_instance, type): - model_class_or_instance = type(model_class_or_instance) - _FUSION_PATTERNS_REGISTRY[model_class_or_instance] = patterns - - -class FusedModuleBase(nn.Module): - def __init__( - self, - modules_to_fuse: list[nn.Module], - source_names: list[str], - fused_module_names: list[str] | None = None, - ): - """ - Args: - modules_to_fuse: The source modules to fuse together. - source_names: The attribute names under which each module lives in its parent - (used to restore them on ``unfuse_modules``). - fused_module_names: The names under which each source module is registered as a - child of this container (i.e. ``self.``). When ``None``, the - ``kernel_layer_name`` attribute of each source module is used. Pass this - explicitly when the source modules do not carry ``@use_kernel_forward_from_hub``. - """ - super().__init__() - if len(modules_to_fuse) == 0: - raise ValueError("At least one module must be provided for fusion.") - if len(modules_to_fuse) != len(source_names): - raise ValueError("Length of modules_to_fuse and source_names must match.") - - self._source_names = source_names - - if fused_module_names is not None: - if len(fused_module_names) != len(modules_to_fuse): - raise ValueError("Length of fused_module_names and modules_to_fuse must match.") - for module, name in zip(modules_to_fuse, fused_module_names): - self.add_module(name, module) - self._fused_module_names = list(fused_module_names) - else: - for module in modules_to_fuse: - attr_name = getattr(module, "kernel_layer_name", None) - if attr_name is None: - raise ValueError( - f"Module {module} does not have a 'kernel_layer_name' attribute. " - f"Either decorate it with @use_kernel_forward_from_hub or provide " - f"explicit names via the inline pattern format: " - f'(("", ""), ...).' - ) - self.add_module(attr_name, module) - self._fused_module_names = [m.kernel_layer_name for m in modules_to_fuse] - - # `kernelize` validates the kernel's forward signature against the class being replaced. - # Since the fused container sits at the position of the first module in the chain, the - # kernel's forward must match that module's signature. We patch the class-level forward - # here (via `functools.wraps`) so the signature is correct when `kernelize` inspects it. - # The body raises because this forward is always replaced by the kernel before any call. - @functools.wraps(type(modules_to_fuse[0]).forward) - def forward(self, *args, **kwargs): - raise NotImplementedError("FusedModule is a placeholder and should not be called directly.") - - self.__class__.forward = forward - - def __repr__(self): - names = ", ".join(self._fused_module_names) - return f"{self.__class__.__name__}(fused=({names}))" - - -@functools.cache -def make_fused_module_class(source_layer_names: tuple[str, ...], kernel_layer_name: str) -> type: - """ - Dynamically create and cache a `FusedModuleBase` subclass for a given fusion combination. - - Args: - source_layer_names (`tuple[str, ...]`): - Ordered tuple of `kernel_layer_name` values of the modules being fused - (e.g. ``("RMSNorm", "MLP")``). Used as the cache key — the same combination - always returns the same class object. - kernel_layer_name (`str`): - The name assigned to the fused class, used by `kernelize` to look up the - kernel in the mapping (e.g. ``"RMSNormMLP"``). - - Returns: - A subclass of `FusedModuleBase` with `kernel_layer_name` set as a class attribute. - """ - return type( - f"Fused_{'_'.join(source_layer_names)}", - (FusedModuleBase,), - {"kernel_layer_name": kernel_layer_name}, - ) - - -def fuse_modules( - model: nn.Module, - module_names_to_fuse: list[str], - kernel_layer_name: str, - source_layer_names: list[str] | None = None, -) -> None: - """ - Fuse a sequence of submodules into a single `FusedModuleBase` subclass in-place. - - For every parent module whose immediate children match all entries in - ``module_names_to_fuse``, the function: - - - replaces the first module with a `FusedModuleBase` subclass instance that holds - all source modules as named children, - - replaces the remaining modules with `nn.Identity()` pass-throughs. - - The fused container's ``forward`` signature is patched to match the first source - module's ``forward``, satisfying the ``kernelize`` signature check. - - Args: - model (`nn.Module`): - The model to modify in-place. - module_names_to_fuse (`list[str]`): - Glob-style paths of the modules to fuse, e.g. - ``["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"]``. - Integer indices are replaced with ``*`` so the same pattern applies to - every repeated block. - kernel_layer_name (`str`): - The ``kernel_layer_name`` assigned to the fused class, used by ``kernelize`` - to look up the kernel in the mapping (e.g. ``"RMSNormMLP"``). - - Example:: - - fuse_modules( - model, - ["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"], - "RMSNormMLP", - ) - """ - pattern = re.compile(r"\d+") - for module_name, module in model.named_modules(): - generic_children = { - re.sub(pattern, "*", f"{module_name}.{n}" if module_name else n): (n, child) - for n, child in module.named_children() - } - if not all(p in generic_children for p in module_names_to_fuse): - continue - - child_names = [generic_children[p][0] for p in module_names_to_fuse] - modules_to_fuse = [generic_children[p][1] for p in module_names_to_fuse] - - if source_layer_names is not None: - resolved_names = tuple(source_layer_names) - else: - resolved_names = tuple(getattr(m, "kernel_layer_name") for m in modules_to_fuse) - FusedClass = make_fused_module_class(resolved_names, kernel_layer_name) - fused_instance = FusedClass(modules_to_fuse, child_names, fused_module_names=list(resolved_names)) - - module.add_module(child_names[0], fused_instance) - for child_name in child_names[1:]: - module.add_module(child_name, nn.Identity()) - - -def unfuse_modules(model: nn.Module) -> None: - """ - Revert a previous `fuse_modules` call in-place, restoring the original modules. - - For each `FusedModuleBase` instance found in the model tree, the function: - - - restores the original first module at the fused container's position, - - restores the remaining original modules at their original positions - (replacing the `nn.Identity()` pass-throughs). - - Args: - model (`nn.Module`): The model to restore in-place. - - Example:: - - fuse_modules(model, ["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"], "RMSNormMLP") - # ... kernelized forward pass ... - unfuse_modules(model) # back to original - """ - for parent in model.modules(): - for name, child in list(parent.named_children()): - if not isinstance(child, FusedModuleBase): - continue - orig_modules = [getattr(child, layer_name) for layer_name in child._fused_module_names] - parent.add_module(name, orig_modules[0]) - for sibling_name, orig_module in zip(child._source_names[1:], orig_modules[1:]): - parent.add_module(sibling_name, orig_module) diff --git a/src/transformers/utils/kernel_config.py b/src/transformers/utils/kernel_config.py index 8f66411480c0..1c7541fbcbc2 100644 --- a/src/transformers/utils/kernel_config.py +++ b/src/transformers/utils/kernel_config.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from ..integrations.hub_kernels import _FUSION_PATTERNS_REGISTRY, fuse_modules from ..utils import PushToHubMixin @@ -122,8 +123,6 @@ def apply_fusions(self, model): registered on the model, fuse the corresponding modules in-place, then replace the tuple key with the resolved kernel layer name so the rest of the pipeline is unchanged. """ - from ..module_fusion import fuse_modules - new_mapping = {} for layer_name, kernel in self.kernel_mapping.items(): if not isinstance(layer_name, tuple): @@ -151,8 +150,6 @@ def apply_fusions(self, model): fuse_modules(model, patterns, kernel_layer_name, source_layer_names=source_names) else: # Legacy format: ("RMSNorm", "MLP") — look up patterns from model class or registry. - from ..module_fusion import _FUSION_PATTERNS_REGISTRY - fusion_patterns = getattr(model, "_kernel_fusion_patterns", None) or _FUSION_PATTERNS_REGISTRY.get( type(model), {} ) From d9d53f08e40d98f9acfbaa6e2814169b983e1651 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Mon, 13 Apr 2026 15:32:00 -0400 Subject: [PATCH 07/49] docs: reformat docstring --- src/transformers/integrations/hub_kernels.py | 74 +++++--------------- 1 file changed, 17 insertions(+), 57 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index fb5f2064917b..d1cc48c26093 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -488,10 +488,6 @@ def allow_all_hub_kernels(): ALLOW_ALL_KERNELS = False -# --------------------------------------------------------------------------- -# Module fusion helpers -# --------------------------------------------------------------------------- - # Model class → {kernel_layer_name → [glob patterns]} # Populated via `register_fusion_patterns` for models that cannot be modified directly. _FUSION_PATTERNS_REGISTRY: dict[type, dict[str, list[str]]] = {} @@ -504,37 +500,15 @@ def register_fusion_patterns( """ Register kernel fusion patterns for a model class without modifying it directly. - This is an alternative to setting ``_kernel_fusion_patterns`` as a class attribute, + This is an alternative to setting `_kernel_fusion_patterns` as a class attribute, useful when the model class is frozen or comes from an external library. Args: model_class_or_instance: The model class (or an instance of it) for which patterns are being registered. patterns (`dict[str, list[str]]`): - Mapping from ``kernel_layer_name`` to a list of glob-style module paths, - identical in format to ``_kernel_fusion_patterns``. For example:: - - { - "RMSNormMLP": [ - "model.layers.*.post_attention_layernorm", - "model.layers.*.mlp", - ] - } - - Example:: - - from transformers.integrations import register_fusion_patterns - from transformers.models.qwen3 import Qwen3ForCausalLM - - register_fusion_patterns( - Qwen3ForCausalLM, - { - "RMSNormMLP": [ - "model.layers.*.post_attention_layernorm", - "model.layers.*.mlp", - ] - }, - ) + Mapping from `kernel_layer_name` to a list of glob-style module paths, + identical in format to `_kernel_fusion_patterns`. """ if not isinstance(model_class_or_instance, type): model_class_or_instance = type(model_class_or_instance) @@ -552,11 +526,11 @@ def __init__( Args: modules_to_fuse: The source modules to fuse together. source_names: The attribute names under which each module lives in its parent - (used to restore them on ``unfuse_modules``). + (used to restore them on `unfuse_modules`). fused_module_names: The names under which each source module is registered as a - child of this container (i.e. ``self.``). When ``None``, the - ``kernel_layer_name`` attribute of each source module is used. Pass this - explicitly when the source modules do not carry ``@use_kernel_forward_from_hub``. + child of this container (i.e. `self.`). When `None`, the + `kernel_layer_name` attribute of each source module is used. Pass this + explicitly when the source modules do not carry `@use_kernel_forward_from_hub`. """ super().__init__() if len(modules_to_fuse) == 0: @@ -609,11 +583,11 @@ def make_fused_module_class(source_layer_names: tuple[str, ...], kernel_layer_na Args: source_layer_names (`tuple[str, ...]`): Ordered tuple of `kernel_layer_name` values of the modules being fused - (e.g. ``("RMSNorm", "MLP")``). Used as the cache key — the same combination + (e.g. `("RMSNorm", "MLP")`). Used as the cache key — the same combination always returns the same class object. kernel_layer_name (`str`): The name assigned to the fused class, used by `kernelize` to look up the - kernel in the mapping (e.g. ``"RMSNormMLP"``). + kernel in the mapping (e.g. `"RMSNormMLP"`). Returns: A subclass of `FusedModuleBase` with `kernel_layer_name` set as a class attribute. @@ -635,38 +609,30 @@ def fuse_modules( Fuse a sequence of submodules into a single `FusedModuleBase` subclass in-place. For every parent module whose immediate children match all entries in - ``module_names_to_fuse``, the function: + `module_names_to_fuse`, the function: - replaces the first module with a `FusedModuleBase` subclass instance that holds all source modules as named children, - replaces the remaining modules with `nn.Identity()` pass-throughs. - The fused container's ``forward`` signature is patched to match the first source - module's ``forward``, satisfying the ``kernelize`` signature check. + The fused container's `forward` signature is patched to match the first source + module's `forward`, satisfying the `kernelize` signature check. Args: model (`nn.Module`): The model to modify in-place. module_names_to_fuse (`list[str]`): Glob-style paths of the modules to fuse, e.g. - ``["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"]``. - Integer indices are replaced with ``*`` so the same pattern applies to + `["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"]`. + Integer indices are replaced with `*` so the same pattern applies to every repeated block. kernel_layer_name (`str`): - The ``kernel_layer_name`` assigned to the fused class, used by ``kernelize`` - to look up the kernel in the mapping (e.g. ``"RMSNormMLP"``). + The `kernel_layer_name` assigned to the fused class, used by `kernelize` + to look up the kernel in the mapping (e.g. `"RMSNormMLP"`). source_layer_names (`list[str]`, *optional*): Explicit names for the child modules inside the fused container - (e.g. ``["RMSNorm", "MLP"]``). When ``None``, the ``kernel_layer_name`` + (e.g. `["RMSNorm", "MLP"]`). When `None`, the `kernel_layer_name` attribute of each source module is used. - - Example:: - - fuse_modules( - model, - ["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"], - "RMSNormMLP", - ) """ pattern = re.compile(r"\d+") for module_name, module in model.named_modules(): @@ -704,12 +670,6 @@ def unfuse_modules(model: "nn.Module") -> None: Args: model (`nn.Module`): The model to restore in-place. - - Example:: - - fuse_modules(model, ["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"], "RMSNormMLP") - # ... kernelized forward pass ... - unfuse_modules(model) # back to original """ for parent in model.modules(): for name, child in list(parent.named_children()): From e1c7f3f0bc4f3f702e2f7af43655f0dcfbab2d0a Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Mon, 13 Apr 2026 15:40:12 -0400 Subject: [PATCH 08/49] refactor: remove comment --- src/transformers/utils/kernel_config.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/transformers/utils/kernel_config.py b/src/transformers/utils/kernel_config.py index 1c7541fbcbc2..3a64a3808279 100644 --- a/src/transformers/utils/kernel_config.py +++ b/src/transformers/utils/kernel_config.py @@ -149,7 +149,6 @@ def apply_fusions(self, model): patterns = [item[1] for item in layer_name] fuse_modules(model, patterns, kernel_layer_name, source_layer_names=source_names) else: - # Legacy format: ("RMSNorm", "MLP") — look up patterns from model class or registry. fusion_patterns = getattr(model, "_kernel_fusion_patterns", None) or _FUSION_PATTERNS_REGISTRY.get( type(model), {} ) From bd640ae368c0ae1058bc09c42fa9984c189d2071 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Mon, 27 Apr 2026 17:35:14 -0400 Subject: [PATCH 09/49] refactor: update example script for testing --- fused_qwen_example.py | 38 ++++++++++++++++++++++++++------------ 1 file changed, 26 insertions(+), 12 deletions(-) diff --git a/fused_qwen_example.py b/fused_qwen_example.py index 0b7abd98439c..ca9bc231dd1f 100644 --- a/fused_qwen_example.py +++ b/fused_qwen_example.py @@ -1,13 +1,27 @@ -import copy import torch from transformers import AutoModelForCausalLM, AutoTokenizer, KernelConfig -from transformers.integrations import unfuse_modules model_id = "michaelbenayoun/qwen3-tiny-4kv-heads-4layers-random" tokenizer = AutoTokenizer.from_pretrained(model_id) +input_ids = tokenizer("Hello, how are you?", return_tensors="pt").input_ids +# --- baseline: plain model, no fusion --- +print("=" * 60) +print("Loading baseline model (no fusion)...") +baseline = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda") +baseline.eval() +input_ids = input_ids.to(baseline.device) + +with torch.no_grad(): + baseline_out = baseline(input_ids).logits +print("Baseline output shape:", baseline_out.shape) +del baseline + +# --- fused model --- +print("=" * 60) +print("Loading fused model...") kernel_config = KernelConfig({ ( ("RMSNorm", "model.layers.*.post_attention_layernorm"), @@ -15,16 +29,16 @@ ): "michaelbenayoun/dummy-rmsnorm-mlp:RMSNormMLP", }) -model = AutoModelForCausalLM.from_pretrained(model_id, use_kernels=True, kernel_config=kernel_config, device_map="cuda") - -input_ids = tokenizer("Hello, how are you?", return_tensors="pt").input_ids.to(model.device) - -original_model = copy.deepcopy(model) -unfuse_modules(original_model) -original_model.eval() +fused_model = AutoModelForCausalLM.from_pretrained( + model_id, use_kernels=True, kernel_config=kernel_config, device_map="cuda" +) +fused_model.eval() +print(fused_model) with torch.no_grad(): - fused_out = model(input_ids).logits - original_out = original_model(input_ids).logits + fused_out = fused_model(input_ids).logits +print("Fused output shape:", fused_out.shape) -print("Max diff fused vs original:", (fused_out - original_out).abs().max().item()) +# --- compare --- +print("=" * 60) +print("Max diff fused vs baseline:", (fused_out - baseline_out).abs().max().item()) From 323b0009dc7ae0f595c7051cf95c81fc0ff48866 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Mon, 27 Apr 2026 17:35:35 -0400 Subject: [PATCH 10/49] wip: remove apply_fusions method --- src/transformers/utils/kernel_config.py | 49 ------------------------- 1 file changed, 49 deletions(-) diff --git a/src/transformers/utils/kernel_config.py b/src/transformers/utils/kernel_config.py index 3a64a3808279..bb4f965ddbf4 100644 --- a/src/transformers/utils/kernel_config.py +++ b/src/transformers/utils/kernel_config.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from ..integrations.hub_kernels import _FUSION_PATTERNS_REGISTRY, fuse_modules from ..utils import PushToHubMixin @@ -117,54 +116,6 @@ def update_kernel(self, repo_id, registered_name, layer_name, device, mode, revi } } - def apply_fusions(self, model): - """ - For each n-to-1 entry (tuple key) in the kernel mapping, find the fusion patterns - registered on the model, fuse the corresponding modules in-place, then replace the - tuple key with the resolved kernel layer name so the rest of the pipeline is unchanged. - """ - new_mapping = {} - for layer_name, kernel in self.kernel_mapping.items(): - if not isinstance(layer_name, tuple): - new_mapping[layer_name] = kernel - continue - - # Parse the target kernel layer name from the repo string (the part after ':') - repo_str = kernel if isinstance(kernel, str) else next(iter(kernel.values())) - if isinstance(repo_str, dict): - repo_str = next(iter(repo_str.values())) - kernel_layer_name = repo_str.split(":")[1] if ":" in repo_str else repo_str.split("/")[-1] - - # Only fuse if a kernel is available for the current device - current_device = infer_device(model) - has_kernel_for_device = isinstance(kernel, str) or current_device in kernel - if not has_kernel_for_device: - continue - - # Detect inline format: tuple of (kernel_layer_name, glob_pattern) pairs. - # e.g. (("RMSNorm", "model.layers.*.post_attention_layernorm"), ("MLP", "model.layers.*.mlp")) - is_inline = all(isinstance(item, tuple) and len(item) == 2 for item in layer_name) - if is_inline: - source_names = [item[0] for item in layer_name] - patterns = [item[1] for item in layer_name] - fuse_modules(model, patterns, kernel_layer_name, source_layer_names=source_names) - else: - fusion_patterns = getattr(model, "_kernel_fusion_patterns", None) or _FUSION_PATTERNS_REGISTRY.get( - type(model), {} - ) - if kernel_layer_name not in fusion_patterns: - raise ValueError( - f"{type(model).__name__} does not define fusion patterns for '{kernel_layer_name}'. " - f'Either add `_kernel_fusion_patterns = {{"{kernel_layer_name}": [...]}}` to the model class, ' - f"call `register_fusion_patterns({type(model).__name__}, ...)` before loading the model, " - f"or use the inline pattern format: " - f'(("{kernel_layer_name}", ""), ...).' - ) - fuse_modules(model, fusion_patterns[kernel_layer_name], kernel_layer_name) - new_mapping[kernel_layer_name] = kernel - - self.kernel_mapping = new_mapping - def store_registered_layer_names(self, model): for name, module in model.named_modules(): if hasattr(module, "kernel_layer_name"): From fe3002d201a74dfa1400ad36992a18f4c3baa89e Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Mon, 27 Apr 2026 17:36:15 -0400 Subject: [PATCH 11/49] wip: add core feature for integration with the current fusing API --- src/transformers/integrations/hub_kernels.py | 227 +++++++++++++------ src/transformers/modeling_utils.py | 9 +- 2 files changed, 165 insertions(+), 71 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index e3ede38fd0eb..506d5a2bce13 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -18,15 +18,25 @@ from collections.abc import Callable from contextlib import contextmanager from types import ModuleType +from typing import TYPE_CHECKING from packaging import version as pkg_version +from ..conversion_mapping import get_checkpoint_conversion_mapping, register_checkpoint_conversion_mapping +from ..core_model_loading import WeightRenaming +from ..monkey_patching import register_patch_mapping from ..utils import ENV_VARS_TRUE_VALUES, logging from ..utils.import_utils import is_kernels_available, is_torch_available from .flash_attention import flash_attention_forward +if TYPE_CHECKING: + from .configuration_utils import PretrainedConfig + from .modeling_utils import PreTrainedModel + from .utils.kernel_config import KernelConfig + if is_torch_available(): + import torch import torch.nn as nn @@ -616,100 +626,183 @@ def make_fused_module_class(source_layer_names: tuple[str, ...], kernel_layer_na ) -def fuse_modules( - model: "nn.Module", - module_names_to_fuse: list[str], +def make_fused_parent_class( + parent_cls: type, + child_names: list[str], + source_names: list[str], kernel_layer_name: str, - source_layer_names: list[str] | None = None, -) -> None: +) -> type: """ - Fuse a sequence of submodules into a single `FusedModuleBase` subclass in-place. - - For every parent module whose immediate children match all entries in - `module_names_to_fuse`, the function: - - - replaces the first module with a `FusedModuleBase` subclass instance that holds - all source modules as named children, - - replaces the remaining modules with `nn.Identity()` pass-throughs. + Return a subclass of `parent_cls` whose __init__ rearranges the named children into a + FusedModuleBase container (keyed by `kernel_layer_name`) and replaces the remaining + children with nn.Identity placeholders. + """ + original_init = parent_cls.__init__ + _child_names = list(child_names) + _source_names = list(source_names) + _kernel_layer_name = kernel_layer_name + + def fused_init(self, *args, **kwargs): + original_init(self, *args, **kwargs) + modules_to_fuse = [getattr(self, name) for name in _child_names] + FusedClass = make_fused_module_class(tuple(_source_names), _kernel_layer_name) + fused = FusedClass(modules_to_fuse, _child_names, fused_module_names=list(_source_names)) + setattr(self, _child_names[0], fused) + for name in _child_names[1:]: + setattr(self, name, nn.Identity()) + + fused_cls = type(f"Fused{parent_cls.__name__}", (parent_cls,), {"__init__": fused_init}) + fused_cls.__qualname__ = f"Fused{parent_cls.__qualname__}" + return fused_cls + + +def infer_kernel_fusion_transforms( + patterns_with_names: list[tuple[str, str]], +) -> list[WeightRenaming]: + """ + Auto-infer WeightRenaming transforms for the path changes caused by FusedModuleBase wrapping. - The fused container's `forward` signature is patched to match the first source - module's `forward`, satisfying the `kernelize` signature check. + For each fusion (source_name, glob_path) pair: + - The first module's weights get source_name inserted between the module path and the weight suffix. + - Each subsequent module's weights are relocated from their original path to under the first module. - Args: - model (`nn.Module`): - The model to modify in-place. - module_names_to_fuse (`list[str]`): - Glob-style paths of the modules to fuse, e.g. - `["model.layers.*.post_attention_layernorm", "model.layers.*.mlp"]`. - Integer indices are replaced with `*` so the same pattern applies to - every repeated block. - kernel_layer_name (`str`): - The `kernel_layer_name` assigned to the fused class, used by `kernelize` - to look up the kernel in the mapping (e.g. `"RMSNormMLP"`). - source_layer_names (`list[str]`, *optional*): - Explicit names for the child modules inside the fused container - (e.g. `["RMSNorm", "MLP"]`). When `None`, the `kernel_layer_name` - attribute of each source module is used. + For example, given [("RMSNorm", "model.layers.*.post_attention_layernorm"), ("MLP", "model.layers.*.mlp")]: + - model.layers.0.post_attention_layernorm.weight -> model.layers.0.post_attention_layernorm.RMSNorm.weight + - model.layers.0.mlp.gate_proj.weight -> model.layers.0.post_attention_layernorm.MLP.gate_proj.weight """ - pattern = re.compile(r"\d+") - for module_name, module in model.named_modules(): - generic_children = { - re.sub(pattern, "*", f"{module_name}.{n}" if module_name else n): (n, child) - for n, child in module.named_children() - } - if not all(p in generic_children for p in module_names_to_fuse): - continue - child_names = [generic_children[p][0] for p in module_names_to_fuse] - modules_to_fuse = [generic_children[p][1] for p in module_names_to_fuse] + def _seg_to_regex(seg: str) -> str: + return r"[^.]+" if seg == "*" else re.escape(seg) - if source_layer_names is not None: - resolved_names = tuple(source_layer_names) - else: - resolved_names = tuple(getattr(m, "kernel_layer_name") for m in modules_to_fuse) - FusedClass = make_fused_module_class(resolved_names, kernel_layer_name) - fused_instance = FusedClass(modules_to_fuse, child_names, fused_module_names=list(resolved_names)) + def _glob_to_regex(glob: str) -> str: + return r"\.".join(_seg_to_regex(s) for s in glob.split(".")) - module.add_module(child_names[0], fused_instance) - for child_name in child_names[1:]: - module.add_module(child_name, nn.Identity()) + first_source_name, first_glob = patterns_with_names[0] + *parent_parts, first_child = first_glob.split(".") + parent_regex = _glob_to_regex(".".join(parent_parts)) + first_child_regex = _seg_to_regex(first_child) + transforms: list[WeightRenaming] = [] -def unfuse_modules(model: "nn.Module") -> None: - """ - Revert a previous `fuse_modules` call in-place, restoring the original modules. + # Module 0: insert source_name between the module path and the weight suffix. + # Capture the full module prefix so \1 can be used in the target for variable layer indices. + transforms.append( + WeightRenaming( + source_patterns=rf"({parent_regex}\.{first_child_regex}\.)", + target_patterns=rf"\1{first_source_name}.", + ) + ) + + # Modules i > 0: relocate weights from their original path to under the first module. + for source_name, glob in patterns_with_names[1:]: + *_, child = glob.split(".") + child_regex = _seg_to_regex(child) + transforms.append( + WeightRenaming( + source_patterns=rf"({parent_regex}\.){child_regex}\.", + target_patterns=rf"\1{first_child}.{source_name}.", + ) + ) - For each `FusedModuleBase` instance found in the model tree, the function: + return transforms - - restores the original first module at the fused container's position, - - restores the remaining original modules at their original positions - (replacing the `nn.Identity()` pass-throughs). - Args: - model (`nn.Module`): The model to restore in-place. +def register_kernel_fusions( + cls: "type[PreTrainedModel]", + config: "PretrainedConfig", + kernel_config: "KernelConfig", +) -> None: """ - for parent in model.modules(): - for name, child in list(parent.named_children()): - if not isinstance(child, FusedModuleBase): + Pre-register hub kernel n-to-1 fusions (tuple keys in KernelConfig) before model instantiation. + + For each inline tuple key `(("Name", "glob.path"), ...)` in `kernel_config.kernel_mapping`: + 1. Meta-instantiate the model to discover parent classes containing all target children. + 2. Register a monkey patch mapping each parent class to a fused subclass whose __init__ + wraps the target children in a FusedModuleBase with the resolved kernel_layer_name. + 3. Auto-infer WeightRenaming transforms for the path changes introduced by the wrapper + and register them via register_checkpoint_conversion_mapping. + 4. Replace the tuple key with the scalar kernel_layer_name in kernel_config so the + downstream pipeline (sanitize, create_compatible_mapping, kernelize) is unchanged. + + Scalar keys are passed through untouched. + """ + + if not hasattr(cls, "config_class") or not hasattr(cls.config_class, "model_type"): + raise ValueError(f"Model {cls.__name__} has no config_class or model_type.") + model_type = cls.config_class.model_type + + new_mapping: dict = {} + for layer_name, hub_repo in kernel_config.kernel_mapping.items(): + # Pass scalar keys through unchanged. + if not isinstance(layer_name, tuple) or not all( + isinstance(item, tuple) and len(item) == 2 for item in layer_name + ): + new_mapping[layer_name] = hub_repo + continue + + source_names = [item[0] for item in layer_name] + glob_patterns = [item[1] for item in layer_name] + child_names = [p.rsplit(".", 1)[1] for p in glob_patterns] + + # Resolve kernel_layer_name from "org/repo:ClassName" or device-nested dict. + repo_str = hub_repo + if isinstance(repo_str, dict): + repo_str = next(iter(repo_str.values())) + if isinstance(repo_str, dict): + repo_str = next(iter(repo_str.values())) + kernel_layer_name = repo_str.split(":")[1] if ":" in repo_str else repo_str.split("/")[-1] + + # 1. Meta-device scan — find parent classes that contain all target children. + with torch.device("meta"): + meta_model = cls(config) + + seen: set[type] = set() + patch_mapping: dict[str, type] = {} + for module in meta_model.modules(): + module_cls = type(module) + if module_cls in seen: continue - orig_modules = [getattr(child, layer_name) for layer_name in child._fused_module_names] - parent.add_module(name, orig_modules[0]) - for sibling_name, orig_module in zip(child._source_names[1:], orig_modules[1:]): - parent.add_module(sibling_name, orig_module) + if not all(hasattr(module, name) for name in child_names): + continue + seen.add(module_cls) + fused_cls = make_fused_parent_class(module_cls, child_names, source_names, kernel_layer_name) + patch_mapping[module_cls.__name__] = fused_cls + + if not patch_mapping: + logger.warning( + f"No parent modules found containing children {child_names} for kernel fusion " + f"'{kernel_layer_name}'. Skipping tuple key." + ) + new_mapping[layer_name] = hub_repo + continue + + # 2. Register class-level monkey patches. + register_patch_mapping(patch_mapping, overwrite=True) + + # 3. Register weight conversion mapping for the path changes introduced by the wrapper. + transforms = infer_kernel_fusion_transforms(list(zip(source_names, glob_patterns))) + existing = get_checkpoint_conversion_mapping(model_type) + if existing is not None: + transforms = existing + transforms + register_checkpoint_conversion_mapping(model_type, transforms, overwrite=True) + + # 4. Resolve tuple key -> scalar key for the downstream pipeline. + new_mapping[kernel_layer_name] = hub_repo + + kernel_config.kernel_mapping = new_mapping __all__ = [ "FusedModuleBase", "LayerRepository", - "fuse_modules", "get_kernel", + "register_kernel_fusions", "lazy_load_kernel", "make_fused_module_class", "register_fusion_patterns", "register_kernel_mapping", "register_kernel_mapping_transformers", "replace_kernel_forward_from_hub", - "unfuse_modules", "use_kernel_forward_from_hub", "use_kernel_func_from_hub", "use_kernelized_func", diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 2b9cd76f4914..8f349a44974f 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -3738,10 +3738,6 @@ def set_use_kernels(self, use_kernels, kernel_config: KernelConfig | None = None register_kernel_mapping_transformers() if kernel_config is not None and isinstance(kernel_config, KernelConfig): - # For n-to-1 entries (tuple keys), fuse the corresponding modules in the model - # before validation so that FusedModuleBase instances are visible to sanitize. - kernel_config.apply_fusions(self) - # This will make sure the mapping is valid, and the layers are registered in the model kernel_config.sanitize_kernel_mapping(self) @@ -4198,6 +4194,11 @@ def from_pretrained( register_fusion_patches(cls, config, fusion_config) + if kernel_config is not None and use_kernels: + from .integrations.hub_kernels import register_kernel_fusions + + register_kernel_fusions(cls, config, kernel_config) + model_init_context = cls.get_init_context(dtype, is_quantized, _is_ds_init_called, allow_all_kernels) config = copy.deepcopy(config) # We do not want to modify the config inplace in from_pretrained. From b541453097ec69f0ae6c8d3498256c161e0c0b64 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 28 Apr 2026 16:07:37 -0400 Subject: [PATCH 12/49] fix: move kernel mapping patching to kernelize --- src/transformers/modeling_utils.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 8f349a44974f..8a99730da7f6 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -1317,6 +1317,9 @@ def __init__(self, config: PreTrainedConfig, *inputs, **kwargs): self.config = config self.name_or_path = config.name_or_path + # Make kernel_config an attribute that can be used by the model. + self.kernel_config = None + # Check the attention implementation is supported, or set it if not yet set (on the internal attr, to avoid # setting it recursively) self.config._attn_implementation_internal = self._check_and_adjust_attn_implementation( @@ -3731,30 +3734,27 @@ def set_use_kernels(self, use_kernels, kernel_config: KernelConfig | None = None raise ValueError( "`use_kernels=True` requires kernels>=0.9.0. Please install the latest version with `pip install -U kernels`" ) - from kernels import use_kernel_mapping - from .integrations.hub_kernels import register_kernel_mapping_transformers register_kernel_mapping_transformers() if kernel_config is not None and isinstance(kernel_config, KernelConfig): + # Since kernel_config is a correct value, set it as an attribute of the model so it can be used. + self.kernel_config = kernel_config + # This will make sure the mapping is valid, and the layers are registered in the model kernel_config.sanitize_kernel_mapping(self) # This will create a compatible mapping for the model with the kernels library kernel_config.create_compatible_mapping(self) - - # This is a context manager to override the default kernel mapping - # We are calling kernelize inside this context manager using the use_kernels setter - # Param inherit_mapping should be False to avoid still loading kernel from remote - inherit_mapping = not kernel_config.use_local_kernel - with use_kernel_mapping(kernel_config.kernel_mapping, inherit_mapping=inherit_mapping): - self.use_kernels = True + self.use_kernels = True # We use the default kernel mapping in .integrations.hub_kernels else: self.use_kernels = True + self.kernel_config = None else: self.use_kernels = False + self.kernel_config = None @classmethod def from_pretrained( @@ -4194,6 +4194,7 @@ def from_pretrained( register_fusion_patches(cls, config, fusion_config) + # Fusion patches related to kernels if kernel_config is not None and use_kernels: from .integrations.hub_kernels import register_kernel_fusions @@ -4589,7 +4590,14 @@ def detach_hidden_kernels(module): self.apply(attach_hidden_kernels) mode = Mode.INFERENCE if not self.training else Mode.TRAINING if mode is None else mode - kernelize(self, device=Device(type=self.device.type), mode=mode) + if self.kernel_config is not None: + from kernels import use_kernel_mapping + + inherit_mapping = not self.kernel_config.use_local_kernel + with use_kernel_mapping(self.kernel_config.kernel_mapping, inherit_mapping=inherit_mapping): + kernelize(self, device=Device(type=self.device.type), mode=mode) + else: + kernelize(self, device=Device(type=self.device.type), mode=mode) self._use_kernels = True finally: From b3d73a7ac7de52e090126a6b70542e5f7c93d0ea Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 28 Apr 2026 16:36:07 -0400 Subject: [PATCH 13/49] wip: update example script --- fused_qwen_example.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/fused_qwen_example.py b/fused_qwen_example.py index ca9bc231dd1f..87a87ee5843c 100644 --- a/fused_qwen_example.py +++ b/fused_qwen_example.py @@ -5,29 +5,31 @@ model_id = "michaelbenayoun/qwen3-tiny-4kv-heads-4layers-random" tokenizer = AutoTokenizer.from_pretrained(model_id) -input_ids = tokenizer("Hello, how are you?", return_tensors="pt").input_ids +inputs = tokenizer("Hello, how are you?", return_tensors="pt") # --- baseline: plain model, no fusion --- print("=" * 60) print("Loading baseline model (no fusion)...") baseline = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda") baseline.eval() -input_ids = input_ids.to(baseline.device) +inputs = {k: v.to(baseline.device) for k, v in inputs.items()} with torch.no_grad(): - baseline_out = baseline(input_ids).logits + baseline_out = baseline(**inputs).logits print("Baseline output shape:", baseline_out.shape) -del baseline +# del baseline # --- fused model --- print("=" * 60) print("Loading fused model...") -kernel_config = KernelConfig({ - ( - ("RMSNorm", "model.layers.*.post_attention_layernorm"), - ("MLP", "model.layers.*.mlp"), - ): "michaelbenayoun/dummy-rmsnorm-mlp:RMSNormMLP", -}) +kernel_config = KernelConfig( + { + ( + ("RMSNorm", "model.layers.*.post_attention_layernorm"), + ("MLP", "model.layers.*.mlp"), + ): "michaelbenayoun/dummy-rmsnorm-mlp:RMSNormMLP", + }, +) fused_model = AutoModelForCausalLM.from_pretrained( model_id, use_kernels=True, kernel_config=kernel_config, device_map="cuda" @@ -36,7 +38,7 @@ print(fused_model) with torch.no_grad(): - fused_out = fused_model(input_ids).logits + fused_out = fused_model(**inputs).logits print("Fused output shape:", fused_out.shape) # --- compare --- From 084522272472cfed8a722b015b66d4e6639597e2 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Fri, 1 May 2026 10:50:59 -0400 Subject: [PATCH 14/49] wip: add transform_model method for WeightTransform --- src/transformers/core_model_loading.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index cd0710649c91..0cb348e81c4e 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -658,6 +658,16 @@ def __init__(self, source_patterns: str | list[str], target_patterns: str | list def __repr__(self): return f"{self.__class__.__name__}(source_patterns={self.source_patterns}, target_patterns={self.target_patterns})" + def transform_model(self, model: "PreTrainedModel") -> None: + """ + Optionally restructure the live model graph before weights are loaded. + + Override this in subclasses that need to add, remove, or replace `nn.Module` + instances (e.g. fusing two `nn.Linear` layers into one) so that the module + structure matches what the weight transform will produce. The default + implementation is a no-op. + """ + def __setattr__(self, name, value): if name in ("source_patterns", "target_patterns"): # We do not allow to re-set the patterns, as they are linked between each other and changing one From ecfd97defc882850b560476e5825b53c9bc4dee4 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Fri, 1 May 2026 16:04:08 -0400 Subject: [PATCH 15/49] wip: conversion_mapping in Kernel --- src/transformers/integrations/hub_kernels.py | 81 +++++++++++++++----- 1 file changed, 62 insertions(+), 19 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 506d5a2bce13..11bb2af173fd 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -46,6 +46,7 @@ from kernels import ( Device, LayerRepository, + LocalLayerRepository, Mode, register_kernel_mapping, replace_kernel_forward_from_hub, @@ -284,6 +285,10 @@ class LayerRepository: def __init__(self, *args, **kwargs): raise RuntimeError("LayerRepository requires `kernels` to be installed. Run `pip install kernels`.") + class LocalLayerRepository: + def __init__(self, *args, **kwargs): + raise RuntimeError("LocalLayerRepository requires `kernels` to be installed. Run `pip install kernels`.") + def replace_kernel_forward_from_hub(*args, **kwargs): raise RuntimeError( "replace_kernel_forward_from_hub requires `kernels` to be installed. Run `pip install kernels`." @@ -631,29 +636,25 @@ def make_fused_parent_class( child_names: list[str], source_names: list[str], kernel_layer_name: str, -) -> type: - """ - Return a subclass of `parent_cls` whose __init__ rearranges the named children into a - FusedModuleBase container (keyed by `kernel_layer_name`) and replaces the remaining - children with nn.Identity placeholders. - """ +) -> tuple[type, type]: original_init = parent_cls.__init__ _child_names = list(child_names) _source_names = list(source_names) _kernel_layer_name = kernel_layer_name + fused_module_cls = make_fused_module_class(tuple(_source_names), _kernel_layer_name) + def fused_init(self, *args, **kwargs): original_init(self, *args, **kwargs) modules_to_fuse = [getattr(self, name) for name in _child_names] - FusedClass = make_fused_module_class(tuple(_source_names), _kernel_layer_name) - fused = FusedClass(modules_to_fuse, _child_names, fused_module_names=list(_source_names)) + fused = fused_module_cls(modules_to_fuse, _child_names, fused_module_names=list(_source_names)) setattr(self, _child_names[0], fused) for name in _child_names[1:]: setattr(self, name, nn.Identity()) fused_cls = type(f"Fused{parent_cls.__name__}", (parent_cls,), {"__init__": fused_init}) fused_cls.__qualname__ = f"Fused{parent_cls.__qualname__}" - return fused_cls + return fused_cls, fused_module_cls def infer_kernel_fusion_transforms( @@ -707,6 +708,35 @@ def _glob_to_regex(glob: str) -> str: return transforms +def _try_load_kernel_class(repo_str: str, use_local: bool = False) -> type | None: + """ + Load a kernel class from a repo string like "org/repo:ClassName" or "/abs/path:ClassName". + + Returns the loaded class, or None if the string is malformed or loading fails. + The kernels library caches the result, so repeated calls are free. + """ + if ":" not in repo_str: + return None + repo_id, _, layer_name = repo_str.rpartition(":") + if not repo_id or not layer_name: + return None + try: + if use_local: + from pathlib import Path + + package_name = repo_id.rstrip("/").split("/")[-1] + repo = LocalLayerRepository( + repo_path=Path(repo_id), + package_name=package_name, + layer_name=layer_name, + ) + else: + repo = LayerRepository(repo_id=repo_id, layer_name=layer_name) + return repo.load() + except Exception: + return None + + def register_kernel_fusions( cls: "type[PreTrainedModel]", config: "PretrainedConfig", @@ -716,13 +746,16 @@ def register_kernel_fusions( Pre-register hub kernel n-to-1 fusions (tuple keys in KernelConfig) before model instantiation. For each inline tuple key `(("Name", "glob.path"), ...)` in `kernel_config.kernel_mapping`: - 1. Meta-instantiate the model to discover parent classes containing all target children. - 2. Register a monkey patch mapping each parent class to a fused subclass whose __init__ - wraps the target children in a FusedModuleBase with the resolved kernel_layer_name. - 3. Auto-infer WeightRenaming transforms for the path changes introduced by the wrapper - and register them via register_checkpoint_conversion_mapping. - 4. Replace the tuple key with the scalar kernel_layer_name in kernel_config so the - downstream pipeline (sanitize, create_compatible_mapping, kernelize) is unchanged. + + 1. Meta-instantiate the model to discover parent classes containing all target children. + + 2. Register a monkey patch mapping each parent class to a fused subclass whose __init__ + wraps the target children in a FusedModuleBase with the resolved kernel_layer_name. + + 3. Load or infer conversion mapping to handle the weight transformations caused by the kernel fusion. + + 4. Replace the tuple key with the scalar kernel_layer_name in kernel_config so the + downstream pipeline (sanitize, create_compatible_mapping, kernelize) is unchanged. Scalar keys are passed through untouched. """ @@ -765,8 +798,8 @@ def register_kernel_fusions( if not all(hasattr(module, name) for name in child_names): continue seen.add(module_cls) - fused_cls = make_fused_parent_class(module_cls, child_names, source_names, kernel_layer_name) - patch_mapping[module_cls.__name__] = fused_cls + fused_parent_cls, fused_cls = make_fused_parent_class(module_cls, child_names, source_names, kernel_layer_name) + patch_mapping[module_cls.__name__] = fused_parent_cls if not patch_mapping: logger.warning( @@ -779,11 +812,20 @@ def register_kernel_fusions( # 2. Register class-level monkey patches. register_patch_mapping(patch_mapping, overwrite=True) - # 3. Register weight conversion mapping for the path changes introduced by the wrapper. + # 3. Register weight conversion mapping. + # Always start with inferred transforms for the basic path renaming caused by fusion + # (e.g. layernorm.weight → layernorm.RMSNorm.weight, mlp.* → layernorm.MLP.*). + # Then append any kernel-provided custom transforms (e.g. weight concatenation). + # The kernels library caches the result, so the later kernelize() call is a free cache hit. transforms = infer_kernel_fusion_transforms(list(zip(source_names, glob_patterns))) + kernel_cls = _try_load_kernel_class(repo_str, use_local=kernel_config.use_local_kernel) + if kernel_cls is not None and hasattr(kernel_cls, "conversion_mapping"): + transforms = transforms + list(kernel_cls.conversion_mapping) + existing = get_checkpoint_conversion_mapping(model_type) if existing is not None: transforms = existing + transforms + register_checkpoint_conversion_mapping(model_type, transforms, overwrite=True) # 4. Resolve tuple key -> scalar key for the downstream pipeline. @@ -795,6 +837,7 @@ def register_kernel_fusions( __all__ = [ "FusedModuleBase", "LayerRepository", + "LocalLayerRepository", "get_kernel", "register_kernel_fusions", "lazy_load_kernel", From 0f0a64b7941df3cbced5faa185345bd484ad6f0a Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Mon, 4 May 2026 16:11:48 -0400 Subject: [PATCH 16/49] wip: remove things from __all__ --- src/transformers/integrations/hub_kernels.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 11bb2af173fd..34b7887aa929 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -837,11 +837,9 @@ def register_kernel_fusions( __all__ = [ "FusedModuleBase", "LayerRepository", - "LocalLayerRepository", "get_kernel", "register_kernel_fusions", "lazy_load_kernel", - "make_fused_module_class", "register_fusion_patterns", "register_kernel_mapping", "register_kernel_mapping_transformers", From 91177ae0b849b052c6198003ff273e7ed0971923 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Mon, 4 May 2026 16:23:16 -0400 Subject: [PATCH 17/49] wip: remove imports --- src/transformers/integrations/__init__.py | 6 ------ src/transformers/integrations/hub_kernels.py | 2 +- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index f74836d890f7..eb783adb64f2 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -69,13 +69,10 @@ "hub_kernels": [ "FusedModuleBase", "LayerRepository", - "fuse_modules", "lazy_load_kernel", - "make_fused_module_class", "register_fusion_patterns", "register_kernel_mapping", "replace_kernel_forward_from_hub", - "unfuse_modules", "use_kernel_forward_from_hub", "use_kernel_func_from_hub", "use_kernelized_func", @@ -228,13 +225,10 @@ from .hub_kernels import ( FusedModuleBase, LayerRepository, - fuse_modules, lazy_load_kernel, - make_fused_module_class, register_fusion_patterns, register_kernel_mapping, replace_kernel_forward_from_hub, - unfuse_modules, use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func, diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 34b7887aa929..00120b39a3d7 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -557,8 +557,8 @@ def __init__( """ Args: modules_to_fuse: The source modules to fuse together. + # TODO: are source_names still required? source_names: The attribute names under which each module lives in its parent - (used to restore them on `unfuse_modules`). fused_module_names: The names under which each source module is registered as a child of this container (i.e. `self.`). When `None`, the `kernel_layer_name` attribute of each source module is used. Pass this From 2636c06acb3924a4bf20956d19d619529cfa56d1 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Mon, 4 May 2026 17:09:08 -0400 Subject: [PATCH 18/49] fix: remove register_fusion_pattern path --- src/transformers/integrations/__init__.py | 2 -- src/transformers/integrations/hub_kernels.py | 28 -------------------- 2 files changed, 30 deletions(-) diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index eb783adb64f2..bf98af1f0275 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -70,7 +70,6 @@ "FusedModuleBase", "LayerRepository", "lazy_load_kernel", - "register_fusion_patterns", "register_kernel_mapping", "replace_kernel_forward_from_hub", "use_kernel_forward_from_hub", @@ -226,7 +225,6 @@ FusedModuleBase, LayerRepository, lazy_load_kernel, - register_fusion_patterns, register_kernel_mapping, replace_kernel_forward_from_hub, use_kernel_forward_from_hub, diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 00120b39a3d7..b0ec0aa00e4a 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -520,33 +520,6 @@ def allow_all_hub_kernels(): ALLOW_ALL_KERNELS = False -# Model class → {kernel_layer_name → [glob patterns]} -# Populated via `register_fusion_patterns` for models that cannot be modified directly. -_FUSION_PATTERNS_REGISTRY: dict[type, dict[str, list[str]]] = {} - - -def register_fusion_patterns( - model_class_or_instance, - patterns: dict[str, list[str]], -) -> None: - """ - Register kernel fusion patterns for a model class without modifying it directly. - - This is an alternative to setting `_kernel_fusion_patterns` as a class attribute, - useful when the model class is frozen or comes from an external library. - - Args: - model_class_or_instance: - The model class (or an instance of it) for which patterns are being registered. - patterns (`dict[str, list[str]]`): - Mapping from `kernel_layer_name` to a list of glob-style module paths, - identical in format to `_kernel_fusion_patterns`. - """ - if not isinstance(model_class_or_instance, type): - model_class_or_instance = type(model_class_or_instance) - _FUSION_PATTERNS_REGISTRY[model_class_or_instance] = patterns - - class FusedModuleBase(nn.Module): def __init__( self, @@ -840,7 +813,6 @@ def register_kernel_fusions( "get_kernel", "register_kernel_fusions", "lazy_load_kernel", - "register_fusion_patterns", "register_kernel_mapping", "register_kernel_mapping_transformers", "replace_kernel_forward_from_hub", From f7c15bd693a70ebf2f82e0fd4a9835dfbf4a5439 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Mon, 4 May 2026 17:32:35 -0400 Subject: [PATCH 19/49] fix: remove unused attribute --- src/transformers/integrations/hub_kernels.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index be76485a02f4..41f11f384e22 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -524,14 +524,11 @@ class FusedModuleBase(nn.Module): def __init__( self, modules_to_fuse: list["nn.Module"], - source_names: list[str], fused_module_names: list[str] | None = None, ): """ Args: modules_to_fuse: The source modules to fuse together. - # TODO: are source_names still required? - source_names: The attribute names under which each module lives in its parent fused_module_names: The names under which each source module is registered as a child of this container (i.e. `self.`). When `None`, the `kernel_layer_name` attribute of each source module is used. Pass this @@ -540,10 +537,6 @@ def __init__( super().__init__() if len(modules_to_fuse) == 0: raise ValueError("At least one module must be provided for fusion.") - if len(modules_to_fuse) != len(source_names): - raise ValueError("Length of modules_to_fuse and source_names must match.") - - self._source_names = source_names if fused_module_names is not None: if len(fused_module_names) != len(modules_to_fuse): @@ -620,7 +613,7 @@ def make_fused_parent_class( def fused_init(self, *args, **kwargs): original_init(self, *args, **kwargs) modules_to_fuse = [getattr(self, name) for name in _child_names] - fused = fused_module_cls(modules_to_fuse, _child_names, fused_module_names=list(_source_names)) + fused = fused_module_cls(modules_to_fuse, fused_module_names=list(_source_names)) setattr(self, _child_names[0], fused) for name in _child_names[1:]: setattr(self, name, nn.Identity()) From f9d42999f52b84d0f0114d003587d53003ed2010 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 5 May 2026 10:15:23 -0400 Subject: [PATCH 20/49] wip: update experimentation script --- fused_qwen_example.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/fused_qwen_example.py b/fused_qwen_example.py index 87a87ee5843c..0be834b6a97d 100644 --- a/fused_qwen_example.py +++ b/fused_qwen_example.py @@ -22,12 +22,15 @@ # --- fused model --- print("=" * 60) print("Loading fused model...") + +kernel_repo_id = "michaelbenayoun/dummy-rmsnorm-mlp:RMSNormMLP" +kernel_repo_id = "michaelbenayoun/dummy-rmsnorm-mlp-with-transformations:RMSNormMLP" kernel_config = KernelConfig( { ( ("RMSNorm", "model.layers.*.post_attention_layernorm"), ("MLP", "model.layers.*.mlp"), - ): "michaelbenayoun/dummy-rmsnorm-mlp:RMSNormMLP", + ): kernel_repo_id, }, ) From 847dbd4ce3db057f7b9d260b85a75831fe4e51d0 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 5 May 2026 14:24:37 -0400 Subject: [PATCH 21/49] refactor: add convert as abstract method --- src/transformers/core_model_loading.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 065691e94d64..1f1a48ea5037 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -19,7 +19,7 @@ import os import re import traceback -from abc import abstractmethod +from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Callable from concurrent.futures import Future, ThreadPoolExecutor @@ -80,7 +80,7 @@ def build_glob_alternation( return alternation, src_group_to_glob, tgt_group_to_glob -class ConversionOps: +class ConversionOps(ABC): """Base class for weight conversion operations.""" def __repr__(self): @@ -774,6 +774,16 @@ def was_used(self) -> bool: """ return self._was_used + @abstractmethod + def convert( + self, + layer_name: str, + model=None, + config=None, + hf_quantizer=None, + loading_info: LoadStateDictInfo | None = None, + ) -> dict[str, torch.Tensor]: ... + class WeightRenaming(WeightTransform): # Special case of WeightTransform that only renames keys without any conversion. From 4443c9a8fe9c234769779a9b9a1248495cdfd0d9 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 5 May 2026 14:39:14 -0400 Subject: [PATCH 22/49] style: reformat hub_kernels.py --- src/transformers/integrations/hub_kernels.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 41f11f384e22..9521cc37505d 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -764,7 +764,9 @@ def register_kernel_fusions( if not all(hasattr(module, name) for name in child_names): continue seen.add(module_cls) - fused_parent_cls, fused_cls = make_fused_parent_class(module_cls, child_names, source_names, kernel_layer_name) + fused_parent_cls, fused_cls = make_fused_parent_class( + module_cls, child_names, source_names, kernel_layer_name + ) patch_mapping[module_cls.__name__] = fused_parent_cls if not patch_mapping: From 51c59c96a38c50ea68abf3ad31f4947abf43778e Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 5 May 2026 14:39:36 -0400 Subject: [PATCH 23/49] wip: transform_model API --- src/transformers/core_model_loading.py | 20 ++++++++++---------- src/transformers/modeling_utils.py | 4 ++++ 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 1f1a48ea5037..29c75583cebd 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -661,16 +661,6 @@ def __init__(self, source_patterns: str | list[str], target_patterns: str | list def __repr__(self): return f"{self.__class__.__name__}(source_patterns={self.source_patterns}, target_patterns={self.target_patterns})" - def transform_model(self, model: "PreTrainedModel") -> None: - """ - Optionally restructure the live model graph before weights are loaded. - - Override this in subclasses that need to add, remove, or replace `nn.Module` - instances (e.g. fusing two `nn.Linear` layers into one) so that the module - structure matches what the weight transform will produce. The default - implementation is a no-op. - """ - def __setattr__(self, name, value): if name in ("source_patterns", "target_patterns"): # We do not allow to re-set the patterns, as they are linked between each other and changing one @@ -784,6 +774,13 @@ def convert( loading_info: LoadStateDictInfo | None = None, ) -> dict[str, torch.Tensor]: ... + @abstractmethod + def transform_model(self, model: PreTrainedModel) -> None: + """ + Optionally restructure the live model graph before weights are loaded. + """ + ... + class WeightRenaming(WeightTransform): # Special case of WeightTransform that only renames keys without any conversion. @@ -954,6 +951,9 @@ def convert( ) return collected_tensors + def transform_model(self, model: PreTrainedModel) -> None: + print(f"[DEBUG] Applying model transformation for {self}...") + # For I/O bound operations (i.e. here reading files), it is better to have fewer threads, e.g. 4 is a good default. # Having too many is actually harming performances quite a lot, i.e. using 16 can sometimes lead to taking TWICE diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index dfdf63b03f3d..4482ca3c8cbc 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4223,6 +4223,10 @@ def from_pretrained( # Obtain the weight conversion mapping for this model if any are registered and apply to all submodels recursively weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer) + # Apply any structural model transforms while parameters are still on the meta device. + for transform in weight_conversions or []: + transform.transform_model(model) + if _torch_distributed_available and device_mesh is not None: # add hooks to nn.Modules: no weights model = distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size) From 4c5850300a27c48888af29ff5445fb3616f22e00 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 5 May 2026 15:01:03 -0400 Subject: [PATCH 24/49] wip: transform_model API, WeightTransform --- src/transformers/core_model_loading.py | 64 ++++++++++++++++++++++++-- src/transformers/modeling_utils.py | 6 +-- 2 files changed, 63 insertions(+), 7 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 29c75583cebd..addf33dca926 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -775,9 +775,12 @@ def convert( ) -> dict[str, torch.Tensor]: ... @abstractmethod - def transform_model(self, model: PreTrainedModel) -> None: + def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: """ - Optionally restructure the live model graph before weights are loaded. + Analogous to `convert` but restructures `model`'s module graph rather than operating on tensor data. + + Called once per target parameter group (identified by `layer_name`) while the model is on the meta device. + Access the associated source parameter names via `self.layer_targets[layer_name]`. """ ... @@ -821,6 +824,29 @@ def convert( return collected_tensors + def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: + source_names = list(self.layer_targets.get(layer_name, [])) + if not source_names: + return + source_name = source_names[0] + if source_name == layer_name: + return + + source_mod_path, _, source_attr = source_name.rpartition(".") + target_mod_path, _, target_attr = layer_name.rpartition(".") + + try: + source_mod = model.get_submodule(source_mod_path) if source_mod_path else model + except AttributeError: + return # already moved + + if source_attr in source_mod._parameters: + obj = source_mod._parameters.pop(source_attr) + model.get_submodule(target_mod_path).register_parameter(target_attr, obj) + elif source_attr in source_mod._buffers: + obj = source_mod._buffers.pop(source_attr) + model.get_submodule(target_mod_path).register_buffer(target_attr, obj) + class PrefixChange(WeightRenaming): """ @@ -951,8 +977,38 @@ def convert( ) return collected_tensors - def transform_model(self, model: PreTrainedModel) -> None: - print(f"[DEBUG] Applying model transformation for {self}...") + def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: + pass + + +def apply_transforms_to_meta_model( + model: PreTrainedModel, + weight_conversions: list[WeightRenaming | WeightConverter], +) -> None: + """ + Restructure `model`'s module graph in-place according to `weight_conversions`, analogously to + `convert_and_load_state_dict_in_model` but operating on module attributes rather than checkpoint tensors. + + Designed to be called while the model is on the meta device, making the operation cost-free. + """ + meta_state_dict = model.state_dict() + + renamings = [x for x in weight_conversions if isinstance(x, WeightRenaming)] + converters = [x for x in weight_conversions if isinstance(x, WeightConverter)] + pattern_to_converter = {k: c for c in converters for k in c.source_patterns} + param_name_to_transform: dict[str, WeightRenaming | WeightConverter] = {} + + for param_name in sorted(meta_state_dict, key=dot_natural_key): + renamed_key, source_pattern = rename_source_key(param_name, renamings, converters) + if source_pattern is not None: + transform = param_name_to_transform.setdefault(renamed_key, deepcopy(pattern_to_converter[source_pattern])) + else: + transform = param_name_to_transform.setdefault(renamed_key, WeightRenaming(param_name, renamed_key)) + source_pattern = param_name + transform.layer_targets[renamed_key].add(param_name) + + for layer_name, transform in param_name_to_transform.items(): + transform.transform_model(model, layer_name) # For I/O bound operations (i.e. here reading files), it is better to have fewer threads, e.g. 4 is a good default. diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 4482ca3c8cbc..287391b8f0c3 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -48,6 +48,7 @@ from .core_model_loading import ( WeightConverter, WeightRenaming, + apply_transforms_to_meta_model, convert_and_load_state_dict_in_model, revert_weight_conversion, ) @@ -4223,9 +4224,8 @@ def from_pretrained( # Obtain the weight conversion mapping for this model if any are registered and apply to all submodels recursively weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer) - # Apply any structural model transforms while parameters are still on the meta device. - for transform in weight_conversions or []: - transform.transform_model(model) + # Apply any structural model transforms while parameters are still on the meta device (free operation). + apply_transforms_to_meta_model(model, weight_conversions) if _torch_distributed_available and device_mesh is not None: # add hooks to nn.Modules: no weights model = distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size) From a7f983f33f68471433f8eb1d5edf734ed98f80e7 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 5 May 2026 17:34:05 -0400 Subject: [PATCH 25/49] wip: transform_model API, WeightConverter --- src/transformers/core_model_loading.py | 96 ++++++++++++++++++++++++-- 1 file changed, 89 insertions(+), 7 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index addf33dca926..a84e68414b6b 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -778,9 +778,6 @@ def convert( def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: """ Analogous to `convert` but restructures `model`'s module graph rather than operating on tensor data. - - Called once per target parameter group (identified by `layer_name`) while the model is on the meta device. - Access the associated source parameter names via `self.layer_targets[layer_name]`. """ ... @@ -842,10 +839,10 @@ def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: if source_attr in source_mod._parameters: obj = source_mod._parameters.pop(source_attr) - model.get_submodule(target_mod_path).register_parameter(target_attr, obj) + _get_submodule_or_create(model, target_mod_path).register_parameter(target_attr, obj) elif source_attr in source_mod._buffers: obj = source_mod._buffers.pop(source_attr) - model.get_submodule(target_mod_path).register_buffer(target_attr, obj) + _get_submodule_or_create(model, target_mod_path).register_buffer(target_attr, obj) class PrefixChange(WeightRenaming): @@ -978,7 +975,91 @@ def convert( return collected_tensors def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: - pass + source_names = list(self.layer_targets.get(layer_name, [])) + if not source_names: + return + + collected_tensors = self.materialize_tensors() + + result: dict[str, Any] = dict(collected_tensors) + for op in self.operations: + result = op.convert(result, source_patterns=self.source_patterns, target_patterns=self.target_patterns) + + # Expand target pattern keys to full parameter names (same logic as convert) + full_name = layer_name.replace(".*.", ".0.") if ".*." in layer_name else layer_name + try: + prefix, _, suffix = next(full_name.partition(k) for k in result if k in full_name) + result = {prefix + k + suffix: v for k, v in result.items()} + except StopIteration: + pass + + # Retrieve one of the source modules to use as original module when creating missing target modules + original_mod: type[torch.nn.Module] | None = None + for sname in source_names: + smod_path = sname.rpartition(".")[0] + try: + original_mod = model.get_submodule(smod_path) if smod_path else model + break + except AttributeError: + pass + + # Remove source parameters/buffers + for source_name in source_names: + mod_path, _, attr = source_name.rpartition(".") + try: + source_mod = model.get_submodule(mod_path) if mod_path else model + except AttributeError: + continue + source_mod._parameters.pop(attr, None) + + # Register the target parameters, creating any missing module at the leaf path. + for target_name, meta_val in result.items(): + meta_tensor = meta_val[0] if isinstance(meta_val, list) else meta_val + t_mod_path, _, t_attr = target_name.rpartition(".") + target_mod = _get_submodule_or_create( + model, t_mod_path, original_module=original_mod, future_weight=meta_tensor + ) + if not isinstance(meta_tensor, torch.nn.Parameter): + meta_tensor = torch.nn.Parameter(meta_tensor, requires_grad=meta_tensor.is_floating_point()) + target_mod.register_parameter(t_attr, meta_tensor) + + +def _get_submodule_or_create( + model: torch.nn.Module, + module_path: str, + original_module: torch.nn.Module | None = None, + future_weight: torch.Tensor | None = None, +) -> torch.nn.Module: + if not module_path: + return model + + if original_module is not None and future_weight is None: + raise ValueError("original_module cannot be provided without future_weight to infer the module class") + parent, _, mod_name = module_path.rpartition(".") + parent_mod = model.get_submodule(parent) if parent else model + if mod_name not in parent_mod._modules: + cls_ = original_module.__class__ if original_module is not None else torch.nn.Module + if cls_ is torch.nn.Linear: + out_features, in_features = future_weight.shape + mod = cls_(in_features=in_features, out_features=out_features, bias=original_module.bias is not None) + elif cls_ is torch.nn.Embedding: + num_embeddings, embedding_dim = future_weight.shape + mod = cls_( + num_embeddings=num_embeddings, + embedding_dim=embedding_dim, + padding_idx=original_module.padding_idx, + max_norm=original_module.max_norm, + norm_type=original_module.norm_type, + scale_grad_by_freq=original_module.scale_grad_by_freq, + sparse=original_module.sparse, + ) + else: + mod = cls_() + parent_mod.add_module(mod_name, mod) + else: + mod = parent_mod._modules[mod_name] + + return mod def apply_transforms_to_meta_model( @@ -1005,7 +1086,8 @@ def apply_transforms_to_meta_model( else: transform = param_name_to_transform.setdefault(renamed_key, WeightRenaming(param_name, renamed_key)) source_pattern = param_name - transform.layer_targets[renamed_key].add(param_name) + + transform.add_tensor(renamed_key, param_name, source_pattern, meta_state_dict[param_name]) for layer_name, transform in param_name_to_transform.items(): transform.transform_model(model, layer_name) From b8d860f7ffc47130e7360a441853b1e44b37e8e6 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 5 May 2026 18:01:44 -0400 Subject: [PATCH 26/49] wip: transform_model API, WeightConverter --- src/transformers/core_model_loading.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index a84e68414b6b..7df502c6ee98 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1023,6 +1023,23 @@ def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: meta_tensor = torch.nn.Parameter(meta_tensor, requires_grad=meta_tensor.is_floating_point()) target_mod.register_parameter(t_attr, meta_tensor) + # Remove source modules if they are empty after the transform. + for source_name in source_names: + mod_path, _, attr = source_name.rpartition(".") + try: + source_mod = model.get_submodule(mod_path) if mod_path else model + except AttributeError: + continue + if ( + all(p is None for p in source_mod._parameters.values()) + and all(b is None for b in source_mod._buffers.values()) + and (m is None for m in source_mod._modules.values()) + ): + # If the module is now empty, we can remove it from its parent to avoid keeping empty modules after conversion + parent_path, _, mod_name = mod_path.rpartition(".") + parent_mod = model.get_submodule(parent_path) if parent_path else model + parent_mod._modules.pop(mod_name, None) + def _get_submodule_or_create( model: torch.nn.Module, From 3d5f353f100e67c04ad23728ffcbd2e8d9ff6298 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Wed, 6 May 2026 12:08:42 -0400 Subject: [PATCH 27/49] wip: make transform_model idempotent --- src/transformers/core_model_loading.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 7df502c6ee98..1bc1386392c7 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -778,6 +778,9 @@ def convert( def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: """ Analogous to `convert` but restructures `model`'s module graph rather than operating on tensor data. + + Must be idempotent, i.e. if called multiple times on the same `model` and `layer_name`, only the first call + should have an effect. """ ... @@ -832,6 +835,16 @@ def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: source_mod_path, _, source_attr = source_name.rpartition(".") target_mod_path, _, target_attr = layer_name.rpartition(".") + # Idempotency guard: if the target already holds the parameter/buffer, the + # model is already in the correct shape (e.g. from monkey-patching before + # from_pretrained). Nothing to move. + try: + target_mod = model.get_submodule(target_mod_path) if target_mod_path else model + if target_attr in target_mod._parameters or target_attr in target_mod._buffers: + return + except AttributeError: + pass + try: source_mod = model.get_submodule(source_mod_path) if source_mod_path else model except AttributeError: @@ -979,6 +992,16 @@ def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: if not source_names: return + # Idempotency guard: if the first target parameter already exists, the model is + # already in the correct shape. Nothing to convert. + t_mod_path, _, t_attr = layer_name.rpartition(".") + try: + t_mod = model.get_submodule(t_mod_path) if t_mod_path else model + if t_attr in t_mod._parameters or t_attr in t_mod._buffers: + return + except AttributeError: + pass + collected_tensors = self.materialize_tensors() result: dict[str, Any] = dict(collected_tensors) From c35b51334e2dbf01d3bb6aa955fcf2df938b7ce2 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Wed, 6 May 2026 15:35:01 -0400 Subject: [PATCH 28/49] refactor: infer_kernel_fusion_transforms --- src/transformers/integrations/hub_kernels.py | 58 +++++++++++--------- 1 file changed, 31 insertions(+), 27 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 9521cc37505d..87769e7ecffa 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -630,44 +630,48 @@ def infer_kernel_fusion_transforms( Auto-infer WeightRenaming transforms for the path changes caused by FusedModuleBase wrapping. For each fusion (source_name, glob_path) pair: - - The first module's weights get source_name inserted between the module path and the weight suffix. - - Each subsequent module's weights are relocated from their original path to under the first module. - - For example, given [("RMSNorm", "model.layers.*.post_attention_layernorm"), ("MLP", "model.layers.*.mlp")]: - - model.layers.0.post_attention_layernorm.weight -> model.layers.0.post_attention_layernorm.RMSNorm.weight - - model.layers.0.mlp.gate_proj.weight -> model.layers.0.post_attention_layernorm.MLP.gate_proj.weight + - The first module's weights get source_name inserted between the module path and the weight suffix. + - Each subsequent module's weights are relocated from their original path to under the first module. + + For example, given + [ + ("RMSNorm", "model.layers.*.post_attention_layernorm"), + ("MLP", "model.layers.*.mlp") + ] + the inferred transforms would be: + - model.layers.0.post_attention_layernorm.weight -> model.layers.0.post_attention_layernorm.RMSNorm.weight + - model.layers.0.mlp.gate_proj.weight -> model.layers.0.post_attention_layernorm.MLP.gate_proj.weight """ - def _seg_to_regex(seg: str) -> str: + def segment_to_regex(seg: str) -> str: + """ + Convert a glob segment to a regex segment. + The wildcard "*" matches any non-empty sequence of characters that does not include a dot. + Otherwise, the segment is escaped. + """ return r"[^.]+" if seg == "*" else re.escape(seg) - def _glob_to_regex(glob: str) -> str: - return r"\.".join(_seg_to_regex(s) for s in glob.split(".")) + _, first_glob = patterns_with_names[0] + parent_path, first_child = first_glob.rsplit(".", 1) - first_source_name, first_glob = patterns_with_names[0] - *parent_parts, first_child = first_glob.split(".") - parent_regex = _glob_to_regex(".".join(parent_parts)) - first_child_regex = _seg_to_regex(first_child) + # Convert parent_path to a regex pattern, capturing it as a group for reuse in the target pattern. + parent_regex = r"\.".join(segment_to_regex(s) for s in parent_path.split(".")) transforms: list[WeightRenaming] = [] - # Module 0: insert source_name between the module path and the weight suffix. - # Capture the full module prefix so \1 can be used in the target for variable layer indices. - transforms.append( - WeightRenaming( - source_patterns=rf"({parent_regex}\.{first_child_regex}\.)", - target_patterns=rf"\1{first_source_name}.", - ) - ) + # We must not match the moved child modules under the first child, so we add a negative lookahead. + child_regexs = "|".join(rf"{name}\." for name, _ in patterns_with_names) + moved_child_lookahead = f"(?!{child_regexs})" - # Modules i > 0: relocate weights from their original path to under the first module. - for source_name, glob in patterns_with_names[1:]: - *_, child = glob.split(".") - child_regex = _seg_to_regex(child) + for i, (source_name, glob) in enumerate(patterns_with_names): + child = glob.rsplit(".", 1)[1] + child_re = segment_to_regex(child) + guard = moved_child_lookahead if i == 0 else "" transforms.append( WeightRenaming( - source_patterns=rf"({parent_regex}\.){child_regex}\.", - target_patterns=rf"\1{first_child}.{source_name}.", + source_patterns=rf"({parent_regex}\.){child_re}\.{guard}", + # \1 refers to the captured parent path + target_patterns=rf"\1{first_child}.{source_name}\.", ) ) From 6e369c0ca9a1abd94e5c8eafc94c75e481bb2884 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Wed, 6 May 2026 15:43:36 -0400 Subject: [PATCH 29/49] style: regexs -> regexes --- src/transformers/integrations/hub_kernels.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 87769e7ecffa..7134813d8f2b 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -660,8 +660,8 @@ def segment_to_regex(seg: str) -> str: transforms: list[WeightRenaming] = [] # We must not match the moved child modules under the first child, so we add a negative lookahead. - child_regexs = "|".join(rf"{name}\." for name, _ in patterns_with_names) - moved_child_lookahead = f"(?!{child_regexs})" + child_regexes = "|".join(rf"{name}\." for name, _ in patterns_with_names) + moved_child_lookahead = f"(?!{child_regexes})" for i, (source_name, glob) in enumerate(patterns_with_names): child = glob.rsplit(".", 1)[1] From 731d0b7208e9c5e6e70e6606806083d06240a30b Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Wed, 6 May 2026 15:52:39 -0400 Subject: [PATCH 30/49] refactor: register_kernel_fusions --- src/transformers/integrations/hub_kernels.py | 32 ++++++-------------- 1 file changed, 10 insertions(+), 22 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 7134813d8f2b..71c9bd314b79 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -602,7 +602,7 @@ def make_fused_parent_class( child_names: list[str], source_names: list[str], kernel_layer_name: str, -) -> tuple[type, type]: +) -> type: original_init = parent_cls.__init__ _child_names = list(child_names) _source_names = list(source_names) @@ -620,7 +620,7 @@ def fused_init(self, *args, **kwargs): fused_cls = type(f"Fused{parent_cls.__name__}", (parent_cls,), {"__init__": fused_init}) fused_cls.__qualname__ = f"Fused{parent_cls.__qualname__}" - return fused_cls, fused_module_cls + return fused_cls def infer_kernel_fusion_transforms( @@ -679,12 +679,6 @@ def segment_to_regex(seg: str) -> str: def _try_load_kernel_class(repo_str: str, use_local: bool = False) -> type | None: - """ - Load a kernel class from a repo string like "org/repo:ClassName" or "/abs/path:ClassName". - - Returns the loaded class, or None if the string is malformed or loading fails. - The kernels library caches the result, so repeated calls are free. - """ if ":" not in repo_str: return None repo_id, _, layer_name = repo_str.rpartition(":") @@ -719,15 +713,13 @@ def register_kernel_fusions( 1. Meta-instantiate the model to discover parent classes containing all target children. - 2. Register a monkey patch mapping each parent class to a fused subclass whose __init__ - wraps the target children in a FusedModuleBase with the resolved kernel_layer_name. + 2. Register a monkey patch mapping each parent class to a fused subclass whose `__init__` + wraps the target children in a `FusedModuleBase` with the resolved `kernel_layer_name`. 3. Load or infer conversion mapping to handle the weight transformations caused by the kernel fusion. 4. Replace the tuple key with the scalar kernel_layer_name in kernel_config so the downstream pipeline (sanitize, create_compatible_mapping, kernelize) is unchanged. - - Scalar keys are passed through untouched. """ if not hasattr(cls, "config_class") or not hasattr(cls.config_class, "model_type"): @@ -736,7 +728,6 @@ def register_kernel_fusions( new_mapping: dict = {} for layer_name, hub_repo in kernel_config.kernel_mapping.items(): - # Pass scalar keys through unchanged. if not isinstance(layer_name, tuple) or not all( isinstance(item, tuple) and len(item) == 2 for item in layer_name ): @@ -768,9 +759,7 @@ def register_kernel_fusions( if not all(hasattr(module, name) for name in child_names): continue seen.add(module_cls) - fused_parent_cls, fused_cls = make_fused_parent_class( - module_cls, child_names, source_names, kernel_layer_name - ) + fused_parent_cls = make_fused_parent_class(module_cls, child_names, source_names, kernel_layer_name) patch_mapping[module_cls.__name__] = fused_parent_cls if not patch_mapping: @@ -784,23 +773,22 @@ def register_kernel_fusions( # 2. Register class-level monkey patches. register_patch_mapping(patch_mapping, overwrite=True) - # 3. Register weight conversion mapping. - # Always start with inferred transforms for the basic path renaming caused by fusion - # (e.g. layernorm.weight → layernorm.RMSNorm.weight, mlp.* → layernorm.MLP.*). - # Then append any kernel-provided custom transforms (e.g. weight concatenation). - # The kernels library caches the result, so the later kernelize() call is a free cache hit. + # 3. Infer weight renaming mapping. transforms = infer_kernel_fusion_transforms(list(zip(source_names, glob_patterns))) + + # 4. Load weight conversion mapping from the kernel class if available. kernel_cls = _try_load_kernel_class(repo_str, use_local=kernel_config.use_local_kernel) if kernel_cls is not None and hasattr(kernel_cls, "conversion_mapping"): transforms = transforms + list(kernel_cls.conversion_mapping) + # 5. Load existing conversion mapping for this model type. existing = get_checkpoint_conversion_mapping(model_type) if existing is not None: transforms = existing + transforms + # 6. Register the combined mappings. register_checkpoint_conversion_mapping(model_type, transforms, overwrite=True) - # 4. Resolve tuple key -> scalar key for the downstream pipeline. new_mapping[kernel_layer_name] = hub_repo kernel_config.kernel_mapping = new_mapping From 4e56ad35eaf55e6df099607d1ab8e07355d90ca9 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Wed, 6 May 2026 16:07:31 -0400 Subject: [PATCH 31/49] refactor: post transformation cleanup --- src/transformers/core_model_loading.py | 43 ++++++++++++++++---------- 1 file changed, 27 insertions(+), 16 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 1bc1386392c7..c21873b0d3df 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -784,6 +784,26 @@ def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: """ ... + def post_transformation_cleanup(self, model: PreTrainedModel, source_names: list[str]) -> None: + """ + After transforming the model's module graph, we may end up with empty modules that held the original weights, + we remove them. + """ + for source_name in source_names: + mod_path, _ = source_name.rsplit(".", 1) if "." in source_name else ("", source_name) + try: + source_mod = model.get_submodule(mod_path) if mod_path else model + except AttributeError: + continue + if ( + all(p is None for p in source_mod._parameters.values()) + and all(b is None for b in source_mod._buffers.values()) + and (m is None for m in source_mod._modules.values()) + ): + parent_path, _, mod_name = mod_path.rpartition(".") + parent_mod = model.get_submodule(parent_path) if parent_path else model + parent_mod._modules.pop(mod_name, None) + class WeightRenaming(WeightTransform): # Special case of WeightTransform that only renames keys without any conversion. @@ -857,6 +877,9 @@ def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: obj = source_mod._buffers.pop(source_attr) _get_submodule_or_create(model, target_mod_path).register_buffer(target_attr, obj) + # Cleanup any empty modules that may be left after moving the parameter/buffer. + self.post_transformation_cleanup(model, source_names) + class PrefixChange(WeightRenaming): """ @@ -1047,21 +1070,7 @@ def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: target_mod.register_parameter(t_attr, meta_tensor) # Remove source modules if they are empty after the transform. - for source_name in source_names: - mod_path, _, attr = source_name.rpartition(".") - try: - source_mod = model.get_submodule(mod_path) if mod_path else model - except AttributeError: - continue - if ( - all(p is None for p in source_mod._parameters.values()) - and all(b is None for b in source_mod._buffers.values()) - and (m is None for m in source_mod._modules.values()) - ): - # If the module is now empty, we can remove it from its parent to avoid keeping empty modules after conversion - parent_path, _, mod_name = mod_path.rpartition(".") - parent_mod = model.get_submodule(parent_path) if parent_path else model - parent_mod._modules.pop(mod_name, None) + self.post_transformation_cleanup(model, source_names) def _get_submodule_or_create( @@ -1094,7 +1103,9 @@ def _get_submodule_or_create( sparse=original_module.sparse, ) else: - mod = cls_() + raise ValueError( + f"Cannot infer instantiation for {cls_} when creating submodule {mod_name} at path {module_path}" + ) parent_mod.add_module(mod_name, mod) else: mod = parent_mod._modules[mod_name] From 105a403427fd619be735f4a663302d8cdc425e28 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Thu, 7 May 2026 10:29:28 -0400 Subject: [PATCH 32/49] style: fix comment --- src/transformers/core_model_loading.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index c21873b0d3df..0b6e623cd138 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1049,7 +1049,7 @@ def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: except AttributeError: pass - # Remove source parameters/buffers + # Remove source parameters for source_name in source_names: mod_path, _, attr = source_name.rpartition(".") try: From b1c9645313116bc1c7586725a65ed65f414eb10b Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Thu, 7 May 2026 12:04:53 -0400 Subject: [PATCH 33/49] test: add TestApplyTransformsToMetaModel tests --- src/transformers/core_model_loading.py | 6 +- tests/utils/test_core_model_loading.py | 350 +++++++++++++++++++++++++ 2 files changed, 354 insertions(+), 2 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 0b6e623cd138..9f15850e543c 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -798,7 +798,7 @@ def post_transformation_cleanup(self, model: PreTrainedModel, source_names: list if ( all(p is None for p in source_mod._parameters.values()) and all(b is None for b in source_mod._buffers.values()) - and (m is None for m in source_mod._modules.values()) + and all(m is None for m in source_mod._modules.values()) ): parent_path, _, mod_name = mod_path.rpartition(".") parent_mod = model.get_submodule(parent_path) if parent_path else model @@ -872,7 +872,9 @@ def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: if source_attr in source_mod._parameters: obj = source_mod._parameters.pop(source_attr) - _get_submodule_or_create(model, target_mod_path).register_parameter(target_attr, obj) + _get_submodule_or_create( + model, target_mod_path, original_module=source_mod, future_weight=obj + ).register_parameter(target_attr, obj) elif source_attr in source_mod._buffers: obj = source_mod._buffers.pop(source_attr) _get_submodule_or_create(model, target_mod_path).register_buffer(target_attr, obj) diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index a358822d19f8..c5546a3dbe1d 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -31,6 +31,7 @@ PrefixChange, WeightConverter, WeightRenaming, + apply_transforms_to_meta_model, build_glob_alternation, convert_and_load_state_dict_in_model, rename_source_key, @@ -1021,5 +1022,354 @@ def test_can_add_prefix_submodule(self): self.assertTrue((v == model_state_dict[k]).all()) +class TestApplyTransformsToMetaModel(unittest.TestCase): + def test_weight_renaming_renames_parameter(self): + class SimpleModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4, bias=False) + + model = SimpleModel() + original_weight = model.linear.weight.data.clone() + self.assertIn("linear.weight", model.state_dict()) + + apply_transforms_to_meta_model(model, [WeightRenaming("linear.weight", "linear.kernel")]) + + state_dict = model.state_dict() + self.assertIn("linear.kernel", state_dict) + self.assertNotIn("linear.weight", state_dict) + torch.testing.assert_close(state_dict["linear.kernel"], original_weight) + + def test_weight_renaming_is_idempotent(self): + class SimpleModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4, bias=False) + + model = SimpleModel() + apply_transforms_to_meta_model(model, [WeightRenaming("linear.weight", "linear.kernel")]) + keys_after_first = set(model.state_dict().keys()) + + apply_transforms_to_meta_model(model, [WeightRenaming("linear.weight", "linear.kernel")]) + self.assertEqual(keys_after_first, set(model.state_dict().keys())) + + def test_weight_renaming_no_match_is_noop(self): + class SimpleModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4, bias=False) + + model = SimpleModel() + original_keys = set(model.state_dict().keys()) + + apply_transforms_to_meta_model(model, [WeightRenaming("non_existent.weight", "something.weight")]) + + self.assertEqual(original_keys, set(model.state_dict().keys())) + + def test_weight_converter_splits_fused_projection(self): + class QKVAttn(nn.Module): + def __init__(self): + super().__init__() + self.qkv_proj = nn.Linear(2, 6, bias=False) + + class QKVModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.self_attn = QKVAttn() + + model = QKVModel() + model.config = PretrainedConfig() + qkv_weight = torch.randn(6, 2) + model.self_attn.qkv_proj.weight.data.copy_(qkv_weight) + + apply_transforms_to_meta_model( + model, + [ + WeightConverter( + "self_attn.qkv_proj.weight", + ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + operations=[Chunk(dim=0)], + ) + ], + ) + + state_dict = model.state_dict() + self.assertIn("self_attn.q_proj.weight", state_dict) + self.assertIn("self_attn.k_proj.weight", state_dict) + self.assertIn("self_attn.v_proj.weight", state_dict) + self.assertNotIn("self_attn.qkv_proj.weight", state_dict) + # The source module must be gone, the target modules must exist + self.assertNotIn("qkv_proj", model.self_attn._modules) + self.assertIn("q_proj", model.self_attn._modules) + self.assertIn("k_proj", model.self_attn._modules) + self.assertIn("v_proj", model.self_attn._modules) + + expected_q, expected_k, expected_v = torch.chunk(qkv_weight, 3, dim=0) + torch.testing.assert_close(state_dict["self_attn.q_proj.weight"], expected_q) + torch.testing.assert_close(state_dict["self_attn.k_proj.weight"], expected_k) + torch.testing.assert_close(state_dict["self_attn.v_proj.weight"], expected_v) + + def test_weight_converter_split_is_idempotent(self): + class QKVAttn(nn.Module): + def __init__(self): + super().__init__() + self.qkv_proj = nn.Linear(2, 6, bias=False) + + class QKVModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.self_attn = QKVAttn() + + def make_mapping(): + return [ + WeightConverter( + "self_attn.qkv_proj.weight", + ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + operations=[Chunk(dim=0)], + ) + ] + + model = QKVModel() + model.config = PretrainedConfig() + + apply_transforms_to_meta_model(model, make_mapping()) + keys_after_first = set(model.state_dict().keys()) + + apply_transforms_to_meta_model(model, make_mapping()) + self.assertEqual(keys_after_first, set(model.state_dict().keys())) + + def test_weight_converter_splits_across_multiple_layers(self): + class QKVAttn(nn.Module): + def __init__(self): + super().__init__() + self.qkv_proj = nn.Linear(2, 6, bias=False) + + class QKVLayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = QKVAttn() + + class QKVModel(nn.Module): + base_model_prefix = "model" + + def __init__(self): + super().__init__() + self.model = nn.ModuleList([QKVLayer(), QKVLayer()]) + + model = QKVModel() + model.config = PretrainedConfig() + + qkv_weights = {} + for i in range(2): + w = torch.randn(6, 2) + model.model[i].self_attn.qkv_proj.weight.data.copy_(w) + qkv_weights[i] = w + + apply_transforms_to_meta_model( + model, + [ + WeightConverter( + "self_attn.qkv_proj.weight", + ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + operations=[Chunk(dim=0)], + ) + ], + ) + + state_dict = model.state_dict() + self.assertNotIn("model.0.self_attn.qkv_proj.weight", state_dict) + self.assertNotIn("model.1.self_attn.qkv_proj.weight", state_dict) + + for i in range(2): + layer_attn = model.model[i].self_attn + self.assertNotIn("qkv_proj", layer_attn._modules) + self.assertIn("q_proj", layer_attn._modules) + self.assertIn("k_proj", layer_attn._modules) + self.assertIn("v_proj", layer_attn._modules) + + expected_q, expected_k, expected_v = torch.chunk(qkv_weights[i], 3, dim=0) + torch.testing.assert_close(state_dict[f"model.{i}.self_attn.q_proj.weight"], expected_q) + torch.testing.assert_close(state_dict[f"model.{i}.self_attn.k_proj.weight"], expected_k) + torch.testing.assert_close(state_dict[f"model.{i}.self_attn.v_proj.weight"], expected_v) + + def test_weight_converter_concatenates_projections(self): + class SeparateAttn(nn.Module): + def __init__(self): + super().__init__() + self.q_proj = nn.Linear(2, 2, bias=False) + self.k_proj = nn.Linear(2, 2, bias=False) + self.v_proj = nn.Linear(2, 2, bias=False) + + class SeparateModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.self_attn = SeparateAttn() + + model = SeparateModel() + model.config = PretrainedConfig() + + q_weight = torch.randn(2, 2) + k_weight = torch.randn(2, 2) + v_weight = torch.randn(2, 2) + model.self_attn.q_proj.weight.data.copy_(q_weight) + model.self_attn.k_proj.weight.data.copy_(k_weight) + model.self_attn.v_proj.weight.data.copy_(v_weight) + + apply_transforms_to_meta_model( + model, + [ + WeightConverter( + ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + "self_attn.qkv_proj.weight", + operations=[Concatenate(dim=0)], + ) + ], + ) + + state_dict = model.state_dict() + self.assertIn("self_attn.qkv_proj.weight", state_dict) + self.assertNotIn("self_attn.q_proj.weight", state_dict) + self.assertNotIn("self_attn.k_proj.weight", state_dict) + self.assertNotIn("self_attn.v_proj.weight", state_dict) + # Source modules must be gone, fused module must exist + self.assertNotIn("q_proj", model.self_attn._modules) + self.assertNotIn("k_proj", model.self_attn._modules) + self.assertNotIn("v_proj", model.self_attn._modules) + self.assertIn("qkv_proj", model.self_attn._modules) + + expected_qkv = torch.cat([q_weight, k_weight, v_weight], dim=0) + torch.testing.assert_close(state_dict["self_attn.qkv_proj.weight"], expected_qkv) + + def test_multiple_independent_converters_in_single_call(self): + class Attn(nn.Module): + def __init__(self): + super().__init__() + self.qkv_proj = nn.Linear(2, 6, bias=False) + + class MLP(nn.Module): + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(2, 4, bias=False) + self.up_proj = nn.Linear(2, 4, bias=False) + + class TwoHeadModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.self_attn = Attn() + self.mlp = MLP() + + model = TwoHeadModel() + model.config = PretrainedConfig() + + qkv_weight = torch.randn(6, 2) + gate_weight = torch.randn(4, 2) + up_weight = torch.randn(4, 2) + model.self_attn.qkv_proj.weight.data.copy_(qkv_weight) + model.mlp.gate_proj.weight.data.copy_(gate_weight) + model.mlp.up_proj.weight.data.copy_(up_weight) + + apply_transforms_to_meta_model( + model, + [ + WeightConverter( + "self_attn.qkv_proj.weight", + ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + operations=[Chunk(dim=0)], + ), + WeightConverter( + ["mlp.gate_proj.weight", "mlp.up_proj.weight"], + "mlp.gate_up_proj.weight", + operations=[Concatenate(dim=0)], + ), + ], + ) + + state_dict = model.state_dict() + + # Attention: fused → split + self.assertNotIn("self_attn.qkv_proj.weight", state_dict) + self.assertIn("self_attn.q_proj.weight", state_dict) + self.assertIn("self_attn.k_proj.weight", state_dict) + self.assertIn("self_attn.v_proj.weight", state_dict) + self.assertNotIn("qkv_proj", model.self_attn._modules) + self.assertIn("q_proj", model.self_attn._modules) + self.assertIn("k_proj", model.self_attn._modules) + self.assertIn("v_proj", model.self_attn._modules) + expected_q, expected_k, expected_v = torch.chunk(qkv_weight, 3, dim=0) + torch.testing.assert_close(state_dict["self_attn.q_proj.weight"], expected_q) + torch.testing.assert_close(state_dict["self_attn.k_proj.weight"], expected_k) + torch.testing.assert_close(state_dict["self_attn.v_proj.weight"], expected_v) + + # MLP: separate → fused + self.assertNotIn("mlp.gate_proj.weight", state_dict) + self.assertNotIn("mlp.up_proj.weight", state_dict) + self.assertIn("mlp.gate_up_proj.weight", state_dict) + self.assertNotIn("gate_proj", model.mlp._modules) + self.assertNotIn("up_proj", model.mlp._modules) + self.assertIn("gate_up_proj", model.mlp._modules) + torch.testing.assert_close(state_dict["mlp.gate_up_proj.weight"], torch.cat([gate_weight, up_weight], dim=0)) + + def test_renaming_happens_before_conversion(self): + class MLPModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.mlp = nn.Module() + self.dummy = nn.Module() + self.mlp.gate_proj = nn.Linear(2, 4, bias=False) + self.dummy.up_proj = nn.Linear(2, 4, bias=False) + + model = MLPModel() + model.config = PretrainedConfig() + + gate_weight = torch.randn(4, 2) + up_weight = torch.randn(4, 2) + model.mlp.gate_proj.weight.data.copy_(gate_weight) + model.dummy.up_proj.weight.data.copy_(up_weight) + + apply_transforms_to_meta_model( + model, + [ + WeightRenaming("dummy.up_proj.weight", "mlp.up_proj.weight"), + WeightConverter( + ["mlp.gate_proj.weight", "mlp.up_proj.weight"], + "mlp.gate_up_proj.weight", + operations=[Concatenate(dim=0)], + ), + ], + ) + + state_dict = model.state_dict() + + # Fused weight exists; both source weights are gone + self.assertIn("mlp.gate_up_proj.weight", state_dict) + self.assertNotIn("mlp.gate_proj.weight", state_dict) + self.assertNotIn("mlp.up_proj.weight", state_dict) + + # Source modules are cleaned up; fused module is present + self.assertNotIn("gate_proj", model.mlp._modules) + self.assertNotIn("up_proj", model.mlp._modules) + self.assertIn("gate_up_proj", model.mlp._modules) + + # Values: gate comes first (as declared in source_patterns), then up + expected = torch.cat([gate_weight, up_weight], dim=0) + torch.testing.assert_close(state_dict["mlp.gate_up_proj.weight"], expected) + + if __name__ == "__main__": unittest.main() From b10b864ae4998d51597f99b437348a7f52aad5e7 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Thu, 7 May 2026 13:51:26 -0400 Subject: [PATCH 34/49] test: add kernels test --- fused_qwen_example.py | 2 +- tests/kernels/test_kernels.py | 186 ++++++++++++++++++++++++++++++++++ 2 files changed, 187 insertions(+), 1 deletion(-) diff --git a/fused_qwen_example.py b/fused_qwen_example.py index 0be834b6a97d..b54061a07cb8 100644 --- a/fused_qwen_example.py +++ b/fused_qwen_example.py @@ -24,7 +24,7 @@ print("Loading fused model...") kernel_repo_id = "michaelbenayoun/dummy-rmsnorm-mlp:RMSNormMLP" -kernel_repo_id = "michaelbenayoun/dummy-rmsnorm-mlp-with-transformations:RMSNormMLP" +#kernel_repo_id = "michaelbenayoun/dummy-rmsnorm-mlp-with-transformations:RMSNormMLP" kernel_config = KernelConfig( { ( diff --git a/tests/kernels/test_kernels.py b/tests/kernels/test_kernels.py index 1bd9a7c79792..54dabed2de3e 100644 --- a/tests/kernels/test_kernels.py +++ b/tests/kernels/test_kernels.py @@ -21,13 +21,22 @@ import types from unittest.mock import MagicMock, patch +import torch +import torch.nn as nn + from transformers import AutoModelForCausalLM, AutoTokenizer, KernelConfig +from transformers.conversion_mapping import get_checkpoint_conversion_mapping +from transformers.core_model_loading import Concatenate, WeightConverter, WeightRenaming, rename_source_key from transformers.integrations.hub_kernels import ( _HUB_KERNEL_MAPPING, _KERNEL_MODULE_MAPPING, + FusedModuleBase, + infer_kernel_fusion_transforms, is_kernel, lazy_load_kernel, load_and_register_attn_kernel, + make_fused_parent_class, + register_kernel_fusions, ) from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS @@ -207,6 +216,54 @@ def test_kernels_mapping(self): del model + def _kernel_fusion_tuple_matches_baseline(self, with_transformations: bool): + model_id = "michaelbenayoun/qwen3-tiny-4kv-heads-4layers-random" + + if with_transformations: + kernel_config = KernelConfig( + { + ( + ("RMSNorm", "model.layers.*.post_attention_layernorm"), + ("MLP", "model.layers.*.mlp"), + ): "michaelbenayoun/dummy-rmsnorm-mlp-with-transformations:RMSNormMLP", + } + ) + else: + kernel_config = KernelConfig( + { + ( + ("RMSNorm", "model.layers.*.post_attention_layernorm"), + ("MLP", "model.layers.*.mlp"), + ): "michaelbenayoun/dummy-rmsnorm-mlp:RMSNormMLP", + } + ) + + tokenizer = AutoTokenizer.from_pretrained(model_id) + inputs = tokenizer("Hello, how are you?", return_tensors="pt") + + baseline = AutoModelForCausalLM.from_pretrained(model_id, device_map=torch_device) + baseline.eval() + inputs = {k: v.to(torch_device) for k, v in inputs.items()} + with torch.no_grad(): + baseline_out = baseline(**inputs).logits + del baseline + + fused = AutoModelForCausalLM.from_pretrained( + model_id, use_kernels=True, kernel_config=kernel_config, device_map=torch_device + ) + fused.eval() + with torch.no_grad(): + fused_out = fused(**inputs).logits + del fused + + torch.testing.assert_close(baseline_out, fused_out, atol=1e-4, rtol=1e-4) + + def test_kernel_fusion_tuple_matches_baseline_without_transformations(self): + self._kernel_fusion_tuple_matches_baseline(with_transformations=False) + + def test_kernel_fusion_tuple_matches_baseline_with_transformations(self): + self._kernel_fusion_tuple_matches_baseline(with_transformations=True) + def test_faulty_kernel_mapping_layer_name(self): kernel_config = KernelConfig(kernel_mapping={"RMSNorm1": "kernels-community/layer_norm:LlamaRMSNorm"}) with self.assertRaises(ValueError): @@ -536,3 +593,132 @@ def test_single_device_mapping_still_works(self): result_mapping = kernel_config.kernel_mapping self.assertIn("RMSNorm", result_mapping, "RMSNorm should be in mapping") + + +class TestKernelFusions(TestCasePlus): + _MODEL_TYPE = "kernel_fusion_test_model" + + def tearDown(self): + import transformers.conversion_mapping as _cm + + if _cm._checkpoint_conversion_mapping_cache is not None: + _cm._checkpoint_conversion_mapping_cache.pop(self._MODEL_TYPE, None) + + def _make_layer_and_model(self): + class Norm(nn.Module): + kernel_layer_name = "RMSNorm" + + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.ones(4)) + + class MLP(nn.Module): + kernel_layer_name = "MLP" + + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(4, 8, bias=False) + self.up_proj = nn.Linear(4, 8, bias=False) + self.down_proj = nn.Linear(8, 4, bias=False) + + class Layer(nn.Module): + def __init__(self): + super().__init__() + self.post_attention_layernorm = Norm() + self.mlp = MLP() + + class Config: + model_type = self._MODEL_TYPE + + class Model(nn.Module): + config_class = Config + + def __init__(self, config=None): + super().__init__() + self.layers = nn.ModuleList([Layer(), Layer()]) + + return Layer, Model, Config + + def test_infer_kernel_fusion_transforms(self): + transforms = infer_kernel_fusion_transforms( + [ + ("RMSNorm", "model.layers.*.post_attention_layernorm"), + ("MLP", "model.layers.*.mlp"), + ] + ) + + first_key = "model.layers.0.post_attention_layernorm.weight" + self.assertEqual( + rename_source_key(first_key, transforms, [])[0], + "model.layers.0.post_attention_layernorm.RMSNorm.weight", + ) + + second_key = "model.layers.2.mlp.gate_proj.weight" + self.assertEqual( + rename_source_key(second_key, transforms, [])[0], + "model.layers.2.post_attention_layernorm.MLP.gate_proj.weight", + ) + + def test_make_fused_parent_class(self): + Layer, _, _ = self._make_layer_and_model() + FusedLayer = make_fused_parent_class( + Layer, + child_names=["post_attention_layernorm", "mlp"], + source_names=["RMSNorm", "MLP"], + kernel_layer_name="RMSNormMLP", + ) + layer = FusedLayer() + + self.assertIsInstance(layer.post_attention_layernorm, FusedModuleBase) + self.assertIsInstance(layer.mlp, nn.Identity) + self.assertIn("RMSNorm", layer.post_attention_layernorm._modules) + self.assertIn("MLP", layer.post_attention_layernorm._modules) + + def test_register_kernel_fusions_replaces_tuple_key(self): + _, Model, Config = self._make_layer_and_model() + kernel_config = KernelConfig( + { + ( + ("RMSNorm", "layers.*.post_attention_layernorm"), + ("MLP", "layers.*.mlp"), + ): "dummy/fake-kernel:RMSNormMLP" + } + ) + + with ( + patch("transformers.integrations.hub_kernels._try_load_kernel_class", return_value=None), + patch("transformers.integrations.hub_kernels.register_patch_mapping"), + ): + register_kernel_fusions(Model, Config(), kernel_config) + + self.assertIn("RMSNormMLP", kernel_config.kernel_mapping) + self.assertFalse(any(isinstance(k, tuple) for k in kernel_config.kernel_mapping)) + + def test_register_kernel_fusions_merges_kernel_conversion_mapping(self): + """WeightRenaming from structure + WeightConverter from kernel.conversion_mapping are both registered.""" + _, Model, Config = self._make_layer_and_model() + + class FakeKernel: + conversion_mapping = [ + WeightConverter(["MLP.gate_proj", "MLP.up_proj"], "MLP.gate_up_proj", [Concatenate(dim=0)]) + ] + + kernel_config = KernelConfig( + { + ( + ("RMSNorm", "layers.*.post_attention_layernorm"), + ("MLP", "layers.*.mlp"), + ): "dummy/fake-kernel:RMSNormMLP" + } + ) + + with ( + patch("transformers.integrations.hub_kernels._try_load_kernel_class", return_value=FakeKernel), + patch("transformers.integrations.hub_kernels.register_patch_mapping"), + ): + register_kernel_fusions(Model, Config(), kernel_config) + + mapping = get_checkpoint_conversion_mapping(self._MODEL_TYPE) + self.assertIsNotNone(mapping) + self.assertTrue(any(isinstance(t, WeightRenaming) for t in mapping)) + self.assertTrue(any(isinstance(t, WeightConverter) for t in mapping)) From fb0a748c7c4f21b6897a8faf18bbfd06c3962366 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 12 May 2026 11:42:52 +0200 Subject: [PATCH 35/49] test: fix hub_kernels package reload --- tests/kernels/test_kernels.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/kernels/test_kernels.py b/tests/kernels/test_kernels.py index 54dabed2de3e..57a37fd0fe8d 100644 --- a/tests/kernels/test_kernels.py +++ b/tests/kernels/test_kernels.py @@ -282,24 +282,24 @@ def test_faulty_kernel_mapping_type(self): @require_kernels class TestKernelsEnv(TestCasePlus): def test_disable_hub_kernels(self): - with patch.dict(os.environ, {"USE_HUB_KERNELS": "OFF"}): - import importlib + import importlib - from transformers.integrations import hub_kernels + from transformers.integrations import hub_kernels + with patch.dict(os.environ, {"USE_HUB_KERNELS": "OFF"}): importlib.reload(hub_kernels) - self.assertFalse(hub_kernels._kernels_enabled) + importlib.reload(hub_kernels) def test_enable_hub_kernels(self): - with patch.dict(os.environ, {"USE_HUB_KERNELS": "ON"}): - import importlib + import importlib - from transformers.integrations import hub_kernels + from transformers.integrations import hub_kernels + with patch.dict(os.environ, {"USE_HUB_KERNELS": "ON"}): importlib.reload(hub_kernels) - self.assertTrue(hub_kernels._kernels_enabled) + importlib.reload(hub_kernels) @require_kernels @@ -669,7 +669,7 @@ def test_make_fused_parent_class(self): ) layer = FusedLayer() - self.assertIsInstance(layer.post_attention_layernorm, FusedModuleBase) + self.assertIsInstance(layer.post_attention_layernorm, hub_kernels_pkg.FusedModuleBase) self.assertIsInstance(layer.mlp, nn.Identity) self.assertIn("RMSNorm", layer.post_attention_layernorm._modules) self.assertIn("MLP", layer.post_attention_layernorm._modules) From ded8b5fc1b4eea61139c53aaf6a4ae5cee4a21ec Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 12 May 2026 11:43:31 +0200 Subject: [PATCH 36/49] style: ruff --- tests/kernels/test_kernels.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/kernels/test_kernels.py b/tests/kernels/test_kernels.py index 57a37fd0fe8d..469f0e597167 100644 --- a/tests/kernels/test_kernels.py +++ b/tests/kernels/test_kernels.py @@ -30,7 +30,6 @@ from transformers.integrations.hub_kernels import ( _HUB_KERNEL_MAPPING, _KERNEL_MODULE_MAPPING, - FusedModuleBase, infer_kernel_fusion_transforms, is_kernel, lazy_load_kernel, From 318553e375843dffe8e413171a2875b6e89662d6 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 12 May 2026 15:53:09 +0200 Subject: [PATCH 37/49] refactor: do not create dynamic classes in test --- tests/kernels/test_kernels.py | 69 ++++++++++++++++------------------- 1 file changed, 31 insertions(+), 38 deletions(-) diff --git a/tests/kernels/test_kernels.py b/tests/kernels/test_kernels.py index 469f0e597167..9f932d040ad5 100644 --- a/tests/kernels/test_kernels.py +++ b/tests/kernels/test_kernels.py @@ -597,46 +597,42 @@ def test_single_device_mapping_still_works(self): class TestKernelFusions(TestCasePlus): _MODEL_TYPE = "kernel_fusion_test_model" - def tearDown(self): - import transformers.conversion_mapping as _cm - - if _cm._checkpoint_conversion_mapping_cache is not None: - _cm._checkpoint_conversion_mapping_cache.pop(self._MODEL_TYPE, None) + # Define dummy model structure to test the kernel fusion utilities. + class Norm(nn.Module): + kernel_layer_name = "RMSNorm" - def _make_layer_and_model(self): - class Norm(nn.Module): - kernel_layer_name = "RMSNorm" + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.ones(4)) - def __init__(self): - super().__init__() - self.weight = nn.Parameter(torch.ones(4)) + class MLP(nn.Module): + kernel_layer_name = "MLP" - class MLP(nn.Module): - kernel_layer_name = "MLP" + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(4, 8, bias=False) + self.up_proj = nn.Linear(4, 8, bias=False) + self.down_proj = nn.Linear(8, 4, bias=False) - def __init__(self): - super().__init__() - self.gate_proj = nn.Linear(4, 8, bias=False) - self.up_proj = nn.Linear(4, 8, bias=False) - self.down_proj = nn.Linear(8, 4, bias=False) + class Layer(nn.Module): + def __init__(self): + super().__init__() + self.post_attention_layernorm = TestKernelFusions.Norm() + self.mlp = TestKernelFusions.MLP() - class Layer(nn.Module): - def __init__(self): - super().__init__() - self.post_attention_layernorm = Norm() - self.mlp = MLP() + class Model(nn.Module): + class config_class: + model_type = "kernel_fusion_test_model" - class Config: - model_type = self._MODEL_TYPE + def __init__(self, config=None): + super().__init__() + self.layers = nn.ModuleList([TestKernelFusions.Layer(), TestKernelFusions.Layer()]) - class Model(nn.Module): - config_class = Config - - def __init__(self, config=None): - super().__init__() - self.layers = nn.ModuleList([Layer(), Layer()]) + def tearDown(self): + import transformers.conversion_mapping as _cm - return Layer, Model, Config + if _cm._checkpoint_conversion_mapping_cache is not None: + _cm._checkpoint_conversion_mapping_cache.pop(self._MODEL_TYPE, None) def test_infer_kernel_fusion_transforms(self): transforms = infer_kernel_fusion_transforms( @@ -659,9 +655,8 @@ def test_infer_kernel_fusion_transforms(self): ) def test_make_fused_parent_class(self): - Layer, _, _ = self._make_layer_and_model() FusedLayer = make_fused_parent_class( - Layer, + TestKernelFusions.Layer, child_names=["post_attention_layernorm", "mlp"], source_names=["RMSNorm", "MLP"], kernel_layer_name="RMSNormMLP", @@ -674,7 +669,6 @@ def test_make_fused_parent_class(self): self.assertIn("MLP", layer.post_attention_layernorm._modules) def test_register_kernel_fusions_replaces_tuple_key(self): - _, Model, Config = self._make_layer_and_model() kernel_config = KernelConfig( { ( @@ -688,14 +682,13 @@ def test_register_kernel_fusions_replaces_tuple_key(self): patch("transformers.integrations.hub_kernels._try_load_kernel_class", return_value=None), patch("transformers.integrations.hub_kernels.register_patch_mapping"), ): - register_kernel_fusions(Model, Config(), kernel_config) + register_kernel_fusions(TestKernelFusions.Model, None, kernel_config) self.assertIn("RMSNormMLP", kernel_config.kernel_mapping) self.assertFalse(any(isinstance(k, tuple) for k in kernel_config.kernel_mapping)) def test_register_kernel_fusions_merges_kernel_conversion_mapping(self): """WeightRenaming from structure + WeightConverter from kernel.conversion_mapping are both registered.""" - _, Model, Config = self._make_layer_and_model() class FakeKernel: conversion_mapping = [ @@ -715,7 +708,7 @@ class FakeKernel: patch("transformers.integrations.hub_kernels._try_load_kernel_class", return_value=FakeKernel), patch("transformers.integrations.hub_kernels.register_patch_mapping"), ): - register_kernel_fusions(Model, Config(), kernel_config) + register_kernel_fusions(TestKernelFusions.Model, None, kernel_config) mapping = get_checkpoint_conversion_mapping(self._MODEL_TYPE) self.assertIsNotNone(mapping) From 2a247602505f47f1b6947e416a21626ba9eea77f Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 12 May 2026 16:00:35 +0200 Subject: [PATCH 38/49] refactor: no dynamic class creation in tests --- tests/utils/test_core_model_loading.py | 86 ++++++++------------------ 1 file changed, 25 insertions(+), 61 deletions(-) diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 646cb436a4ca..e174b70182c4 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -1254,15 +1254,27 @@ class _TstRootSame(PreTrainedModel): ... class TestApplyTransformsToMetaModel(unittest.TestCase): - def test_weight_renaming_renames_parameter(self): - class SimpleModel(nn.Module): - base_model_prefix = "" + class SimpleModel(nn.Module): + base_model_prefix = "" - def __init__(self): - super().__init__() - self.linear = nn.Linear(4, 4, bias=False) + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4, bias=False) + + class QKVAttn(nn.Module): + def __init__(self): + super().__init__() + self.qkv_proj = nn.Linear(2, 6, bias=False) + + class QKVModel(nn.Module): + base_model_prefix = "" - model = SimpleModel() + def __init__(self): + super().__init__() + self.self_attn = TestApplyTransformsToMetaModel.QKVAttn() + + def test_weight_renaming_renames_parameter(self): + model = TestApplyTransformsToMetaModel.SimpleModel() original_weight = model.linear.weight.data.clone() self.assertIn("linear.weight", model.state_dict()) @@ -1274,14 +1286,7 @@ def __init__(self): torch.testing.assert_close(state_dict["linear.kernel"], original_weight) def test_weight_renaming_is_idempotent(self): - class SimpleModel(nn.Module): - base_model_prefix = "" - - def __init__(self): - super().__init__() - self.linear = nn.Linear(4, 4, bias=False) - - model = SimpleModel() + model = TestApplyTransformsToMetaModel.SimpleModel() apply_transforms_to_meta_model(model, [WeightRenaming("linear.weight", "linear.kernel")]) keys_after_first = set(model.state_dict().keys()) @@ -1289,14 +1294,7 @@ def __init__(self): self.assertEqual(keys_after_first, set(model.state_dict().keys())) def test_weight_renaming_no_match_is_noop(self): - class SimpleModel(nn.Module): - base_model_prefix = "" - - def __init__(self): - super().__init__() - self.linear = nn.Linear(4, 4, bias=False) - - model = SimpleModel() + model = TestApplyTransformsToMetaModel.SimpleModel() original_keys = set(model.state_dict().keys()) apply_transforms_to_meta_model(model, [WeightRenaming("non_existent.weight", "something.weight")]) @@ -1304,19 +1302,7 @@ def __init__(self): self.assertEqual(original_keys, set(model.state_dict().keys())) def test_weight_converter_splits_fused_projection(self): - class QKVAttn(nn.Module): - def __init__(self): - super().__init__() - self.qkv_proj = nn.Linear(2, 6, bias=False) - - class QKVModel(nn.Module): - base_model_prefix = "" - - def __init__(self): - super().__init__() - self.self_attn = QKVAttn() - - model = QKVModel() + model = TestApplyTransformsToMetaModel.QKVModel() model.config = PretrainedConfig() qkv_weight = torch.randn(6, 2) model.self_attn.qkv_proj.weight.data.copy_(qkv_weight) @@ -1349,18 +1335,6 @@ def __init__(self): torch.testing.assert_close(state_dict["self_attn.v_proj.weight"], expected_v) def test_weight_converter_split_is_idempotent(self): - class QKVAttn(nn.Module): - def __init__(self): - super().__init__() - self.qkv_proj = nn.Linear(2, 6, bias=False) - - class QKVModel(nn.Module): - base_model_prefix = "" - - def __init__(self): - super().__init__() - self.self_attn = QKVAttn() - def make_mapping(): return [ WeightConverter( @@ -1370,7 +1344,7 @@ def make_mapping(): ) ] - model = QKVModel() + model = TestApplyTransformsToMetaModel.QKVModel() model.config = PretrainedConfig() apply_transforms_to_meta_model(model, make_mapping()) @@ -1380,15 +1354,10 @@ def make_mapping(): self.assertEqual(keys_after_first, set(model.state_dict().keys())) def test_weight_converter_splits_across_multiple_layers(self): - class QKVAttn(nn.Module): - def __init__(self): - super().__init__() - self.qkv_proj = nn.Linear(2, 6, bias=False) - class QKVLayer(nn.Module): def __init__(self): super().__init__() - self.self_attn = QKVAttn() + self.self_attn = TestApplyTransformsToMetaModel.QKVAttn() class QKVModel(nn.Module): base_model_prefix = "model" @@ -1484,11 +1453,6 @@ def __init__(self): torch.testing.assert_close(state_dict["self_attn.qkv_proj.weight"], expected_qkv) def test_multiple_independent_converters_in_single_call(self): - class Attn(nn.Module): - def __init__(self): - super().__init__() - self.qkv_proj = nn.Linear(2, 6, bias=False) - class MLP(nn.Module): def __init__(self): super().__init__() @@ -1500,7 +1464,7 @@ class TwoHeadModel(nn.Module): def __init__(self): super().__init__() - self.self_attn = Attn() + self.self_attn = TestApplyTransformsToMetaModel.QKVAttn() self.mlp = MLP() model = TwoHeadModel() From 81367813568bc0a4a15e7997f563addd2b755445 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Wed, 13 May 2026 10:42:44 +0200 Subject: [PATCH 39/49] refactor: test --- tests/kernels/test_kernels.py | 3 ++ tests/utils/test_core_model_loading.py | 38 ++++++++++++++++++++++++-- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/tests/kernels/test_kernels.py b/tests/kernels/test_kernels.py index 9f932d040ad5..85cd7a37853b 100644 --- a/tests/kernels/test_kernels.py +++ b/tests/kernels/test_kernels.py @@ -594,6 +594,7 @@ def test_single_device_mapping_still_works(self): self.assertIn("RMSNorm", result_mapping, "RMSNorm should be in mapping") +@require_kernels class TestKernelFusions(TestCasePlus): _MODEL_TYPE = "kernel_fusion_test_model" @@ -630,9 +631,11 @@ def __init__(self, config=None): def tearDown(self): import transformers.conversion_mapping as _cm + from transformers.integrations.hub_kernels import make_fused_module_class if _cm._checkpoint_conversion_mapping_cache is not None: _cm._checkpoint_conversion_mapping_cache.pop(self._MODEL_TYPE, None) + make_fused_module_class.cache_clear() def test_infer_kernel_fusion_transforms(self): transforms = infer_kernel_fusion_transforms( diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index e174b70182c4..cdda7a169ff3 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -1359,14 +1359,14 @@ def __init__(self): super().__init__() self.self_attn = TestApplyTransformsToMetaModel.QKVAttn() - class QKVModel(nn.Module): + class QKVLayeredModel(nn.Module): base_model_prefix = "model" def __init__(self): super().__init__() self.model = nn.ModuleList([QKVLayer(), QKVLayer()]) - model = QKVModel() + model = QKVLayeredModel() model.config = PretrainedConfig() qkv_weights = {} @@ -1402,6 +1402,40 @@ def __init__(self): torch.testing.assert_close(state_dict[f"model.{i}.self_attn.k_proj.weight"], expected_k) torch.testing.assert_close(state_dict[f"model.{i}.self_attn.v_proj.weight"], expected_v) + def test_weight_converter_glob_merges_modulelist(self): + class ExpertModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.experts = nn.ModuleList([nn.Linear(2, 4, bias=False), nn.Linear(2, 4, bias=False)]) + # Pre-create empty target so _get_submodule_or_create finds it without + # needing to infer a module type from the 3-D stacked tensor. + self.merged = nn.Module() + + model = ExpertModel() + model.config = PretrainedConfig() + + weight_0 = model.experts[0].weight.data.clone() + weight_1 = model.experts[1].weight.data.clone() + + apply_transforms_to_meta_model( + model, + [ + WeightConverter( + "experts.*.weight", + "merged.weight", + operations=[MergeModulelist(dim=0)], + ) + ], + ) + + state_dict = model.state_dict() + self.assertIn("merged.weight", state_dict) + self.assertNotIn("experts.0.weight", state_dict) + self.assertNotIn("experts.1.weight", state_dict) + torch.testing.assert_close(state_dict["merged.weight"], torch.stack([weight_0, weight_1], dim=0)) + def test_weight_converter_concatenates_projections(self): class SeparateAttn(nn.Module): def __init__(self): From 17f4a9e9e6bec5260307cffb7645211e67b70c8f Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Wed, 13 May 2026 15:04:40 +0200 Subject: [PATCH 40/49] fix: TYPE_CHECKING imports were broken --- src/transformers/integrations/hub_kernels.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 71c9bd314b79..fdd9ef7df475 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -31,9 +31,9 @@ if TYPE_CHECKING: - from .configuration_utils import PretrainedConfig - from .modeling_utils import PreTrainedModel - from .utils.kernel_config import KernelConfig + from ..configuration_utils import PretrainedConfig + from ..modeling_utils import PreTrainedModel + from ..utils.kernel_config import KernelConfig if is_torch_available(): import torch From fabadcadf4c0be3a9c7a079f402e4a2c07af1a42 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 19 May 2026 11:58:36 +0200 Subject: [PATCH 41/49] wip: get rid of transform_model methods --- src/transformers/core_model_loading.py | 211 ------------------- src/transformers/integrations/hub_kernels.py | 132 ++++++++++++ src/transformers/modeling_utils.py | 8 +- tests/utils/test_core_model_loading.py | 2 +- 4 files changed, 138 insertions(+), 215 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 6427ce76a42e..dcf16f97f6c9 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -791,46 +791,6 @@ def was_used(self) -> bool: """ return self._was_used - @abstractmethod - def convert( - self, - layer_name: str, - model=None, - config=None, - hf_quantizer=None, - loading_info: LoadStateDictInfo | None = None, - ) -> dict[str, torch.Tensor]: ... - - @abstractmethod - def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: - """ - Analogous to `convert` but restructures `model`'s module graph rather than operating on tensor data. - - Must be idempotent, i.e. if called multiple times on the same `model` and `layer_name`, only the first call - should have an effect. - """ - ... - - def post_transformation_cleanup(self, model: PreTrainedModel, source_names: list[str]) -> None: - """ - After transforming the model's module graph, we may end up with empty modules that held the original weights, - we remove them. - """ - for source_name in source_names: - mod_path, _ = source_name.rsplit(".", 1) if "." in source_name else ("", source_name) - try: - source_mod = model.get_submodule(mod_path) if mod_path else model - except AttributeError: - continue - if ( - all(p is None for p in source_mod._parameters.values()) - and all(b is None for b in source_mod._buffers.values()) - and all(m is None for m in source_mod._modules.values()) - ): - parent_path, _, mod_name = mod_path.rpartition(".") - parent_mod = model.get_submodule(parent_path) if parent_path else model - parent_mod._modules.pop(mod_name, None) - class WeightRenaming(WeightTransform): # Special case of WeightTransform that only renames keys without any conversion. @@ -871,44 +831,6 @@ def convert( return collected_tensors - def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: - source_names = list(self.layer_targets.get(layer_name, [])) - if not source_names: - return - source_name = source_names[0] - if source_name == layer_name: - return - - source_mod_path, _, source_attr = source_name.rpartition(".") - target_mod_path, _, target_attr = layer_name.rpartition(".") - - # Idempotency guard: if the target already holds the parameter/buffer, the - # model is already in the correct shape (e.g. from monkey-patching before - # from_pretrained). Nothing to move. - try: - target_mod = model.get_submodule(target_mod_path) if target_mod_path else model - if target_attr in target_mod._parameters or target_attr in target_mod._buffers: - return - except AttributeError: - pass - - try: - source_mod = model.get_submodule(source_mod_path) if source_mod_path else model - except AttributeError: - return # already moved - - if source_attr in source_mod._parameters: - obj = source_mod._parameters.pop(source_attr) - _get_submodule_or_create( - model, target_mod_path, original_module=source_mod, future_weight=obj - ).register_parameter(target_attr, obj) - elif source_attr in source_mod._buffers: - obj = source_mod._buffers.pop(source_attr) - _get_submodule_or_create(model, target_mod_path).register_buffer(target_attr, obj) - - # Cleanup any empty modules that may be left after moving the parameter/buffer. - self.post_transformation_cleanup(model, source_names) - class PrefixChange(WeightRenaming): """ @@ -1035,139 +957,6 @@ def convert( ) return collected_tensors - def transform_model(self, model: PreTrainedModel, layer_name: str) -> None: - source_names = list(self.layer_targets.get(layer_name, [])) - if not source_names: - return - - # Idempotency guard: if the first target parameter already exists, the model is - # already in the correct shape. Nothing to convert. - t_mod_path, _, t_attr = layer_name.rpartition(".") - try: - t_mod = model.get_submodule(t_mod_path) if t_mod_path else model - if t_attr in t_mod._parameters or t_attr in t_mod._buffers: - return - except AttributeError: - pass - - collected_tensors = self.materialize_tensors() - - result: dict[str, Any] = dict(collected_tensors) - for op in self.operations: - result = op.convert(result, source_patterns=self.source_patterns, target_patterns=self.target_patterns) - - # Expand target pattern keys to full parameter names (same logic as convert) - full_name = layer_name.replace(".*.", ".0.") if ".*." in layer_name else layer_name - try: - prefix, _, suffix = next(full_name.partition(k) for k in result if k in full_name) - result = {prefix + k + suffix: v for k, v in result.items()} - except StopIteration: - pass - - # Retrieve one of the source modules to use as original module when creating missing target modules - original_mod: type[torch.nn.Module] | None = None - for sname in source_names: - smod_path = sname.rpartition(".")[0] - try: - original_mod = model.get_submodule(smod_path) if smod_path else model - break - except AttributeError: - pass - - # Remove source parameters - for source_name in source_names: - mod_path, _, attr = source_name.rpartition(".") - try: - source_mod = model.get_submodule(mod_path) if mod_path else model - except AttributeError: - continue - source_mod._parameters.pop(attr, None) - - # Register the target parameters, creating any missing module at the leaf path. - for target_name, meta_val in result.items(): - meta_tensor = meta_val[0] if isinstance(meta_val, list) else meta_val - t_mod_path, _, t_attr = target_name.rpartition(".") - target_mod = _get_submodule_or_create( - model, t_mod_path, original_module=original_mod, future_weight=meta_tensor - ) - if not isinstance(meta_tensor, torch.nn.Parameter): - meta_tensor = torch.nn.Parameter(meta_tensor, requires_grad=meta_tensor.is_floating_point()) - target_mod.register_parameter(t_attr, meta_tensor) - - # Remove source modules if they are empty after the transform. - self.post_transformation_cleanup(model, source_names) - - -def _get_submodule_or_create( - model: torch.nn.Module, - module_path: str, - original_module: torch.nn.Module | None = None, - future_weight: torch.Tensor | None = None, -) -> torch.nn.Module: - if not module_path: - return model - - if original_module is not None and future_weight is None: - raise ValueError("original_module cannot be provided without future_weight to infer the module class") - parent, _, mod_name = module_path.rpartition(".") - parent_mod = model.get_submodule(parent) if parent else model - if mod_name not in parent_mod._modules: - cls_ = original_module.__class__ if original_module is not None else torch.nn.Module - if cls_ is torch.nn.Linear: - out_features, in_features = future_weight.shape - mod = cls_(in_features=in_features, out_features=out_features, bias=original_module.bias is not None) - elif cls_ is torch.nn.Embedding: - num_embeddings, embedding_dim = future_weight.shape - mod = cls_( - num_embeddings=num_embeddings, - embedding_dim=embedding_dim, - padding_idx=original_module.padding_idx, - max_norm=original_module.max_norm, - norm_type=original_module.norm_type, - scale_grad_by_freq=original_module.scale_grad_by_freq, - sparse=original_module.sparse, - ) - else: - raise ValueError( - f"Cannot infer instantiation for {cls_} when creating submodule {mod_name} at path {module_path}" - ) - parent_mod.add_module(mod_name, mod) - else: - mod = parent_mod._modules[mod_name] - - return mod - - -def apply_transforms_to_meta_model( - model: PreTrainedModel, - weight_conversions: list[WeightRenaming | WeightConverter], -) -> None: - """ - Restructure `model`'s module graph in-place according to `weight_conversions`, analogously to - `convert_and_load_state_dict_in_model` but operating on module attributes rather than checkpoint tensors. - - Designed to be called while the model is on the meta device, making the operation cost-free. - """ - meta_state_dict = model.state_dict() - - renamings = [x for x in weight_conversions if isinstance(x, WeightRenaming)] - converters = [x for x in weight_conversions if isinstance(x, WeightConverter)] - pattern_to_converter = {k: c for c in converters for k in c.source_patterns} - param_name_to_transform: dict[str, WeightRenaming | WeightConverter] = {} - - for param_name in sorted(meta_state_dict, key=dot_natural_key): - renamed_key, source_pattern = rename_source_key(param_name, renamings, converters) - if source_pattern is not None: - transform = param_name_to_transform.setdefault(renamed_key, deepcopy(pattern_to_converter[source_pattern])) - else: - transform = param_name_to_transform.setdefault(renamed_key, WeightRenaming(param_name, renamed_key)) - source_pattern = param_name - - transform.add_tensor(renamed_key, param_name, source_pattern, meta_state_dict[param_name]) - - for layer_name, transform in param_name_to_transform.items(): - transform.transform_model(model, layer_name) - # For I/O bound operations (i.e. here reading files), it is better to have fewer threads, e.g. 4 is a good default. # Having too many is actually harming performances quite a lot, i.e. using 16 can sometimes lead to taking TWICE diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index fdd9ef7df475..9837c1980d10 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -17,6 +17,7 @@ import re from collections.abc import Callable from contextlib import contextmanager +from copy import deepcopy from types import ModuleType from typing import TYPE_CHECKING @@ -520,6 +521,136 @@ def allow_all_hub_kernels(): ALLOW_ALL_KERNELS = False +def _get_submodule_or_create( + model: "nn.Module", + module_path: str, + original_module: "nn.Module | None" = None, + future_weight: "torch.Tensor | None" = None, +) -> "nn.Module": + if not module_path: + return model + if original_module is not None and future_weight is None: + raise ValueError("original_module cannot be provided without future_weight to infer the module class") + parent, _, mod_name = module_path.rpartition(".") + parent_mod = model.get_submodule(parent) if parent else model + if mod_name not in parent_mod._modules: + cls_ = original_module.__class__ if original_module is not None else nn.Module + if cls_ is nn.Linear: + out_features, in_features = future_weight.shape + mod = cls_(in_features=in_features, out_features=out_features, bias=original_module.bias is not None) + elif cls_ is nn.Embedding: + num_embeddings, embedding_dim = future_weight.shape + mod = cls_( + num_embeddings=num_embeddings, + embedding_dim=embedding_dim, + padding_idx=original_module.padding_idx, + max_norm=original_module.max_norm, + norm_type=original_module.norm_type, + scale_grad_by_freq=original_module.scale_grad_by_freq, + sparse=original_module.sparse, + ) + else: + raise ValueError( + f"Cannot infer instantiation for {cls_} when creating submodule {mod_name} at path {module_path}" + ) + parent_mod.add_module(mod_name, mod) + else: + mod = parent_mod._modules[mod_name] + return mod + + +def _post_transformation_cleanup(model: "nn.Module", source_names: list[str]) -> None: + for source_name in source_names: + mod_path, _ = source_name.rsplit(".", 1) if "." in source_name else ("", source_name) + try: + source_mod = model.get_submodule(mod_path) if mod_path else model + except AttributeError: + continue + if ( + all(p is None for p in source_mod._parameters.values()) + and all(b is None for b in source_mod._buffers.values()) + and all(m is None for m in source_mod._modules.values()) + ): + parent_path, _, mod_name = mod_path.rpartition(".") + parent_mod = model.get_submodule(parent_path) if parent_path else model + parent_mod._modules.pop(mod_name, None) + + +def apply_kernel_structural_transforms( + model: "PreTrainedModel", + weight_conversions: list, +) -> None: + """ + Restructure `model`'s module graph according to `weight_conversions` while parameters are on the + meta device (cost-free). Called after model init and before the loader, so the loader sees the + correct parameter structure. + + Kept in hub_kernels rather than core_model_loading to leave the loading machinery clean. + """ + from ..core_model_loading import WeightConverter, WeightRenaming, dot_natural_key, rename_source_key + + meta_state_dict = model.state_dict() + + renamings = [x for x in weight_conversions if isinstance(x, WeightRenaming)] + converters = [x for x in weight_conversions if isinstance(x, WeightConverter)] + pattern_to_converter = {k: c for c in converters for k in c.source_patterns} + param_name_to_transform: dict = {} + + for param_name in sorted(meta_state_dict, key=dot_natural_key): + renamed_key, source_pattern = rename_source_key(param_name, renamings, converters) + if source_pattern is not None: + transform = param_name_to_transform.setdefault(renamed_key, deepcopy(pattern_to_converter[source_pattern])) + else: + transform = param_name_to_transform.setdefault(renamed_key, WeightRenaming(param_name, renamed_key)) + source_pattern = param_name + transform.add_tensor(renamed_key, param_name, source_pattern, meta_state_dict[param_name]) + + for layer_name, transform in param_name_to_transform.items(): + source_names = list(transform.layer_targets.get(layer_name, [])) + if not source_names: + continue + + # Skip identity renames: parameter is already at the correct path. + if len(source_names) == 1 and source_names[0] == layer_name: + continue + + # Find a source module for type inference when _get_submodule_or_create must create a new node. + original_mod = None + for sname in source_names: + smod_path = sname.rpartition(".")[0] + try: + original_mod = model.get_submodule(smod_path) if smod_path else model + break + except AttributeError: + pass + + # Use convert() on meta tensors — correct shapes/keys, cost-free on meta device. + result = transform.convert(layer_name) + + # Remove source parameters from the model graph. + for sname in source_names: + mod_path, _, attr = sname.rpartition(".") + try: + source_mod = model.get_submodule(mod_path) if mod_path else model + except AttributeError: + continue + source_mod._parameters.pop(attr, None) + source_mod._buffers.pop(attr, None) + + # Register target parameters, creating intermediate module nodes as needed. + for target_key, meta_val in result.items(): + meta_tensor = meta_val[0] if isinstance(meta_val, list) else meta_val + t_mod_path, _, t_attr = target_key.rpartition(".") + target_mod = _get_submodule_or_create( + model, t_mod_path, original_module=original_mod, future_weight=meta_tensor + ) + if not isinstance(meta_tensor, nn.Parameter): + meta_tensor = nn.Parameter(meta_tensor, requires_grad=meta_tensor.is_floating_point()) + target_mod.register_parameter(t_attr, meta_tensor) + + _post_transformation_cleanup(model, source_names) + + class FusedModuleBase(nn.Module): def __init__( self, @@ -797,6 +928,7 @@ def register_kernel_fusions( __all__ = [ "FusedModuleBase", "LayerRepository", + "apply_kernel_structural_transforms", "get_kernel", "register_kernel_fusions", "lazy_load_kernel", diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 78af178fe776..9db2df09a707 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -48,7 +48,6 @@ from .core_model_loading import ( WeightConverter, WeightRenaming, - apply_transforms_to_meta_model, convert_and_load_state_dict_in_model, revert_weight_conversion, ) @@ -4253,8 +4252,11 @@ def from_pretrained( # Obtain the weight conversion mapping for this model if any are registered and apply to all submodels recursively weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer) - # Apply any structural model transforms while parameters are still on the meta device (free operation). - apply_transforms_to_meta_model(model, weight_conversions) + # Apply kernel-specific structural transforms before loading (module graph changes, cost-free on meta device). + if use_kernels: + from .integrations.hub_kernels import apply_kernel_structural_transforms + + apply_kernel_structural_transforms(model, weight_conversions) if _torch_distributed_available and device_mesh is not None: # add hooks to nn.Modules: no weights model = distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size) diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index cdda7a169ff3..6299d634ad47 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -35,12 +35,12 @@ PrefixChange, WeightConverter, WeightRenaming, - apply_transforms_to_meta_model, build_glob_alternation, convert_and_load_state_dict_in_model, rename_source_key, revert_weight_conversion, ) +from transformers.integrations.hub_kernels import apply_kernel_structural_transforms as apply_transforms_to_meta_model from transformers.modeling_utils import LoadStateDictConfig from transformers.utils.import_utils import is_triton_available From e1f3a833a08ee61cc0cdc8f73dff1caee3252fe6 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 19 May 2026 15:09:52 +0200 Subject: [PATCH 42/49] wip: move tests --- tests/kernels/test_kernels.py | 350 ++++++++++++++++++++++++- tests/utils/test_core_model_loading.py | 348 ------------------------ 2 files changed, 348 insertions(+), 350 deletions(-) diff --git a/tests/kernels/test_kernels.py b/tests/kernels/test_kernels.py index 85cd7a37853b..6bbafb433fcd 100644 --- a/tests/kernels/test_kernels.py +++ b/tests/kernels/test_kernels.py @@ -19,17 +19,26 @@ import os import tempfile import types +import unittest from unittest.mock import MagicMock, patch import torch import torch.nn as nn -from transformers import AutoModelForCausalLM, AutoTokenizer, KernelConfig +from transformers import AutoModelForCausalLM, AutoTokenizer, KernelConfig, PretrainedConfig from transformers.conversion_mapping import get_checkpoint_conversion_mapping -from transformers.core_model_loading import Concatenate, WeightConverter, WeightRenaming, rename_source_key +from transformers.core_model_loading import ( + Chunk, + Concatenate, + MergeModulelist, + WeightConverter, + WeightRenaming, + rename_source_key, +) from transformers.integrations.hub_kernels import ( _HUB_KERNEL_MAPPING, _KERNEL_MODULE_MAPPING, + apply_kernel_structural_transforms, infer_kernel_fusion_transforms, is_kernel, lazy_load_kernel, @@ -717,3 +726,340 @@ class FakeKernel: self.assertIsNotNone(mapping) self.assertTrue(any(isinstance(t, WeightRenaming) for t in mapping)) self.assertTrue(any(isinstance(t, WeightConverter) for t in mapping)) + + +class TestApplyKernelStructuralTransforms(unittest.TestCase): + class SimpleModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4, bias=False) + + class QKVAttn(nn.Module): + def __init__(self): + super().__init__() + self.qkv_proj = nn.Linear(2, 6, bias=False) + + class QKVModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.self_attn = TestApplyKernelStructuralTransforms.QKVAttn() + + def test_weight_renaming_renames_parameter(self): + model = TestApplyKernelStructuralTransforms.SimpleModel() + original_weight = model.linear.weight.data.clone() + self.assertIn("linear.weight", model.state_dict()) + + apply_kernel_structural_transforms(model, [WeightRenaming("linear.weight", "linear.kernel")]) + + state_dict = model.state_dict() + self.assertIn("linear.kernel", state_dict) + self.assertNotIn("linear.weight", state_dict) + torch.testing.assert_close(state_dict["linear.kernel"], original_weight) + + def test_weight_renaming_is_idempotent(self): + model = TestApplyKernelStructuralTransforms.SimpleModel() + apply_kernel_structural_transforms(model, [WeightRenaming("linear.weight", "linear.kernel")]) + keys_after_first = set(model.state_dict().keys()) + + apply_kernel_structural_transforms(model, [WeightRenaming("linear.weight", "linear.kernel")]) + self.assertEqual(keys_after_first, set(model.state_dict().keys())) + + def test_weight_renaming_no_match_is_noop(self): + model = TestApplyKernelStructuralTransforms.SimpleModel() + original_keys = set(model.state_dict().keys()) + + apply_kernel_structural_transforms(model, [WeightRenaming("non_existent.weight", "something.weight")]) + + self.assertEqual(original_keys, set(model.state_dict().keys())) + + def test_weight_converter_splits_fused_projection(self): + model = TestApplyKernelStructuralTransforms.QKVModel() + model.config = PretrainedConfig() + qkv_weight = torch.randn(6, 2) + model.self_attn.qkv_proj.weight.data.copy_(qkv_weight) + + apply_kernel_structural_transforms( + model, + [ + WeightConverter( + "self_attn.qkv_proj.weight", + ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + operations=[Chunk(dim=0)], + ) + ], + ) + + state_dict = model.state_dict() + self.assertIn("self_attn.q_proj.weight", state_dict) + self.assertIn("self_attn.k_proj.weight", state_dict) + self.assertIn("self_attn.v_proj.weight", state_dict) + self.assertNotIn("self_attn.qkv_proj.weight", state_dict) + self.assertNotIn("qkv_proj", model.self_attn._modules) + self.assertIn("q_proj", model.self_attn._modules) + self.assertIn("k_proj", model.self_attn._modules) + self.assertIn("v_proj", model.self_attn._modules) + + expected_q, expected_k, expected_v = torch.chunk(qkv_weight, 3, dim=0) + torch.testing.assert_close(state_dict["self_attn.q_proj.weight"], expected_q) + torch.testing.assert_close(state_dict["self_attn.k_proj.weight"], expected_k) + torch.testing.assert_close(state_dict["self_attn.v_proj.weight"], expected_v) + + def test_weight_converter_split_is_idempotent(self): + def make_mapping(): + return [ + WeightConverter( + "self_attn.qkv_proj.weight", + ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + operations=[Chunk(dim=0)], + ) + ] + + model = TestApplyKernelStructuralTransforms.QKVModel() + model.config = PretrainedConfig() + + apply_kernel_structural_transforms(model, make_mapping()) + keys_after_first = set(model.state_dict().keys()) + + apply_kernel_structural_transforms(model, make_mapping()) + self.assertEqual(keys_after_first, set(model.state_dict().keys())) + + def test_weight_converter_splits_across_multiple_layers(self): + class QKVLayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = TestApplyKernelStructuralTransforms.QKVAttn() + + class QKVLayeredModel(nn.Module): + base_model_prefix = "model" + + def __init__(self): + super().__init__() + self.model = nn.ModuleList([QKVLayer(), QKVLayer()]) + + model = QKVLayeredModel() + model.config = PretrainedConfig() + + qkv_weights = {} + for i in range(2): + w = torch.randn(6, 2) + model.model[i].self_attn.qkv_proj.weight.data.copy_(w) + qkv_weights[i] = w + + apply_kernel_structural_transforms( + model, + [ + WeightConverter( + "self_attn.qkv_proj.weight", + ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + operations=[Chunk(dim=0)], + ) + ], + ) + + state_dict = model.state_dict() + self.assertNotIn("model.0.self_attn.qkv_proj.weight", state_dict) + self.assertNotIn("model.1.self_attn.qkv_proj.weight", state_dict) + + for i in range(2): + layer_attn = model.model[i].self_attn + self.assertNotIn("qkv_proj", layer_attn._modules) + self.assertIn("q_proj", layer_attn._modules) + self.assertIn("k_proj", layer_attn._modules) + self.assertIn("v_proj", layer_attn._modules) + + expected_q, expected_k, expected_v = torch.chunk(qkv_weights[i], 3, dim=0) + torch.testing.assert_close(state_dict[f"model.{i}.self_attn.q_proj.weight"], expected_q) + torch.testing.assert_close(state_dict[f"model.{i}.self_attn.k_proj.weight"], expected_k) + torch.testing.assert_close(state_dict[f"model.{i}.self_attn.v_proj.weight"], expected_v) + + def test_weight_converter_glob_merges_modulelist(self): + class ExpertModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.experts = nn.ModuleList([nn.Linear(2, 4, bias=False), nn.Linear(2, 4, bias=False)]) + self.merged = nn.Module() + + model = ExpertModel() + model.config = PretrainedConfig() + + weight_0 = model.experts[0].weight.data.clone() + weight_1 = model.experts[1].weight.data.clone() + + apply_kernel_structural_transforms( + model, + [ + WeightConverter( + "experts.*.weight", + "merged.weight", + operations=[MergeModulelist(dim=0)], + ) + ], + ) + + state_dict = model.state_dict() + self.assertIn("merged.weight", state_dict) + self.assertNotIn("experts.0.weight", state_dict) + self.assertNotIn("experts.1.weight", state_dict) + torch.testing.assert_close(state_dict["merged.weight"], torch.stack([weight_0, weight_1], dim=0)) + + def test_weight_converter_concatenates_projections(self): + class SeparateAttn(nn.Module): + def __init__(self): + super().__init__() + self.q_proj = nn.Linear(2, 2, bias=False) + self.k_proj = nn.Linear(2, 2, bias=False) + self.v_proj = nn.Linear(2, 2, bias=False) + + class SeparateModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.self_attn = SeparateAttn() + + model = SeparateModel() + model.config = PretrainedConfig() + + q_weight = torch.randn(2, 2) + k_weight = torch.randn(2, 2) + v_weight = torch.randn(2, 2) + model.self_attn.q_proj.weight.data.copy_(q_weight) + model.self_attn.k_proj.weight.data.copy_(k_weight) + model.self_attn.v_proj.weight.data.copy_(v_weight) + + apply_kernel_structural_transforms( + model, + [ + WeightConverter( + ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + "self_attn.qkv_proj.weight", + operations=[Concatenate(dim=0)], + ) + ], + ) + + state_dict = model.state_dict() + self.assertIn("self_attn.qkv_proj.weight", state_dict) + self.assertNotIn("self_attn.q_proj.weight", state_dict) + self.assertNotIn("self_attn.k_proj.weight", state_dict) + self.assertNotIn("self_attn.v_proj.weight", state_dict) + self.assertNotIn("q_proj", model.self_attn._modules) + self.assertNotIn("k_proj", model.self_attn._modules) + self.assertNotIn("v_proj", model.self_attn._modules) + self.assertIn("qkv_proj", model.self_attn._modules) + + expected_qkv = torch.cat([q_weight, k_weight, v_weight], dim=0) + torch.testing.assert_close(state_dict["self_attn.qkv_proj.weight"], expected_qkv) + + def test_multiple_independent_converters_in_single_call(self): + class MLP(nn.Module): + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(2, 4, bias=False) + self.up_proj = nn.Linear(2, 4, bias=False) + + class TwoHeadModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.self_attn = TestApplyKernelStructuralTransforms.QKVAttn() + self.mlp = MLP() + + model = TwoHeadModel() + model.config = PretrainedConfig() + + qkv_weight = torch.randn(6, 2) + gate_weight = torch.randn(4, 2) + up_weight = torch.randn(4, 2) + model.self_attn.qkv_proj.weight.data.copy_(qkv_weight) + model.mlp.gate_proj.weight.data.copy_(gate_weight) + model.mlp.up_proj.weight.data.copy_(up_weight) + + apply_kernel_structural_transforms( + model, + [ + WeightConverter( + "self_attn.qkv_proj.weight", + ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + operations=[Chunk(dim=0)], + ), + WeightConverter( + ["mlp.gate_proj.weight", "mlp.up_proj.weight"], + "mlp.gate_up_proj.weight", + operations=[Concatenate(dim=0)], + ), + ], + ) + + state_dict = model.state_dict() + + self.assertNotIn("self_attn.qkv_proj.weight", state_dict) + self.assertIn("self_attn.q_proj.weight", state_dict) + self.assertIn("self_attn.k_proj.weight", state_dict) + self.assertIn("self_attn.v_proj.weight", state_dict) + self.assertNotIn("qkv_proj", model.self_attn._modules) + self.assertIn("q_proj", model.self_attn._modules) + self.assertIn("k_proj", model.self_attn._modules) + self.assertIn("v_proj", model.self_attn._modules) + expected_q, expected_k, expected_v = torch.chunk(qkv_weight, 3, dim=0) + torch.testing.assert_close(state_dict["self_attn.q_proj.weight"], expected_q) + torch.testing.assert_close(state_dict["self_attn.k_proj.weight"], expected_k) + torch.testing.assert_close(state_dict["self_attn.v_proj.weight"], expected_v) + + self.assertNotIn("mlp.gate_proj.weight", state_dict) + self.assertNotIn("mlp.up_proj.weight", state_dict) + self.assertIn("mlp.gate_up_proj.weight", state_dict) + self.assertNotIn("gate_proj", model.mlp._modules) + self.assertNotIn("up_proj", model.mlp._modules) + self.assertIn("gate_up_proj", model.mlp._modules) + torch.testing.assert_close(state_dict["mlp.gate_up_proj.weight"], torch.cat([gate_weight, up_weight], dim=0)) + + def test_renaming_happens_before_conversion(self): + class MLPModel(nn.Module): + base_model_prefix = "" + + def __init__(self): + super().__init__() + self.mlp = nn.Module() + self.dummy = nn.Module() + self.mlp.gate_proj = nn.Linear(2, 4, bias=False) + self.dummy.up_proj = nn.Linear(2, 4, bias=False) + + model = MLPModel() + model.config = PretrainedConfig() + + gate_weight = torch.randn(4, 2) + up_weight = torch.randn(4, 2) + model.mlp.gate_proj.weight.data.copy_(gate_weight) + model.dummy.up_proj.weight.data.copy_(up_weight) + + apply_kernel_structural_transforms( + model, + [ + WeightRenaming("dummy.up_proj.weight", "mlp.up_proj.weight"), + WeightConverter( + ["mlp.gate_proj.weight", "mlp.up_proj.weight"], + "mlp.gate_up_proj.weight", + operations=[Concatenate(dim=0)], + ), + ], + ) + + state_dict = model.state_dict() + + self.assertIn("mlp.gate_up_proj.weight", state_dict) + self.assertNotIn("mlp.gate_proj.weight", state_dict) + self.assertNotIn("mlp.up_proj.weight", state_dict) + self.assertNotIn("gate_proj", model.mlp._modules) + self.assertNotIn("up_proj", model.mlp._modules) + self.assertIn("gate_up_proj", model.mlp._modules) + + expected = torch.cat([gate_weight, up_weight], dim=0) + torch.testing.assert_close(state_dict["mlp.gate_up_proj.weight"], expected) diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 6299d634ad47..d35d1f6ad9b8 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -40,7 +40,6 @@ rename_source_key, revert_weight_conversion, ) -from transformers.integrations.hub_kernels import apply_kernel_structural_transforms as apply_transforms_to_meta_model from transformers.modeling_utils import LoadStateDictConfig from transformers.utils.import_utils import is_triton_available @@ -1253,352 +1252,5 @@ class _TstRootSame(PreTrainedModel): ... self.assertIsNone(transforms[0].scope_prefix) -class TestApplyTransformsToMetaModel(unittest.TestCase): - class SimpleModel(nn.Module): - base_model_prefix = "" - - def __init__(self): - super().__init__() - self.linear = nn.Linear(4, 4, bias=False) - - class QKVAttn(nn.Module): - def __init__(self): - super().__init__() - self.qkv_proj = nn.Linear(2, 6, bias=False) - - class QKVModel(nn.Module): - base_model_prefix = "" - - def __init__(self): - super().__init__() - self.self_attn = TestApplyTransformsToMetaModel.QKVAttn() - - def test_weight_renaming_renames_parameter(self): - model = TestApplyTransformsToMetaModel.SimpleModel() - original_weight = model.linear.weight.data.clone() - self.assertIn("linear.weight", model.state_dict()) - - apply_transforms_to_meta_model(model, [WeightRenaming("linear.weight", "linear.kernel")]) - - state_dict = model.state_dict() - self.assertIn("linear.kernel", state_dict) - self.assertNotIn("linear.weight", state_dict) - torch.testing.assert_close(state_dict["linear.kernel"], original_weight) - - def test_weight_renaming_is_idempotent(self): - model = TestApplyTransformsToMetaModel.SimpleModel() - apply_transforms_to_meta_model(model, [WeightRenaming("linear.weight", "linear.kernel")]) - keys_after_first = set(model.state_dict().keys()) - - apply_transforms_to_meta_model(model, [WeightRenaming("linear.weight", "linear.kernel")]) - self.assertEqual(keys_after_first, set(model.state_dict().keys())) - - def test_weight_renaming_no_match_is_noop(self): - model = TestApplyTransformsToMetaModel.SimpleModel() - original_keys = set(model.state_dict().keys()) - - apply_transforms_to_meta_model(model, [WeightRenaming("non_existent.weight", "something.weight")]) - - self.assertEqual(original_keys, set(model.state_dict().keys())) - - def test_weight_converter_splits_fused_projection(self): - model = TestApplyTransformsToMetaModel.QKVModel() - model.config = PretrainedConfig() - qkv_weight = torch.randn(6, 2) - model.self_attn.qkv_proj.weight.data.copy_(qkv_weight) - - apply_transforms_to_meta_model( - model, - [ - WeightConverter( - "self_attn.qkv_proj.weight", - ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], - operations=[Chunk(dim=0)], - ) - ], - ) - - state_dict = model.state_dict() - self.assertIn("self_attn.q_proj.weight", state_dict) - self.assertIn("self_attn.k_proj.weight", state_dict) - self.assertIn("self_attn.v_proj.weight", state_dict) - self.assertNotIn("self_attn.qkv_proj.weight", state_dict) - # The source module must be gone, the target modules must exist - self.assertNotIn("qkv_proj", model.self_attn._modules) - self.assertIn("q_proj", model.self_attn._modules) - self.assertIn("k_proj", model.self_attn._modules) - self.assertIn("v_proj", model.self_attn._modules) - - expected_q, expected_k, expected_v = torch.chunk(qkv_weight, 3, dim=0) - torch.testing.assert_close(state_dict["self_attn.q_proj.weight"], expected_q) - torch.testing.assert_close(state_dict["self_attn.k_proj.weight"], expected_k) - torch.testing.assert_close(state_dict["self_attn.v_proj.weight"], expected_v) - - def test_weight_converter_split_is_idempotent(self): - def make_mapping(): - return [ - WeightConverter( - "self_attn.qkv_proj.weight", - ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], - operations=[Chunk(dim=0)], - ) - ] - - model = TestApplyTransformsToMetaModel.QKVModel() - model.config = PretrainedConfig() - - apply_transforms_to_meta_model(model, make_mapping()) - keys_after_first = set(model.state_dict().keys()) - - apply_transforms_to_meta_model(model, make_mapping()) - self.assertEqual(keys_after_first, set(model.state_dict().keys())) - - def test_weight_converter_splits_across_multiple_layers(self): - class QKVLayer(nn.Module): - def __init__(self): - super().__init__() - self.self_attn = TestApplyTransformsToMetaModel.QKVAttn() - - class QKVLayeredModel(nn.Module): - base_model_prefix = "model" - - def __init__(self): - super().__init__() - self.model = nn.ModuleList([QKVLayer(), QKVLayer()]) - - model = QKVLayeredModel() - model.config = PretrainedConfig() - - qkv_weights = {} - for i in range(2): - w = torch.randn(6, 2) - model.model[i].self_attn.qkv_proj.weight.data.copy_(w) - qkv_weights[i] = w - - apply_transforms_to_meta_model( - model, - [ - WeightConverter( - "self_attn.qkv_proj.weight", - ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], - operations=[Chunk(dim=0)], - ) - ], - ) - - state_dict = model.state_dict() - self.assertNotIn("model.0.self_attn.qkv_proj.weight", state_dict) - self.assertNotIn("model.1.self_attn.qkv_proj.weight", state_dict) - - for i in range(2): - layer_attn = model.model[i].self_attn - self.assertNotIn("qkv_proj", layer_attn._modules) - self.assertIn("q_proj", layer_attn._modules) - self.assertIn("k_proj", layer_attn._modules) - self.assertIn("v_proj", layer_attn._modules) - - expected_q, expected_k, expected_v = torch.chunk(qkv_weights[i], 3, dim=0) - torch.testing.assert_close(state_dict[f"model.{i}.self_attn.q_proj.weight"], expected_q) - torch.testing.assert_close(state_dict[f"model.{i}.self_attn.k_proj.weight"], expected_k) - torch.testing.assert_close(state_dict[f"model.{i}.self_attn.v_proj.weight"], expected_v) - - def test_weight_converter_glob_merges_modulelist(self): - class ExpertModel(nn.Module): - base_model_prefix = "" - - def __init__(self): - super().__init__() - self.experts = nn.ModuleList([nn.Linear(2, 4, bias=False), nn.Linear(2, 4, bias=False)]) - # Pre-create empty target so _get_submodule_or_create finds it without - # needing to infer a module type from the 3-D stacked tensor. - self.merged = nn.Module() - - model = ExpertModel() - model.config = PretrainedConfig() - - weight_0 = model.experts[0].weight.data.clone() - weight_1 = model.experts[1].weight.data.clone() - - apply_transforms_to_meta_model( - model, - [ - WeightConverter( - "experts.*.weight", - "merged.weight", - operations=[MergeModulelist(dim=0)], - ) - ], - ) - - state_dict = model.state_dict() - self.assertIn("merged.weight", state_dict) - self.assertNotIn("experts.0.weight", state_dict) - self.assertNotIn("experts.1.weight", state_dict) - torch.testing.assert_close(state_dict["merged.weight"], torch.stack([weight_0, weight_1], dim=0)) - - def test_weight_converter_concatenates_projections(self): - class SeparateAttn(nn.Module): - def __init__(self): - super().__init__() - self.q_proj = nn.Linear(2, 2, bias=False) - self.k_proj = nn.Linear(2, 2, bias=False) - self.v_proj = nn.Linear(2, 2, bias=False) - - class SeparateModel(nn.Module): - base_model_prefix = "" - - def __init__(self): - super().__init__() - self.self_attn = SeparateAttn() - - model = SeparateModel() - model.config = PretrainedConfig() - - q_weight = torch.randn(2, 2) - k_weight = torch.randn(2, 2) - v_weight = torch.randn(2, 2) - model.self_attn.q_proj.weight.data.copy_(q_weight) - model.self_attn.k_proj.weight.data.copy_(k_weight) - model.self_attn.v_proj.weight.data.copy_(v_weight) - - apply_transforms_to_meta_model( - model, - [ - WeightConverter( - ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], - "self_attn.qkv_proj.weight", - operations=[Concatenate(dim=0)], - ) - ], - ) - - state_dict = model.state_dict() - self.assertIn("self_attn.qkv_proj.weight", state_dict) - self.assertNotIn("self_attn.q_proj.weight", state_dict) - self.assertNotIn("self_attn.k_proj.weight", state_dict) - self.assertNotIn("self_attn.v_proj.weight", state_dict) - # Source modules must be gone, fused module must exist - self.assertNotIn("q_proj", model.self_attn._modules) - self.assertNotIn("k_proj", model.self_attn._modules) - self.assertNotIn("v_proj", model.self_attn._modules) - self.assertIn("qkv_proj", model.self_attn._modules) - - expected_qkv = torch.cat([q_weight, k_weight, v_weight], dim=0) - torch.testing.assert_close(state_dict["self_attn.qkv_proj.weight"], expected_qkv) - - def test_multiple_independent_converters_in_single_call(self): - class MLP(nn.Module): - def __init__(self): - super().__init__() - self.gate_proj = nn.Linear(2, 4, bias=False) - self.up_proj = nn.Linear(2, 4, bias=False) - - class TwoHeadModel(nn.Module): - base_model_prefix = "" - - def __init__(self): - super().__init__() - self.self_attn = TestApplyTransformsToMetaModel.QKVAttn() - self.mlp = MLP() - - model = TwoHeadModel() - model.config = PretrainedConfig() - - qkv_weight = torch.randn(6, 2) - gate_weight = torch.randn(4, 2) - up_weight = torch.randn(4, 2) - model.self_attn.qkv_proj.weight.data.copy_(qkv_weight) - model.mlp.gate_proj.weight.data.copy_(gate_weight) - model.mlp.up_proj.weight.data.copy_(up_weight) - - apply_transforms_to_meta_model( - model, - [ - WeightConverter( - "self_attn.qkv_proj.weight", - ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], - operations=[Chunk(dim=0)], - ), - WeightConverter( - ["mlp.gate_proj.weight", "mlp.up_proj.weight"], - "mlp.gate_up_proj.weight", - operations=[Concatenate(dim=0)], - ), - ], - ) - - state_dict = model.state_dict() - - # Attention: fused → split - self.assertNotIn("self_attn.qkv_proj.weight", state_dict) - self.assertIn("self_attn.q_proj.weight", state_dict) - self.assertIn("self_attn.k_proj.weight", state_dict) - self.assertIn("self_attn.v_proj.weight", state_dict) - self.assertNotIn("qkv_proj", model.self_attn._modules) - self.assertIn("q_proj", model.self_attn._modules) - self.assertIn("k_proj", model.self_attn._modules) - self.assertIn("v_proj", model.self_attn._modules) - expected_q, expected_k, expected_v = torch.chunk(qkv_weight, 3, dim=0) - torch.testing.assert_close(state_dict["self_attn.q_proj.weight"], expected_q) - torch.testing.assert_close(state_dict["self_attn.k_proj.weight"], expected_k) - torch.testing.assert_close(state_dict["self_attn.v_proj.weight"], expected_v) - - # MLP: separate → fused - self.assertNotIn("mlp.gate_proj.weight", state_dict) - self.assertNotIn("mlp.up_proj.weight", state_dict) - self.assertIn("mlp.gate_up_proj.weight", state_dict) - self.assertNotIn("gate_proj", model.mlp._modules) - self.assertNotIn("up_proj", model.mlp._modules) - self.assertIn("gate_up_proj", model.mlp._modules) - torch.testing.assert_close(state_dict["mlp.gate_up_proj.weight"], torch.cat([gate_weight, up_weight], dim=0)) - - def test_renaming_happens_before_conversion(self): - class MLPModel(nn.Module): - base_model_prefix = "" - - def __init__(self): - super().__init__() - self.mlp = nn.Module() - self.dummy = nn.Module() - self.mlp.gate_proj = nn.Linear(2, 4, bias=False) - self.dummy.up_proj = nn.Linear(2, 4, bias=False) - - model = MLPModel() - model.config = PretrainedConfig() - - gate_weight = torch.randn(4, 2) - up_weight = torch.randn(4, 2) - model.mlp.gate_proj.weight.data.copy_(gate_weight) - model.dummy.up_proj.weight.data.copy_(up_weight) - - apply_transforms_to_meta_model( - model, - [ - WeightRenaming("dummy.up_proj.weight", "mlp.up_proj.weight"), - WeightConverter( - ["mlp.gate_proj.weight", "mlp.up_proj.weight"], - "mlp.gate_up_proj.weight", - operations=[Concatenate(dim=0)], - ), - ], - ) - - state_dict = model.state_dict() - - # Fused weight exists; both source weights are gone - self.assertIn("mlp.gate_up_proj.weight", state_dict) - self.assertNotIn("mlp.gate_proj.weight", state_dict) - self.assertNotIn("mlp.up_proj.weight", state_dict) - - # Source modules are cleaned up; fused module is present - self.assertNotIn("gate_proj", model.mlp._modules) - self.assertNotIn("up_proj", model.mlp._modules) - self.assertIn("gate_up_proj", model.mlp._modules) - - # Values: gate comes first (as declared in source_patterns), then up - expected = torch.cat([gate_weight, up_weight], dim=0) - torch.testing.assert_close(state_dict["mlp.gate_up_proj.weight"], expected) - - if __name__ == "__main__": unittest.main() From 68c36591c1eaa9dba5b381315d93de7b089c1cd2 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 19 May 2026 15:53:11 +0200 Subject: [PATCH 43/49] wip: make conversion happen before fused module instantiation --- src/transformers/integrations/hub_kernels.py | 33 ++++++++++++++------ src/transformers/modeling_utils.py | 6 ---- 2 files changed, 24 insertions(+), 15 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 9837c1980d10..e65473e0afe8 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -733,11 +733,13 @@ def make_fused_parent_class( child_names: list[str], source_names: list[str], kernel_layer_name: str, + conversion_mapping: list | None = None, ) -> type: original_init = parent_cls.__init__ _child_names = list(child_names) _source_names = list(source_names) _kernel_layer_name = kernel_layer_name + _conversion_mapping = list(conversion_mapping) if conversion_mapping else [] fused_module_cls = make_fused_module_class(tuple(_source_names), _kernel_layer_name) @@ -745,6 +747,8 @@ def fused_init(self, *args, **kwargs): original_init(self, *args, **kwargs) modules_to_fuse = [getattr(self, name) for name in _child_names] fused = fused_module_cls(modules_to_fuse, fused_module_names=list(_source_names)) + if _conversion_mapping: + apply_kernel_structural_transforms(fused, _conversion_mapping) setattr(self, _child_names[0], fused) for name in _child_names[1:]: setattr(self, name, nn.Identity()) @@ -877,7 +881,15 @@ def register_kernel_fusions( repo_str = next(iter(repo_str.values())) kernel_layer_name = repo_str.split(":")[1] if ":" in repo_str else repo_str.split("/")[-1] - # 1. Meta-device scan — find parent classes that contain all target children. + # 1. Load the kernel class to get any conversion_mapping it declares. + kernel_cls = _try_load_kernel_class(repo_str, use_local=kernel_config.use_local_kernel) + kernel_conversion_mapping = ( + list(kernel_cls.conversion_mapping) + if kernel_cls is not None and hasattr(kernel_cls, "conversion_mapping") + else [] + ) + + # 2. Meta-device scan — find parent classes that contain all target children. with torch.device("meta"): meta_model = cls(config) @@ -890,7 +902,13 @@ def register_kernel_fusions( if not all(hasattr(module, name) for name in child_names): continue seen.add(module_cls) - fused_parent_cls = make_fused_parent_class(module_cls, child_names, source_names, kernel_layer_name) + fused_parent_cls = make_fused_parent_class( + module_cls, + child_names, + source_names, + kernel_layer_name, + conversion_mapping=kernel_conversion_mapping, + ) patch_mapping[module_cls.__name__] = fused_parent_cls if not patch_mapping: @@ -901,16 +919,13 @@ def register_kernel_fusions( new_mapping[layer_name] = hub_repo continue - # 2. Register class-level monkey patches. + # 3. Register class-level monkey patches. register_patch_mapping(patch_mapping, overwrite=True) - # 3. Infer weight renaming mapping. + # 4. Infer weight renaming mapping for the loader. transforms = infer_kernel_fusion_transforms(list(zip(source_names, glob_patterns))) - - # 4. Load weight conversion mapping from the kernel class if available. - kernel_cls = _try_load_kernel_class(repo_str, use_local=kernel_config.use_local_kernel) - if kernel_cls is not None and hasattr(kernel_cls, "conversion_mapping"): - transforms = transforms + list(kernel_cls.conversion_mapping) + if kernel_conversion_mapping: + transforms = transforms + kernel_conversion_mapping # 5. Load existing conversion mapping for this model type. existing = get_checkpoint_conversion_mapping(model_type) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 9db2df09a707..12b6a65183dc 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -4252,12 +4252,6 @@ def from_pretrained( # Obtain the weight conversion mapping for this model if any are registered and apply to all submodels recursively weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer) - # Apply kernel-specific structural transforms before loading (module graph changes, cost-free on meta device). - if use_kernels: - from .integrations.hub_kernels import apply_kernel_structural_transforms - - apply_kernel_structural_transforms(model, weight_conversions) - if _torch_distributed_available and device_mesh is not None: # add hooks to nn.Modules: no weights model = distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size) From 041182fa63de34acaa046c4db7ab965d30483826 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 19 May 2026 16:30:47 +0200 Subject: [PATCH 44/49] refactor --- src/transformers/integrations/hub_kernels.py | 235 +++++------ tests/kernels/test_kernels.py | 393 ++++++------------- 2 files changed, 231 insertions(+), 397 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index e65473e0afe8..7a1c6a6e3e88 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -521,136 +521,6 @@ def allow_all_hub_kernels(): ALLOW_ALL_KERNELS = False -def _get_submodule_or_create( - model: "nn.Module", - module_path: str, - original_module: "nn.Module | None" = None, - future_weight: "torch.Tensor | None" = None, -) -> "nn.Module": - if not module_path: - return model - if original_module is not None and future_weight is None: - raise ValueError("original_module cannot be provided without future_weight to infer the module class") - parent, _, mod_name = module_path.rpartition(".") - parent_mod = model.get_submodule(parent) if parent else model - if mod_name not in parent_mod._modules: - cls_ = original_module.__class__ if original_module is not None else nn.Module - if cls_ is nn.Linear: - out_features, in_features = future_weight.shape - mod = cls_(in_features=in_features, out_features=out_features, bias=original_module.bias is not None) - elif cls_ is nn.Embedding: - num_embeddings, embedding_dim = future_weight.shape - mod = cls_( - num_embeddings=num_embeddings, - embedding_dim=embedding_dim, - padding_idx=original_module.padding_idx, - max_norm=original_module.max_norm, - norm_type=original_module.norm_type, - scale_grad_by_freq=original_module.scale_grad_by_freq, - sparse=original_module.sparse, - ) - else: - raise ValueError( - f"Cannot infer instantiation for {cls_} when creating submodule {mod_name} at path {module_path}" - ) - parent_mod.add_module(mod_name, mod) - else: - mod = parent_mod._modules[mod_name] - return mod - - -def _post_transformation_cleanup(model: "nn.Module", source_names: list[str]) -> None: - for source_name in source_names: - mod_path, _ = source_name.rsplit(".", 1) if "." in source_name else ("", source_name) - try: - source_mod = model.get_submodule(mod_path) if mod_path else model - except AttributeError: - continue - if ( - all(p is None for p in source_mod._parameters.values()) - and all(b is None for b in source_mod._buffers.values()) - and all(m is None for m in source_mod._modules.values()) - ): - parent_path, _, mod_name = mod_path.rpartition(".") - parent_mod = model.get_submodule(parent_path) if parent_path else model - parent_mod._modules.pop(mod_name, None) - - -def apply_kernel_structural_transforms( - model: "PreTrainedModel", - weight_conversions: list, -) -> None: - """ - Restructure `model`'s module graph according to `weight_conversions` while parameters are on the - meta device (cost-free). Called after model init and before the loader, so the loader sees the - correct parameter structure. - - Kept in hub_kernels rather than core_model_loading to leave the loading machinery clean. - """ - from ..core_model_loading import WeightConverter, WeightRenaming, dot_natural_key, rename_source_key - - meta_state_dict = model.state_dict() - - renamings = [x for x in weight_conversions if isinstance(x, WeightRenaming)] - converters = [x for x in weight_conversions if isinstance(x, WeightConverter)] - pattern_to_converter = {k: c for c in converters for k in c.source_patterns} - param_name_to_transform: dict = {} - - for param_name in sorted(meta_state_dict, key=dot_natural_key): - renamed_key, source_pattern = rename_source_key(param_name, renamings, converters) - if source_pattern is not None: - transform = param_name_to_transform.setdefault(renamed_key, deepcopy(pattern_to_converter[source_pattern])) - else: - transform = param_name_to_transform.setdefault(renamed_key, WeightRenaming(param_name, renamed_key)) - source_pattern = param_name - transform.add_tensor(renamed_key, param_name, source_pattern, meta_state_dict[param_name]) - - for layer_name, transform in param_name_to_transform.items(): - source_names = list(transform.layer_targets.get(layer_name, [])) - if not source_names: - continue - - # Skip identity renames: parameter is already at the correct path. - if len(source_names) == 1 and source_names[0] == layer_name: - continue - - # Find a source module for type inference when _get_submodule_or_create must create a new node. - original_mod = None - for sname in source_names: - smod_path = sname.rpartition(".")[0] - try: - original_mod = model.get_submodule(smod_path) if smod_path else model - break - except AttributeError: - pass - - # Use convert() on meta tensors — correct shapes/keys, cost-free on meta device. - result = transform.convert(layer_name) - - # Remove source parameters from the model graph. - for sname in source_names: - mod_path, _, attr = sname.rpartition(".") - try: - source_mod = model.get_submodule(mod_path) if mod_path else model - except AttributeError: - continue - source_mod._parameters.pop(attr, None) - source_mod._buffers.pop(attr, None) - - # Register target parameters, creating intermediate module nodes as needed. - for target_key, meta_val in result.items(): - meta_tensor = meta_val[0] if isinstance(meta_val, list) else meta_val - t_mod_path, _, t_attr = target_key.rpartition(".") - target_mod = _get_submodule_or_create( - model, t_mod_path, original_module=original_mod, future_weight=meta_tensor - ) - if not isinstance(meta_tensor, nn.Parameter): - meta_tensor = nn.Parameter(meta_tensor, requires_grad=meta_tensor.is_floating_point()) - target_mod.register_parameter(t_attr, meta_tensor) - - _post_transformation_cleanup(model, source_names) - - class FusedModuleBase(nn.Module): def __init__( self, @@ -703,6 +573,107 @@ def __repr__(self): names = ", ".join(self._fused_module_names) return f"{self.__class__.__name__}(fused=({names}))" + def _get_or_create_submodule( + self, + module_path: str, + original_module: "nn.Module | None", + future_weight: "torch.Tensor", + ) -> "nn.Module": + if not module_path: + return self + parent, _, mod_name = module_path.rpartition(".") + parent_mod = self.get_submodule(parent) if parent else self + if mod_name not in parent_mod._modules: + cls_ = original_module.__class__ if original_module is not None else nn.Module + if cls_ is nn.Linear: + out_features, in_features = future_weight.shape + mod = cls_(in_features=in_features, out_features=out_features, bias=original_module.bias is not None) + elif cls_ is nn.Embedding: + num_embeddings, embedding_dim = future_weight.shape + mod = cls_( + num_embeddings=num_embeddings, + embedding_dim=embedding_dim, + padding_idx=original_module.padding_idx, + max_norm=original_module.max_norm, + norm_type=original_module.norm_type, + scale_grad_by_freq=original_module.scale_grad_by_freq, + sparse=original_module.sparse, + ) + else: + raise ValueError( + f"Cannot create submodule {mod_name!r} at {module_path!r}: unsupported class {cls_.__name__!r}" + ) + parent_mod.add_module(mod_name, mod) + else: + mod = parent_mod._modules[mod_name] + return mod + + def _apply_weight_conversions(self, conversion_mapping: list) -> None: + """Apply WeightConverter transforms to this module's graph (cost-free on meta device).""" + + from ..core_model_loading import dot_natural_key, rename_source_key + + state_dict = self.state_dict() + pattern_to_converter = {k: c for c in conversion_mapping for k in c.source_patterns} + param_name_to_transform: dict = {} + + for param_name in sorted(state_dict, key=dot_natural_key): + renamed_key, source_pattern = rename_source_key(param_name, [], conversion_mapping) + if source_pattern is None: + continue + transform = param_name_to_transform.setdefault(renamed_key, deepcopy(pattern_to_converter[source_pattern])) + transform.add_tensor(renamed_key, param_name, source_pattern, state_dict[param_name]) + + for layer_name, transform in param_name_to_transform.items(): + source_names = list(transform.layer_targets.get(layer_name, [])) + if not source_names: + continue + + original_mod = None + for sname in source_names: + smod_path = sname.rpartition(".")[0] + try: + original_mod = self.get_submodule(smod_path) if smod_path else self + break + except AttributeError: + pass + + result = transform.convert(layer_name) + + for sname in source_names: + mod_path, _, attr = sname.rpartition(".") + try: + source_mod = self.get_submodule(mod_path) if mod_path else self + except AttributeError: + continue + source_mod._parameters.pop(attr, None) + source_mod._buffers.pop(attr, None) + + for target_key, meta_val in result.items(): + meta_tensor = meta_val[0] if isinstance(meta_val, list) else meta_val + t_mod_path, _, t_attr = target_key.rpartition(".") + target_mod = self._get_or_create_submodule(t_mod_path, original_mod, meta_tensor) + if not isinstance(meta_tensor, nn.Parameter): + meta_tensor = nn.Parameter(meta_tensor, requires_grad=meta_tensor.is_floating_point()) + target_mod.register_parameter(t_attr, meta_tensor) + + for source_name in source_names: + mod_path = source_name.rpartition(".")[0] + if not mod_path: + continue + try: + source_mod = self.get_submodule(mod_path) + except AttributeError: + continue + if ( + all(p is None for p in source_mod._parameters.values()) + and all(b is None for b in source_mod._buffers.values()) + and all(m is None for m in source_mod._modules.values()) + ): + parent_path, _, mod_name = mod_path.rpartition(".") + parent_mod = self.get_submodule(parent_path) if parent_path else self + parent_mod._modules.pop(mod_name, None) + @functools.cache def make_fused_module_class(source_layer_names: tuple[str, ...], kernel_layer_name: str) -> type: @@ -748,7 +719,7 @@ def fused_init(self, *args, **kwargs): modules_to_fuse = [getattr(self, name) for name in _child_names] fused = fused_module_cls(modules_to_fuse, fused_module_names=list(_source_names)) if _conversion_mapping: - apply_kernel_structural_transforms(fused, _conversion_mapping) + fused._apply_weight_conversions(_conversion_mapping) setattr(self, _child_names[0], fused) for name in _child_names[1:]: setattr(self, name, nn.Identity()) @@ -943,8 +914,8 @@ def register_kernel_fusions( __all__ = [ "FusedModuleBase", "LayerRepository", - "apply_kernel_structural_transforms", "get_kernel", + "make_fused_module_class", "register_kernel_fusions", "lazy_load_kernel", "register_kernel_mapping", diff --git a/tests/kernels/test_kernels.py b/tests/kernels/test_kernels.py index 6bbafb433fcd..141f1e4becea 100644 --- a/tests/kernels/test_kernels.py +++ b/tests/kernels/test_kernels.py @@ -25,7 +25,7 @@ import torch import torch.nn as nn -from transformers import AutoModelForCausalLM, AutoTokenizer, KernelConfig, PretrainedConfig +from transformers import AutoModelForCausalLM, AutoTokenizer, KernelConfig from transformers.conversion_mapping import get_checkpoint_conversion_mapping from transformers.core_model_loading import ( Chunk, @@ -38,11 +38,11 @@ from transformers.integrations.hub_kernels import ( _HUB_KERNEL_MAPPING, _KERNEL_MODULE_MAPPING, - apply_kernel_structural_transforms, infer_kernel_fusion_transforms, is_kernel, lazy_load_kernel, load_and_register_attn_kernel, + make_fused_module_class, make_fused_parent_class, register_kernel_fusions, ) @@ -728,338 +728,201 @@ class FakeKernel: self.assertTrue(any(isinstance(t, WeightConverter) for t in mapping)) -class TestApplyKernelStructuralTransforms(unittest.TestCase): - class SimpleModel(nn.Module): - base_model_prefix = "" +class TestFusedModuleWeightConversions(unittest.TestCase): + def _make_fused(self, *modules, names=None): + if names is None: + names = [type(m).__name__ for m in modules] + FusedCls = make_fused_module_class(tuple(names), "kernel") + return FusedCls(list(modules), fused_module_names=list(names)) - def __init__(self): - super().__init__() - self.linear = nn.Linear(4, 4, bias=False) - - class QKVAttn(nn.Module): - def __init__(self): - super().__init__() - self.qkv_proj = nn.Linear(2, 6, bias=False) - - class QKVModel(nn.Module): - base_model_prefix = "" - - def __init__(self): - super().__init__() - self.self_attn = TestApplyKernelStructuralTransforms.QKVAttn() - - def test_weight_renaming_renames_parameter(self): - model = TestApplyKernelStructuralTransforms.SimpleModel() - original_weight = model.linear.weight.data.clone() - self.assertIn("linear.weight", model.state_dict()) + def test_weight_converter_concatenates_projections(self): + class SeparateAttn(nn.Module): + def __init__(self): + super().__init__() + self.q_proj = nn.Linear(2, 2, bias=False) + self.k_proj = nn.Linear(2, 2, bias=False) + self.v_proj = nn.Linear(2, 2, bias=False) - apply_kernel_structural_transforms(model, [WeightRenaming("linear.weight", "linear.kernel")]) + fused = self._make_fused(SeparateAttn(), names=["Attn"]) - state_dict = model.state_dict() - self.assertIn("linear.kernel", state_dict) - self.assertNotIn("linear.weight", state_dict) - torch.testing.assert_close(state_dict["linear.kernel"], original_weight) + q_weight = torch.randn(2, 2) + k_weight = torch.randn(2, 2) + v_weight = torch.randn(2, 2) + fused.Attn.q_proj.weight.data.copy_(q_weight) + fused.Attn.k_proj.weight.data.copy_(k_weight) + fused.Attn.v_proj.weight.data.copy_(v_weight) - def test_weight_renaming_is_idempotent(self): - model = TestApplyKernelStructuralTransforms.SimpleModel() - apply_kernel_structural_transforms(model, [WeightRenaming("linear.weight", "linear.kernel")]) - keys_after_first = set(model.state_dict().keys()) + fused._apply_weight_conversions( + [ + WeightConverter( + ["Attn.q_proj.weight", "Attn.k_proj.weight", "Attn.v_proj.weight"], + "Attn.qkv_proj.weight", + operations=[Concatenate(dim=0)], + ) + ] + ) - apply_kernel_structural_transforms(model, [WeightRenaming("linear.weight", "linear.kernel")]) - self.assertEqual(keys_after_first, set(model.state_dict().keys())) + state_dict = fused.state_dict() + self.assertIn("Attn.qkv_proj.weight", state_dict) + self.assertNotIn("Attn.q_proj.weight", state_dict) + self.assertNotIn("Attn.k_proj.weight", state_dict) + self.assertNotIn("Attn.v_proj.weight", state_dict) + self.assertIn("qkv_proj", fused.Attn._modules) + self.assertNotIn("q_proj", fused.Attn._modules) + self.assertNotIn("k_proj", fused.Attn._modules) + self.assertNotIn("v_proj", fused.Attn._modules) - def test_weight_renaming_no_match_is_noop(self): - model = TestApplyKernelStructuralTransforms.SimpleModel() - original_keys = set(model.state_dict().keys()) + expected_qkv = torch.cat([q_weight, k_weight, v_weight], dim=0) + torch.testing.assert_close(state_dict["Attn.qkv_proj.weight"], expected_qkv) - apply_kernel_structural_transforms(model, [WeightRenaming("non_existent.weight", "something.weight")]) + def test_weight_converter_splits_fused_projection(self): + class QKVAttn(nn.Module): + def __init__(self): + super().__init__() + self.qkv_proj = nn.Linear(2, 6, bias=False) - self.assertEqual(original_keys, set(model.state_dict().keys())) + fused = self._make_fused(QKVAttn(), names=["QKV"]) - def test_weight_converter_splits_fused_projection(self): - model = TestApplyKernelStructuralTransforms.QKVModel() - model.config = PretrainedConfig() qkv_weight = torch.randn(6, 2) - model.self_attn.qkv_proj.weight.data.copy_(qkv_weight) + fused.QKV.qkv_proj.weight.data.copy_(qkv_weight) - apply_kernel_structural_transforms( - model, + fused._apply_weight_conversions( [ WeightConverter( - "self_attn.qkv_proj.weight", - ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + "QKV.qkv_proj.weight", + ["QKV.q_proj.weight", "QKV.k_proj.weight", "QKV.v_proj.weight"], operations=[Chunk(dim=0)], ) - ], + ] ) - state_dict = model.state_dict() - self.assertIn("self_attn.q_proj.weight", state_dict) - self.assertIn("self_attn.k_proj.weight", state_dict) - self.assertIn("self_attn.v_proj.weight", state_dict) - self.assertNotIn("self_attn.qkv_proj.weight", state_dict) - self.assertNotIn("qkv_proj", model.self_attn._modules) - self.assertIn("q_proj", model.self_attn._modules) - self.assertIn("k_proj", model.self_attn._modules) - self.assertIn("v_proj", model.self_attn._modules) + state_dict = fused.state_dict() + self.assertIn("QKV.q_proj.weight", state_dict) + self.assertIn("QKV.k_proj.weight", state_dict) + self.assertIn("QKV.v_proj.weight", state_dict) + self.assertNotIn("QKV.qkv_proj.weight", state_dict) + self.assertNotIn("qkv_proj", fused.QKV._modules) + self.assertIn("q_proj", fused.QKV._modules) + self.assertIn("k_proj", fused.QKV._modules) + self.assertIn("v_proj", fused.QKV._modules) expected_q, expected_k, expected_v = torch.chunk(qkv_weight, 3, dim=0) - torch.testing.assert_close(state_dict["self_attn.q_proj.weight"], expected_q) - torch.testing.assert_close(state_dict["self_attn.k_proj.weight"], expected_k) - torch.testing.assert_close(state_dict["self_attn.v_proj.weight"], expected_v) + torch.testing.assert_close(state_dict["QKV.q_proj.weight"], expected_q) + torch.testing.assert_close(state_dict["QKV.k_proj.weight"], expected_k) + torch.testing.assert_close(state_dict["QKV.v_proj.weight"], expected_v) def test_weight_converter_split_is_idempotent(self): - def make_mapping(): - return [ - WeightConverter( - "self_attn.qkv_proj.weight", - ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], - operations=[Chunk(dim=0)], - ) - ] - - model = TestApplyKernelStructuralTransforms.QKVModel() - model.config = PretrainedConfig() - - apply_kernel_structural_transforms(model, make_mapping()) - keys_after_first = set(model.state_dict().keys()) - - apply_kernel_structural_transforms(model, make_mapping()) - self.assertEqual(keys_after_first, set(model.state_dict().keys())) - - def test_weight_converter_splits_across_multiple_layers(self): - class QKVLayer(nn.Module): + class QKVAttn(nn.Module): def __init__(self): super().__init__() - self.self_attn = TestApplyKernelStructuralTransforms.QKVAttn() - - class QKVLayeredModel(nn.Module): - base_model_prefix = "model" - - def __init__(self): - super().__init__() - self.model = nn.ModuleList([QKVLayer(), QKVLayer()]) - - model = QKVLayeredModel() - model.config = PretrainedConfig() - - qkv_weights = {} - for i in range(2): - w = torch.randn(6, 2) - model.model[i].self_attn.qkv_proj.weight.data.copy_(w) - qkv_weights[i] = w + self.qkv_proj = nn.Linear(2, 6, bias=False) - apply_kernel_structural_transforms( - model, - [ + def make_mapping(): + return [ WeightConverter( - "self_attn.qkv_proj.weight", - ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + "QKV.qkv_proj.weight", + ["QKV.q_proj.weight", "QKV.k_proj.weight", "QKV.v_proj.weight"], operations=[Chunk(dim=0)], ) - ], - ) - - state_dict = model.state_dict() - self.assertNotIn("model.0.self_attn.qkv_proj.weight", state_dict) - self.assertNotIn("model.1.self_attn.qkv_proj.weight", state_dict) + ] - for i in range(2): - layer_attn = model.model[i].self_attn - self.assertNotIn("qkv_proj", layer_attn._modules) - self.assertIn("q_proj", layer_attn._modules) - self.assertIn("k_proj", layer_attn._modules) - self.assertIn("v_proj", layer_attn._modules) + fused = self._make_fused(QKVAttn(), names=["QKV"]) + fused._apply_weight_conversions(make_mapping()) + keys_after_first = set(fused.state_dict().keys()) - expected_q, expected_k, expected_v = torch.chunk(qkv_weights[i], 3, dim=0) - torch.testing.assert_close(state_dict[f"model.{i}.self_attn.q_proj.weight"], expected_q) - torch.testing.assert_close(state_dict[f"model.{i}.self_attn.k_proj.weight"], expected_k) - torch.testing.assert_close(state_dict[f"model.{i}.self_attn.v_proj.weight"], expected_v) + fused._apply_weight_conversions(make_mapping()) + self.assertEqual(keys_after_first, set(fused.state_dict().keys())) def test_weight_converter_glob_merges_modulelist(self): - class ExpertModel(nn.Module): - base_model_prefix = "" - + class ExpertModule(nn.Module): def __init__(self): super().__init__() self.experts = nn.ModuleList([nn.Linear(2, 4, bias=False), nn.Linear(2, 4, bias=False)]) self.merged = nn.Module() - model = ExpertModel() - model.config = PretrainedConfig() + fused = self._make_fused(ExpertModule(), names=["M"]) - weight_0 = model.experts[0].weight.data.clone() - weight_1 = model.experts[1].weight.data.clone() + weight_0 = fused.M.experts[0].weight.data.clone() + weight_1 = fused.M.experts[1].weight.data.clone() - apply_kernel_structural_transforms( - model, - [ - WeightConverter( - "experts.*.weight", - "merged.weight", - operations=[MergeModulelist(dim=0)], - ) - ], + fused._apply_weight_conversions( + [WeightConverter("M.experts.*.weight", "M.merged.weight", operations=[MergeModulelist(dim=0)])] ) - state_dict = model.state_dict() - self.assertIn("merged.weight", state_dict) - self.assertNotIn("experts.0.weight", state_dict) - self.assertNotIn("experts.1.weight", state_dict) - torch.testing.assert_close(state_dict["merged.weight"], torch.stack([weight_0, weight_1], dim=0)) - - def test_weight_converter_concatenates_projections(self): - class SeparateAttn(nn.Module): - def __init__(self): - super().__init__() - self.q_proj = nn.Linear(2, 2, bias=False) - self.k_proj = nn.Linear(2, 2, bias=False) - self.v_proj = nn.Linear(2, 2, bias=False) - - class SeparateModel(nn.Module): - base_model_prefix = "" + state_dict = fused.state_dict() + self.assertIn("M.merged.weight", state_dict) + self.assertNotIn("M.experts.0.weight", state_dict) + self.assertNotIn("M.experts.1.weight", state_dict) + torch.testing.assert_close(state_dict["M.merged.weight"], torch.stack([weight_0, weight_1], dim=0)) + def test_multiple_independent_converters(self): + class QKVAttn(nn.Module): def __init__(self): super().__init__() - self.self_attn = SeparateAttn() - - model = SeparateModel() - model.config = PretrainedConfig() - - q_weight = torch.randn(2, 2) - k_weight = torch.randn(2, 2) - v_weight = torch.randn(2, 2) - model.self_attn.q_proj.weight.data.copy_(q_weight) - model.self_attn.k_proj.weight.data.copy_(k_weight) - model.self_attn.v_proj.weight.data.copy_(v_weight) + self.qkv_proj = nn.Linear(2, 6, bias=False) - apply_kernel_structural_transforms( - model, - [ - WeightConverter( - ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], - "self_attn.qkv_proj.weight", - operations=[Concatenate(dim=0)], - ) - ], - ) - - state_dict = model.state_dict() - self.assertIn("self_attn.qkv_proj.weight", state_dict) - self.assertNotIn("self_attn.q_proj.weight", state_dict) - self.assertNotIn("self_attn.k_proj.weight", state_dict) - self.assertNotIn("self_attn.v_proj.weight", state_dict) - self.assertNotIn("q_proj", model.self_attn._modules) - self.assertNotIn("k_proj", model.self_attn._modules) - self.assertNotIn("v_proj", model.self_attn._modules) - self.assertIn("qkv_proj", model.self_attn._modules) - - expected_qkv = torch.cat([q_weight, k_weight, v_weight], dim=0) - torch.testing.assert_close(state_dict["self_attn.qkv_proj.weight"], expected_qkv) - - def test_multiple_independent_converters_in_single_call(self): class MLP(nn.Module): def __init__(self): super().__init__() self.gate_proj = nn.Linear(2, 4, bias=False) self.up_proj = nn.Linear(2, 4, bias=False) - class TwoHeadModel(nn.Module): - base_model_prefix = "" - - def __init__(self): - super().__init__() - self.self_attn = TestApplyKernelStructuralTransforms.QKVAttn() - self.mlp = MLP() - - model = TwoHeadModel() - model.config = PretrainedConfig() + fused = self._make_fused(QKVAttn(), MLP(), names=["Attn", "MLP"]) qkv_weight = torch.randn(6, 2) gate_weight = torch.randn(4, 2) up_weight = torch.randn(4, 2) - model.self_attn.qkv_proj.weight.data.copy_(qkv_weight) - model.mlp.gate_proj.weight.data.copy_(gate_weight) - model.mlp.up_proj.weight.data.copy_(up_weight) + fused.Attn.qkv_proj.weight.data.copy_(qkv_weight) + fused.MLP.gate_proj.weight.data.copy_(gate_weight) + fused.MLP.up_proj.weight.data.copy_(up_weight) - apply_kernel_structural_transforms( - model, + fused._apply_weight_conversions( [ WeightConverter( - "self_attn.qkv_proj.weight", - ["self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"], + "Attn.qkv_proj.weight", + ["Attn.q_proj.weight", "Attn.k_proj.weight", "Attn.v_proj.weight"], operations=[Chunk(dim=0)], ), WeightConverter( - ["mlp.gate_proj.weight", "mlp.up_proj.weight"], - "mlp.gate_up_proj.weight", + ["MLP.gate_proj.weight", "MLP.up_proj.weight"], + "MLP.gate_up_proj.weight", operations=[Concatenate(dim=0)], ), - ], + ] ) - state_dict = model.state_dict() + state_dict = fused.state_dict() - self.assertNotIn("self_attn.qkv_proj.weight", state_dict) - self.assertIn("self_attn.q_proj.weight", state_dict) - self.assertIn("self_attn.k_proj.weight", state_dict) - self.assertIn("self_attn.v_proj.weight", state_dict) - self.assertNotIn("qkv_proj", model.self_attn._modules) - self.assertIn("q_proj", model.self_attn._modules) - self.assertIn("k_proj", model.self_attn._modules) - self.assertIn("v_proj", model.self_attn._modules) + self.assertNotIn("Attn.qkv_proj.weight", state_dict) + self.assertIn("Attn.q_proj.weight", state_dict) + self.assertIn("Attn.k_proj.weight", state_dict) + self.assertIn("Attn.v_proj.weight", state_dict) + self.assertNotIn("qkv_proj", fused.Attn._modules) + self.assertIn("q_proj", fused.Attn._modules) + self.assertIn("k_proj", fused.Attn._modules) + self.assertIn("v_proj", fused.Attn._modules) expected_q, expected_k, expected_v = torch.chunk(qkv_weight, 3, dim=0) - torch.testing.assert_close(state_dict["self_attn.q_proj.weight"], expected_q) - torch.testing.assert_close(state_dict["self_attn.k_proj.weight"], expected_k) - torch.testing.assert_close(state_dict["self_attn.v_proj.weight"], expected_v) - - self.assertNotIn("mlp.gate_proj.weight", state_dict) - self.assertNotIn("mlp.up_proj.weight", state_dict) - self.assertIn("mlp.gate_up_proj.weight", state_dict) - self.assertNotIn("gate_proj", model.mlp._modules) - self.assertNotIn("up_proj", model.mlp._modules) - self.assertIn("gate_up_proj", model.mlp._modules) - torch.testing.assert_close(state_dict["mlp.gate_up_proj.weight"], torch.cat([gate_weight, up_weight], dim=0)) - - def test_renaming_happens_before_conversion(self): - class MLPModel(nn.Module): - base_model_prefix = "" - + torch.testing.assert_close(state_dict["Attn.q_proj.weight"], expected_q) + + self.assertNotIn("MLP.gate_proj.weight", state_dict) + self.assertNotIn("MLP.up_proj.weight", state_dict) + self.assertIn("MLP.gate_up_proj.weight", state_dict) + self.assertNotIn("gate_proj", fused.MLP._modules) + self.assertNotIn("up_proj", fused.MLP._modules) + self.assertIn("gate_up_proj", fused.MLP._modules) + torch.testing.assert_close(state_dict["MLP.gate_up_proj.weight"], torch.cat([gate_weight, up_weight], dim=0)) + + def test_no_match_is_noop(self): + class SimpleModule(nn.Module): def __init__(self): super().__init__() - self.mlp = nn.Module() - self.dummy = nn.Module() - self.mlp.gate_proj = nn.Linear(2, 4, bias=False) - self.dummy.up_proj = nn.Linear(2, 4, bias=False) + self.linear = nn.Linear(4, 4, bias=False) - model = MLPModel() - model.config = PretrainedConfig() - - gate_weight = torch.randn(4, 2) - up_weight = torch.randn(4, 2) - model.mlp.gate_proj.weight.data.copy_(gate_weight) - model.dummy.up_proj.weight.data.copy_(up_weight) + fused = self._make_fused(SimpleModule(), names=["M"]) + original_keys = set(fused.state_dict().keys()) - apply_kernel_structural_transforms( - model, - [ - WeightRenaming("dummy.up_proj.weight", "mlp.up_proj.weight"), - WeightConverter( - ["mlp.gate_proj.weight", "mlp.up_proj.weight"], - "mlp.gate_up_proj.weight", - operations=[Concatenate(dim=0)], - ), - ], + fused._apply_weight_conversions( + [WeightConverter("M.non_existent.weight", "M.something.weight", operations=[Concatenate(dim=0)])] ) - state_dict = model.state_dict() - - self.assertIn("mlp.gate_up_proj.weight", state_dict) - self.assertNotIn("mlp.gate_proj.weight", state_dict) - self.assertNotIn("mlp.up_proj.weight", state_dict) - self.assertNotIn("gate_proj", model.mlp._modules) - self.assertNotIn("up_proj", model.mlp._modules) - self.assertIn("gate_up_proj", model.mlp._modules) - - expected = torch.cat([gate_weight, up_weight], dim=0) - torch.testing.assert_close(state_dict["mlp.gate_up_proj.weight"], expected) + self.assertEqual(original_keys, set(fused.state_dict().keys())) From 519e673c7efb8c7e87145931184694cfa331cb97 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 19 May 2026 18:23:06 +0200 Subject: [PATCH 45/49] wip: move conversion_mapping inside the init --- src/transformers/integrations/hub_kernels.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 7a1c6a6e3e88..f1e592ce8c83 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -526,6 +526,7 @@ def __init__( self, modules_to_fuse: list["nn.Module"], fused_module_names: list[str] | None = None, + conversion_mapping: list | None = None, ): """ Args: @@ -534,6 +535,9 @@ def __init__( child of this container (i.e. `self.`). When `None`, the `kernel_layer_name` attribute of each source module is used. Pass this explicitly when the source modules do not carry `@use_kernel_forward_from_hub`. + conversion_mapping: Optional list of WeightConverter transforms to apply immediately + after the child modules are registered, reshaping parameters to the structure + the kernel expects (e.g. concatenating q/k/v into qkv). Cost-free on meta device. """ super().__init__() if len(modules_to_fuse) == 0: @@ -558,6 +562,9 @@ def __init__( self.add_module(attr_name, module) self._fused_module_names = [m.kernel_layer_name for m in modules_to_fuse] + if conversion_mapping: + self._apply_weight_conversions(conversion_mapping) + # `kernelize` validates the kernel's forward signature against the class being replaced. # Since the fused container sits at the position of the first module in the chain, the # kernel's forward must match that module's signature. We patch the class-level forward @@ -717,9 +724,11 @@ def make_fused_parent_class( def fused_init(self, *args, **kwargs): original_init(self, *args, **kwargs) modules_to_fuse = [getattr(self, name) for name in _child_names] - fused = fused_module_cls(modules_to_fuse, fused_module_names=list(_source_names)) - if _conversion_mapping: - fused._apply_weight_conversions(_conversion_mapping) + fused = fused_module_cls( + modules_to_fuse, + fused_module_names=list(_source_names), + conversion_mapping=_conversion_mapping or None, + ) setattr(self, _child_names[0], fused) for name in _child_names[1:]: setattr(self, name, nn.Identity()) From 88e0aee841a3ed7da56a5fea973fc9444b89eae0 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 26 May 2026 12:18:17 +0200 Subject: [PATCH 46/49] wip: without any transform_model --- src/transformers/integrations/hub_kernels.py | 176 +++++++++++++++++-- tests/kernels/test_kernels.py | 96 ++++++++++ 2 files changed, 254 insertions(+), 18 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index f1e592ce8c83..82b1b51b028d 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -526,7 +526,6 @@ def __init__( self, modules_to_fuse: list["nn.Module"], fused_module_names: list[str] | None = None, - conversion_mapping: list | None = None, ): """ Args: @@ -535,9 +534,6 @@ def __init__( child of this container (i.e. `self.`). When `None`, the `kernel_layer_name` attribute of each source module is used. Pass this explicitly when the source modules do not carry `@use_kernel_forward_from_hub`. - conversion_mapping: Optional list of WeightConverter transforms to apply immediately - after the child modules are registered, reshaping parameters to the structure - the kernel expects (e.g. concatenating q/k/v into qkv). Cost-free on meta device. """ super().__init__() if len(modules_to_fuse) == 0: @@ -562,9 +558,6 @@ def __init__( self.add_module(attr_name, module) self._fused_module_names = [m.kernel_layer_name for m in modules_to_fuse] - if conversion_mapping: - self._apply_weight_conversions(conversion_mapping) - # `kernelize` validates the kernel's forward signature against the class being replaced. # Since the fused container sits at the position of the first module in the chain, the # kernel's forward must match that module's signature. We patch the class-level forward @@ -659,7 +652,7 @@ def _apply_weight_conversions(self, conversion_mapping: list) -> None: for target_key, meta_val in result.items(): meta_tensor = meta_val[0] if isinstance(meta_val, list) else meta_val t_mod_path, _, t_attr = target_key.rpartition(".") - target_mod = self._get_or_create_submodule(t_mod_path, original_mod, meta_tensor) + target_mod = FusedModuleBase._get_or_create_submodule(self, t_mod_path, original_mod, meta_tensor) if not isinstance(meta_tensor, nn.Parameter): meta_tensor = nn.Parameter(meta_tensor, requires_grad=meta_tensor.is_floating_point()) target_mod.register_parameter(t_attr, meta_tensor) @@ -706,18 +699,154 @@ def make_fused_module_class(source_layer_names: tuple[str, ...], kernel_layer_na ) +def _strip_converter_prefix(converter, prefix: str): + """Return a new WeightConverter with `prefix.` stripped from all source and target patterns.""" + from ..core_model_loading import WeightConverter + + prefix_dot = prefix + "." + new_src = [p.removeprefix(prefix_dot) for p in converter._original_source_patterns] + new_tgt = [p.removeprefix(prefix_dot) for p in converter._original_target_patterns] + return WeightConverter(new_src, new_tgt, operations=converter.operations) + + +def _apply_converter_to_module(module: "nn.Module", converter) -> None: + """ + Apply a single WeightConverter to any nn.Module. + + Uses the converter's _original_source_patterns / _original_target_patterns to locate + source submodules directly (no state-dict walk), infers output shape by running the + converter operations on the existing meta tensors, then creates the target submodule + with the proper typed class (nn.Linear, nn.Embedding, …) and removes the source + submodules that become empty afterwards. + """ + from copy import deepcopy + + from ..core_model_loading import dot_natural_key, rename_source_key + + src_patterns = converter._original_source_patterns + + # Collect source tensors directly from the module — bail if any are missing. + src_tensors: dict[str, torch.Tensor] = {} + for pattern in src_patterns: + mod_path, _, attr = pattern.rpartition(".") + try: + src_mod = module.get_submodule(mod_path) if mod_path else module + except AttributeError: + return + t = src_mod._parameters.get(attr) + if t is None: + t = src_mod._buffers.get(attr) + if t is None: + return + src_tensors[pattern] = t + + # Feed the source tensors into a converter copy and run the operation to get the + # result meta tensor (shape inference is free on meta device). + conv = deepcopy(converter) + for param_name in sorted(src_tensors, key=dot_natural_key): + renamed_key, source_pattern = rename_source_key(param_name, [], [conv]) + if source_pattern is not None: + conv.add_tensor(renamed_key, param_name, source_pattern, src_tensors[param_name]) + + target_keys = list(conv.layer_targets) + if not target_keys: + return + target_key = target_keys[0] + result = conv.convert(target_key) + + # Infer the typed module class from the first source submodule. + first_src_path = src_patterns[0].rpartition(".")[0] + try: + first_src_mod = module.get_submodule(first_src_path) if first_src_path else module + except AttributeError: + first_src_mod = None + + # Remove source parameters before creating the target so shape queries stay clean. + for pattern in src_patterns: + mod_path, _, attr = pattern.rpartition(".") + try: + src_mod = module.get_submodule(mod_path) if mod_path else module + except AttributeError: + continue + src_mod._parameters.pop(attr, None) + src_mod._buffers.pop(attr, None) + + # Create the target submodule with the correct shape inferred from the result. + # The module is initialised with the right in/out dimensions — no parameter replacement needed. + for result_key, meta_val in result.items(): + meta_tensor = meta_val[0] if isinstance(meta_val, list) else meta_val + t_mod_path, _, _ = result_key.rpartition(".") + + if t_mod_path: + parent_path, _, mod_name = t_mod_path.rpartition(".") + parent_mod = module.get_submodule(parent_path) if parent_path else module + if mod_name not in parent_mod._modules: + parent_mod.add_module(mod_name, _create_typed_module(first_src_mod, meta_tensor)) + + # Prune source submodules that are now completely empty. + for pattern in src_patterns: + mod_path = pattern.rpartition(".")[0] + if not mod_path: + continue + try: + src_mod = module.get_submodule(mod_path) + except AttributeError: + continue + if ( + all(p is None for p in src_mod._parameters.values()) + and all(b is None for b in src_mod._buffers.values()) + and all(m is None for m in src_mod._modules.values()) + ): + parent_path, _, mod_name = mod_path.rpartition(".") + parent_mod = module.get_submodule(parent_path) if parent_path else module + parent_mod._modules.pop(mod_name, None) + + +def _create_typed_module(source_mod: "nn.Module | None", weight: "torch.Tensor") -> "nn.Module": + """Instantiate a properly-typed nn.Module whose shape matches weight, inferred from source_mod.""" + cls_ = type(source_mod) if source_mod is not None else nn.Module + if cls_ is nn.Linear: + out_features, in_features = weight.shape + return cls_(in_features=in_features, out_features=out_features, bias=source_mod.bias is not None) + if cls_ is nn.Embedding: + num_embeddings, embedding_dim = weight.shape + return cls_( + num_embeddings=num_embeddings, + embedding_dim=embedding_dim, + padding_idx=source_mod.padding_idx, + max_norm=source_mod.max_norm, + norm_type=source_mod.norm_type, + scale_grad_by_freq=source_mod.scale_grad_by_freq, + sparse=source_mod.sparse, + ) + raise ValueError(f"Cannot create target module: unsupported source class {cls_.__name__!r}") + + +def _make_converted_child_class(child_cls: type, stripped_converters: list) -> type: + """Return a subclass of child_cls whose __init__ applies weight conversions after construction.""" + original_init = child_cls.__init__ + _converters = list(stripped_converters) + + def converted_init(self, *args, **kwargs): + original_init(self, *args, **kwargs) + for converter in _converters: + _apply_converter_to_module(self, converter) + + converted_cls = type(f"Converted{child_cls.__name__}", (child_cls,), {"__init__": converted_init}) + converted_cls.__qualname__ = f"Converted{child_cls.__qualname__}" + return converted_cls + + def make_fused_parent_class( parent_cls: type, child_names: list[str], source_names: list[str], kernel_layer_name: str, - conversion_mapping: list | None = None, ) -> type: original_init = parent_cls.__init__ _child_names = list(child_names) _source_names = list(source_names) _kernel_layer_name = kernel_layer_name - _conversion_mapping = list(conversion_mapping) if conversion_mapping else [] fused_module_cls = make_fused_module_class(tuple(_source_names), _kernel_layer_name) @@ -727,7 +856,6 @@ def fused_init(self, *args, **kwargs): fused = fused_module_cls( modules_to_fuse, fused_module_names=list(_source_names), - conversion_mapping=_conversion_mapping or None, ) setattr(self, _child_names[0], fused) for name in _child_names[1:]: @@ -874,6 +1002,7 @@ def register_kernel_fusions( meta_model = cls(config) seen: set[type] = set() + seen_children: set[type] = set() patch_mapping: dict[str, type] = {} for module in meta_model.modules(): module_cls = type(module) @@ -882,13 +1011,24 @@ def register_kernel_fusions( if not all(hasattr(module, name) for name in child_names): continue seen.add(module_cls) - fused_parent_cls = make_fused_parent_class( - module_cls, - child_names, - source_names, - kernel_layer_name, - conversion_mapping=kernel_conversion_mapping, - ) + + if kernel_conversion_mapping: + for source_name, child_name in zip(source_names, child_names): + child_cls = type(getattr(module, child_name)) + if child_cls in seen_children: + continue + seen_children.add(child_cls) + prefix_dot = source_name + "." + relevant = [ + c + for c in kernel_conversion_mapping + if any(p.startswith(prefix_dot) for p in c._original_source_patterns) + ] + if relevant: + stripped = [_strip_converter_prefix(c, source_name) for c in relevant] + patch_mapping[child_cls.__name__] = _make_converted_child_class(child_cls, stripped) + + fused_parent_cls = make_fused_parent_class(module_cls, child_names, source_names, kernel_layer_name) patch_mapping[module_cls.__name__] = fused_parent_cls if not patch_mapping: diff --git a/tests/kernels/test_kernels.py b/tests/kernels/test_kernels.py index 141f1e4becea..ca664acefd62 100644 --- a/tests/kernels/test_kernels.py +++ b/tests/kernels/test_kernels.py @@ -38,6 +38,9 @@ from transformers.integrations.hub_kernels import ( _HUB_KERNEL_MAPPING, _KERNEL_MODULE_MAPPING, + _apply_converter_to_module, + _make_converted_child_class, + _strip_converter_prefix, infer_kernel_fusion_transforms, is_kernel, lazy_load_kernel, @@ -926,3 +929,96 @@ def __init__(self): ) self.assertEqual(original_keys, set(fused.state_dict().keys())) + + +class TestConvertedChildClass(unittest.TestCase): + def test_strip_converter_prefix_removes_prefix(self): + converter = WeightConverter( + ["MLP.gate_proj.weight", "MLP.up_proj.weight"], + "MLP.gate_up_proj.weight", + operations=[Concatenate(dim=0)], + ) + stripped = _strip_converter_prefix(converter, "MLP") + self.assertEqual(stripped._original_source_patterns, ["gate_proj.weight", "up_proj.weight"]) + self.assertEqual(stripped._original_target_patterns, ["gate_up_proj.weight"]) + + def test_strip_converter_prefix_leaves_non_matching_unchanged(self): + converter = WeightConverter( + ["gate_proj.weight", "up_proj.weight"], + "gate_up_proj.weight", + operations=[Concatenate(dim=0)], + ) + stripped = _strip_converter_prefix(converter, "MLP") + self.assertEqual(stripped._original_source_patterns, ["gate_proj.weight", "up_proj.weight"]) + self.assertEqual(stripped._original_target_patterns, ["gate_up_proj.weight"]) + + def test_apply_converter_to_module_creates_linear(self): + class MLP(nn.Module): + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(4, 8, bias=False) + self.up_proj = nn.Linear(4, 8, bias=False) + + mlp = MLP() + converter = WeightConverter( + ["gate_proj.weight", "up_proj.weight"], + "gate_up_proj.weight", + operations=[Concatenate(dim=0)], + ) + _apply_converter_to_module(mlp, converter) + state_dict = mlp.state_dict() + self.assertIn("gate_up_proj.weight", state_dict) + self.assertNotIn("gate_proj.weight", state_dict) + self.assertNotIn("up_proj.weight", state_dict) + self.assertEqual(state_dict["gate_up_proj.weight"].shape, (16, 4)) + self.assertIsInstance(mlp.gate_up_proj, nn.Linear) + self.assertFalse(hasattr(mlp, "gate_proj")) + self.assertFalse(hasattr(mlp, "up_proj")) + + def test_apply_converter_to_module_preserves_bias(self): + class MLP(nn.Module): + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(4, 8, bias=True) + self.up_proj = nn.Linear(4, 8, bias=True) + + mlp = MLP() + converter = WeightConverter( + ["gate_proj.weight", "up_proj.weight"], + "gate_up_proj.weight", + operations=[Concatenate(dim=0)], + ) + _apply_converter_to_module(mlp, converter) + self.assertIsInstance(mlp.gate_up_proj, nn.Linear) + self.assertIsNotNone(mlp.gate_up_proj.bias) + + def test_make_converted_child_class_applies_conversion(self): + class MLP(nn.Module): + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(4, 8, bias=False) + self.up_proj = nn.Linear(4, 8, bias=False) + + converter = WeightConverter( + ["gate_proj.weight", "up_proj.weight"], + "gate_up_proj.weight", + operations=[Concatenate(dim=0)], + ) + ConvertedMLP = _make_converted_child_class(MLP, [converter]) + mlp = ConvertedMLP() + state_dict = mlp.state_dict() + self.assertIn("gate_up_proj.weight", state_dict) + self.assertNotIn("gate_proj.weight", state_dict) + self.assertNotIn("up_proj.weight", state_dict) + self.assertEqual(state_dict["gate_up_proj.weight"].shape, (16, 4)) + self.assertIsInstance(mlp.gate_up_proj, nn.Linear) + + def test_make_converted_child_class_is_subclass(self): + class Base(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(2, 2, bias=False) + + Converted = _make_converted_child_class(Base, []) + self.assertTrue(issubclass(Converted, Base)) + self.assertEqual(Converted.__name__, "ConvertedBase") From ad0e24e970eb13ac46f11fce866b350118a28a74 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 26 May 2026 15:43:49 +0200 Subject: [PATCH 47/49] wip: remove dead code --- src/transformers/integrations/hub_kernels.py | 201 +++++++++--------- tests/kernels/test_kernels.py | 203 ++++++++----------- 2 files changed, 180 insertions(+), 224 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 82b1b51b028d..86375222bc1f 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -573,107 +573,6 @@ def __repr__(self): names = ", ".join(self._fused_module_names) return f"{self.__class__.__name__}(fused=({names}))" - def _get_or_create_submodule( - self, - module_path: str, - original_module: "nn.Module | None", - future_weight: "torch.Tensor", - ) -> "nn.Module": - if not module_path: - return self - parent, _, mod_name = module_path.rpartition(".") - parent_mod = self.get_submodule(parent) if parent else self - if mod_name not in parent_mod._modules: - cls_ = original_module.__class__ if original_module is not None else nn.Module - if cls_ is nn.Linear: - out_features, in_features = future_weight.shape - mod = cls_(in_features=in_features, out_features=out_features, bias=original_module.bias is not None) - elif cls_ is nn.Embedding: - num_embeddings, embedding_dim = future_weight.shape - mod = cls_( - num_embeddings=num_embeddings, - embedding_dim=embedding_dim, - padding_idx=original_module.padding_idx, - max_norm=original_module.max_norm, - norm_type=original_module.norm_type, - scale_grad_by_freq=original_module.scale_grad_by_freq, - sparse=original_module.sparse, - ) - else: - raise ValueError( - f"Cannot create submodule {mod_name!r} at {module_path!r}: unsupported class {cls_.__name__!r}" - ) - parent_mod.add_module(mod_name, mod) - else: - mod = parent_mod._modules[mod_name] - return mod - - def _apply_weight_conversions(self, conversion_mapping: list) -> None: - """Apply WeightConverter transforms to this module's graph (cost-free on meta device).""" - - from ..core_model_loading import dot_natural_key, rename_source_key - - state_dict = self.state_dict() - pattern_to_converter = {k: c for c in conversion_mapping for k in c.source_patterns} - param_name_to_transform: dict = {} - - for param_name in sorted(state_dict, key=dot_natural_key): - renamed_key, source_pattern = rename_source_key(param_name, [], conversion_mapping) - if source_pattern is None: - continue - transform = param_name_to_transform.setdefault(renamed_key, deepcopy(pattern_to_converter[source_pattern])) - transform.add_tensor(renamed_key, param_name, source_pattern, state_dict[param_name]) - - for layer_name, transform in param_name_to_transform.items(): - source_names = list(transform.layer_targets.get(layer_name, [])) - if not source_names: - continue - - original_mod = None - for sname in source_names: - smod_path = sname.rpartition(".")[0] - try: - original_mod = self.get_submodule(smod_path) if smod_path else self - break - except AttributeError: - pass - - result = transform.convert(layer_name) - - for sname in source_names: - mod_path, _, attr = sname.rpartition(".") - try: - source_mod = self.get_submodule(mod_path) if mod_path else self - except AttributeError: - continue - source_mod._parameters.pop(attr, None) - source_mod._buffers.pop(attr, None) - - for target_key, meta_val in result.items(): - meta_tensor = meta_val[0] if isinstance(meta_val, list) else meta_val - t_mod_path, _, t_attr = target_key.rpartition(".") - target_mod = FusedModuleBase._get_or_create_submodule(self, t_mod_path, original_mod, meta_tensor) - if not isinstance(meta_tensor, nn.Parameter): - meta_tensor = nn.Parameter(meta_tensor, requires_grad=meta_tensor.is_floating_point()) - target_mod.register_parameter(t_attr, meta_tensor) - - for source_name in source_names: - mod_path = source_name.rpartition(".")[0] - if not mod_path: - continue - try: - source_mod = self.get_submodule(mod_path) - except AttributeError: - continue - if ( - all(p is None for p in source_mod._parameters.values()) - and all(b is None for b in source_mod._buffers.values()) - and all(m is None for m in source_mod._modules.values()) - ): - parent_path, _, mod_name = mod_path.rpartition(".") - parent_mod = self.get_submodule(parent_path) if parent_path else self - parent_mod._modules.pop(mod_name, None) - @functools.cache def make_fused_module_class(source_layer_names: tuple[str, ...], kernel_layer_name: str) -> type: @@ -699,6 +598,87 @@ def make_fused_module_class(source_layer_names: tuple[str, ...], kernel_layer_na ) +def _get_or_create_submodule( + root: "nn.Module", + mod_path: str, + original_module: "nn.Module | None", + future_weight: "torch.Tensor", +) -> "nn.Module": + if not mod_path: + return root + parent, _, mod_name = mod_path.rpartition(".") + parent_mod = root.get_submodule(parent) if parent else root + if mod_name not in parent_mod._modules: + parent_mod.add_module(mod_name, _create_typed_module(original_module, future_weight)) + return parent_mod._modules[mod_name] + + +def _apply_weight_conversions(module: "nn.Module", conversion_mapping: list) -> None: + """Apply WeightConverter transforms to a module via state-dict walk. Handles glob patterns.""" + from ..core_model_loading import dot_natural_key, rename_source_key + + state_dict = module.state_dict() + pattern_to_converter = {k: c for c in conversion_mapping for k in c.source_patterns} + param_name_to_transform: dict = {} + + for param_name in sorted(state_dict, key=dot_natural_key): + renamed_key, source_pattern = rename_source_key(param_name, [], conversion_mapping) + if source_pattern is None: + continue + transform = param_name_to_transform.setdefault(renamed_key, deepcopy(pattern_to_converter[source_pattern])) + transform.add_tensor(renamed_key, param_name, source_pattern, state_dict[param_name]) + + for layer_name, transform in param_name_to_transform.items(): + source_names = list(transform.layer_targets.get(layer_name, [])) + if not source_names: + continue + + original_mod = None + for sname in source_names: + smod_path = sname.rpartition(".")[0] + try: + original_mod = module.get_submodule(smod_path) if smod_path else module + break + except AttributeError: + pass + + result = transform.convert(layer_name) + + for sname in source_names: + mod_path, _, attr = sname.rpartition(".") + try: + source_mod = module.get_submodule(mod_path) if mod_path else module + except AttributeError: + continue + source_mod._parameters.pop(attr, None) + source_mod._buffers.pop(attr, None) + + for target_key, meta_val in result.items(): + meta_tensor = meta_val[0] if isinstance(meta_val, list) else meta_val + t_mod_path, _, t_attr = target_key.rpartition(".") + target_mod = _get_or_create_submodule(module, t_mod_path, original_mod, meta_tensor) + if not isinstance(meta_tensor, nn.Parameter): + meta_tensor = nn.Parameter(meta_tensor, requires_grad=meta_tensor.is_floating_point()) + target_mod.register_parameter(t_attr, meta_tensor) + + for source_name in source_names: + mod_path = source_name.rpartition(".")[0] + if not mod_path: + continue + try: + source_mod = module.get_submodule(mod_path) + except AttributeError: + continue + if ( + all(p is None for p in source_mod._parameters.values()) + and all(b is None for b in source_mod._buffers.values()) + and all(m is None for m in source_mod._modules.values()) + ): + parent_path, _, mod_name = mod_path.rpartition(".") + parent_mod = module.get_submodule(parent_path) if parent_path else module + parent_mod._modules.pop(mod_name, None) + + def _strip_converter_prefix(converter, prefix: str): """Return a new WeightConverter with `prefix.` stripped from all source and target patterns.""" from ..core_model_loading import WeightConverter @@ -713,18 +693,23 @@ def _apply_converter_to_module(module: "nn.Module", converter) -> None: """ Apply a single WeightConverter to any nn.Module. - Uses the converter's _original_source_patterns / _original_target_patterns to locate - source submodules directly (no state-dict walk), infers output shape by running the - converter operations on the existing meta tensors, then creates the target submodule - with the proper typed class (nn.Linear, nn.Embedding, …) and removes the source - submodules that become empty afterwards. + For literal (non-glob) patterns: locates source submodules directly by path, runs the + converter operations to infer the output shape, then creates the target submodule with + the proper typed class (nn.Linear, nn.Embedding, …) and prunes empty source submodules. + For glob patterns (containing *): falls back to _apply_weight_conversions which does a + full state-dict walk and registers the result parameter explicitly. """ + src_patterns = converter._original_source_patterns + + # Glob patterns require a state-dict walk to expand wildcards. + if any("*" in p for p in src_patterns): + _apply_weight_conversions(module, [converter]) + return + from copy import deepcopy from ..core_model_loading import dot_natural_key, rename_source_key - src_patterns = converter._original_source_patterns - # Collect source tensors directly from the module — bail if any are missing. src_tensors: dict[str, torch.Tensor] = {} for pattern in src_patterns: diff --git a/tests/kernels/test_kernels.py b/tests/kernels/test_kernels.py index ca664acefd62..1dcbb0776718 100644 --- a/tests/kernels/test_kernels.py +++ b/tests/kernels/test_kernels.py @@ -39,6 +39,7 @@ _HUB_KERNEL_MAPPING, _KERNEL_MODULE_MAPPING, _apply_converter_to_module, + _apply_weight_conversions, _make_converted_child_class, _strip_converter_prefix, infer_kernel_fusion_transforms, @@ -643,7 +644,6 @@ def __init__(self, config=None): def tearDown(self): import transformers.conversion_mapping as _cm - from transformers.integrations.hub_kernels import make_fused_module_class if _cm._checkpoint_conversion_mapping_cache is not None: _cm._checkpoint_conversion_mapping_cache.pop(self._MODEL_TYPE, None) @@ -732,12 +732,6 @@ class FakeKernel: class TestFusedModuleWeightConversions(unittest.TestCase): - def _make_fused(self, *modules, names=None): - if names is None: - names = [type(m).__name__ for m in modules] - FusedCls = make_fused_module_class(tuple(names), "kernel") - return FusedCls(list(modules), fused_module_names=list(names)) - def test_weight_converter_concatenates_projections(self): class SeparateAttn(nn.Module): def __init__(self): @@ -746,37 +740,32 @@ def __init__(self): self.k_proj = nn.Linear(2, 2, bias=False) self.v_proj = nn.Linear(2, 2, bias=False) - fused = self._make_fused(SeparateAttn(), names=["Attn"]) - - q_weight = torch.randn(2, 2) - k_weight = torch.randn(2, 2) - v_weight = torch.randn(2, 2) - fused.Attn.q_proj.weight.data.copy_(q_weight) - fused.Attn.k_proj.weight.data.copy_(k_weight) - fused.Attn.v_proj.weight.data.copy_(v_weight) + attn = SeparateAttn() + q_weight = attn.q_proj.weight.data.clone() + k_weight = attn.k_proj.weight.data.clone() + v_weight = attn.v_proj.weight.data.clone() - fused._apply_weight_conversions( + _apply_weight_conversions( + attn, [ WeightConverter( - ["Attn.q_proj.weight", "Attn.k_proj.weight", "Attn.v_proj.weight"], - "Attn.qkv_proj.weight", + ["q_proj.weight", "k_proj.weight", "v_proj.weight"], + "qkv_proj.weight", operations=[Concatenate(dim=0)], ) - ] + ], ) - state_dict = fused.state_dict() - self.assertIn("Attn.qkv_proj.weight", state_dict) - self.assertNotIn("Attn.q_proj.weight", state_dict) - self.assertNotIn("Attn.k_proj.weight", state_dict) - self.assertNotIn("Attn.v_proj.weight", state_dict) - self.assertIn("qkv_proj", fused.Attn._modules) - self.assertNotIn("q_proj", fused.Attn._modules) - self.assertNotIn("k_proj", fused.Attn._modules) - self.assertNotIn("v_proj", fused.Attn._modules) - - expected_qkv = torch.cat([q_weight, k_weight, v_weight], dim=0) - torch.testing.assert_close(state_dict["Attn.qkv_proj.weight"], expected_qkv) + state_dict = attn.state_dict() + self.assertIn("qkv_proj.weight", state_dict) + self.assertNotIn("q_proj.weight", state_dict) + self.assertNotIn("k_proj.weight", state_dict) + self.assertNotIn("v_proj.weight", state_dict) + self.assertIn("qkv_proj", attn._modules) + self.assertNotIn("q_proj", attn._modules) + self.assertNotIn("k_proj", attn._modules) + self.assertNotIn("v_proj", attn._modules) + torch.testing.assert_close(state_dict["qkv_proj.weight"], torch.cat([q_weight, k_weight, v_weight], dim=0)) def test_weight_converter_splits_fused_projection(self): class QKVAttn(nn.Module): @@ -784,35 +773,31 @@ def __init__(self): super().__init__() self.qkv_proj = nn.Linear(2, 6, bias=False) - fused = self._make_fused(QKVAttn(), names=["QKV"]) - - qkv_weight = torch.randn(6, 2) - fused.QKV.qkv_proj.weight.data.copy_(qkv_weight) + attn = QKVAttn() + qkv_weight = attn.qkv_proj.weight.data.clone() - fused._apply_weight_conversions( + _apply_weight_conversions( + attn, [ WeightConverter( - "QKV.qkv_proj.weight", - ["QKV.q_proj.weight", "QKV.k_proj.weight", "QKV.v_proj.weight"], - operations=[Chunk(dim=0)], + "qkv_proj.weight", ["q_proj.weight", "k_proj.weight", "v_proj.weight"], operations=[Chunk(dim=0)] ) - ] + ], ) - state_dict = fused.state_dict() - self.assertIn("QKV.q_proj.weight", state_dict) - self.assertIn("QKV.k_proj.weight", state_dict) - self.assertIn("QKV.v_proj.weight", state_dict) - self.assertNotIn("QKV.qkv_proj.weight", state_dict) - self.assertNotIn("qkv_proj", fused.QKV._modules) - self.assertIn("q_proj", fused.QKV._modules) - self.assertIn("k_proj", fused.QKV._modules) - self.assertIn("v_proj", fused.QKV._modules) - + state_dict = attn.state_dict() + self.assertIn("q_proj.weight", state_dict) + self.assertIn("k_proj.weight", state_dict) + self.assertIn("v_proj.weight", state_dict) + self.assertNotIn("qkv_proj.weight", state_dict) + self.assertNotIn("qkv_proj", attn._modules) + self.assertIn("q_proj", attn._modules) + self.assertIn("k_proj", attn._modules) + self.assertIn("v_proj", attn._modules) expected_q, expected_k, expected_v = torch.chunk(qkv_weight, 3, dim=0) - torch.testing.assert_close(state_dict["QKV.q_proj.weight"], expected_q) - torch.testing.assert_close(state_dict["QKV.k_proj.weight"], expected_k) - torch.testing.assert_close(state_dict["QKV.v_proj.weight"], expected_v) + torch.testing.assert_close(state_dict["q_proj.weight"], expected_q) + torch.testing.assert_close(state_dict["k_proj.weight"], expected_k) + torch.testing.assert_close(state_dict["v_proj.weight"], expected_v) def test_weight_converter_split_is_idempotent(self): class QKVAttn(nn.Module): @@ -823,18 +808,15 @@ def __init__(self): def make_mapping(): return [ WeightConverter( - "QKV.qkv_proj.weight", - ["QKV.q_proj.weight", "QKV.k_proj.weight", "QKV.v_proj.weight"], - operations=[Chunk(dim=0)], + "qkv_proj.weight", ["q_proj.weight", "k_proj.weight", "v_proj.weight"], operations=[Chunk(dim=0)] ) ] - fused = self._make_fused(QKVAttn(), names=["QKV"]) - fused._apply_weight_conversions(make_mapping()) - keys_after_first = set(fused.state_dict().keys()) - - fused._apply_weight_conversions(make_mapping()) - self.assertEqual(keys_after_first, set(fused.state_dict().keys())) + attn = QKVAttn() + _apply_weight_conversions(attn, make_mapping()) + keys_after_first = set(attn.state_dict().keys()) + _apply_weight_conversions(attn, make_mapping()) + self.assertEqual(keys_after_first, set(attn.state_dict().keys())) def test_weight_converter_glob_merges_modulelist(self): class ExpertModule(nn.Module): @@ -843,20 +825,19 @@ def __init__(self): self.experts = nn.ModuleList([nn.Linear(2, 4, bias=False), nn.Linear(2, 4, bias=False)]) self.merged = nn.Module() - fused = self._make_fused(ExpertModule(), names=["M"]) + mod = ExpertModule() + weight_0 = mod.experts[0].weight.data.clone() + weight_1 = mod.experts[1].weight.data.clone() - weight_0 = fused.M.experts[0].weight.data.clone() - weight_1 = fused.M.experts[1].weight.data.clone() - - fused._apply_weight_conversions( - [WeightConverter("M.experts.*.weight", "M.merged.weight", operations=[MergeModulelist(dim=0)])] + _apply_weight_conversions( + mod, [WeightConverter("experts.*.weight", "merged.weight", operations=[MergeModulelist(dim=0)])] ) - state_dict = fused.state_dict() - self.assertIn("M.merged.weight", state_dict) - self.assertNotIn("M.experts.0.weight", state_dict) - self.assertNotIn("M.experts.1.weight", state_dict) - torch.testing.assert_close(state_dict["M.merged.weight"], torch.stack([weight_0, weight_1], dim=0)) + state_dict = mod.state_dict() + self.assertIn("merged.weight", state_dict) + self.assertNotIn("experts.0.weight", state_dict) + self.assertNotIn("experts.1.weight", state_dict) + torch.testing.assert_close(state_dict["merged.weight"], torch.stack([weight_0, weight_1], dim=0)) def test_multiple_independent_converters(self): class QKVAttn(nn.Module): @@ -870,50 +851,42 @@ def __init__(self): self.gate_proj = nn.Linear(2, 4, bias=False) self.up_proj = nn.Linear(2, 4, bias=False) - fused = self._make_fused(QKVAttn(), MLP(), names=["Attn", "MLP"]) - - qkv_weight = torch.randn(6, 2) - gate_weight = torch.randn(4, 2) - up_weight = torch.randn(4, 2) - fused.Attn.qkv_proj.weight.data.copy_(qkv_weight) - fused.MLP.gate_proj.weight.data.copy_(gate_weight) - fused.MLP.up_proj.weight.data.copy_(up_weight) + attn = QKVAttn() + mlp = MLP() + qkv_weight = attn.qkv_proj.weight.data.clone() + gate_weight = mlp.gate_proj.weight.data.clone() + up_weight = mlp.up_proj.weight.data.clone() - fused._apply_weight_conversions( + _apply_weight_conversions( + attn, [ WeightConverter( - "Attn.qkv_proj.weight", - ["Attn.q_proj.weight", "Attn.k_proj.weight", "Attn.v_proj.weight"], - operations=[Chunk(dim=0)], - ), + "qkv_proj.weight", ["q_proj.weight", "k_proj.weight", "v_proj.weight"], operations=[Chunk(dim=0)] + ) + ], + ) + _apply_weight_conversions( + mlp, + [ WeightConverter( - ["MLP.gate_proj.weight", "MLP.up_proj.weight"], - "MLP.gate_up_proj.weight", - operations=[Concatenate(dim=0)], - ), - ] + ["gate_proj.weight", "up_proj.weight"], "gate_up_proj.weight", operations=[Concatenate(dim=0)] + ) + ], ) - state_dict = fused.state_dict() - - self.assertNotIn("Attn.qkv_proj.weight", state_dict) - self.assertIn("Attn.q_proj.weight", state_dict) - self.assertIn("Attn.k_proj.weight", state_dict) - self.assertIn("Attn.v_proj.weight", state_dict) - self.assertNotIn("qkv_proj", fused.Attn._modules) - self.assertIn("q_proj", fused.Attn._modules) - self.assertIn("k_proj", fused.Attn._modules) - self.assertIn("v_proj", fused.Attn._modules) + attn_sd = attn.state_dict() + self.assertNotIn("qkv_proj.weight", attn_sd) + self.assertIn("q_proj.weight", attn_sd) + self.assertIn("k_proj.weight", attn_sd) + self.assertIn("v_proj.weight", attn_sd) expected_q, expected_k, expected_v = torch.chunk(qkv_weight, 3, dim=0) - torch.testing.assert_close(state_dict["Attn.q_proj.weight"], expected_q) + torch.testing.assert_close(attn_sd["q_proj.weight"], expected_q) - self.assertNotIn("MLP.gate_proj.weight", state_dict) - self.assertNotIn("MLP.up_proj.weight", state_dict) - self.assertIn("MLP.gate_up_proj.weight", state_dict) - self.assertNotIn("gate_proj", fused.MLP._modules) - self.assertNotIn("up_proj", fused.MLP._modules) - self.assertIn("gate_up_proj", fused.MLP._modules) - torch.testing.assert_close(state_dict["MLP.gate_up_proj.weight"], torch.cat([gate_weight, up_weight], dim=0)) + mlp_sd = mlp.state_dict() + self.assertNotIn("gate_proj.weight", mlp_sd) + self.assertNotIn("up_proj.weight", mlp_sd) + self.assertIn("gate_up_proj.weight", mlp_sd) + torch.testing.assert_close(mlp_sd["gate_up_proj.weight"], torch.cat([gate_weight, up_weight], dim=0)) def test_no_match_is_noop(self): class SimpleModule(nn.Module): @@ -921,14 +894,12 @@ def __init__(self): super().__init__() self.linear = nn.Linear(4, 4, bias=False) - fused = self._make_fused(SimpleModule(), names=["M"]) - original_keys = set(fused.state_dict().keys()) - - fused._apply_weight_conversions( - [WeightConverter("M.non_existent.weight", "M.something.weight", operations=[Concatenate(dim=0)])] + mod = SimpleModule() + original_keys = set(mod.state_dict().keys()) + _apply_weight_conversions( + mod, [WeightConverter("non_existent.weight", "something.weight", operations=[Concatenate(dim=0)])] ) - - self.assertEqual(original_keys, set(fused.state_dict().keys())) + self.assertEqual(original_keys, set(mod.state_dict().keys())) class TestConvertedChildClass(unittest.TestCase): From 3924cf358eb7eba234bbfd0cae51314f1d61e156 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 26 May 2026 16:01:41 +0200 Subject: [PATCH 48/49] wip: api imrpovement --- src/transformers/integrations/hub_kernels.py | 21 ++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 86375222bc1f..443e97f4a3c8 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -672,7 +672,7 @@ def _apply_weight_conversions(module: "nn.Module", conversion_mapping: list) -> if ( all(p is None for p in source_mod._parameters.values()) and all(b is None for b in source_mod._buffers.values()) - and all(m is None for m in source_mod._modules.values()) + and not source_mod._modules ): parent_path, _, mod_name = mod_path.rpartition(".") parent_mod = module.get_submodule(parent_path) if parent_path else module @@ -736,8 +736,11 @@ def _apply_converter_to_module(module: "nn.Module", converter) -> None: target_keys = list(conv.layer_targets) if not target_keys: return - target_key = target_keys[0] - result = conv.convert(target_key) + + # Collect results from all target keys; a converter may produce more than one output tensor. + result: dict = {} + for target_key in target_keys: + result.update(conv.convert(target_key)) # Infer the typed module class from the first source submodule. first_src_path = src_patterns[0].rpartition(".")[0] @@ -777,10 +780,13 @@ def _apply_converter_to_module(module: "nn.Module", converter) -> None: src_mod = module.get_submodule(mod_path) except AttributeError: continue + # _parameters and _buffers may hold None-valued entries (e.g. bias=False in nn.Linear + # registers _parameters['bias'] = None), so check all-None rather than emptiness. + # _modules entries are always non-None module objects, so use emptiness there. if ( all(p is None for p in src_mod._parameters.values()) and all(b is None for b in src_mod._buffers.values()) - and all(m is None for m in src_mod._modules.values()) + and not src_mod._modules ): parent_path, _, mod_name = mod_path.rpartition(".") parent_mod = module.get_submodule(parent_path) if parent_path else module @@ -790,10 +796,10 @@ def _apply_converter_to_module(module: "nn.Module", converter) -> None: def _create_typed_module(source_mod: "nn.Module | None", weight: "torch.Tensor") -> "nn.Module": """Instantiate a properly-typed nn.Module whose shape matches weight, inferred from source_mod.""" cls_ = type(source_mod) if source_mod is not None else nn.Module - if cls_ is nn.Linear: + if issubclass(cls_, nn.Linear): out_features, in_features = weight.shape return cls_(in_features=in_features, out_features=out_features, bias=source_mod.bias is not None) - if cls_ is nn.Embedding: + if issubclass(cls_, nn.Embedding): num_embeddings, embedding_dim = weight.shape return cls_( num_embeddings=num_embeddings, @@ -831,9 +837,8 @@ def make_fused_parent_class( original_init = parent_cls.__init__ _child_names = list(child_names) _source_names = list(source_names) - _kernel_layer_name = kernel_layer_name - fused_module_cls = make_fused_module_class(tuple(_source_names), _kernel_layer_name) + fused_module_cls = make_fused_module_class(tuple(_source_names), kernel_layer_name) def fused_init(self, *args, **kwargs): original_init(self, *args, **kwargs) From c948834232157a6d825509a6d83c2495e071fb75 Mon Sep 17 00:00:00 2001 From: Michael Benayoun Date: Tue, 26 May 2026 16:31:03 +0200 Subject: [PATCH 49/49] wip: refactor --- src/transformers/integrations/hub_kernels.py | 23 +++++++++++++++----- 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 443e97f4a3c8..fc9a39043a31 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -911,6 +911,18 @@ def segment_to_regex(seg: str) -> str: return transforms +def _first_str_leaf(obj) -> str | None: + """Recursively extract the first string leaf from a potentially nested dict (device → mode → str).""" + if isinstance(obj, str): + return obj + if isinstance(obj, dict): + for v in obj.values(): + result = _first_str_leaf(v) + if result is not None: + return result + return None + + def _try_load_kernel_class(repo_str: str, use_local: bool = False) -> type | None: if ":" not in repo_str: return None @@ -971,12 +983,11 @@ def register_kernel_fusions( glob_patterns = [item[1] for item in layer_name] child_names = [p.rsplit(".", 1)[1] for p in glob_patterns] - # Resolve kernel_layer_name from "org/repo:ClassName" or device-nested dict. - repo_str = hub_repo - if isinstance(repo_str, dict): - repo_str = next(iter(repo_str.values())) - if isinstance(repo_str, dict): - repo_str = next(iter(repo_str.values())) + # Resolve kernel_layer_name from "org/repo:ClassName" or a device/mode-nested dict. + # All leaf strings in the nested dict point to the same kernel class, so any will do. + repo_str = _first_str_leaf(hub_repo) + if repo_str is None: + raise ValueError(f"Cannot resolve a repo string from hub_repo={hub_repo!r}") kernel_layer_name = repo_str.split(":")[1] if ":" in repo_str else repo_str.split("/")[-1] # 1. Load the kernel class to get any conversion_mapping it declares.