fix: improve expert parallel implementation and refactors

tests/parallel_layer_test.py

@@ -0,0 +1,94 @@
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import os
+
+
+def test_megablocks_moe_mlp_import():
+    from megablocks.layers import MegaBlocksMoeMLP
+
+    assert MegaBlocksMoeMLP is not None, "MegaBlocksMoeMLP import failed."
+
+
+def run_distributed_test(rank, world_size):
+    from megablocks.layers import MegaBlocksMoeMLP
+
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = "12355"
+    os.environ["RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = str(world_size)
+
+    dist.init_process_group(
+        backend="gloo",
+        rank=rank,
+        world_size=world_size,
+    )
+
+    expert_parallel_group = torch.distributed.new_group(
+        range(torch.distributed.get_world_size())
+    )
+
+    model = MegaBlocksMoeMLP()
+    model.expert_parallel_group = expert_parallel_group
+
+    class Experts:
+        def __init__(self):
+            self.gate_up_proj = None
+            self.gate_up_proj_bias = None
+            self.down_proj = None
+            self.down_proj_bias = None
+            self.hidden_size = None
+
+    model.experts = Experts()
+
+    num_experts = 128
+    hidden_size = 1152
+    intermediate_size = 3072
+
+    ne, hs, isz = num_experts, hidden_size, intermediate_size
+
+    experts_per_rank = ne // world_size
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    model.router = torch.nn.Linear(hs, ne).to(device)
+    model.router.weight.data.fill_(1)
+
+    e = model.experts
+    e.gate_up_proj = torch.nn.Parameter(
+        torch.ones(experts_per_rank, hs, isz, device=device)
+    )
+    e.gate_up_proj_bias = torch.nn.Parameter(
+        torch.zeros(experts_per_rank, isz, device=device)
+    )
+    e.down_proj = torch.nn.Parameter(
+        torch.ones(experts_per_rank, 1536, hs, device=device)
+    )
+    e.down_proj_bias = torch.nn.Parameter(
+        torch.zeros(experts_per_rank, hs, device=device)
+    )
+    e.hidden_size = hs
+
+    x = torch.randn(1, 1, 1152).to(device)
+    output, expert_weights_out = model(x)
+
+    assert output.shape == (1, 1, 1152), f"Output shape mismatch on rank {rank}."
+
+    print(f"Rank {rank}: Test passed! Output shape: {output.shape}")
+
+    dist.destroy_process_group()
+
+
+def test_megablocks_moe_mlp_functionality():
+    world_size = 2
+
+    mp.spawn(run_distributed_test, args=(world_size,), nprocs=world_size, join=True)
+
+    print("Multi-process test completed successfully!")
+
+
+if __name__ == "__main__":
+    test_megablocks_moe_mlp_import()
+    print("Import test passed!")
+
+    test_megablocks_moe_mlp_functionality()

torch-ext/megablocks/layers.py

@@ -333,7 +333,6 @@ def permute_and_compute(
     gradient_scale,
     alpha,
 ):
-    """Permute tokens and compute expert outputs."""
     # Route tokens to experts
     x = x.view(-1, x.shape[-1])
 
@@ -367,6 +366,7 @@ def forward_once(
     expert_parallel_group: int = None,
     moe_capacity_factor: float = 1.0,
     moe_expert_model_parallelism: bool = False,
+    mlp_impl: Optional[str] = None,
 ):
     # x: [sl, bs, hs]
     # expert_weights: [sl * bs, top-k]
@@ -430,11 +430,15 @@ def parallel_forward_once(
     moe_capacity_factor: float = 1.0,
     moe_expert_model_parallelism: bool = True,
     hidden_size: int = 1152,
+    mlp_impl: Optional[str] = "grouped",
 ):
     # Flatten inputs
     expert_weights = expert_weights.flatten()
     top_experts = top_experts.flatten()
 
+    # TODO: remove debugging var
+    # my_rank = dist.get_rank(expert_parallel_group) if expert_parallel_group else 0
+
     with torch.no_grad():
         # Step 1: Local permutation setup
         indices, bin_ids, bins, tokens_per_expert = indices_and_bins(
@@ -455,9 +459,7 @@ def parallel_forward_once(
 
         # 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,
@@ -493,20 +495,13 @@ def parallel_forward_once(
     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
+    parallel_x, parallel_x_handle = _layers.all_to_all.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 = ops.inclusive_cumsum(parallel_tokens_per_expert.flatten(), 0)
         replicate_bins = (
             replicate_bins.view(1) if not len(replicate_bins.size()) else replicate_bins
         )
@@ -528,7 +523,7 @@ def parallel_forward_once(
 
         # Sort tokens by expert assignment
         parallel_bin_ids, parallel_indices = ops.sort(
-            parallel_top_expert, 
+            parallel_top_expert,
             sort_end_bit,
         )
 
@@ -536,10 +531,7 @@ def parallel_forward_once(
         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 = ops.inclusive_cumsum(parallel_tokens_per_expert, 0)
         parallel_bins = (
             parallel_bins.view(1) if not len(parallel_bins.size()) else parallel_bins
         )
@@ -558,10 +550,7 @@ def parallel_forward_once(
 
     # 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:
+    if mlp_impl == "grouped":
         # 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.
@@ -591,7 +580,9 @@ def parallel_forward_once(
     )
 
     # Step 6: Reverse communication - send results back
-    x, _ = ops.all_to_all(parallel_x, send_counts, recv_counts, expert_parallel_group)
+    x, _ = _layers.all_to_all.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)
@@ -603,139 +594,135 @@ def parallel_forward_once(
     return x, tokens_per_expert.flatten()
 
 
-class MyReplacementLayer(torch.nn.Module):
-    def forward(
-        x: torch.Tensor,
-        router_weight: torch.Tensor,
-        moe_top_k: int,
-        moe_num_experts: int,
-        moe_jitter_eps: float = None,
-        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,
-        w2_bias: torch.Tensor = None,
-        gradient_scale: Optional[float] = None,
-        alpha: float = 1.702,
-        sort_end_bit: int = 0,
-        expert_parallel_group: torch.distributed.ProcessGroup = None,
-        moe_capacity_factor: float = 1.0,
-        moe_expert_model_parallelism: bool = False,
-        forward_fn: Any = None,
-        hidden_size: int = None,  # Required for parallel forward
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-
-        # Route tokens to experts
-        logits, expert_weights, expert_indices = route_tokens(
-            x,
-            router_weight,
-            moe_top_k,
-            moe_num_experts,
-            moe_jitter_eps,
-            moe_normalize_expert_weights,
-            uniform_expert_assignment,
-            training,
-        )
+def moe_forward(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    moe_top_k: int,
+    moe_num_experts: int,
+    moe_jitter_eps: float = None,
+    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,
+    w2_bias: torch.Tensor = None,
+    gradient_scale: Optional[float] = None,
+    alpha: float = 1.702,
+    sort_end_bit: int = 0,
+    expert_parallel_group: torch.distributed.ProcessGroup = None,
+    moe_capacity_factor: float = 1.0,
+    moe_expert_model_parallelism: bool = False,
+    forward_fn: Any = None,
+    hidden_size: int = None,
+    mlp_impl: str = "grouped",
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
 
-        # Create router scores for output
-        router_scores = (
-            torch.zeros_like(logits)
-            .scatter_(1, expert_indices, expert_weights)
-            .transpose(0, 1)
-        )
+    # Route tokens to experts
+    logits, expert_weights, expert_indices = route_tokens(
+        x,
+        router_weight,
+        moe_top_k,
+        moe_num_experts,
+        moe_jitter_eps,
+        moe_normalize_expert_weights,
+        uniform_expert_assignment,
+        training,
+    )
 
-        in_shape = x.size()
-
-        # Prepare forward function arguments
-        forward_args = {
-            "x": x,
-            "expert_weights": expert_weights,
-            "top_experts": expert_indices,
-            "w1": w1,
-            "w2": w2,
-            "w1_bias": w1_bias,
-            "w2_bias": w2_bias,
-            "gradient_scale": gradient_scale,
-            "alpha": alpha,
-            "sort_end_bit": sort_end_bit,
-            "top_k": moe_top_k,
-            "num_experts": moe_num_experts,
-            "expert_parallel_group": expert_parallel_group,
-            "moe_capacity_factor": moe_capacity_factor,
-            "moe_expert_model_parallelism": moe_expert_model_parallelism,
-        }
-
-        # Add hidden_size for parallel forward
-        if moe_expert_model_parallelism and hidden_size is not None:
-            forward_args["hidden_size"] = hidden_size
-        elif moe_expert_model_parallelism and hidden_size is None:
-            # Infer hidden_size from input shape
-            forward_args["hidden_size"] = x.shape[-1]
-
-        # Compute expert outputs
-        x, tokens_per_expert = forward_fn(**forward_args)
-
-        # Save load balancing loss if needed
-        moe_loss_weight = 0.0  # Can be made configurable
-        if training and moe_loss_weight > 0:
-            save_load_balancing_loss((tokens_per_expert, logits))
-
-        # Restore original shape
-        x = x.view(in_shape)
-
-        return x, expert_weights, router_scores
+    # Create router scores for output
+    router_scores = (
+        torch.zeros_like(logits)
+        .scatter_(1, expert_indices, expert_weights)
+        .transpose(0, 1)
+    )
+
+    in_shape = x.size()
+
+    # Prepare forward function arguments
+    forward_args = {
+        "x": x,
+        "expert_weights": expert_weights,
+        "top_experts": expert_indices,
+        "w1": w1,
+        "w2": w2,
+        "w1_bias": w1_bias,
+        "w2_bias": w2_bias,
+        "gradient_scale": gradient_scale,
+        "alpha": alpha,
+        "sort_end_bit": sort_end_bit,
+        "top_k": moe_top_k,
+        "num_experts": moe_num_experts,
+        "expert_parallel_group": expert_parallel_group,
+        "moe_capacity_factor": moe_capacity_factor,
+        "moe_expert_model_parallelism": moe_expert_model_parallelism,
+        "mlp_impl": mlp_impl,
+    }
+
+    # Add hidden_size for parallel forward
+    if moe_expert_model_parallelism and hidden_size is not None:
+        forward_args["hidden_size"] = hidden_size
+    elif moe_expert_model_parallelism and hidden_size is None:
+        # Infer hidden_size from input shape
+        forward_args["hidden_size"] = x.shape[-1]
+
+    # Compute expert outputs
+    x, tokens_per_expert = forward_fn(**forward_args)
+
+    # Save load balancing loss if needed
+    moe_loss_weight = 0.0  # Can be made configurable
+    if training and moe_loss_weight > 0:
+        save_load_balancing_loss((tokens_per_expert, logits))
+
+    # Restore original shape
+    x = x.view(in_shape)
+
+    return x, expert_weights, router_scores
 
 
 class MegaBlocksMoeMLP(torch.nn.Module):
 
-    def forward(
-        self,
-        x: torch.Tensor,
-    ) -> torch.Tensor:
-        router_weight = self.router.weight
-        moe_top_k = 4
-        moe_num_experts = 128
-        w1 = self.experts.gate_up_proj.data
-        w2 = self.experts.down_proj.data
-        w1_bias = self.experts.gate_up_proj_bias.data
-        w2_bias = self.experts.down_proj_bias.data
-
-        # 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
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        moe_top_k = getattr(self, "moe_top_k", 4)
+        moe_num_experts = getattr(self, "moe_num_experts", 128)
+        gradient_scale = getattr(self, "gradient_scale", None)
+        alpha = getattr(self, "alpha", 1.702)
+        moe_capacity_factor = getattr(self, "moe_capacity_factor", 1.0)
+        moe_jitter_eps = getattr(self, "moe_jitter_eps", None)
+        moe_normalize_expert_weights = getattr(
+            self, "moe_normalize_expert_weights", None
+        )
+        uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
 
+        has_parallel = hasattr(self, "expert_parallel_group")
+        expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        forward_fn = parallel_forward_once if has_parallel else 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(
+        mlp_impl = getattr(self, "mlp_impl", "grouped")  # or sparse
+
+        output, expert_weights_out, _ = moe_forward(
             x=x,
-            router_weight=router_weight,
+            router_weight=self.router.weight,
             moe_top_k=moe_top_k,
             moe_num_experts=moe_num_experts,
-            moe_jitter_eps=None,
-            moe_normalize_expert_weights=None,
-            uniform_expert_assignment=False,
-            training=False,
-            w1=w1,
-            w2=w2,
-            w1_bias=w1_bias,
-            w2_bias=w2_bias,
-            gradient_scale=None,
-            alpha=1.702,
+            moe_jitter_eps=moe_jitter_eps,
+            moe_normalize_expert_weights=moe_normalize_expert_weights,
+            uniform_expert_assignment=uniform_expert_assignment,
+            training=self.training,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            w1_bias=self.experts.gate_up_proj_bias,
+            w2_bias=self.experts.down_proj_bias,
+            gradient_scale=gradient_scale,
+            alpha=alpha,
             sort_end_bit=sort_end_bit,
             expert_parallel_group=expert_parallel_group,
-            moe_capacity_factor=1.0,
-            moe_expert_model_parallelism=moe_expert_model_parallelism,
+            moe_capacity_factor=moe_capacity_factor,
+            moe_expert_model_parallelism=has_parallel,
             forward_fn=forward_fn,
-            hidden_size=hidden_size,
+            hidden_size=self.experts.hidden_size,
+            mlp_impl=mlp_impl,
         )
         return output, expert_weights_out