diff --git a/test/prototype/moe_training/test_mxfp8_backward_override.py b/test/prototype/moe_training/test_mxfp8_backward_override.py new file mode 100644 index 0000000000..365f594802 --- /dev/null +++ b/test/prototype/moe_training/test_mxfp8_backward_override.py @@ -0,0 +1,239 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import pytest +import torch + +from torchao.utils import is_sm_version + +if not (torch.cuda.is_available() and is_sm_version(10, 0)): + pytest.skip( + "MXFP8 grouped GEMM backward overrides require SM100", + allow_module_level=True, + ) + +pytest.importorskip("triton", reason="Triton required to run this test") + +from torchao.prototype.moe_training.config import MXFP8TrainingOpConfig +from torchao.prototype.moe_training.mxfp8_grouped_mm import ( + _SM100_KERNELS_AVAILABLE, + _compute_wgrad_sm100, + _to_mxfp8_then_scaled_grouped_mm, +) +from torchao.prototype.moe_training.utils import _quantize_then_scaled_grouped_mm +from torchao.prototype.mx_formats.config import ScaleCalculationMode + +if not _SM100_KERNELS_AVAILABLE: + pytest.skip( + "SM100 MXFP8 kernels (CUDA + Triton) unavailable", + allow_module_level=True, + ) + +# Group sizes must be multiples of 128 for the CuTe DSL 1x32 quantizer used +# by the stock forward, so every test allocates 128 rows per expert. +_ROWS_PER_EXPERT = 128 + + +def _make_inputs(num_experts=4, K=256, N=512, seed=0, tail_rows=0): + torch.manual_seed(seed) + M = num_experts * _ROWS_PER_EXPERT + A = torch.randn(M + tail_rows, K, dtype=torch.bfloat16, device="cuda") + if tail_rows: + A[M:] = 0 + B = torch.randn(num_experts, N, K, dtype=torch.bfloat16, device="cuda") * 0.1 + offs = torch.arange( + _ROWS_PER_EXPERT, + M + 1, + _ROWS_PER_EXPERT, + dtype=torch.int32, + device="cuda", + ) + return A, B, offs + + +def _reference_grads(A, B, offs, grad_output): + """bf16 reference gradients from the same two-grouped-GEMM formulation.""" + grad_output = grad_output.contiguous() + ref_grad_A = torch._grouped_mm(grad_output, B, offs=offs, out_dtype=torch.bfloat16) + ref_grad_B = torch._grouped_mm( + grad_output.transpose(-2, -1), A, offs=offs, out_dtype=torch.bfloat16 + ) + return ref_grad_A, ref_grad_B + + +def _run_override(A, B, offs, backward_override, grad_output=None): + """Run forward+backward through the op; returns (out, A.grad, B.grad).""" + A = A.detach().clone().requires_grad_(True) + B = B.detach().clone().requires_grad_(True) + out = _to_mxfp8_then_scaled_grouped_mm( + A, + B.transpose(-2, -1), + offs=offs, + backward_override=backward_override, + ) + if grad_output is None: + grad_output = torch.randn_like(out) + out.backward(grad_output) + return out, A.grad, B.grad + + +def test_forward_bitwise_identical_across_overrides(): + """The override must not change the forward: the Triton dim0 quantizer is + bitwise identical to the CuTe DSL 1x32 rceil kernel, so all three arms + consume identical GEMM operands.""" + A, B, offs = _make_inputs() + outputs = {} + for override in (None, "high_precision", "dequantized"): + out = _to_mxfp8_then_scaled_grouped_mm( + A.detach(), + B.detach().transpose(-2, -1), + offs=offs, + backward_override=override, + ) + outputs[override] = out + assert torch.equal(outputs[None], outputs["high_precision"]) + assert torch.equal(outputs[None], outputs["dequantized"]) + # "quantized" is an accepted alias for the default backward. + out_quantized = _to_mxfp8_then_scaled_grouped_mm( + A.detach(), + B.detach().transpose(-2, -1), + offs=offs, + backward_override="quantized", + ) + assert torch.equal(outputs[None], out_quantized) + + +def test_high_precision_backward_matches_grouped_mm_reference(): + A, B, offs = _make_inputs() + grad_output = torch.randn( + A.shape[0], B.shape[1], dtype=torch.bfloat16, device="cuda" + ) + _, grad_A, grad_B = _run_override(A, B, offs, "high_precision", grad_output) + ref_grad_A, ref_grad_B = _reference_grads(A, B, offs, grad_output) + # Same two torch._grouped_mm calls on the same operands: bitwise. + assert torch.equal(grad_A, ref_grad_A) + assert torch.equal(grad_B, ref_grad_B) + + +def test_dequantized_backward_close_to_reference_and_distinct(): + A, B, offs = _make_inputs() + grad_output = torch.randn( + A.shape[0], B.shape[1], dtype=torch.bfloat16, device="cuda" + ) + _, grad_A, grad_B = _run_override(A, B, offs, "dequantized", grad_output) + assert torch.isfinite(grad_A).all() + assert torch.isfinite(grad_B).all() + + ref_grad_A, ref_grad_B = _reference_grads(A, B, offs, grad_output) + # Loose tolerance: the operands are round-tripped through MXFP8 (e4m3 + # data with e8m0 block scales), so gradients carry one quantization + # error per operand — observed ~5% max relative error at these shapes; + # 0.15 is a deliberately loose bound to avoid seed sensitivity. + for grad, ref in ((grad_A, ref_grad_A), (grad_B, ref_grad_B)): + rel_err = ((grad - ref).abs().max() / ref.abs().max()).item() + assert rel_err < 0.15, f"relative error too large: {rel_err}" + + # Sanity: the dequantized backward consumed quantized operands, so its + # gradients differ from the high-precision ones. + _, hp_grad_A, hp_grad_B = _run_override(A, B, offs, "high_precision", grad_output) + assert not torch.equal(grad_A, hp_grad_A) + assert not torch.equal(grad_B, hp_grad_B) + + +@pytest.mark.parametrize("backward_override", ["high_precision", "dequantized"]) +def test_tail_rows_past_final_offset_tolerated(backward_override): + """Token dispatchers may over-allocate activation rows past offs[-1]; + the grouped GEMMs only read rows covered by the offsets.""" + tail_rows = 128 + A, B, offs = _make_inputs(tail_rows=tail_rows) + M_logical = int(offs[-1]) + grad_output = torch.randn( + A.shape[0], B.shape[1], dtype=torch.bfloat16, device="cuda" + ) + out, grad_A, grad_B = _run_override(A, B, offs, backward_override, grad_output) + assert out.shape[0] == A.shape[0] + ref_grad_A, ref_grad_B = _reference_grads( + A[:M_logical], B, offs, grad_output[:M_logical] + ) + # Rows past offs[-1] belong to no group and may be uninitialized, so + # only the logical rows are checked. + assert torch.isfinite(grad_A[:M_logical]).all() + assert torch.isfinite(grad_B).all() + if backward_override == "high_precision": + assert torch.equal(grad_A[:M_logical], ref_grad_A) + assert torch.equal(grad_B, ref_grad_B) + else: + for grad, ref in ((grad_A[:M_logical], ref_grad_A), (grad_B, ref_grad_B)): + rel_err = ((grad - ref).abs().max() / ref.abs().max()).item() + assert rel_err < 0.15, f"relative error too large: {rel_err}" + + +def test_noncontiguous_grad_output_wgrad(): + """The CUDA dim1 cast in the quantized wgrad requires contiguous inputs; + expanded (stride-0) and other non-contiguous gradient views must not + crash it.""" + A, B, offs = _make_inputs() + N = B.shape[1] + M = A.shape[0] + + # Direct wgrad call with an expanded (stride-0) grad_output. + grad_output_expanded = torch.randn( + M, 1, dtype=torch.bfloat16, device="cuda" + ).expand(M, N) + grad_weight_t = _compute_wgrad_sm100( + grad_output_expanded, + A, + offs, + 32, + torch.bfloat16, + ScaleCalculationMode.RCEIL, + wgrad_with_hp=False, + ) + assert torch.isfinite(grad_weight_t).all() + + # End-to-end: sum().backward() feeds an expanded ones gradient into the + # default quantized backward. + A_leaf = A.detach().clone().requires_grad_(True) + B_leaf = B.detach().clone().requires_grad_(True) + out = _to_mxfp8_then_scaled_grouped_mm(A_leaf, B_leaf.transpose(-2, -1), offs=offs) + out.sum().backward() + assert torch.isfinite(A_leaf.grad).all() + assert torch.isfinite(B_leaf.grad).all() + + +def test_invalid_backward_override_rejected(): + A, B, offs = _make_inputs(num_experts=1, K=128, N=128) + with pytest.raises(AssertionError, match="backward_override"): + _to_mxfp8_then_scaled_grouped_mm( + A.detach(), + B.detach().transpose(-2, -1), + offs=offs, + backward_override="hp", + ) + + +def test_config_threads_backward_override(): + """backward_override must participate in config equality/hashing (no + silent aliasing) and reach the grouped GEMM through the op config.""" + default_config = MXFP8TrainingOpConfig() + hp_config = MXFP8TrainingOpConfig(backward_override="high_precision") + assert default_config != hp_config + assert hash(default_config) != hash(hp_config) + assert hp_config == MXFP8TrainingOpConfig(backward_override="high_precision") + + A, B, offs = _make_inputs() + grad_output = torch.randn( + A.shape[0], B.shape[1], dtype=torch.bfloat16, device="cuda" + ) + A_leaf = A.detach().clone().requires_grad_(True) + B_leaf = B.detach().clone().requires_grad_(True) + out = _quantize_then_scaled_grouped_mm( + A_leaf, B_leaf.transpose(-2, -1), config=hp_config, offs=offs + ) + out.backward(grad_output) + ref_grad_A, ref_grad_B = _reference_grads(A, B, offs, grad_output) + assert torch.equal(A_leaf.grad, ref_grad_A) + assert torch.equal(B_leaf.grad, ref_grad_B) diff --git a/torchao/prototype/moe_training/config.py b/torchao/prototype/moe_training/config.py index e60c7f7682..356c7e7890 100644 --- a/torchao/prototype/moe_training/config.py +++ b/torchao/prototype/moe_training/config.py @@ -170,6 +170,12 @@ class MXFP8TrainingOpConfig(TrainingOpBaseConfig): # Whether to pad the token group sizes to multiples of 32 (MXFP8 scaling block size). pad_token_groups_for_grouped_mm: bool = False + # Backward computation override for the grouped GEMM. None or "quantized" uses the + # quantized MXFP8 backward (default). "high_precision" computes both gradients with + # plain grouped GEMMs on the saved high-precision operands. "dequantized" computes + # them from the dequantized forward operands. + backward_override: Optional[str] = None + @classmethod def from_recipe( cls, @@ -212,6 +218,7 @@ def __eq__(self, other): and self.scale_calculation_mode == other.scale_calculation_mode and self.pad_token_groups_for_grouped_mm == other.pad_token_groups_for_grouped_mm + and self.backward_override == other.backward_override ) return NotImplemented @@ -223,6 +230,7 @@ def __hash__(self): self.wgrad_with_hp, self.scale_calculation_mode, self.pad_token_groups_for_grouped_mm, + self.backward_override, ) ) diff --git a/torchao/prototype/moe_training/mxfp8_grouped_mm.py b/torchao/prototype/moe_training/mxfp8_grouped_mm.py index eaf3a1c5af..d6567e7df7 100644 --- a/torchao/prototype/moe_training/mxfp8_grouped_mm.py +++ b/torchao/prototype/moe_training/mxfp8_grouped_mm.py @@ -17,6 +17,7 @@ mxfp8_quantize_2d_1x32_cutedsl, mxfp8_quantize_cuda_3d, triton_mx_block_rearrange_2d_K_groups, + triton_mx_block_rearrange_2d_M_groups, triton_mx_block_rearrange_per_group_3d, ) from torchao.prototype.moe_training.utils import ( @@ -63,6 +64,7 @@ def _to_mxfp8_then_scaled_grouped_mm( wgrad_with_hp: bool = False, scale_calculation_mode: ScaleCalculationMode = ScaleCalculationMode.RCEIL, pad_token_groups_for_grouped_mm: bool = False, + backward_override: Optional[str] = None, ) -> torch.Tensor: """ Differentiable mxfp8 grouped gemm with dynamic mxfp8 quantization. @@ -82,6 +84,10 @@ def _to_mxfp8_then_scaled_grouped_mm( wgrad_with_hp (bool): Whether to compute weight gradient in high precision. Defaults to False. scale_calculation_mode (ScaleCalculationMode): Mode for scale calculation (RCEIL, FLOOR, etc.). Defaults to ScaleCalculationMode.RCEIL. pad_token_groups_for_grouped_mm (bool): Whether to pad token groups to the next multiple of 32 (requirement for MXFP8 grouped GEMM). If your tokens are already padded, set to False. + backward_override (Optional[str]): Backward computation override. None or "quantized" uses the + quantized MXFP8 backward (default). "high_precision" saves the high-precision operands and + computes both gradients with plain grouped GEMMs. "dequantized" saves the quantized forward + operands and computes both gradients from their dequantized values. Returns: out (torch.Tensor): The result of the mxfp8 scaled grouped gemm. @@ -98,6 +104,7 @@ def _to_mxfp8_then_scaled_grouped_mm( wgrad_with_hp, scale_calculation_mode, pad_token_groups_for_grouped_mm, + backward_override, ) # add bias outside the autograd function so that autograd @@ -128,6 +135,7 @@ def forward( wgrad_with_hp: bool = False, scale_calculation_mode: ScaleCalculationMode = ScaleCalculationMode.RCEIL, pad_token_groups_for_grouped_mm: bool = False, + backward_override: Optional[str] = None, ) -> torch.Tensor: """ Forward pass: Quantize inputs and perform grouped GEMM. @@ -141,6 +149,9 @@ def forward( wgrad_with_hp: Compute weight gradient in high precision scale_calculation_mode: Mode for scale calculation (RCEIL, FLOOR, etc.) pad_token_groups_for_grouped_mm: Whether to pad token groups to the next multiple of 32 + backward_override: None/"quantized" for the quantized backward, "high_precision" + to compute gradients from the saved high-precision operands, or "dequantized" + to compute them from the dequantized forward operands Returns: Output tensor, shape (M, N) @@ -153,6 +164,37 @@ def forward( KernelPreference.EMULATED, ), "kernel_preference must be AUTO or EMULATED" + assert backward_override in ( + None, + "quantized", + "high_precision", + "dequantized", + ), ( + "backward_override must be None, 'quantized', 'high_precision', or " + f"'dequantized', got {backward_override!r}" + ) + # None and "quantized" both select the quantized backward default. + if backward_override == "quantized": + backward_override = None + if backward_override is not None: + assert not isinstance(input_act, MXTensor), ( + "backward_override requires high-precision input activations" + ) + # The override backwards quantize saved operands with the bf16-only + # dim0 Triton cast, so reject other dtypes here rather than deep + # inside the kernel. + assert ( + input_act.dtype == torch.bfloat16 and weight_t.dtype == torch.bfloat16 + ), ( + "backward_override supports bfloat16 operands only, got " + f"{input_act.dtype} and {weight_t.dtype}" + ) + if backward_override == "dequantized": + assert kernel_preference != KernelPreference.EMULATED, ( + "the dequantized backward saves quantized operands from the SM100 " + "forward and does not support kernel_preference=EMULATED" + ) + # Validate SM100 kernels are available if not using emulated mode if kernel_preference != KernelPreference.EMULATED: assert _SM100_KERNELS_AVAILABLE, ( @@ -198,15 +240,33 @@ def forward( padded_group_end_offsets = group_end_offsets # Perform forward computation using appropriate path - output = _compute_fwd( - padded_input_act, - weight_t, - padded_group_end_offsets, - block_size, - out_dtype, - scale_calculation_mode, - kernel_preference, - ) + if backward_override == "dequantized": + # The dequantized backward needs the quantized operands with + # logical scales, which this forward variant also returns. + ( + output, + input_act_e4m3, + input_act_scales, + weight_e4m3, + weight_scales, + ) = _compute_fwd_sm100_logical_scales( + padded_input_act, + weight_t, + padded_group_end_offsets, + block_size, + out_dtype, + scale_calculation_mode, + ) + else: + output = _compute_fwd( + padded_input_act, + weight_t, + padded_group_end_offsets, + block_size, + out_dtype, + scale_calculation_mode, + kernel_preference, + ) # Unpad output if padding was used if pad_token_groups_for_grouped_mm: @@ -220,13 +280,31 @@ def forward( ) # Save tensors and config for backward - ctx.save_for_backward( - padded_input_act, - weight_t, - group_end_offsets, - padded_group_start_offsets, - padded_group_end_offsets, - ) + if backward_override == "high_precision": + # The high-precision backward only needs the unpadded operands. + ctx.save_for_backward(input_act, weight_t, group_end_offsets) + elif backward_override == "dequantized": + if padded_group_start_offsets is None: + padded_group_start_offsets = group_end_offsets.new_zeros(0) + ctx.save_for_backward( + input_act_e4m3, + input_act_scales, + weight_e4m3, + weight_scales, + group_end_offsets, + padded_group_start_offsets, + ) + ctx.input_hp_dtype = padded_input_act.dtype + ctx.weight_hp_dtype = weight_t.dtype + else: + ctx.save_for_backward( + padded_input_act, + weight_t, + group_end_offsets, + padded_group_start_offsets, + padded_group_end_offsets, + ) + ctx.backward_override = backward_override ctx.out_dtype = out_dtype ctx.kernel_preference = kernel_preference ctx.wgrad_with_hp = wgrad_with_hp @@ -249,6 +327,65 @@ def backward(ctx, grad_output: torch.Tensor): Returns: tuple: (grad_input, grad_weight_t, None, ...) matching forward args """ + # block_size is always 32 for MXFP8 + block_size = 32 + + # The backward overrides compute both gradients with plain grouped + # GEMMs on high-precision operands instead of the quantized backward. + if ctx.backward_override == "high_precision": + input_act, weight_t, group_end_offsets = ctx.saved_tensors + grad_input, grad_weight_t = _compute_grads_from_hp_operands( + grad_output, + input_act, + weight_t.transpose(-2, -1), + group_end_offsets, + ctx.out_dtype, + ) + return _backward_override_grads_tuple(grad_input, grad_weight_t) + elif ctx.backward_override == "dequantized": + ( + input_act_e4m3, + input_act_scales, + weight_e4m3, + weight_scales, + group_end_offsets, + padded_group_start_offsets, + ) = ctx.saved_tensors + input_act = triton_mxfp8_dequant_dim0( + input_act_e4m3, + input_act_scales.view( + torch.uint8 + ), # Triton can't handle e8m0 directly yet + out_dtype=ctx.input_hp_dtype, + scale_block_size=block_size, + ) + if ctx.pad_token_groups_for_grouped_mm: + input_act = unpad_token_groups( + input_act, + group_end_offsets, + padded_group_start_offsets, + ctx.num_tokens, + alignment_size=block_size, + kernel_preference=ctx.kernel_preference, + ) + num_experts, weight_n, weight_k = weight_e4m3.shape + weight = triton_mxfp8_dequant_dim0( + weight_e4m3.reshape(num_experts * weight_n, weight_k), + weight_scales.view(torch.uint8).reshape( + num_experts * weight_n, weight_k // block_size + ), + out_dtype=ctx.weight_hp_dtype, + scale_block_size=block_size, + ).reshape(num_experts, weight_n, weight_k) + grad_input, grad_weight_t = _compute_grads_from_hp_operands( + grad_output, + input_act, + weight, + group_end_offsets, + ctx.out_dtype, + ) + return _backward_override_grads_tuple(grad_input, grad_weight_t) + # Retrieve saved tensors and config ( padded_input_act, @@ -257,9 +394,6 @@ def backward(ctx, grad_output: torch.Tensor): padded_group_start_offsets, padded_group_end_offsets, ) = ctx.saved_tensors - - # block_size is always 32 for MXFP8 - block_size = 32 out_dtype = ctx.out_dtype kernel_preference = ctx.kernel_preference wgrad_with_hp = ctx.wgrad_with_hp @@ -324,9 +458,69 @@ def backward(ctx, grad_output: torch.Tensor): None, # wgrad_with_hp None, # scale_calculation_mode None, # pad_token_groups_for_grouped_mm + None, # backward_override ) +def _backward_override_grads_tuple( + grad_input: torch.Tensor, + grad_weight_t: torch.Tensor, +) -> tuple: + """Pad the override gradients with Nones for forward's non-tensor args.""" + return ( + grad_input, + grad_weight_t, + None, # group_end_offsets + None, # out_dtype + None, # kernel_preference + None, # wgrad_with_hp + None, # scale_calculation_mode + None, # pad_token_groups_for_grouped_mm + None, # backward_override + ) + + +def _compute_grads_from_hp_operands( + grad_output: torch.Tensor, + input_act: torch.Tensor, + weight: torch.Tensor, + group_end_offsets: torch.Tensor, + out_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Compute both gradients with plain grouped GEMMs on high-precision operands. + + Shared by the "high_precision" and "dequantized" backward overrides. The + grouped GEMMs only read rows covered by group_end_offsets, so + over-allocated tail rows past the final offset are tolerated. + + Args: + grad_output: Gradient output, shape (M, N) + input_act: High-precision input activations, shape (M, K) + weight: High-precision expert weights, shape (E, N, K) + group_end_offsets: End index of each token group, shape (E,) + out_dtype: Output dtype for the gradients + + Returns: + tuple: (grad_input, grad_weight_t) with shapes (M, K) and (E, K, N) + """ + grad_output = grad_output.contiguous() + grad_input = torch._grouped_mm( + grad_output, + weight, + offs=group_end_offsets, + out_dtype=out_dtype, + ) + grad_weight = torch._grouped_mm( + grad_output.transpose(-2, -1), + input_act, + offs=group_end_offsets, + out_dtype=out_dtype, + ) + # Transpose to match weight_t shape in forward: (E, N, K) -> (E, K, N) + return grad_input, grad_weight.transpose(-2, -1) + + def _compute_fwd( padded_input_act: torch.Tensor, weight_t: torch.Tensor, @@ -549,6 +743,65 @@ def _compute_fwd_sm100( return output +def _compute_fwd_sm100_logical_scales( + padded_input_act: torch.Tensor, + weight_t: torch.Tensor, + padded_group_end_offsets: torch.Tensor, + block_size: int, + out_dtype: torch.dtype, + scale_calculation_mode: ScaleCalculationMode, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Forward computation using AUTO path (SM100 kernels), also returning the + quantized operands with logical (pre-blocking) scales. + + Used when the backward needs to dequantize the forward operands. The + Triton dim0 quantizer produces qdata bitwise identical to the CuTe DSL + 1x32 rceil kernel used by `_compute_fwd_sm100`, but returns logical + scales (the CuTe DSL kernel returns blocked-only scales and no + unswizzler exists), so this path quantizes input activations with Triton + and blocks the scales separately for the grouped GEMM. + + Args: + padded_input_act: Input activations (possibly padded), shape (M, K) + weight_t: Expert weights transposed, shape (E, K, N) + padded_group_end_offsets: Group offsets (possibly padded) + block_size: Block size for quantization + out_dtype: Output dtype + scale_calculation_mode: Mode for scale calculation + + Returns: + tuple: (output, input_act_e4m3, input_act_scales, weight_e4m3, weight_scales) + - output: shape (M, N) + - input_act_e4m3: shape (M, K) with logical scales (M, K//block_size) + - weight_e4m3: shape (E, N, K) with logical scales (E, N, K//block_size) + """ + # Quantize input activations along dim0, keeping the logical scales + input_act_e4m3, input_act_scales = triton_to_mxfp8_dim0( + padded_input_act, block_size, scale_calculation_mode.value.lower() + ) + input_act_scales_blocked = triton_mx_block_rearrange_2d_M_groups( + input_act_scales, padded_group_end_offsets + ) + + # Quantize weights along dim0 (after transposing from (E, K, N) to (E, N, K)) + weight_e4m3, weight_scales = triton_to_mxfp8_dim0( + weight_t.transpose(-2, -1), block_size, scale_calculation_mode.value.lower() + ) + weight_scales_blocked = triton_mx_block_rearrange_per_group_3d(weight_scales) + + # Compute output using SM100 kernel + output = torch._scaled_grouped_mm( + input_act_e4m3, + weight_e4m3.transpose(-2, -1), # Transpose back to (E, K, N) + input_act_scales_blocked, + weight_scales_blocked, + offs=padded_group_end_offsets, + out_dtype=out_dtype, + ) + return output, input_act_e4m3, input_act_scales, weight_e4m3, weight_scales + + def _compute_fwd_emulated( padded_input_act: torch.Tensor, weight_t: torch.Tensor, @@ -625,8 +878,12 @@ def _compute_dgrad_sm100( grad_output.scales, group_end_offsets ) else: + # As in the wgrad path: autograd can hand backward an expanded or + # otherwise non-contiguous grad_output view (sum().backward()'s + # stride-0 ones), and the CuTe DSL quantize kernel's handling of such + # strides is unverified — materialize first (no-op when contiguous). grad_out_e4m3, grad_output_scales_blocked = mxfp8_quantize_2d_1x32_cutedsl( - grad_output, + grad_output.contiguous(), scaling_mode=scale_calculation_mode.value.lower(), offs=group_end_offsets, ) @@ -747,6 +1004,11 @@ def _compute_wgrad_sm100( return grad_weight.transpose(-2, -1) # Use CUDA kernel for dim1 quant + # The kernel requires contiguous inputs; autograd can hand backward an + # expanded or otherwise non-contiguous grad_output view (e.g. from + # sum().backward()), so make both operands contiguous first. + grad_output = grad_output.contiguous() + input_act = input_act.contiguous() grad_output_t_mx = _to_mxfp8_dim1_kernel_wrapper( grad_output, block_size, diff --git a/torchao/prototype/moe_training/utils.py b/torchao/prototype/moe_training/utils.py index 9ba6fe78a3..32c5c13f97 100644 --- a/torchao/prototype/moe_training/utils.py +++ b/torchao/prototype/moe_training/utils.py @@ -396,6 +396,7 @@ def _quantize_then_scaled_grouped_mm( "wgrad_with_hp": config.wgrad_with_hp, "scale_calculation_mode": config.scale_calculation_mode, "pad_token_groups_for_grouped_mm": config.pad_token_groups_for_grouped_mm, + "backward_override": config.backward_override, } if bias is not None: kwargs["bias"] = bias