🚨 [Kernels] Fix kernel function registration

src/transformers/integrations/hub_kernels.py

@@ -225,9 +225,12 @@ def use_kernel_func_from_hub(func_name: str):
                 )
             },
             "cuda": {
+                Mode.TRAINING: FuncRepository(
+                    repo_id="kernels-community/rotary", func_name="apply_rotary_transformers"
+                ),
                 Mode.INFERENCE: FuncRepository(
                     repo_id="kernels-community/rotary", func_name="apply_rotary_transformers"
-                )
+                ),
             },
         }
 
@@ -438,11 +441,13 @@ def get_kernel(
 
 def use_kernelized_func(module_names: list[Callable] | Callable):
     """
-    This decorator attaches the target function as an attribute of the module.
-    The function must already be decorated with @use_kernel_func_from_hub
-    this decorator then wraps it as an nn.Module internally.
-    When kernelize is later applied to the full model, the function can be accessed as a regular module attribute and kernelized just like any other layer.
-    The kernelization is performed in place, modifying the module directly.
+    This decorator attaches the target function within the module as a plain attribute (not as a submodule).
+    Keep in mind that this registration is only meant for `kernelize` to recognize its target modules (i.e.
+    function exchanged for a weightless `nn.Module` with the same forward) to then exchange to the kernel
+    variation (in-place) if the conditions are met.
+
+    We cache each of these function-based registrations: After proper registration and exchange it is removed
+    from the module's `_modules` dict as it does not really act as `nn.Module` but a base function.
     """
     if isinstance(module_names, Callable):
         module_names = [module_names]
@@ -452,18 +457,20 @@ def decorator(cls):
 
         def new_init(self, *args, **kwargs):
             orig_init(self, *args, **kwargs)
-            # Skip attaching the kernelized submodule under DeepSpeed ZeRO-3: the coordinator traces
-            # the module graph at init time, and a child `nn.Module` that is not actually invoked
-            # during forward (e.g. when the model keeps calling the plain Python `apply_rotary_pos_emb`)
-            # breaks the parameter fetch trace and raises `IndexError: pop from an empty deque`.
-            # See https://github.com/huggingface/transformers/issues/45137
-            from .deepspeed import is_deepspeed_zero3_enabled
-
-            if is_deepspeed_zero3_enabled():
-                return
+
+            # Register new function as non-submodule within the modules dict
+            hidden_kernels = self.__dict__.setdefault("_hidden_kernels", {})
             for fn in module_names:
-                # we hardcode the name of the function to "rotary_fn" for now
-                setattr(self, "rotary_fn", fn)
+                name = (
+                    getattr(fn, "__name__", None)
+                    or getattr(fn, "kernel_layer_name", None)
+                    or getattr(fn, "func_name", None)
+                )
+                if name is None:
+                    raise ValueError(f"Could not infer kernel function name for {fn!r}")
+
+                # Do not register as submodule! Hide it behind a dict to be removed later after registering it
+                hidden_kernels[name] = fn
 
         cls.__init__ = new_init
         return cls

src/transformers/modeling_utils.py

@@ -4466,15 +4466,31 @@ def loss_function(self, value):
         self._loss_function = value
 
     def kernelize(self, mode=None):
+        """Temporarily register hidden kernel wrappers so `kernelize` can discover and replace them."""
         if not is_kernels_available():
             raise ValueError(
-                "Kernels are not available. To use kernels, please install kernels using `pip install kernels`"
+                "Kernels are not available. To use kernels, please install kernels using `pip install -U kernels`"
             )
         from kernels import Device, Mode, kernelize
 
-        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)
-        self._use_kernels = True
+        def attach_hidden_kernels(module):
+            for name, fn in getattr(module, "_hidden_kernels", {}).items():
+                if name not in dict(module.named_children()):
+                    module.register_module(name, fn)
+
+        def detach_hidden_kernels(module):
+            for name in getattr(module, "_hidden_kernels", {}):
+                delattr(module, name)
+
+        try:
+            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)
+            self._use_kernels = True
+
+        finally:
+            self.apply(detach_hidden_kernels)
 
     @property
     def use_kernels(self) -> bool: