From a6bab0663e38dac514d1f80095eb56c22fc8136a Mon Sep 17 00:00:00 2001 From: Animesh Jain Date: Fri, 28 Aug 2026 12:22:57 -0700 Subject: [PATCH] Pass the expert weight, not its transpose, to the grouped-mm seam GroupedExperts._grouped_mm is the override point low-precision variants use to swap in a scaled grouped GEMM. It currently receives B_t, the already transposed weight, which is enough for a variant that quantizes on every call but not for one whose quantized representation is owned by the weight's FSDP unshard lifetime: that representation is keyed off the stored (experts, out_features, in_features) orientation, and the transpose has erased which parameter it came from. Pass weight_ENK and move the transpose inside the seam. Behavior is unchanged; the base implementation transposes exactly as before. Update every caller: the common, gpt_oss, and kimi_k3 experts, the fused SwiGLU override, and the float8 converter's override. The fused SwiGLU case reshapes its 4D w13 parameter to the (E, 2F, D) grouped weight first, which the seam then transposes. --- torchtitan/components/quantization/float8.py | 7 ++-- torchtitan/models/common/moe.py | 36 +++++++++----------- torchtitan/models/gpt_oss/moe.py | 6 ++-- torchtitan/models/kimi_k3/moe.py | 6 ++-- torchtitan/overrides/fused_swiglu.py | 10 +++--- 5 files changed, 32 insertions(+), 33 deletions(-) diff --git a/torchtitan/components/quantization/float8.py b/torchtitan/components/quantization/float8.py index 99770ed0c9..9d8339c649 100644 --- a/torchtitan/components/quantization/float8.py +++ b/torchtitan/components/quantization/float8.py @@ -198,13 +198,16 @@ def __init__(self, config: Config): self._float8_op_config = Float8TrainingOpConfig() - def _grouped_mm(self, *, A, B_t, offs): + def _grouped_mm(self, *, A, weight_EOI, offs): from torchao.prototype.moe_training.utils import ( _quantize_then_scaled_grouped_mm, ) return _quantize_then_scaled_grouped_mm( - A, B_t, config=self._float8_op_config, offs=offs + A, + weight_EOI.bfloat16().transpose(-2, -1), + config=self._float8_op_config, + offs=offs, ) Float8GroupedExperts.__name__ = f"Float8{parent_cls.__name__}" diff --git a/torchtitan/models/common/moe.py b/torchtitan/models/common/moe.py index 1a1930f143..63bc5a8859 100644 --- a/torchtitan/models/common/moe.py +++ b/torchtitan/models/common/moe.py @@ -92,32 +92,30 @@ def forward( spmd.mutate_type(offsets_E, axis, src=spmd.P, dst=spmd.V) h_RF = F.silu( - self._grouped_mm( - A=x_RD.bfloat16(), - B_t=w1_EFD.bfloat16().transpose(-2, -1), - offs=offsets_E, - ) + self._grouped_mm(A=x_RD.bfloat16(), weight_EOI=w1_EFD, offs=offsets_E) ) h_RF = h_RF * self._grouped_mm( - A=x_RD.bfloat16(), - B_t=w3_EFD.bfloat16().transpose(-2, -1), - offs=offsets_E, + A=x_RD.bfloat16(), weight_EOI=w3_EFD, offs=offsets_E ) - return self._grouped_mm( - A=h_RF, B_t=w2_EDF.bfloat16().transpose(-2, -1), offs=offsets_E - ).type_as(x_RD) + return self._grouped_mm(A=h_RF, weight_EOI=w2_EDF, offs=offsets_E).type_as(x_RD) def _grouped_mm( - self, *, A: torch.Tensor, B_t: torch.Tensor, offs: torch.Tensor + self, *, A: torch.Tensor, weight_EOI: torch.Tensor, offs: torch.Tensor ) -> torch.Tensor: - """Grouped matmul of ``A @ B_t`` with per-expert token offsets. - - Overridable seam for low-precision variants (e.g. the MXFP8 converter - swaps this for a dynamically-quantized scaled grouped GEMM). Keeping the - op here -- rather than behind a tensor-subclass ``__torch_function__`` -- - means it is captured by FX tracers such as graph_trainer's make_fx path. + """Grouped matmul of ``A @ weight_EOI.transpose(-2, -1)``. + + ``weight_EOI`` is the grouped expert weight in its stored + ``(experts, out_features, in_features)`` orientation; the transpose to + the grouped-GEMM right operand happens here. Overridable seam for + low-precision variants (e.g. the MXFP8 converter swaps this for a + scaled grouped GEMM). Variants receive the weight rather than its + transpose because a quantized representation may be owned by the + weight's FSDP unshard lifetime and is keyed off the stored orientation. + Keeping the op here -- rather than behind a tensor-subclass + ``__torch_function__`` -- means it is captured by FX tracers such as + graph_trainer's make_fx path. """ - return torch._grouped_mm(A, B_t, offs=offs) + return torch._grouped_mm(A, weight_EOI.bfloat16().transpose(-2, -1), offs=offs) class RoutedExperts(Module): diff --git a/torchtitan/models/gpt_oss/moe.py b/torchtitan/models/gpt_oss/moe.py index c6a274c171..5150c5eab4 100644 --- a/torchtitan/models/gpt_oss/moe.py +++ b/torchtitan/models/gpt_oss/moe.py @@ -162,7 +162,7 @@ def forward( # G = gate+up dimension (2*F) h_RG = self._grouped_mm( A=x_RD.bfloat16(), - B_t=mlp1_weight_EGD.transpose(-2, -1).bfloat16(), + weight_EOI=mlp1_weight_EGD, offs=offsets_E, ) @@ -175,9 +175,7 @@ def forward( h_RG = h_RG + b1_RG.to(h_RG.dtype) h_RF = swiglu(h_RG, limit=self.swiglu_limit) - h_RD = self._grouped_mm( - A=h_RF, B_t=mlp2_weight_EDF.transpose(-2, -1).bfloat16(), offs=offsets_E - ) + h_RD = self._grouped_mm(A=h_RF, weight_EOI=mlp2_weight_EDF, offs=offsets_E) # Apply custom autograd function to scale bias in forward but not in backward b2 = torch.cat( diff --git a/torchtitan/models/kimi_k3/moe.py b/torchtitan/models/kimi_k3/moe.py index 990104a9d8..a5856d93a5 100644 --- a/torchtitan/models/kimi_k3/moe.py +++ b/torchtitan/models/kimi_k3/moe.py @@ -89,12 +89,12 @@ def forward( gate_RF = self._grouped_mm( A=x_RD.bfloat16(), - B_t=w1_EFD.bfloat16().transpose(-2, -1), + weight_EOI=w1_EFD, offs=offsets_E, ) up_RF = self._grouped_mm( A=x_RD.bfloat16(), - B_t=w3_EFD.bfloat16().transpose(-2, -1), + weight_EOI=w3_EFD, offs=offsets_E, ) @@ -102,7 +102,7 @@ def forward( return self._grouped_mm( A=h_RF, - B_t=w2_EDF.bfloat16().transpose(-2, -1), + weight_EOI=w2_EDF, offs=offsets_E, ).type_as(x_RD) diff --git a/torchtitan/overrides/fused_swiglu.py b/torchtitan/overrides/fused_swiglu.py index 9b1cd25a80..26138ccbac 100644 --- a/torchtitan/overrides/fused_swiglu.py +++ b/torchtitan/overrides/fused_swiglu.py @@ -647,15 +647,15 @@ def forward( E, F, _, D = w13.shape offsets_E = torch.cumsum(num_tokens_per_expert_E, dim=0, dtype=torch.int32) - w13_E_D_2F = w13.bfloat16().reshape(E, F * 2, D).transpose(-2, -1) + # The fused parameter stores gate and up interleaved as (E, F, 2, D); + # the grouped GEMM consumes them as one (E, 2F, D) expert weight. + w13_E_2F_D = w13.bfloat16().reshape(E, F * 2, D) gate_up_R2F = self._grouped_mm( - A=x_RD.bfloat16(), B_t=w13_E_D_2F, offs=offsets_E + A=x_RD.bfloat16(), weight_EOI=w13_E_2F_D, offs=offsets_E ) gate_RF, up_RF = gate_up_R2F.reshape(-1, F, 2).unbind(-1) h_RF = silu_and_mul_op(gate_RF, up_RF, offsets_E) - return self._grouped_mm( - A=h_RF, B_t=w2_EDF.bfloat16().transpose(-2, -1), offs=offsets_E - ).type_as(x_RD) + return self._grouped_mm(A=h_RF, weight_EOI=w2_EDF, offs=offsets_E).type_as(x_RD) @staticmethod def _split_w13_on_save(module, state_dict, prefix, local_metadata) -> None: