fix: add parallel forward functional logic

torch-ext/megablocks/layers.py

@@ -121,7 +121,15 @@ def scale_grad(
 
 
 # Forward pass for the MLP layer
-def mlp_forward(x, w1, w2, w1_bias, w2_bias, gradient_scale=None, alpha: float = 1.702):
+def mlp_forward(
+    x: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_bias: torch.Tensor,
+    w2_bias: torch.Tensor,
+    gradient_scale: Optional[float] = None,
+    alpha: float = 1.702,
+):
     # Scale weights
     w1 = scale_grad(w1, gradient_scale)
     w2 = scale_grad(w2, gradient_scale)
@@ -144,8 +152,6 @@ def mlp_forward(x, w1, w2, w1_bias, w2_bias, gradient_scale=None, alpha: float =
     return torch.bmm(x, w2) + w2_bias[..., None, :]
 
 
-## START: Load Balancing Loss (unused at the moment)
-
 # Global variable to store load balancing loss
 _LOAD_BALANCING_LOSS = []
 
@@ -234,9 +240,6 @@ def batched_load_balancing_loss(args):
     return scale * torch.dot(tokens_per_expert, expert_scores)
 
 
-## END Load Balancing Loss
-
-
 # Calculate the expert capacity based on tokens, top_k, number of experts,
 # expert parallel group, capacity factor, and whether expert model parallelism is used.
 def expert_capacity(
@@ -410,7 +413,6 @@ def forward_once(
     return x, tokens_per_expert
 
 
-# TODO: replace with functional logic once aligned with ref
 def parallel_forward_once(
     x: torch.Tensor,
     expert_weights: torch.Tensor,
@@ -429,15 +431,180 @@ def parallel_forward_once(
     moe_expert_model_parallelism: bool = True,
     hidden_size: int = 1152,
 ):
-    pass
+    # Flatten inputs
+    expert_weights = expert_weights.flatten()
+    top_experts = top_experts.flatten()
+
+    with torch.no_grad():
+        # Step 1: Local permutation setup
+        indices, bin_ids, bins, tokens_per_expert = indices_and_bins(
+            top_experts, sort_end_bit, num_experts
+        )
 
+        # Calculate sharding parameters
+        world_size = dist.get_world_size(expert_parallel_group)
+        hidden_sharding_deg = hidden_sharding_degree(
+            world_size, num_experts, hidden_size
+        )
+        experts_per_rank_val = experts_per_rank(num_experts, world_size)
 
-class MyReplacementLayer(torch.nn.Module):
-    # def __init__(self):
-    #     super().__init__()
+        # Replicate token counts for hidden sharding
+        repeated_tokens_per_expert = ops.repeat(
+            tokens_per_expert, (hidden_sharding_deg,)
+        )
+
+        # Exchange token counts across devices
+        parallel_tokens_per_expert = torch.empty_like(repeated_tokens_per_expert)
+        # print("world_size:", world_size)
+        # print("experts_per_rank_val:", experts_per_rank_val)
+        
+        # Ensure CUB knows which device to use
+        tpe_handle = dist.all_to_all_single(
+            parallel_tokens_per_expert,
+            repeated_tokens_per_expert,
+            group=expert_parallel_group,
+            async_op=True,
+        )
+
+    # Step 2: Local permutation - group tokens by target device
+    x = x.view(-1, x.shape[-1])  # [sl * bs, hs]
+    x = ops.gather(x, indices, bin_ids, bins, top_k)
+
+    # Step 3: Compute communication counts and exchange tokens
+    with torch.no_grad():
+        tpe_handle.wait()
+
+        # Reshape for per-device calculations
+        repeated_tokens_per_expert = repeated_tokens_per_expert.view(
+            world_size, experts_per_rank_val
+        )
+        parallel_tokens_per_expert = parallel_tokens_per_expert.view(
+            world_size, experts_per_rank_val
+        )
+
+        # Calculate send/recv counts
+        send_counts = repeated_tokens_per_expert.cpu().sum(dim=-1).tolist()
+        # recv_counts = parallel_tokens_per_expert.cpu().sum(dim=-1).tolist()
+        parallel_tokens_per_expert_cpu = parallel_tokens_per_expert.cpu()
+        recv_counts = parallel_tokens_per_expert_cpu.sum(dim=-1).tolist()
+        tokens_received = sum(recv_counts)
+
+    # Replicate for hidden sharding
+    x = ops.repeat(x, (hidden_sharding_deg, 1))
+
+    # Cross-device token exchange
+    parallel_x, parallel_x_handle = ops.all_to_all(
+        x,
+        recv_counts,
+        send_counts,
+        expert_parallel_group,
+        async_op=True
+    )
 
+    with torch.no_grad():
+        # Step 4: Setup for local expert computation
+        replicate_bins = ops.inclusive_cumsum(
+            parallel_tokens_per_expert.flatten(), 
+            0
+        )
+        replicate_bins = (
+            replicate_bins.view(1) if not len(replicate_bins.size()) else replicate_bins
+        )
+
+        # Create expert indices for received tokens
+        parallel_top_expert = torch.remainder(
+            torch.arange(
+                num_experts * hidden_sharding_deg,
+                dtype=torch.int32,
+                device=indices.device,
+            ),
+            experts_per_rank_val,
+        )
+        parallel_top_expert = ops.replicate(
+            parallel_top_expert.unsqueeze(dim=0),
+            replicate_bins,
+            tokens_received,
+        ).flatten()
+
+        # Sort tokens by expert assignment
+        parallel_bin_ids, parallel_indices = ops.sort(
+            parallel_top_expert, 
+            sort_end_bit,
+        )
+
+        # Calculate bins for local experts
+        parallel_tokens_per_expert = parallel_tokens_per_expert.sum(
+            dim=0, dtype=torch.int
+        )
+        parallel_bins = ops.inclusive_cumsum(
+            parallel_tokens_per_expert, 
+            0
+        )
+        parallel_bins = (
+            parallel_bins.view(1) if not len(parallel_bins.size()) else parallel_bins
+        )
+
+        # Calculate expert capacity
+        expert_capacity = expert_capacity_fn(
+            tokens_received,
+            top_k,
+            experts_per_rank_val,
+            expert_parallel_group,
+            moe_capacity_factor,
+            moe_expert_model_parallelism,
+        )
+        if expert_capacity == 0:
+            expert_capacity = torch.max(parallel_tokens_per_expert).item()
+
+    # Locally permute the tokens and perform the expert computation.
+    # Block to make sure that the cross-device permutation is complete.
+    # if self.args.mlp_impl == 'grouped':
+
+    # TODO: dont always assume grouped MLP
+    if True:
+        # GroupedMLP requires counts on CPU. We can use the tensor already
+        # moved to CPU for the prior all_to_all, which avoids an extra
+        # device synchronization.
+        parallel_tokens_per_expert = parallel_tokens_per_expert_cpu.sum(
+            dim=0,
+            dtype=torch.int,
+        )
+
+    # Step 5: Expert computation
+    parallel_x_handle.wait()
+
+    parallel_x = permute_and_compute(
+        parallel_x,
+        parallel_tokens_per_expert,
+        parallel_indices,
+        parallel_bin_ids,
+        None,  # expert_weights
+        parallel_bins,
+        expert_capacity,
+        top_k=1,
+        w1=w1,
+        w2=w2,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        gradient_scale=gradient_scale,
+        alpha=alpha,
+    )
+
+    # Step 6: Reverse communication - send results back
+    x, _ = ops.all_to_all(parallel_x, send_counts, recv_counts, expert_parallel_group)
+
+    # Step 7: Reduce across hidden sharding dimension
+    shape = (hidden_sharding_deg, -1, hidden_size)
+    x = x.view(shape).sum(dim=0)
+
+    # Step 8: Final local unpermutation
+    x = ops.scatter(x, indices, bin_ids, expert_weights, bins, top_k)
+
+    return x, tokens_per_expert.flatten()
+
+
+class MyReplacementLayer(torch.nn.Module):
     def forward(
-        # self,
         x: torch.Tensor,
         router_weight: torch.Tensor,
         moe_top_k: int,
@@ -446,7 +613,6 @@ class MyReplacementLayer(torch.nn.Module):
         moe_normalize_expert_weights: int = None,
         uniform_expert_assignment: bool = False,
         training: bool = False,
-        #
         w1: torch.Tensor = None,
         w2: torch.Tensor = None,
         w1_bias: torch.Tensor = None,
@@ -522,7 +688,6 @@ class MyReplacementLayer(torch.nn.Module):
         return x, expert_weights, router_scores
 
 
-
 class MegaBlocksMoeMLP(torch.nn.Module):
 
     def forward(
@@ -536,11 +701,21 @@ class MegaBlocksMoeMLP(torch.nn.Module):
         w2 = self.experts.down_proj.data
         w1_bias = self.experts.gate_up_proj_bias.data
         w2_bias = self.experts.down_proj_bias.data
-        expert_parallel_group = None
 
-        sort_end_bit = max(int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1)
+        # check if the expert_parallel_group attribute is set
+        if hasattr(self, "expert_parallel_group"):
+            expert_parallel_group = self.expert_parallel_group
+            moe_expert_model_parallelism = True
+            forward_fn = parallel_forward_once
+        else:
+            expert_parallel_group = None
+            moe_expert_model_parallelism = False
+            forward_fn = forward_once
+
+        sort_end_bit = max(
+            int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1
+        )
         hidden_size = self.experts.hidden_size
-
         output, expert_weights_out, router_scores = MyReplacementLayer.forward(
             x=x,
             router_weight=router_weight,
@@ -559,8 +734,8 @@ class MegaBlocksMoeMLP(torch.nn.Module):
             sort_end_bit=sort_end_bit,
             expert_parallel_group=expert_parallel_group,
             moe_capacity_factor=1.0,
-            moe_expert_model_parallelism=False,
-            forward_fn=forward_once,
+            moe_expert_model_parallelism=moe_expert_model_parallelism,
+            forward_fn=forward_fn,
             hidden_size=hidden_size,
         )
-        return output, expert_weights_out
+        return output, expert_weights_out