Allow for registered experts from kernels hub

src/transformers/integrations/finegrained_fp8.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 import torch
 import torch.nn as nn
-import triton
 from torch.nn import functional as F
 
 from ..activations import ACT2FN
@@ -159,6 +158,11 @@ def _load_deepgemm_kernel():
     _deepgemm_available = True
 
 
+def _cdiv(a: int, b: int) -> int:
+    """Ceiling division."""
+    return (a + b - 1) // b
+
+
 def w8a8_fp8_matmul(
     A: torch.Tensor,
     B: torch.Tensor,
@@ -603,26 +607,26 @@ def __init__(
         if self.has_gate:
             gu_proj_out, gu_proj_in = 2 * self.intermediate_dim, self.hidden_dim
             self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, gu_proj_out, gu_proj_in, dtype=dtype))
-            gu_scale_out = triton.cdiv(gu_proj_out, self.block_size[0]) if self.block_size is not None else 1
-            gu_scale_in = triton.cdiv(gu_proj_in, self.block_size[1]) if self.block_size is not None else 1
+            gu_scale_out = _cdiv(gu_proj_out, self.block_size[0]) if self.block_size is not None else 1
+            gu_scale_in = _cdiv(gu_proj_in, self.block_size[1]) if self.block_size is not None else 1
             self.gate_up_proj_scale_inv = nn.Parameter(
                 torch.empty(self.num_experts, gu_scale_out, gu_scale_in, dtype=torch.float32)
             )
             self.register_parameter("gate_up_proj_bias", None)
         else:
             u_proj_out, u_proj_in = self.intermediate_dim, self.hidden_dim
             self.up_proj = nn.Parameter(torch.empty(self.num_experts, u_proj_out, u_proj_in, dtype=dtype))
-            u_scale_out = triton.cdiv(u_proj_out, self.block_size[0]) if self.block_size is not None else 1
-            u_scale_in = triton.cdiv(u_proj_in, self.block_size[1]) if self.block_size is not None else 1
+            u_scale_out = _cdiv(u_proj_out, self.block_size[0]) if self.block_size is not None else 1
+            u_scale_in = _cdiv(u_proj_in, self.block_size[1]) if self.block_size is not None else 1
             self.up_proj_scale_inv = nn.Parameter(
                 torch.empty(self.num_experts, u_scale_out, u_scale_in, dtype=torch.float32)
             )
             self.register_parameter("up_proj_bias", None)
 
         d_proj_out, d_proj_in = self.hidden_dim, self.intermediate_dim
         self.down_proj = nn.Parameter(torch.empty(self.num_experts, d_proj_out, d_proj_in, dtype=dtype))
-        d_scale_out = triton.cdiv(d_proj_out, self.block_size[0]) if self.block_size is not None else 1
-        d_scale_in = triton.cdiv(d_proj_in, self.block_size[1]) if self.block_size is not None else 1
+        d_scale_out = _cdiv(d_proj_out, self.block_size[0]) if self.block_size is not None else 1
+        d_scale_in = _cdiv(d_proj_in, self.block_size[1]) if self.block_size is not None else 1
         self.down_proj_scale_inv = nn.Parameter(
             torch.empty(self.num_experts, d_scale_out, d_scale_in, dtype=torch.float32)
         )

src/transformers/modeling_utils.py

@@ -66,10 +66,12 @@
 )
 from .integrations.deepspeed import _load_state_dict_into_zero3_model
 from .integrations.eager_paged import eager_paged_attention_forward
+from .integrations.finegrained_fp8 import ALL_FP8_EXPERTS_FUNCTIONS
 from .integrations.flash_attention import flash_attention_forward
 from .integrations.flash_paged import paged_attention_forward
 from .integrations.flex_attention import flex_attention_forward
 from .integrations.hub_kernels import allow_all_hub_kernels, is_kernel
+from .integrations.moe import ALL_EXPERTS_FUNCTIONS
 from .integrations.peft import maybe_load_adapters
 from .integrations.sdpa_attention import sdpa_attention_forward
 from .integrations.sdpa_paged import sdpa_attention_paged_forward
@@ -1969,11 +1971,14 @@ def get_correct_attn_implementation(self, requested_attention: str | None, is_in
 
     def get_correct_experts_implementation(self, requested_experts: str | None) -> str:
         applicable_experts = "grouped_mm" if requested_experts is None else requested_experts
-        if applicable_experts not in ["eager", "grouped_mm", "batched_mm", "deepgemm"]:
+        base_experts_fns = ["eager"] + list(set(ALL_EXPERTS_FUNCTIONS.keys()) | set(ALL_FP8_EXPERTS_FUNCTIONS.keys()))
+        valid_experts_str_list = [f'`experts_implementation="{fn}"`' for fn in base_experts_fns]
+        valid_experts_str_list[-1] = "and " + valid_experts_str_list[-1]
+        valid_experts_str = ", ".join(valid_experts_str_list)
+        if applicable_experts not in base_experts_fns:
             message = (
                 f'Specified `experts_implementation="{applicable_experts}"` is not supported. The only possible arguments are '
-                '`experts_implementation="eager"`, `"experts_implementation=grouped_mm"`, `"experts_implementation=batched_mm"` '
-                'and `"experts_implementation=deepgemm"`.'
+                f"{valid_experts_str}."
             )
             raise ValueError(message)
 

tests/utils/test_modeling_utils.py

@@ -2823,6 +2823,20 @@ def test_error_wrong_attn_implementation(self):
 
         self.assertTrue('The only possible arguments are `attn_implementation="eager"' in str(cm.exception))
 
+    def test_registered_experts_implementation_is_valid(self):
+        from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS
+
+        def custom_experts_forward(*args, **kwargs):
+            pass
+
+        experts_implementation = "custom_experts"
+        model = BaseModel(PreTrainedConfig())
+
+        with patch.dict(ALL_EXPERTS_FUNCTIONS._global_mapping, {}, clear=False):
+            ALL_EXPERTS_FUNCTIONS.register(experts_implementation, custom_experts_forward)
+
+            self.assertEqual(model.get_correct_experts_implementation(experts_implementation), experts_implementation)
+
     def test_not_available_flash(self):
         if is_flash_attn_2_available():
             self.skipTest(reason="Please uninstall flash-attn package to run test_not_available_flash")