Upload folder using huggingface_hub

build/torch-cuda/__init__.py

@@ -1,36 +1,43 @@
 import torch
 import torch.nn as nn
 
-from transformers.core_model_loading import WeightRenaming
+from transformers.conversion_mapping import WeightRenaming
 
-conversion_mapping = [
-    WeightRenaming(
-        source_patterns=r"(.*(?:input_layernorm|post_attention_layernorm|norm)\.)weight",
-        target_patterns=r"\1scale",
-    ),
-]
 
-
-class CustomRMSNorm(nn.Module):
+class CustomRMSNormLayout(nn.Module):
     """
-    Drop-in replacement for LlamaRMSNorm with a different parameter layout.
+    Layout class for the CustomRMSNorm kernel.
 
-    Stores the learnable gain as ``self.scale`` instead of ``self.weight``.
-    The ``conversion_mapping`` class attribute tells the loader to rename
-    checkpoint keys so that ``*.weight`` → ``*.scale`` for any RMSNorm path.
+    Transformers instantiates this in place of the original RMSNorm, passing
+    the same constructor arguments (hidden_size, eps).  The weight is stored
+    as ``self.scale`` instead of ``self.weight``; the conversion_mapping below
+    handles the checkpoint key rename at load time.
 
-    This kernel exists to test the ``register_kernel_replacements`` path:
-    single-layer replacement where the kernel defines ``__init__`` and
-    carries a ``conversion_mapping``.
+    kernelize then replaces this module's forward with CustomRMSNorm.forward.
     """
 
-    conversion_mapping = conversion_mapping
+    conversion_mapping = [
+        WeightRenaming(
+            source_patterns=r"(.*(?:input_layernorm|post_attention_layernorm|norm)\.)weight",
+            target_patterns=r"\1scale",
+        ),
+    ]
 
     def __init__(self, hidden_size: int, eps: float = 1e-6):
         super().__init__()
         self.scale = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps
 
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        pass  # replaced at runtime by kernelize
+
+
+class CustomRMSNorm(nn.Module):
+    """
+    Forward-only kernel class.  kernelize discovers this via kernel_layer_name
+    set on CustomRMSNormLayout and replaces the layout instance's forward at runtime.
+    """
+
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         print("This dummy kernel is used")
         input_dtype = hidden_states.dtype