Make token group alignment size configurable (#1503)

- For mxfp8, token group sizes must be multiples of "block_size" because
in the backward pass for `grad_weight = grad_output_t @ input`, the "M"
(token) dimension is the contracting dimension, and each token group is
a logically distinct subtensor, so we scale them separately. This means
token groups contracting dimension must be divisible by the mxfp8
block_size (default 32). Here is a diagram showing the problem:
https://www.internalfb.com/excalidraw/EX521879
- To solve this, this PR makes the token group M aligment configurable.

- Integration test with torchao passes:
pytorch/ao#2642
- Did manual test run with llama4 debug model using bf16

Author: danielvegamyhre

test.py

@@ -0,0 +1,62 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+from dataclasses import dataclass
+
+
+@dataclass
+class Config:
+    num_experts=2
+    intermediate_size=1024
+    dim=2048
+
+class MoE(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.w1 = nn.Parameter( # num exp, expert_dim, hidden_dim
+            torch.empty(config.num_experts, config.intermediate_size, config.dim)
+        )  # E, I, D
+        self.w2 = nn.Parameter(
+            torch.empty(config.num_experts, config.dim, config.intermediate_size)
+        )  # E, D, I
+        self.w3 = nn.Parameter(
+            torch.empty(config.num_experts, config.intermediate_size, config.dim)
+        )  # E, I, D
+        self.w2 = torch.nn.Parameter(self.w2.transpose(-2, -1).contiguous().transpose(-2, -1))
+
+        nn.init.normal_(self.w1, std=0.02)
+        nn.init.normal_(self.w2, std=0.02)
+        nn.init.normal_(self.w3, std=0.02)
+        self.w13 = torch.cat((self.w1, self.w3), dim=1)
+        self.w13 = torch.nn.Parameter(self.w13.transpose(-2,-1).contiguous().transpose(-2,-1))
+
+    def forward(self, ordered_inputs, optim=True):
+        M = ordered_inputs.size(0)
+        group_size = M // self.config.num_experts
+        offs = torch.arange(group_size, M + 1, group_size, device="cuda", dtype=torch.int32)
+
+        # Since we do do grouped gemms with same size:
+        # - x @ w1 => (M, K) @ (E, K, D)
+        # - x @ w3 => (M, K) @ (E, K, D)
+        # We can concatenate w1 and w3 along the K dim so we have:
+        # (M, K) @ (E, 2*K, D) -> (M, E, D)
+        # and K dim of inputs is broadcasted to 2*K for the gemm.
+        breakpoint()
+        x13 = torch._grouped_mm(ordered_inputs, self.w13.transpose(-2, -1), offs)
+        x1, x3 = x13.split(self.config.intermediate_size, dim=1)
+        y1 = F.silu(x1) * x3
+        ordered_outs1 = torch._grouped_mm(y1, self.w2.transpose(-2,-1), offs)
+
+        x1 = F.silu(torch._grouped_mm(ordered_inputs, self.w1.transpose(-2,-1), offs))
+        x3 = torch._grouped_mm(ordered_inputs, self.w3.transpose(-2,-1), offs)
+        y1 = x1 * x3
+        ordered_outs2 = torch._grouped_mm(y1, self.w2.transpose(-2,-1), offs)        
+
+        assert torch.equal(ordered_outs1, ordered_outs2)
+        return ordered_outs1
+
+config = Config()
+m = MoE(config).cuda().bfloat16()
+ordered_inputs = torch.randn(256, config.intermediate_size, device="cuda", dtype=torch.bfloat16)
+m(ordered_inputs, optim=True)

torchtitan/experiments/llama4/infra/expert_parallel.py

@@ -272,24 +272,22 @@ def expert_parallel(func: Callable) -> Callable:
     """
 
     def wrapper(
-        w1: torch.Tensor,
+        w13: torch.Tensor,
         w2: torch.Tensor,
-        w3: torch.Tensor,
         x: torch.Tensor,
         num_tokens_per_expert: torch.Tensor | None = None,
     ) -> torch.Tensor:
         global TOKEN_GROUP_ALIGN_SIZE_M
-        if isinstance(w1, DTensor):
-            w1 = w1.to_local()
+        if isinstance(w13, DTensor):
+            w13 = w13.to_local()
             w2 = w2.to_local()
-            w3 = w3.to_local()
 
         if num_tokens_per_expert is not None:
             from torchtitan.experiments.kernels.moe.indices import (
                 generate_permute_indices,
             )
 
-            experts_per_ep_rank = w1.shape[0]
+            experts_per_ep_rank = w13.shape[0]
             num_ep_ranks = num_tokens_per_expert.shape[0] // experts_per_ep_rank
 
             with torch.no_grad():
@@ -309,7 +307,7 @@ def wrapper(
             input_shape = x.shape
             x = x[permuted_indices, :]
 
-        out = func(w1, w2, w3, x, num_tokens_per_expert)
+        out = func(w13, w2, x, num_tokens_per_expert)
 
         if num_tokens_per_expert is not None:
             out_unpermuted = out.new_empty(input_shape)

torchtitan/experiments/llama4/model/moe.py

@@ -23,76 +23,27 @@ def __init__(
     ):
         super().__init__()
         self.num_experts = num_experts
-        self.w1 = nn.Parameter(torch.empty(num_experts, dim, hidden_dim))
-        self.w2 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim))
-        self.w3 = nn.Parameter(torch.empty(num_experts, dim, hidden_dim))
+        # Combine w1 and w3 into a single tensor to perform so we can combine
+        # `x @ w1` and `x @ w3` into a single grouped mm.
+        self.w13 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim * 2))
+        self.w2 = nn.Parameter(torch.empty(num_experts, dim, hidden_dim))
         self.use_grouped_mm = use_grouped_mm
 
     def forward(
         self,
         x: torch.Tensor,
         num_tokens_per_expert: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        if self.use_grouped_mm:
-            return GroupedExperts._run_experts_grouped_mm(
-                self.w1, self.w2, self.w3, x, num_tokens_per_expert
-            )
-        else:
-            return GroupedExperts._run_experts_for_loop(
-                self.w1, self.w2, self.w3, x, num_tokens_per_expert
-            )
-
-    # TODO: keeping this for-loop implementation for comparison
-    #       and readability, may remove later
-    @expert_parallel
-    @staticmethod
-    def _run_experts_for_loop(
-        w1: torch.Tensor,
-        w2: torch.Tensor,
-        w3: torch.Tensor,
-        x: torch.Tensor,
-        num_tokens_per_expert: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        if num_tokens_per_expert is not None:
-            # NOTE: this would incur a synchronization between device and host
-            num_tokens_per_expert = num_tokens_per_expert.tolist()
-
-            # side-effect code due to the usage of generate_permute_indices
-            num_padding = x.shape[0] - sum(num_tokens_per_expert)
-
-            # a tuple of tensors indexed by experts
-            # each with shape (tokens_per_expert(varying), dim)
-            x = torch.split(
-                x[: sum(num_tokens_per_expert)],
-                split_size_or_sections=num_tokens_per_expert,
-                dim=0,
-            )
-            out_experts_splits = []
-            for expert_idx, x_expert in enumerate(x):
-                h = F.silu(torch.matmul(x_expert, w1[expert_idx]))
-                h = h * torch.matmul(x_expert, w3[expert_idx])
-                h = torch.matmul(h, w2[expert_idx])
-                # h shape (tokens_per_expert(varying), dim)
-                out_experts_splits.append(h)
-            out = torch.cat(out_experts_splits, dim=0)
-
-            # side-effect code due to the usage of generate_permute_indices
-            out = torch.vstack((out, out.new_zeros((num_padding, out.shape[-1]))))
-        else:
-            # x shape (num_experts, tokens_per_expert, dim)
-            h = F.silu(torch.bmm(x, w1))
-            h = h * torch.bmm(x, w3)
-            # out shape (num_experts, tokens_per_expert, dim)
-            out = torch.bmm(h, w2)
+        return GroupedExperts._run_experts_grouped_mm(
+            self.w13, self.w2, x, num_tokens_per_expert
+        )
 
-        return out
 
     @expert_parallel
     @staticmethod
     def _run_experts_grouped_mm(
-        w1: torch.Tensor,
+        w13: torch.Tensor,
         w2: torch.Tensor,
-        w3: torch.Tensor,
         x: torch.Tensor,
         num_tokens_per_expert: torch.Tensor | None = None,
     ) -> torch.Tensor:
@@ -105,16 +56,14 @@ def _run_experts_grouped_mm(
             # fall back to regular bmm between 3D tensors
             assert x.dim() == 3
 
-        h = F.silu(torch._grouped_mm(x.bfloat16(), w1.bfloat16(), offs=offsets))
-        h = h * torch._grouped_mm(x.bfloat16(), w3.bfloat16(), offs=offsets)
-        out = torch._grouped_mm(h, w2.bfloat16(), offs=offsets).type_as(x)
-
+        x1, x3 = torch._grouped_mm(x, w13.transpose(-2, -1), offs=offsets).chunk(2, dim=-1)
+        y = F.silu(x1) * x3
+        out = torch._grouped_mm(y, w2.transpose(-2, -1), offs=offsets).type_as(x)
         return out
 
     def init_weights(self, init_std: float):
-        nn.init.trunc_normal_(self.w1, mean=0.0, std=0.02)
+        nn.init.trunc_normal_(self.w13, mean=0.0, std=0.02)
         nn.init.trunc_normal_(self.w2, mean=0.0, std=init_std)
-        nn.init.trunc_normal_(self.w3, mean=0.0, std=init_std)
 
 
 class TokenChoiceTopKRouter(nn.Module):
@@ -299,7 +248,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         # shared expert
         if self.shared_expert is not None:
-            out = self.shared_expert(x.reshape(1, bs * slen, dim)).reshape(
+            out = self.shared_expert(x.reshape(1, bs * slen, dim))
+            out = out.reshape(
                 bs * slen, dim
             )
         else: