diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index d0484b7f7f6b..4a989eac267b 100644 --- a/src/transformers/modeling_utils.py +++ b/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) diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py index 55381a7e3c21..484b071f01bc 100644 --- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/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) diff --git a/src/transformers/models/gpt_oss/modular_gpt_oss.py b/src/transformers/models/gpt_oss/modular_gpt_oss.py index 3354acef2196..cc9c42364511 100644 --- a/src/transformers/models/gpt_oss/modular_gpt_oss.py +++ b/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) diff --git a/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py index 4082a3ee9bf4..e76bd3c35c1b 100644 --- a/src/transformers/models/openai_privacy_filter/modeling_openai_privacy_filter.py +++ b/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) diff --git a/tests/models/gpt_oss/test_modeling_gpt_oss.py b/tests/models/gpt_oss/test_modeling_gpt_oss.py index 45b6b0e87b37..fa39405791f4 100644 --- a/tests/models/gpt_oss/test_modeling_gpt_oss.py +++ b/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