Fix kernelize() crash for gpt_oss: missing @use_kernel_func_from_hub on apply_rotary_pos_emb

src/transformers/modeling_utils.py

@@ -4580,11 +4580,20 @@ def kernelize(self, mode=None):
         def attach_hidden_kernels(module):
             for name, fn in getattr(module, "_hidden_kernels", {}).items():
                 if name not in dict(module.named_children()):
+                    if not isinstance(fn, nn.Module):
+                        raise ValueError(
+                            f"Attempted to register a kernel for {name}, but it was not a `torch.nn.Module`. "
+                            "This means the underlying function needs to be decorated with `@use_kernel_func_from_hub`. "
+                            "Please submit and issue to the transformers repo: `https://github.com/huggingface/transformers/issues`."
+                        )
                     module.register_module(name, fn)
 
         def detach_hidden_kernels(module):
             for name in getattr(module, "_hidden_kernels", {}):
-                delattr(module, name)
+                # Skip deregistering if it failed to properly register,
+                # i.e. `ValueError` will be raised afterwards
+                if hasattr(module, name):
+                    delattr(module, name)
 
         try:
             self.apply(attach_hidden_kernels)

src/transformers/models/gpt_oss/modeling_gpt_oss.py

@@ -27,7 +27,12 @@
 from ... import initialization as init
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func
+from ...integrations import (
+    use_experts_implementation,
+    use_kernel_forward_from_hub,
+    use_kernel_func_from_hub,
+    use_kernelized_func,
+)
 from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
 from ...modeling_layers import (
     GenericForSequenceClassification,
@@ -234,6 +239,7 @@ def _apply_rotary_emb(
     return torch.cat((first_, second_), dim=-1)
 
 
+@use_kernel_func_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     cos = cos.unsqueeze(unsqueeze_dim)
     sin = sin.unsqueeze(unsqueeze_dim)

src/transformers/models/gpt_oss/modular_gpt_oss.py

@@ -19,7 +19,7 @@
 
 from ... import initialization as init
 from ...cache_utils import Cache, DynamicCache
-from ...integrations import use_experts_implementation, use_kernel_forward_from_hub
+from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub
 from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
 from ...modeling_outputs import (
     MoeModelOutputWithPast,
@@ -171,6 +171,7 @@ def _apply_rotary_emb(
     return torch.cat((first_, second_), dim=-1)
 
 
+@use_kernel_func_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     cos = cos.unsqueeze(unsqueeze_dim)
     sin = sin.unsqueeze(unsqueeze_dim)

src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py

@@ -26,7 +26,12 @@
 from torch.nn import functional as F
 
 from ... import initialization as init
-from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func
+from ...integrations import (
+    use_experts_implementation,
+    use_kernel_forward_from_hub,
+    use_kernel_func_from_hub,
+    use_kernelized_func,
+)
 from ...masking_utils import create_bidirectional_sliding_window_mask
 from ...modeling_layers import GenericForTokenClassification, GradientCheckpointingLayer
 from ...modeling_outputs import BaseModelOutput
@@ -137,6 +142,7 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
     return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
 
 
+@use_kernel_func_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     cos = cos.unsqueeze(unsqueeze_dim)
     sin = sin.unsqueeze(unsqueeze_dim)

tests/models/gpt_oss/test_modeling_gpt_oss.py

@@ -69,6 +69,14 @@ class GptOssModelTest(CausalLMModelTest, unittest.TestCase):
     model_split_percents = [0.5, 0.6]
     model_tester_class = GptOssModelTester
 
+    @require_kernels
+    def test_kernelize_does_not_crash(self):
+        """Regression test #45799: `kernelize` should not crash with `use_kernelized_func` + `use_kernel_func_from_hub`."""
+        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
+        model = GptOssModel(config).to(device=torch_device)
+        # This used to raise TypeError because apply_rotary_pos_emb was not wrapped as nn.Module
+        model.use_kernels = True
+
     @require_kernels
     @pytest.mark.flash_attn_test
     @require_torch_gpu