GptOss experts implementation

src/transformers/integrations/moe.py

@@ -16,11 +16,13 @@
 
 from ..utils.generic import GeneralInterface
 from ..utils.import_utils import is_torch_available
+from ..utils.logging import get_logger
 
 
 if is_torch_available():
     import torch
 
+logger = get_logger(__name__)
 
 # Examples of experts class with its eager mm implementation
 # class Experts(nn.Module):
@@ -99,18 +101,29 @@ def batched_mm_experts_forward(
     selected_down = self.down_proj[expert_ids]  # (S, hidden_dim, intermediate_dim)
 
     # --- Up projection per expert (batched) ---
-    gate_up_out = torch.bmm(selected_gate_up, current_hidden_states.unsqueeze(-1)).squeeze(-1)
+    if getattr(self, "experts_are_transposed", False):
+        gate_up_out = torch.bmm(current_hidden_states.unsqueeze(1), selected_gate_up).squeeze(1)
+    else:
+        gate_up_out = torch.bmm(selected_gate_up, current_hidden_states.unsqueeze(-1)).squeeze(-1)
+
     if hasattr(self, "gate_up_proj_bias") and self.gate_up_proj_bias is not None:
         gate_up_out = gate_up_out + self.gate_up_proj_bias[expert_ids]
 
-    # Split into gate and up components
-    gate, up = gate_up_out.chunk(2, dim=-1)  # both have shape (S, intermediate_dim)
-
-    # Apply activation
-    hidden_after_activation = self.act_fn(gate) * up  # (S, intermediate_dim)
+    # Apply gating
+    if hasattr(self, "_apply_gate"):
+        # Custom gating if defined
+        hidden_after_activation = self._apply_gate(gate_up_out)  # (S, intermediate_dim)
+    else:
+        # Default gating mechanism
+        gate, up = gate_up_out.chunk(2, dim=-1)  # (S, intermediate_dim)
+        hidden_after_activation = self.act_fn(gate) * up  # (S, intermediate_dim)
 
     # --- Down projection per expert (batched) ---
-    out_per_sample = torch.bmm(selected_down, hidden_after_activation.unsqueeze(-1)).squeeze(-1)
+    if getattr(self, "experts_are_transposed", False):
+        out_per_sample = torch.bmm(hidden_after_activation.unsqueeze(1), selected_down).squeeze(1)
+    else:
+        out_per_sample = torch.bmm(selected_down, hidden_after_activation.unsqueeze(-1)).squeeze(-1)
+
     if hasattr(self, "down_proj_bias") and self.down_proj_bias is not None:
         out_per_sample = out_per_sample + self.down_proj_bias[expert_ids]
 
@@ -176,19 +189,30 @@ def grouped_mm_experts_forward(
     offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32)
 
     # --- Up projection per expert (grouped_mm) ---
-    gate_up_out = torch._grouped_mm(current_states_g, self.gate_up_proj.transpose(-2, -1), offs=offsets)
+    if getattr(self, "experts_are_transposed", False):
+        gate_up_out = torch._grouped_mm(current_states_g, self.gate_up_proj, offs=offsets)
+    else:
+        gate_up_out = torch._grouped_mm(current_states_g, self.gate_up_proj.transpose(-2, -1), offs=offsets)
+
     if hasattr(self, "gate_up_proj_bias") and self.gate_up_proj_bias is not None:
         # we should be able to pass bias to the grouped_mm call, but it's still not fully supported
         gate_up_out = gate_up_out + self.gate_up_proj_bias[expert_ids_g]
 
-    # Split into gate and up components
-    gate, up = gate_up_out.chunk(2, dim=-1)  # both have shape (S, intermediate_dim)
-
-    # Apply activation
-    hidden_after_activation = self.act_fn(gate) * up  # (S, intermediate_dim)
+    # Apply gating
+    if hasattr(self, "_apply_gate"):
+        # Custom gating if defined
+        hidden_after_activation = self._apply_gate(gate_up_out)  # (S, intermediate_dim)
+    else:
+        # Default gating mechanism
+        gate, up = gate_up_out.chunk(2, dim=-1)  # (S, intermediate_dim)
+        hidden_after_activation = self.act_fn(gate) * up  # (S, intermediate_dim)
 
     # --- Down projection per expert (grouped_mm) ---
-    out_per_sample_g = torch._grouped_mm(hidden_after_activation, self.down_proj.transpose(-2, -1), offs=offsets)
+    if getattr(self, "experts_are_transposed", False):
+        out_per_sample_g = torch._grouped_mm(hidden_after_activation, self.down_proj, offs=offsets)
+    else:
+        out_per_sample_g = torch._grouped_mm(hidden_after_activation, self.down_proj.transpose(-2, -1), offs=offsets)
+
     if hasattr(self, "down_proj_bias") and self.down_proj_bias is not None:
         # we should be able to pass bias to the grouped_mm call, but it's still not fully supported
         out_per_sample_g = out_per_sample_g + self.down_proj_bias[expert_ids_g]
@@ -230,6 +254,15 @@ def __init__(self, config, *args, **kwargs):
     def forward(self, *args, **kwargs):
         experts_forward = original_forward
 
+        if self.config._experts_implementation == "grouped_mm":
+            if self.gate_up_proj.data_ptr() % 16 != 0 or self.down_proj.data_ptr() % 16 != 0:
+                logger.warning(
+                    "'grouped_mm' experts implementation requires 16-byte aligned expert weights. "
+                    "We will fall back to 'eager' implementation to avoid potential crashes. "
+                    "Please re-initialize the expert weights with 16-byte alignment."
+                )
+                self.config._experts_implementation = "eager"
+
         if self.config._experts_implementation != "eager":
             experts_forward = ALL_EXPERTS_FUNCTIONS[self.config._experts_implementation]
 

src/transformers/models/gpt_oss/modeling_gpt_oss.py

@@ -27,7 +27,7 @@
 from ... import initialization as init
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
+from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func
 from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
 from ...modeling_layers import (
     GenericForSequenceClassification,
@@ -64,86 +64,54 @@ def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
 
 
+@use_experts_implementation(transposed_experts=True)
 class GptOssExperts(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.intermediate_size = config.intermediate_size
         self.num_experts = config.num_local_experts
         self.hidden_size = config.hidden_size
-        self.expert_dim = self.intermediate_size
-        self.gate_up_proj = nn.Parameter(torch.zeros(self.num_experts, self.hidden_size, 2 * self.expert_dim))
-        self.gate_up_proj_bias = nn.Parameter(torch.zeros(self.num_experts, 2 * self.expert_dim))
-        self.down_proj = nn.Parameter(torch.empty((self.num_experts, self.expert_dim, self.hidden_size)))
-        self.down_proj_bias = nn.Parameter(torch.zeros(self.num_experts, self.hidden_size))
+        self.experts_are_transposed = True
+        self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, 2 * self.intermediate_size))
+        self.gate_up_proj_bias = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_size))
+        self.down_proj = nn.Parameter(torch.empty((self.num_experts, self.intermediate_size, self.hidden_size)))
+        self.down_proj_bias = nn.Parameter(torch.empty(self.num_experts, self.hidden_size))
         self.alpha = 1.702
         self.limit = 7.0
 
