GPT-OSS MoE expert routing optimization for native pytorch

unsloth_zoo/temporary_patches/gpt_oss.py

@@ -895,20 +895,31 @@ def __init__(self, config):
     def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None) -> torch.Tensor:
         batch_size = hidden_states.shape[0]
         hidden_states = hidden_states.reshape(-1, self.hidden_size)
+        num_tokens = hidden_states.shape[0]
         num_experts = routing_weights.shape[1]
-
+        top_k = router_indices.shape[1]
+
         if self.training:
+            with torch.no_grad():
+                flat_experts = router_indices.flatten()  # [tokens * topk]
+                token_ids = torch.arange(num_tokens, device=hidden_states.device).repeat_interleave(top_k)
+
+                sorted_idx = flat_experts.argsort(stable=True)
+                sorted_tokens = token_ids[sorted_idx]
+
+                counts = torch.bincount(flat_experts, minlength=num_experts).tolist()
+
             next_states = torch.zeros_like(hidden_states, dtype=torch.float32, device=hidden_states.device)
-            # with torch.no_grad():
-                # expert_mask = torch.nn.functional.one_hot(router_indices, num_classes=num_experts)
-                # expert_mask = expert_mask.permute(2, 1, 0)
-                # expert_hitted = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
-            # for expert_idx in expert_hitted[:]:
+            offset = 0
+
             for expert_idx in range(num_experts):
-                with torch.no_grad():
-                    # _, token_idx = torch.where(expert_mask[expert_idx[0]])
-                    token_idx, _ = torch.where(router_indices == expert_idx)
+                count = counts[expert_idx]
+                if count == 0:
+                    continue
+                # Use pre-computed indices (no torch.where needed)
+                token_idx = sorted_tokens[offset:offset + count]
                 current_state = hidden_states[token_idx]
+
                 gate_up = self.gate_up_projs[expert_idx](current_state)
                 gated_output = swiglu_torch_forward(gate_up, self.alpha, self.limit)
                 # gate, up = gate_up[..., ::2], gate_up[..., 1::2]
@@ -917,8 +928,12 @@ def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weig
                 # glu = gate * torch.sigmoid(gate * self.alpha)
                 # gated_output = (up + 1) * glu
                 out = self.down_projs[expert_idx](gated_output)
+
                 weighted_output = out * routing_weights[token_idx, expert_idx, None].to(torch.float32)
                 next_states.index_add_(0, token_idx, weighted_output)
+
+                offset += count
+
             next_states = next_states.view(batch_size, -1, self.hidden_size)
             return next_states.to(hidden_states.dtype)
         else:
@@ -1731,35 +1746,45 @@ def torch_native_forward(
 
     batch_size = hidden_states.shape[0]
     hidden_states = hidden_states.reshape(-1, self.hidden_size)
+    num_tokens = hidden_states.shape[0]
     num_experts = routing_weights.shape[1]
+    top_k = router_indices.shape[1]
+
     if self.training:
+        with torch.no_grad():
+            flat_experts = router_indices.flatten()  # [tokens * topk]
+            token_ids = torch.arange(num_tokens, device=hidden_states.device).repeat_interleave(top_k)
+
+            sorted_idx = flat_experts.argsort(stable=True)
+            sorted_tokens = token_ids[sorted_idx]
+
+            counts = torch.bincount(flat_experts, minlength=num_experts).tolist()
+
         next_states = torch.zeros_like(hidden_states, dtype=torch.float32, device=hidden_states.device)
-        # with torch.no_grad():
-        #     expert_mask = torch.nn.functional.one_hot(router_indices, num_classes=num_experts)
-        #     expert_mask = expert_mask.permute(2, 1, 0)
-        #     expert_hitted = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
-        # for expert_idx in expert_hitted[:]:
+        offset = 0
+
         for expert_idx in range(num_experts):
-            with torch.no_grad():
-                # _, token_idx = torch.where(expert_mask[expert_idx[0]])
-                token_idx, _ = torch.where(router_indices == expert_idx)
+            count = counts[expert_idx]
+            if count == 0:
+                continue
+
+            # Use pre-computed indices (no torch.where needed)
+            token_idx = sorted_tokens[offset:offset + count]
             current_state = hidden_states[token_idx]
+
             gate_up = self.gate_up_projs[expert_idx](current_state)
             down_proj = self.down_projs[expert_idx]
             gated_output = swiglu_torch_forward(gate_up, self.alpha, self.limit, dtype = torch.float32)
-            # 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
 
-            # Force float32 matrix multiply on some down projection modules
             gated_output = gated_output.to(torch.float32)
             device_type = gated_output.device.type if isinstance(gated_output.device.type, str) and gated_output.device.type != "mps" else "cpu"
             with torch.autocast(device_type=device_type, enabled=False): # Force float32
                 out = down_proj(gated_output)
+
             weighted_output = out.to(torch.float32) * routing_weights[token_idx, expert_idx, None].to(torch.float32)
             next_states.index_add_(0, token_idx, weighted_output)
+
+            offset += count
         next_states = next_states.view(batch_size, -1, self.hidden_size)
         return next_states.to(torch.float32)
     else: