Fix backward return count mismatch in _Float8GroupedMM (#3956)

xiaobochen-amd · web-flow · commit 9bdc0ca87c11 · 2026-02-26T21:12:59.000-08:00
diff --git a/torchao/prototype/moe_training/fp8_grouped_mm.py b/torchao/prototype/moe_training/fp8_grouped_mm.py
@@ -254,4 +254,4 @@ def backward(ctx, grad_output: torch.Tensor):
             out_dtype=out_dtype,
             use_fast_accum=True,
         )
-        return grad_A, grad_B.transpose(-2, -1), None, None, None, None
+        return grad_A, grad_B.transpose(-2, -1), None, None, None

Original file line number	Diff line number	Diff line change
`@@ -254,4 +254,4 @@ def backward(ctx, grad_output: torch.Tensor):`
`254`	`254`	`out_dtype=out_dtype,`
`255`	`255`	`use_fast_accum=True,`
`256`	`256`	`)`
`257`		`- return grad_A, grad_B.transpose(-2, -1), None, None, None, None`
	`257`	`+ return grad_A, grad_B.transpose(-2, -1), None, None, None`