-    def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None) -> torch.Tensor:
-        """
-        When training it is more efficient to just loop over the experts and compute the output for each expert
-        as otherwise the memory would explode.
+    def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor:
+        gate, up = gate_up[..., ::2], gate_up[..., 1::2]
+        gate = gate.clamp(min=None, max=self.limit)
+        up = up.clamp(min=-self.limit, max=self.limit)
+        glu = gate * torch.sigmoid(gate * self.alpha)
+        gated_output = (up + 1) * glu
+        return gated_output
 
-        For inference we can sacrifice some memory and compute the output for all experts at once. By repeating the inputs.
-
-        Args:
-            hidden_states (torch.Tensor): (batch_size, seq_len, hidden_size)
-            selected_experts (torch.Tensor): (batch_size * seq_len, top_k)
-            routing_weights (torch.Tensor): (batch_size * seq_len, top_k)
-        Returns:
-            torch.Tensor
-        """
-        batch_size = hidden_states.shape[0]
-        hidden_states = hidden_states.reshape(-1, self.hidden_size)  # (num_tokens, hidden_size)
-        if hidden_states.device.type == "cpu" or self.training:
-            next_states = torch.zeros_like(hidden_states, dtype=hidden_states.dtype, device=hidden_states.device)
+    def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None) -> torch.Tensor:
+        next_states = torch.zeros_like(hidden_states, dtype=hidden_states.dtype, device=hidden_states.device)
+        with torch.no_grad():
+            expert_mask = torch.nn.functional.one_hot(
+                router_indices, num_classes=self.num_experts
+            )  # masking is also a class
+            expert_mask = expert_mask.permute(2, 1, 0)
+            # we sum on the top_k and on the sequence length to get which experts
+            # are hit this time around
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+        for expert_idx in expert_hit:
+            # expert_idx only have 1 element, so we can use scale for fast indexing
+            expert_idx = expert_idx[0]
+            # skip masking index
+            if expert_idx == self.num_experts:
+                continue
             with torch.no_grad():
-                expert_mask = torch.nn.functional.one_hot(
-                    router_indices, num_classes=self.num_experts
-                )  # masking is also a class
-                expert_mask = expert_mask.permute(2, 1, 0)
-                # we sum on the top_k and on the sequence length to get which experts
-                # are hit this time around
-                expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
-            for expert_idx in expert_hit[:]:
-                # expert_idx only have 1 element, so we can use scale for fast indexing
-                expert_idx = expert_idx[0]
-                # skip masking index
-                if expert_idx == self.num_experts:
-                    continue
-                with torch.no_grad():
-                    top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
-                current_state = hidden_states[token_idx]
-                gate_up = current_state @ self.gate_up_proj[expert_idx] + self.gate_up_proj_bias[expert_idx]
-                gate, up = gate_up[..., ::2], gate_up[..., 1::2]
-                gate = gate.clamp(min=None, max=self.limit)
-                up = up.clamp(min=-self.limit, max=self.limit)
-                glu = gate * torch.sigmoid(gate * self.alpha)
-                gated_output = (up + 1) * glu
-                out = gated_output @ self.down_proj[expert_idx] + self.down_proj_bias[expert_idx]
-                weighted_output = out * routing_weights[token_idx, top_k_pos, None]
-                next_states.index_add_(0, token_idx, weighted_output.to(hidden_states.dtype))
-            next_states = next_states.view(batch_size, -1, self.hidden_size)
-        else:
-            num_tokens = hidden_states.shape[0]
-            hidden_states = hidden_states.repeat(self.num_experts, 1)
-            hidden_states = hidden_states.view(self.num_experts, -1, self.hidden_size)
-            gate_up = torch.bmm(hidden_states, self.gate_up_proj) + self.gate_up_proj_bias[..., None, :]
-            gate, up = gate_up[..., ::2], gate_up[..., 1::2]
-            gate = gate.clamp(min=None, max=self.limit)
-            up = up.clamp(min=-self.limit, max=self.limit)
-            glu = gate * torch.sigmoid(gate * self.alpha)
-            next_states = torch.bmm(((up + 1) * glu), self.down_proj)
-            next_states = next_states + self.down_proj_bias[..., None, :]
-            next_states = next_states.view(self.num_experts, batch_size, -1, self.hidden_size)
-
-            full_routing_weights = torch.zeros(
-                num_tokens, self.num_experts, device=routing_weights.device, dtype=routing_weights.dtype
-            )
-            full_routing_weights.scatter_(1, router_indices, routing_weights)
-            full_routing_weights = full_routing_weights.transpose(0, 1).view(self.num_experts, batch_size, -1, 1)
+                top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+            current_state = hidden_states[token_idx]
+            gate_up = current_state @ self.gate_up_proj[expert_idx] + self.gate_up_proj_bias[expert_idx]
+            gated_output = self._apply_gate(gate_up)
+            out = gated_output @ self.down_proj[expert_idx] + self.down_proj_bias[expert_idx]
+            weighted_output = out * routing_weights[token_idx, top_k_pos, None]
+            next_states.index_add_(0, token_idx, weighted_output.to(hidden_states.dtype))
 
-            next_states = next_states * full_routing_weights
-            next_states = next_states.sum(dim=0)
         return next_states
 
 
@@ -157,7 +125,6 @@ def __init__(self, config):
         self.bias = nn.Parameter(torch.zeros(self.num_experts))
 
     def forward(self, hidden_states):
-        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
         router_logits = F.linear(hidden_states, self.weight, self.bias)  # (num_tokens, num_experts)
         router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1)  # (num_tokens, top_k)
         router_top_value = torch.nn.functional.softmax(router_top_value, dim=1, dtype=router_top_value.dtype)
@@ -173,9 +140,12 @@ def __init__(self, config):
         self.experts = GptOssExperts(config)
 
     def forward(self, hidden_states):
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
         _, router_scores, router_indices = self.router(hidden_states)
-        routed_out = self.experts(hidden_states, router_indices, router_scores)
-        return routed_out, router_scores
+        hidden_states = self.experts(hidden_states, router_indices, router_scores)
+        hidden_states = hidden_states.view(batch_size, sequence_length, hidden_dim)
+        return hidden_states, router_scores
 
 
 class GptOssRotaryEmbedding(nn.Module):