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 a16051e96c..6f4b919ed5 100644 --- a/torchtitan/models/common/moe.py +++ b/torchtitan/models/common/moe.py @@ -96,32 +96,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: