feat: bump build

build/torch26-cxx11-cu118-x86_64-linux/megablocks/{_megablocks_63599de.abi3.so → _megablocks_13afbbe_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9b35f3f60e0cbf0ce9e84e1224754d353f9de646cf30df5828168222889d312f
-size 10517576
+oid sha256:5683ac8b3e98fc8b8ab19f964b0dbfb9a980b6135220b0a0c1b50180665ce341
+size 10517608

build/torch26-cxx11-cu118-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_63599de
-ops = torch.ops._megablocks_63599de
+from . import _megablocks_13afbbe_dirty
+ops = torch.ops._megablocks_13afbbe_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_63599de::{op_name}"
+    return f"_megablocks_13afbbe_dirty::{op_name}"

build/torch26-cxx11-cu118-x86_64-linux/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

build/torch26-cxx11-cu118-x86_64-linux/megablocks/ops/all_to_all_benchmark.py

@@ -7,28 +7,126 @@ import torch.distributed as dist
 # from megablocks import benchmark_util
 # from megablocks.layers.all_to_all import all_to_all
 
-from .. import benchmark_util
-from .._layers.all_to_all import all_to_all
+# from .. import benchmark_util
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import torch
+
+
+def log_benchmark(name, arguments, time, std):
+    print("=" * 60)
+    print(f"{name} Benchmark")
+    print("Benchmark Parameters:")
+    for key, value in arguments.items():
+        print(f"{key} = {value}")
+    print("Results:")
+    print("mean time = {:.3f}ms, std time = {:.3f}ms".format(time, std))
+    print("=" * 60)
+
+
+def benchmark_function(fn, iterations=100, warmup=10):
+    print(f"Benchmarking {fn.__name__} with {iterations} iterations and {warmup} warmup iterations")
+    # Warmup iterations.
+    for _ in range(warmup):
+        fn()
+
+    times = []
+    print(f"Running {iterations} iterations...")
+    for i in range(iterations):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+
+        start.record()
+        fn()
+        end.record()
+
+        torch.cuda.synchronize()
+        times.append(start.elapsed_time(end))
+    return np.mean(times), np.std(times)
+
+
+# from .._layers.all_to_all import all_to_all
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.distributed as dist
+
+
+class AllToAllOp(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, output_split_sizes, input_split_sizes, group, async_op):
+        out = torch.empty(
+            (sum(output_split_sizes),) + x.shape[1:], device=x.device, dtype=x.dtype
+        )
+
+        ctx.input_shape = x.shape
+        ctx.output_split_sizes = output_split_sizes
+        ctx.input_split_sizes = input_split_sizes
+        ctx.group = group
+        handle = dist.all_to_all_single(
+            out,
+            x,
+            output_split_sizes=output_split_sizes,
+            input_split_sizes=input_split_sizes,
+            group=group,
+            async_op=async_op,
+        )
+        return out, handle
+
+    @staticmethod
+    def backward(ctx, grad, _):
+        if ctx.needs_input_grad[0]:
+            out = torch.empty(
+                ctx.input_shape,
+                device=grad.device,
+                dtype=grad.dtype,
+            )
+            dist.all_to_all_single(
+                out,
+                grad,
+                output_split_sizes=ctx.input_split_sizes,
+                input_split_sizes=ctx.output_split_sizes,
+                group=ctx.group,
+            )
+            return out, None, None, None, None
+        return None, None, None, None, None
+
+
+def all_to_all(x, output_split_sizes, input_split_sizes, group, async_op=False):
+    return AllToAllOp.apply(
+        x,
+        output_split_sizes,
+        input_split_sizes,
+        group,
+        async_op,
+    )
+
 
 _ALL_TO_ALL_BENCHMARK = (
     (8, 1024),
-    (16, 1024),
-    (32, 1024),
-    (64, 1024),
-    (128, 1024),
-    (256, 1024),
-    (512, 1024),
-    (1024, 1024),
-    (2 * 1024, 1024),
-    (4 * 1024, 1024),
-    (8 * 1024, 1024),
-    (16 * 1024, 1024),
-    (32 * 1024, 1024),
-    (64 * 1024, 1024),
-    (128 * 1024, 1024),
-    (256 * 1024, 1024),
-    (512 * 1024, 1024),
-    (1024 * 1024, 1024),
+    # (16, 1024),
+    # (32, 1024),
+    # (64, 1024),
+    # (128, 1024),
+    # (256, 1024),
+    # (512, 1024),
+    # (1024, 1024),
+    # (2 * 1024, 1024),
+    # (4 * 1024, 1024),
+    # (8 * 1024, 1024),
+    # (16 * 1024, 1024),
+    # (32 * 1024, 1024),
+    # (64 * 1024, 1024),
+    # (128 * 1024, 1024),
+    # (256 * 1024, 1024),
+    # (512 * 1024, 1024),
+    # (1024 * 1024, 1024),
 )
 
 
@@ -47,10 +145,12 @@ def benchmark_all_to_all(group, sl, hs):
     def benchmark():
         return all_to_all(x, send_recv_sizes, send_recv_sizes, group)
 
-    time, std = benchmark_util.benchmark_function(benchmark)
+    # time, std = benchmark_util.benchmark_function(benchmark)
+    time, std = benchmark_function(benchmark)
 
     if dist.get_rank(group) == 0:
-        benchmark_util.log_benchmark('All-To-All', details, time, std)
+        log_benchmark('All-To-All', details, time, std)
+        # benchmark_util.log_benchmark('All-To-All', details, time, std)
 
 
 if __name__ == '__main__':

build/torch26-cxx11-cu124-x86_64-linux/megablocks/{_megablocks_63599de.abi3.so → _megablocks_13afbbe_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:05d38f81524501b75940bfad8686f4f502b5c6af1de85fb1fe5b20da765d4c3c
-size 11869392
+oid sha256:b55d6ee3d41404603fdb75ad9a2949aa92e0224f7056fdbeb4c66934035ebd4b
+size 11869424

build/torch26-cxx11-cu124-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_63599de
-ops = torch.ops._megablocks_63599de
+from . import _megablocks_13afbbe_dirty
+ops = torch.ops._megablocks_13afbbe_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_63599de::{op_name}"
+    return f"_megablocks_13afbbe_dirty::{op_name}"

build/torch26-cxx11-cu124-x86_64-linux/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

build/torch26-cxx11-cu124-x86_64-linux/megablocks/ops/all_to_all_benchmark.py

@@ -7,28 +7,126 @@ import torch.distributed as dist
 # from megablocks import benchmark_util
 # from megablocks.layers.all_to_all import all_to_all
 
-from .. import benchmark_util
-from .._layers.all_to_all import all_to_all
+# from .. import benchmark_util
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import torch
+
+
+def log_benchmark(name, arguments, time, std):
+    print("=" * 60)
+    print(f"{name} Benchmark")
+    print("Benchmark Parameters:")
+    for key, value in arguments.items():
+        print(f"{key} = {value}")
+    print("Results:")
+    print("mean time = {:.3f}ms, std time = {:.3f}ms".format(time, std))
+    print("=" * 60)
+
+
+def benchmark_function(fn, iterations=100, warmup=10):
+    print(f"Benchmarking {fn.__name__} with {iterations} iterations and {warmup} warmup iterations")
+    # Warmup iterations.
+    for _ in range(warmup):
+        fn()
+
+    times = []
+    print(f"Running {iterations} iterations...")
+    for i in range(iterations):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+
+        start.record()
+        fn()
+        end.record()
+
+        torch.cuda.synchronize()
+        times.append(start.elapsed_time(end))
+    return np.mean(times), np.std(times)
+
+
+# from .._layers.all_to_all import all_to_all
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.distributed as dist
+
+
+class AllToAllOp(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, output_split_sizes, input_split_sizes, group, async_op):
+        out = torch.empty(
+            (sum(output_split_sizes),) + x.shape[1:], device=x.device, dtype=x.dtype
+        )
+
+        ctx.input_shape = x.shape
+        ctx.output_split_sizes = output_split_sizes
+        ctx.input_split_sizes = input_split_sizes
+        ctx.group = group
+        handle = dist.all_to_all_single(
+            out,
+            x,
+            output_split_sizes=output_split_sizes,
+            input_split_sizes=input_split_sizes,
+            group=group,
+            async_op=async_op,
+        )
+        return out, handle
+
+    @staticmethod
+    def backward(ctx, grad, _):
+        if ctx.needs_input_grad[0]:
+            out = torch.empty(
+                ctx.input_shape,
+                device=grad.device,
+                dtype=grad.dtype,
+            )
+            dist.all_to_all_single(
+                out,
+                grad,
+                output_split_sizes=ctx.input_split_sizes,
+                input_split_sizes=ctx.output_split_sizes,
+                group=ctx.group,
+            )
+            return out, None, None, None, None
+        return None, None, None, None, None
+
+
+def all_to_all(x, output_split_sizes, input_split_sizes, group, async_op=False):
+    return AllToAllOp.apply(
+        x,
+        output_split_sizes,
+        input_split_sizes,
+        group,
+        async_op,
+    )
+
 
 _ALL_TO_ALL_BENCHMARK = (
     (8, 1024),
-    (16, 1024),
-    (32, 1024),
-    (64, 1024),
-    (128, 1024),
-    (256, 1024),
-    (512, 1024),
-    (1024, 1024),
-    (2 * 1024, 1024),
-    (4 * 1024, 1024),
-    (8 * 1024, 1024),
-    (16 * 1024, 1024),
-    (32 * 1024, 1024),
-    (64 * 1024, 1024),
-    (128 * 1024, 1024),
-    (256 * 1024, 1024),
-    (512 * 1024, 1024),
-    (1024 * 1024, 1024),
+    # (16, 1024),
+    # (32, 1024),
+    # (64, 1024),
+    # (128, 1024),
+    # (256, 1024),
+    # (512, 1024),
+    # (1024, 1024),
+    # (2 * 1024, 1024),
+    # (4 * 1024, 1024),
+    # (8 * 1024, 1024),
+    # (16 * 1024, 1024),
+    # (32 * 1024, 1024),
+    # (64 * 1024, 1024),
+    # (128 * 1024, 1024),
+    # (256 * 1024, 1024),
+    # (512 * 1024, 1024),
+    # (1024 * 1024, 1024),
 )
 
 
@@ -47,10 +145,12 @@ def benchmark_all_to_all(group, sl, hs):
     def benchmark():
         return all_to_all(x, send_recv_sizes, send_recv_sizes, group)
 
-    time, std = benchmark_util.benchmark_function(benchmark)
+    # time, std = benchmark_util.benchmark_function(benchmark)
+    time, std = benchmark_function(benchmark)
 
     if dist.get_rank(group) == 0:
-        benchmark_util.log_benchmark('All-To-All', details, time, std)
+        log_benchmark('All-To-All', details, time, std)
+        # benchmark_util.log_benchmark('All-To-All', details, time, std)
 
 
 if __name__ == '__main__':

build/{torch27-cxx11-cu126-x86_64-linux/megablocks/_megablocks_63599de.abi3.so → torch26-cxx11-cu126-x86_64-linux/megablocks/_megablocks_13afbbe_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3a243e51490184fb48e02dbc1115545ea69313a3d63058f8423c0c493e90bc5a
+oid sha256:516c5026180d4a8d013c500ed284a60ecbed4bc6c9dc084b838913f40327d1a6
 size 11931080

build/torch26-cxx11-cu126-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_63599de
-ops = torch.ops._megablocks_63599de
+from . import _megablocks_13afbbe_dirty
+ops = torch.ops._megablocks_13afbbe_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_63599de::{op_name}"
+    return f"_megablocks_13afbbe_dirty::{op_name}"

build/torch26-cxx11-cu126-x86_64-linux/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

build/torch26-cxx11-cu126-x86_64-linux/megablocks/ops/all_to_all_benchmark.py

@@ -7,28 +7,126 @@ import torch.distributed as dist
 # from megablocks import benchmark_util
 # from megablocks.layers.all_to_all import all_to_all
 
-from .. import benchmark_util
-from .._layers.all_to_all import all_to_all
+# from .. import benchmark_util
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import torch
+
+
+def log_benchmark(name, arguments, time, std):
+    print("=" * 60)
+    print(f"{name} Benchmark")
+    print("Benchmark Parameters:")
+    for key, value in arguments.items():
+        print(f"{key} = {value}")
+    print("Results:")
+    print("mean time = {:.3f}ms, std time = {:.3f}ms".format(time, std))
+    print("=" * 60)
+
+
+def benchmark_function(fn, iterations=100, warmup=10):
+    print(f"Benchmarking {fn.__name__} with {iterations} iterations and {warmup} warmup iterations")
+    # Warmup iterations.
+    for _ in range(warmup):
+        fn()
+
+    times = []
+    print(f"Running {iterations} iterations...")
+    for i in range(iterations):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+
+        start.record()
+        fn()
+        end.record()
+
+        torch.cuda.synchronize()
+        times.append(start.elapsed_time(end))
+    return np.mean(times), np.std(times)
+
+
+# from .._layers.all_to_all import all_to_all
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.distributed as dist
+
+
+class AllToAllOp(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, output_split_sizes, input_split_sizes, group, async_op):
+        out = torch.empty(
+            (sum(output_split_sizes),) + x.shape[1:], device=x.device, dtype=x.dtype
+        )
+
+        ctx.input_shape = x.shape
+        ctx.output_split_sizes = output_split_sizes
+        ctx.input_split_sizes = input_split_sizes
+        ctx.group = group
+        handle = dist.all_to_all_single(
+            out,
+            x,
+            output_split_sizes=output_split_sizes,
+            input_split_sizes=input_split_sizes,
+            group=group,
+            async_op=async_op,
+        )
+        return out, handle
+
+    @staticmethod
+    def backward(ctx, grad, _):
+        if ctx.needs_input_grad[0]:
+            out = torch.empty(
+                ctx.input_shape,
+                device=grad.device,
+                dtype=grad.dtype,
+            )
+            dist.all_to_all_single(
+                out,
+                grad,
+                output_split_sizes=ctx.input_split_sizes,
+                input_split_sizes=ctx.output_split_sizes,
+                group=ctx.group,
+            )
+            return out, None, None, None, None
+        return None, None, None, None, None
+
+
+def all_to_all(x, output_split_sizes, input_split_sizes, group, async_op=False):
+    return AllToAllOp.apply(
+        x,
+        output_split_sizes,
+        input_split_sizes,
+        group,
+        async_op,
+    )
+
 
 _ALL_TO_ALL_BENCHMARK = (
     (8, 1024),
-    (16, 1024),
-    (32, 1024),
-    (64, 1024),
-    (128, 1024),
-    (256, 1024),
-    (512, 1024),
-    (1024, 1024),
-    (2 * 1024, 1024),
-    (4 * 1024, 1024),
-    (8 * 1024, 1024),
-    (16 * 1024, 1024),
-    (32 * 1024, 1024),
-    (64 * 1024, 1024),
-    (128 * 1024, 1024),
-    (256 * 1024, 1024),
-    (512 * 1024, 1024),
-    (1024 * 1024, 1024),
+    # (16, 1024),
+    # (32, 1024),
+    # (64, 1024),
+    # (128, 1024),
+    # (256, 1024),
+    # (512, 1024),
+    # (1024, 1024),
+    # (2 * 1024, 1024),
+    # (4 * 1024, 1024),
+    # (8 * 1024, 1024),
+    # (16 * 1024, 1024),
+    # (32 * 1024, 1024),
+    # (64 * 1024, 1024),
+    # (128 * 1024, 1024),
+    # (256 * 1024, 1024),
+    # (512 * 1024, 1024),
+    # (1024 * 1024, 1024),
 )
 
 
@@ -47,10 +145,12 @@ def benchmark_all_to_all(group, sl, hs):
     def benchmark():
         return all_to_all(x, send_recv_sizes, send_recv_sizes, group)
 
-    time, std = benchmark_util.benchmark_function(benchmark)
+    # time, std = benchmark_util.benchmark_function(benchmark)
+    time, std = benchmark_function(benchmark)
 
     if dist.get_rank(group) == 0:
-        benchmark_util.log_benchmark('All-To-All', details, time, std)
+        log_benchmark('All-To-All', details, time, std)
+        # benchmark_util.log_benchmark('All-To-All', details, time, std)
 
 
 if __name__ == '__main__':

build/{torch26-cxx11-cu126-x86_64-linux/megablocks/_megablocks_63599de.abi3.so → torch26-cxx98-cu118-x86_64-linux/megablocks/_megablocks_13afbbe_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0e9e392427d3157216b82014570075137082c5ec5c5bd6b63c1458d509ed4ff3
-size 11931048
+oid sha256:a5c8c1b700d297741dd86e8c388e03913a30769ceb51b7c12a01245fbdf30128
+size 10510072

build/torch26-cxx98-cu118-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_63599de
-ops = torch.ops._megablocks_63599de
+from . import _megablocks_13afbbe_dirty
+ops = torch.ops._megablocks_13afbbe_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_63599de::{op_name}"
+    return f"_megablocks_13afbbe_dirty::{op_name}"

build/torch26-cxx98-cu118-x86_64-linux/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

build/torch26-cxx98-cu118-x86_64-linux/megablocks/ops/all_to_all_benchmark.py

@@ -7,28 +7,126 @@ import torch.distributed as dist
 # from megablocks import benchmark_util
 # from megablocks.layers.all_to_all import all_to_all
 
-from .. import benchmark_util
-from .._layers.all_to_all import all_to_all
+# from .. import benchmark_util
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import torch
+
+
+def log_benchmark(name, arguments, time, std):
+    print("=" * 60)
+    print(f"{name} Benchmark")
+    print("Benchmark Parameters:")
+    for key, value in arguments.items():
+        print(f"{key} = {value}")
+    print("Results:")
+    print("mean time = {:.3f}ms, std time = {:.3f}ms".format(time, std))
+    print("=" * 60)
+
+
+def benchmark_function(fn, iterations=100, warmup=10):
+    print(f"Benchmarking {fn.__name__} with {iterations} iterations and {warmup} warmup iterations")
+    # Warmup iterations.
+    for _ in range(warmup):
+        fn()
+
+    times = []
+    print(f"Running {iterations} iterations...")
+    for i in range(iterations):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+
+        start.record()
+        fn()
+        end.record()
+
+        torch.cuda.synchronize()
+        times.append(start.elapsed_time(end))
+    return np.mean(times), np.std(times)
+
+
+# from .._layers.all_to_all import all_to_all
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.distributed as dist
+
+
+class AllToAllOp(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, output_split_sizes, input_split_sizes, group, async_op):
+        out = torch.empty(
+            (sum(output_split_sizes),) + x.shape[1:], device=x.device, dtype=x.dtype
+        )
+
+        ctx.input_shape = x.shape
+        ctx.output_split_sizes = output_split_sizes
+        ctx.input_split_sizes = input_split_sizes
+        ctx.group = group
+        handle = dist.all_to_all_single(
+            out,
+            x,
+            output_split_sizes=output_split_sizes,
+            input_split_sizes=input_split_sizes,
+            group=group,
+            async_op=async_op,
+        )
+        return out, handle
+
+    @staticmethod
+    def backward(ctx, grad, _):
+        if ctx.needs_input_grad[0]:
+            out = torch.empty(
+                ctx.input_shape,
+                device=grad.device,
+                dtype=grad.dtype,
+            )
+            dist.all_to_all_single(
+                out,
+                grad,
+                output_split_sizes=ctx.input_split_sizes,
+                input_split_sizes=ctx.output_split_sizes,
+                group=ctx.group,
+            )
+            return out, None, None, None, None
+        return None, None, None, None, None
+
+
+def all_to_all(x, output_split_sizes, input_split_sizes, group, async_op=False):
+    return AllToAllOp.apply(
+        x,
+        output_split_sizes,
+        input_split_sizes,
+        group,
+        async_op,
+    )
+
 
 _ALL_TO_ALL_BENCHMARK = (
     (8, 1024),
-    (16, 1024),
-    (32, 1024),
-    (64, 1024),
-    (128, 1024),
-    (256, 1024),
-    (512, 1024),
-    (1024, 1024),
-    (2 * 1024, 1024),
-    (4 * 1024, 1024),
-    (8 * 1024, 1024),
-    (16 * 1024, 1024),
-    (32 * 1024, 1024),
-    (64 * 1024, 1024),
-    (128 * 1024, 1024),
-    (256 * 1024, 1024),
-    (512 * 1024, 1024),
-    (1024 * 1024, 1024),
+    # (16, 1024),
+    # (32, 1024),
+    # (64, 1024),
+    # (128, 1024),
+    # (256, 1024),
+    # (512, 1024),
+    # (1024, 1024),
+    # (2 * 1024, 1024),
+    # (4 * 1024, 1024),
+    # (8 * 1024, 1024),
+    # (16 * 1024, 1024),
+    # (32 * 1024, 1024),
+    # (64 * 1024, 1024),
+    # (128 * 1024, 1024),
+    # (256 * 1024, 1024),
+    # (512 * 1024, 1024),
+    # (1024 * 1024, 1024),
 )
 
 
@@ -47,10 +145,12 @@ def benchmark_all_to_all(group, sl, hs):
     def benchmark():
         return all_to_all(x, send_recv_sizes, send_recv_sizes, group)
 
-    time, std = benchmark_util.benchmark_function(benchmark)
+    # time, std = benchmark_util.benchmark_function(benchmark)
+    time, std = benchmark_function(benchmark)
 
     if dist.get_rank(group) == 0:
-        benchmark_util.log_benchmark('All-To-All', details, time, std)
+        log_benchmark('All-To-All', details, time, std)
+        # benchmark_util.log_benchmark('All-To-All', details, time, std)
 
 
 if __name__ == '__main__':

build/{torch26-cxx98-cu118-x86_64-linux/megablocks/_megablocks_63599de.abi3.so → torch26-cxx98-cu124-x86_64-linux/megablocks/_megablocks_13afbbe_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b2451173cb1d000c6d270b59b2aaab1aa0e54025422ba81b1ee990621c90a823
-size 10510040
+oid sha256:d915db521f8d37fb887ed8334db60165e5923f8dce817d69f6441c5ba2d210d6
+size 11857952

build/torch26-cxx98-cu124-x86_64-linux/megablocks/_megablocks_63599de.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:f8bfaaeb2a5e226a80403463d15f2c762ac8cb70ca7a44d2156aadfac63ab0d1
-size 11857920

build/torch26-cxx98-cu124-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_63599de
-ops = torch.ops._megablocks_63599de
+from . import _megablocks_13afbbe_dirty
+ops = torch.ops._megablocks_13afbbe_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_63599de::{op_name}"
+    return f"_megablocks_13afbbe_dirty::{op_name}"

build/torch26-cxx98-cu124-x86_64-linux/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

build/torch26-cxx98-cu124-x86_64-linux/megablocks/ops/all_to_all_benchmark.py

@@ -7,28 +7,126 @@ import torch.distributed as dist
 # from megablocks import benchmark_util
 # from megablocks.layers.all_to_all import all_to_all
 
-from .. import benchmark_util
-from .._layers.all_to_all import all_to_all
+# from .. import benchmark_util
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import torch
+
+
+def log_benchmark(name, arguments, time, std):
+    print("=" * 60)
+    print(f"{name} Benchmark")
+    print("Benchmark Parameters:")
+    for key, value in arguments.items():
+        print(f"{key} = {value}")
+    print("Results:")
+    print("mean time = {:.3f}ms, std time = {:.3f}ms".format(time, std))
+    print("=" * 60)
+
+
+def benchmark_function(fn, iterations=100, warmup=10):
+    print(f"Benchmarking {fn.__name__} with {iterations} iterations and {warmup} warmup iterations")
+    # Warmup iterations.
+    for _ in range(warmup):
+        fn()
+
+    times = []
+    print(f"Running {iterations} iterations...")
+    for i in range(iterations):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+
+        start.record()
+        fn()
+        end.record()
+
+        torch.cuda.synchronize()
+        times.append(start.elapsed_time(end))
+    return np.mean(times), np.std(times)
+
+
+# from .._layers.all_to_all import all_to_all
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.distributed as dist
+
+
+class AllToAllOp(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, output_split_sizes, input_split_sizes, group, async_op):
+        out = torch.empty(
+            (sum(output_split_sizes),) + x.shape[1:], device=x.device, dtype=x.dtype
+        )
+
+        ctx.input_shape = x.shape
+        ctx.output_split_sizes = output_split_sizes
+        ctx.input_split_sizes = input_split_sizes
+        ctx.group = group
+        handle = dist.all_to_all_single(
+            out,
+            x,
+            output_split_sizes=output_split_sizes,
+            input_split_sizes=input_split_sizes,
+            group=group,
+            async_op=async_op,
+        )
+        return out, handle
+
+    @staticmethod
+    def backward(ctx, grad, _):
+        if ctx.needs_input_grad[0]:
+            out = torch.empty(
+                ctx.input_shape,
+                device=grad.device,
+                dtype=grad.dtype,
+            )
+            dist.all_to_all_single(
+                out,
+                grad,
+                output_split_sizes=ctx.input_split_sizes,
+                input_split_sizes=ctx.output_split_sizes,
+                group=ctx.group,
+            )
+            return out, None, None, None, None
+        return None, None, None, None, None
+
+
+def all_to_all(x, output_split_sizes, input_split_sizes, group, async_op=False):
+    return AllToAllOp.apply(
+        x,
+        output_split_sizes,
+        input_split_sizes,
+        group,
+        async_op,
+    )
+
 
 _ALL_TO_ALL_BENCHMARK = (
     (8, 1024),
-    (16, 1024),
-    (32, 1024),
-    (64, 1024),
-    (128, 1024),
-    (256, 1024),
-    (512, 1024),
-    (1024, 1024),
-    (2 * 1024, 1024),
-    (4 * 1024, 1024),
-    (8 * 1024, 1024),
-    (16 * 1024, 1024),
-    (32 * 1024, 1024),
-    (64 * 1024, 1024),
-    (128 * 1024, 1024),
-    (256 * 1024, 1024),
-    (512 * 1024, 1024),
-    (1024 * 1024, 1024),
+    # (16, 1024),
+    # (32, 1024),
+    # (64, 1024),
+    # (128, 1024),
+    # (256, 1024),
+    # (512, 1024),
+    # (1024, 1024),
+    # (2 * 1024, 1024),
+    # (4 * 1024, 1024),
+    # (8 * 1024, 1024),
+    # (16 * 1024, 1024),
+    # (32 * 1024, 1024),
+    # (64 * 1024, 1024),
+    # (128 * 1024, 1024),
+    # (256 * 1024, 1024),
+    # (512 * 1024, 1024),
+    # (1024 * 1024, 1024),
 )
 
 
@@ -47,10 +145,12 @@ def benchmark_all_to_all(group, sl, hs):
     def benchmark():
         return all_to_all(x, send_recv_sizes, send_recv_sizes, group)
 
-    time, std = benchmark_util.benchmark_function(benchmark)
+    # time, std = benchmark_util.benchmark_function(benchmark)
+    time, std = benchmark_function(benchmark)
 
     if dist.get_rank(group) == 0:
-        benchmark_util.log_benchmark('All-To-All', details, time, std)
+        log_benchmark('All-To-All', details, time, std)
+        # benchmark_util.log_benchmark('All-To-All', details, time, std)
 
 
 if __name__ == '__main__':

build/torch26-cxx98-cu126-x86_64-linux/megablocks/_megablocks_13afbbe_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94a9a3bb426adceab66b39fe9d179b73e4524167aeb63bed5a67cd7734d31b24
+size 11923704

build/torch26-cxx98-cu126-x86_64-linux/megablocks/_megablocks_63599de.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:637a8c7ef51b1d35911546ef7456854f1ee7cc3278565d2e144e16f733487148
-size 11923672

build/torch26-cxx98-cu126-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_63599de
-ops = torch.ops._megablocks_63599de
+from . import _megablocks_13afbbe_dirty
+ops = torch.ops._megablocks_13afbbe_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_63599de::{op_name}"
+    return f"_megablocks_13afbbe_dirty::{op_name}"

build/torch26-cxx98-cu126-x86_64-linux/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

build/torch26-cxx98-cu126-x86_64-linux/megablocks/ops/all_to_all_benchmark.py

@@ -7,28 +7,126 @@ import torch.distributed as dist
 # from megablocks import benchmark_util
 # from megablocks.layers.all_to_all import all_to_all
 
-from .. import benchmark_util
-from .._layers.all_to_all import all_to_all
+# from .. import benchmark_util
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import torch
+
+
+def log_benchmark(name, arguments, time, std):
+    print("=" * 60)
+    print(f"{name} Benchmark")
+    print("Benchmark Parameters:")
+    for key, value in arguments.items():
+        print(f"{key} = {value}")
+    print("Results:")
+    print("mean time = {:.3f}ms, std time = {:.3f}ms".format(time, std))
+    print("=" * 60)
+
+
+def benchmark_function(fn, iterations=100, warmup=10):
+    print(f"Benchmarking {fn.__name__} with {iterations} iterations and {warmup} warmup iterations")
+    # Warmup iterations.
+    for _ in range(warmup):
+        fn()
+
+    times = []
+    print(f"Running {iterations} iterations...")
+    for i in range(iterations):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+
+        start.record()
+        fn()
+        end.record()
+
+        torch.cuda.synchronize()
+        times.append(start.elapsed_time(end))
+    return np.mean(times), np.std(times)
+
+
+# from .._layers.all_to_all import all_to_all
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.distributed as dist
+
+
+class AllToAllOp(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, output_split_sizes, input_split_sizes, group, async_op):
+        out = torch.empty(
+            (sum(output_split_sizes),) + x.shape[1:], device=x.device, dtype=x.dtype
+        )
+
+        ctx.input_shape = x.shape
+        ctx.output_split_sizes = output_split_sizes
+        ctx.input_split_sizes = input_split_sizes
+        ctx.group = group
+        handle = dist.all_to_all_single(
+            out,
+            x,
+            output_split_sizes=output_split_sizes,
+            input_split_sizes=input_split_sizes,
+            group=group,
+            async_op=async_op,
+        )
+        return out, handle
+
+    @staticmethod
+    def backward(ctx, grad, _):
+        if ctx.needs_input_grad[0]:
+            out = torch.empty(
+                ctx.input_shape,
+                device=grad.device,
+                dtype=grad.dtype,
+            )
+            dist.all_to_all_single(
+                out,
+                grad,
+                output_split_sizes=ctx.input_split_sizes,
+                input_split_sizes=ctx.output_split_sizes,
+                group=ctx.group,
+            )
+            return out, None, None, None, None
+        return None, None, None, None, None
+
+
+def all_to_all(x, output_split_sizes, input_split_sizes, group, async_op=False):
+    return AllToAllOp.apply(
+        x,
+        output_split_sizes,
+        input_split_sizes,
+        group,
+        async_op,
+    )
+
 
 _ALL_TO_ALL_BENCHMARK = (
     (8, 1024),
-    (16, 1024),
-    (32, 1024),
-    (64, 1024),
-    (128, 1024),
-    (256, 1024),
-    (512, 1024),
-    (1024, 1024),
-    (2 * 1024, 1024),
-    (4 * 1024, 1024),
-    (8 * 1024, 1024),
-    (16 * 1024, 1024),
-    (32 * 1024, 1024),
-    (64 * 1024, 1024),
-    (128 * 1024, 1024),
-    (256 * 1024, 1024),
-    (512 * 1024, 1024),
-    (1024 * 1024, 1024),
+    # (16, 1024),
+    # (32, 1024),
+    # (64, 1024),
+    # (128, 1024),
+    # (256, 1024),
+    # (512, 1024),
+    # (1024, 1024),
+    # (2 * 1024, 1024),
+    # (4 * 1024, 1024),
+    # (8 * 1024, 1024),
+    # (16 * 1024, 1024),
+    # (32 * 1024, 1024),
+    # (64 * 1024, 1024),
+    # (128 * 1024, 1024),
+    # (256 * 1024, 1024),
+    # (512 * 1024, 1024),
+    # (1024 * 1024, 1024),
 )
 
 
@@ -47,10 +145,12 @@ def benchmark_all_to_all(group, sl, hs):
     def benchmark():
         return all_to_all(x, send_recv_sizes, send_recv_sizes, group)
 
-    time, std = benchmark_util.benchmark_function(benchmark)
+    # time, std = benchmark_util.benchmark_function(benchmark)
+    time, std = benchmark_function(benchmark)
 
     if dist.get_rank(group) == 0:
-        benchmark_util.log_benchmark('All-To-All', details, time, std)
+        log_benchmark('All-To-All', details, time, std)
+        # benchmark_util.log_benchmark('All-To-All', details, time, std)
 
 
 if __name__ == '__main__':

build/torch27-cxx11-cu118-x86_64-linux/megablocks/_megablocks_13afbbe_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa9d1964e47ec6ff3c4ec77947f6a2a19868b03cec3618daf0555e011f69924d
+size 10517848

build/torch27-cxx11-cu118-x86_64-linux/megablocks/_megablocks_63599de.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:002a58b415ed9e0f6418b103368c4f57f17fa86a851a02f594a33b097b33da09
-size 10517816

build/torch27-cxx11-cu118-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_63599de
-ops = torch.ops._megablocks_63599de
+from . import _megablocks_13afbbe_dirty
+ops = torch.ops._megablocks_13afbbe_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_63599de::{op_name}"
+    return f"_megablocks_13afbbe_dirty::{op_name}"

build/torch27-cxx11-cu118-x86_64-linux/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

build/torch27-cxx11-cu118-x86_64-linux/megablocks/ops/all_to_all_benchmark.py

@@ -7,28 +7,126 @@ import torch.distributed as dist
 # from megablocks import benchmark_util
 # from megablocks.layers.all_to_all import all_to_all
 
-from .. import benchmark_util
-from .._layers.all_to_all import all_to_all
+# from .. import benchmark_util
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import torch
+
+
+def log_benchmark(name, arguments, time, std):
+    print("=" * 60)
+    print(f"{name} Benchmark")
+    print("Benchmark Parameters:")
+    for key, value in arguments.items():
+        print(f"{key} = {value}")
+    print("Results:")
+    print("mean time = {:.3f}ms, std time = {:.3f}ms".format(time, std))
+    print("=" * 60)
+
+
+def benchmark_function(fn, iterations=100, warmup=10):
+    print(f"Benchmarking {fn.__name__} with {iterations} iterations and {warmup} warmup iterations")
+    # Warmup iterations.
+    for _ in range(warmup):
+        fn()
+
+    times = []
+    print(f"Running {iterations} iterations...")
+    for i in range(iterations):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+
+        start.record()
+        fn()
+        end.record()
+
+        torch.cuda.synchronize()
+        times.append(start.elapsed_time(end))
+    return np.mean(times), np.std(times)
+
+
+# from .._layers.all_to_all import all_to_all
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.distributed as dist
+
+
+class AllToAllOp(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, output_split_sizes, input_split_sizes, group, async_op):
+        out = torch.empty(
+            (sum(output_split_sizes),) + x.shape[1:], device=x.device, dtype=x.dtype
+        )
+
+        ctx.input_shape = x.shape
+        ctx.output_split_sizes = output_split_sizes
+        ctx.input_split_sizes = input_split_sizes
+        ctx.group = group
+        handle = dist.all_to_all_single(
+            out,
+            x,
+            output_split_sizes=output_split_sizes,
+            input_split_sizes=input_split_sizes,
+            group=group,
+            async_op=async_op,
+        )
+        return out, handle
+
+    @staticmethod
+    def backward(ctx, grad, _):
+        if ctx.needs_input_grad[0]:
+            out = torch.empty(
+                ctx.input_shape,
+                device=grad.device,
+                dtype=grad.dtype,
+            )
+            dist.all_to_all_single(
+                out,
+                grad,
+                output_split_sizes=ctx.input_split_sizes,
+                input_split_sizes=ctx.output_split_sizes,
+                group=ctx.group,
+            )
+            return out, None, None, None, None
+        return None, None, None, None, None
+
+
+def all_to_all(x, output_split_sizes, input_split_sizes, group, async_op=False):
+    return AllToAllOp.apply(
+        x,
+        output_split_sizes,
+        input_split_sizes,
+        group,
+        async_op,
+    )
+
 
 _ALL_TO_ALL_BENCHMARK = (
     (8, 1024),
-    (16, 1024),
-    (32, 1024),
-    (64, 1024),
-    (128, 1024),
-    (256, 1024),
-    (512, 1024),
-    (1024, 1024),
-    (2 * 1024, 1024),
-    (4 * 1024, 1024),
-    (8 * 1024, 1024),
-    (16 * 1024, 1024),
-    (32 * 1024, 1024),
-    (64 * 1024, 1024),
-    (128 * 1024, 1024),
-    (256 * 1024, 1024),
-    (512 * 1024, 1024),
-    (1024 * 1024, 1024),
+    # (16, 1024),
+    # (32, 1024),
+    # (64, 1024),
+    # (128, 1024),
+    # (256, 1024),
+    # (512, 1024),
+    # (1024, 1024),
+    # (2 * 1024, 1024),
+    # (4 * 1024, 1024),
+    # (8 * 1024, 1024),
+    # (16 * 1024, 1024),
+    # (32 * 1024, 1024),
+    # (64 * 1024, 1024),
+    # (128 * 1024, 1024),
+    # (256 * 1024, 1024),
+    # (512 * 1024, 1024),
+    # (1024 * 1024, 1024),
 )
 
 
@@ -47,10 +145,12 @@ def benchmark_all_to_all(group, sl, hs):
     def benchmark():
         return all_to_all(x, send_recv_sizes, send_recv_sizes, group)
 
-    time, std = benchmark_util.benchmark_function(benchmark)
+    # time, std = benchmark_util.benchmark_function(benchmark)
+    time, std = benchmark_function(benchmark)
 
     if dist.get_rank(group) == 0:
-        benchmark_util.log_benchmark('All-To-All', details, time, std)
+        log_benchmark('All-To-All', details, time, std)
+        # benchmark_util.log_benchmark('All-To-All', details, time, std)
 
 
 if __name__ == '__main__':

build/torch27-cxx11-cu126-x86_64-linux/megablocks/_megablocks_13afbbe_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b204da58db0f8be45dda62abd98b74a8e60f1f983bfc6a128c74ff66f67cf502
+size 11931112

build/torch27-cxx11-cu126-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_63599de
-ops = torch.ops._megablocks_63599de
+from . import _megablocks_13afbbe_dirty
+ops = torch.ops._megablocks_13afbbe_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_63599de::{op_name}"
+    return f"_megablocks_13afbbe_dirty::{op_name}"

build/torch27-cxx11-cu126-x86_64-linux/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

build/torch27-cxx11-cu126-x86_64-linux/megablocks/ops/all_to_all_benchmark.py

@@ -7,28 +7,126 @@ import torch.distributed as dist
 # from megablocks import benchmark_util
 # from megablocks.layers.all_to_all import all_to_all
 
-from .. import benchmark_util
-from .._layers.all_to_all import all_to_all
+# from .. import benchmark_util
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import torch
+
+
+def log_benchmark(name, arguments, time, std):
+    print("=" * 60)
+    print(f"{name} Benchmark")
+    print("Benchmark Parameters:")
+    for key, value in arguments.items():
+        print(f"{key} = {value}")
+    print("Results:")
+    print("mean time = {:.3f}ms, std time = {:.3f}ms".format(time, std))
+    print("=" * 60)
+
+
+def benchmark_function(fn, iterations=100, warmup=10):
+    print(f"Benchmarking {fn.__name__} with {iterations} iterations and {warmup} warmup iterations")
+    # Warmup iterations.
+    for _ in range(warmup):
+        fn()
+
+    times = []
+    print(f"Running {iterations} iterations...")
+    for i in range(iterations):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+
+        start.record()
+        fn()
+        end.record()
+
+        torch.cuda.synchronize()
+        times.append(start.elapsed_time(end))
+    return np.mean(times), np.std(times)
+
+
+# from .._layers.all_to_all import all_to_all
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.distributed as dist
+
+
+class AllToAllOp(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, output_split_sizes, input_split_sizes, group, async_op):
+        out = torch.empty(
+            (sum(output_split_sizes),) + x.shape[1:], device=x.device, dtype=x.dtype
+        )
+
+        ctx.input_shape = x.shape
+        ctx.output_split_sizes = output_split_sizes
+        ctx.input_split_sizes = input_split_sizes
+        ctx.group = group
+        handle = dist.all_to_all_single(
+            out,
+            x,
+            output_split_sizes=output_split_sizes,
+            input_split_sizes=input_split_sizes,
+            group=group,
+            async_op=async_op,
+        )
+        return out, handle
+
+    @staticmethod
+    def backward(ctx, grad, _):
+        if ctx.needs_input_grad[0]:
+            out = torch.empty(
+                ctx.input_shape,
+                device=grad.device,
+                dtype=grad.dtype,
+            )
+            dist.all_to_all_single(
+                out,
+                grad,
+                output_split_sizes=ctx.input_split_sizes,
+                input_split_sizes=ctx.output_split_sizes,
+                group=ctx.group,
+            )
+            return out, None, None, None, None
+        return None, None, None, None, None
+
+
+def all_to_all(x, output_split_sizes, input_split_sizes, group, async_op=False):
+    return AllToAllOp.apply(
+        x,
+        output_split_sizes,
+        input_split_sizes,
+        group,
+        async_op,
+    )
+
 
 _ALL_TO_ALL_BENCHMARK = (
     (8, 1024),
-    (16, 1024),
-    (32, 1024),
-    (64, 1024),
-    (128, 1024),
-    (256, 1024),
-    (512, 1024),
-    (1024, 1024),
-    (2 * 1024, 1024),
-    (4 * 1024, 1024),
-    (8 * 1024, 1024),
-    (16 * 1024, 1024),
-    (32 * 1024, 1024),
-    (64 * 1024, 1024),
-    (128 * 1024, 1024),
-    (256 * 1024, 1024),
-    (512 * 1024, 1024),
-    (1024 * 1024, 1024),
+    # (16, 1024),
+    # (32, 1024),
+    # (64, 1024),
+    # (128, 1024),
+    # (256, 1024),
+    # (512, 1024),
+    # (1024, 1024),
+    # (2 * 1024, 1024),
+    # (4 * 1024, 1024),
+    # (8 * 1024, 1024),
+    # (16 * 1024, 1024),
+    # (32 * 1024, 1024),
+    # (64 * 1024, 1024),
+    # (128 * 1024, 1024),
+    # (256 * 1024, 1024),
+    # (512 * 1024, 1024),
+    # (1024 * 1024, 1024),
 )
 
 
@@ -47,10 +145,12 @@ def benchmark_all_to_all(group, sl, hs):
     def benchmark():
         return all_to_all(x, send_recv_sizes, send_recv_sizes, group)
 
-    time, std = benchmark_util.benchmark_function(benchmark)
+    # time, std = benchmark_util.benchmark_function(benchmark)
+    time, std = benchmark_function(benchmark)
 
     if dist.get_rank(group) == 0:
-        benchmark_util.log_benchmark('All-To-All', details, time, std)
+        log_benchmark('All-To-All', details, time, std)
+        # benchmark_util.log_benchmark('All-To-All', details, time, std)
 
 
 if __name__ == '__main__':

build/torch27-cxx11-cu128-x86_64-linux/megablocks/_megablocks_13afbbe_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f861a8bffedbbf14341d39355f3f43a7c24fee2b99bb9ea7b3a2b9ad21c7ee28
+size 17892656

build/torch27-cxx11-cu128-x86_64-linux/megablocks/_megablocks_63599de.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:dadccc59929c2fdbdf3b153f564d223013924c7b617d1eb2b3ecdc04470a4a60
-size 17892624

build/torch27-cxx11-cu128-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_63599de
-ops = torch.ops._megablocks_63599de
+from . import _megablocks_13afbbe_dirty
+ops = torch.ops._megablocks_13afbbe_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_63599de::{op_name}"
+    return f"_megablocks_13afbbe_dirty::{op_name}"

build/torch27-cxx11-cu128-x86_64-linux/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

build/torch27-cxx11-cu128-x86_64-linux/megablocks/ops/all_to_all_benchmark.py

@@ -7,28 +7,126 @@ import torch.distributed as dist
 # from megablocks import benchmark_util
 # from megablocks.layers.all_to_all import all_to_all
 
-from .. import benchmark_util
-from .._layers.all_to_all import all_to_all
+# from .. import benchmark_util
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import torch
+
+
+def log_benchmark(name, arguments, time, std):
+    print("=" * 60)
+    print(f"{name} Benchmark")
+    print("Benchmark Parameters:")
+    for key, value in arguments.items():
+        print(f"{key} = {value}")
+    print("Results:")
+    print("mean time = {:.3f}ms, std time = {:.3f}ms".format(time, std))
+    print("=" * 60)
+
+
+def benchmark_function(fn, iterations=100, warmup=10):
+    print(f"Benchmarking {fn.__name__} with {iterations} iterations and {warmup} warmup iterations")
+    # Warmup iterations.
+    for _ in range(warmup):
+        fn()
+
+    times = []
+    print(f"Running {iterations} iterations...")
+    for i in range(iterations):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+
+        start.record()
+        fn()
+        end.record()
+
+        torch.cuda.synchronize()
+        times.append(start.elapsed_time(end))
+    return np.mean(times), np.std(times)
+
+
+# from .._layers.all_to_all import all_to_all
+
+# Copyright 2024 Databricks
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.distributed as dist
+
+
+class AllToAllOp(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, output_split_sizes, input_split_sizes, group, async_op):
+        out = torch.empty(
+            (sum(output_split_sizes),) + x.shape[1:], device=x.device, dtype=x.dtype
+        )
+
+        ctx.input_shape = x.shape
+        ctx.output_split_sizes = output_split_sizes
+        ctx.input_split_sizes = input_split_sizes
+        ctx.group = group
+        handle = dist.all_to_all_single(
+            out,
+            x,
+            output_split_sizes=output_split_sizes,
+            input_split_sizes=input_split_sizes,
+            group=group,
+            async_op=async_op,
+        )
+        return out, handle
+
+    @staticmethod
+    def backward(ctx, grad, _):
+        if ctx.needs_input_grad[0]:
+            out = torch.empty(
+                ctx.input_shape,
+                device=grad.device,
+                dtype=grad.dtype,
+            )
+            dist.all_to_all_single(
+                out,
+                grad,
+                output_split_sizes=ctx.input_split_sizes,
+                input_split_sizes=ctx.output_split_sizes,
+                group=ctx.group,
+            )
+            return out, None, None, None, None
+        return None, None, None, None, None
+
+
+def all_to_all(x, output_split_sizes, input_split_sizes, group, async_op=False):
+    return AllToAllOp.apply(
+        x,
+        output_split_sizes,
+        input_split_sizes,
+        group,
+        async_op,
+    )
+
 
 _ALL_TO_ALL_BENCHMARK = (
     (8, 1024),
-    (16, 1024),
-    (32, 1024),
-    (64, 1024),
-    (128, 1024),
-    (256, 1024),
-    (512, 1024),
-    (1024, 1024),
-    (2 * 1024, 1024),
-    (4 * 1024, 1024),
-    (8 * 1024, 1024),
-    (16 * 1024, 1024),
-    (32 * 1024, 1024),
-    (64 * 1024, 1024),
-    (128 * 1024, 1024),
-    (256 * 1024, 1024),
-    (512 * 1024, 1024),
-    (1024 * 1024, 1024),
+    # (16, 1024),
+    # (32, 1024),
+    # (64, 1024),
+    # (128, 1024),
+    # (256, 1024),
+    # (512, 1024),
+    # (1024, 1024),
+    # (2 * 1024, 1024),
+    # (4 * 1024, 1024),
+    # (8 * 1024, 1024),
+    # (16 * 1024, 1024),
+    # (32 * 1024, 1024),
+    # (64 * 1024, 1024),
+    # (128 * 1024, 1024),
+    # (256 * 1024, 1024),
+    # (512 * 1024, 1024),
+    # (1024 * 1024, 1024),
 )
 
 
@@ -47,10 +145,12 @@ def benchmark_all_to_all(group, sl, hs):
     def benchmark():
         return all_to_all(x, send_recv_sizes, send_recv_sizes, group)
 
-    time, std = benchmark_util.benchmark_function(benchmark)
+    # time, std = benchmark_util.benchmark_function(benchmark)
+    time, std = benchmark_function(benchmark)
 
     if dist.get_rank(group) == 0:
-        benchmark_util.log_benchmark('All-To-All', details, time, std)
+        log_benchmark('All-To-All', details, time, std)
+        # benchmark_util.log_benchmark('All-To-All', details, time, std)
 
 
 if __name__ == '__main__':