Fix backward return count mismatch in _Float8GroupedMM

[xiaobochen-amd]

No description provided.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3956

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 84d572b with merge base 4e18d87 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/quantization/pt2e/test_x86inductor_quantizer.py::TestQuantizePT2EX86Inductor::test_set_module_name_with_mixed_configs

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[danielvegamyhre]

surprised 1xh100 CI tests didn't catch this

[xiaobochen-amd]

Before my previous PR, the counts here were mismatched.

[danielvegamyhre]

i think i see the issue, we need to unskip this test:

ao/test/prototype/moe_training/test_scaled_grouped_mm.py

Line 58 in 4e18d87

@pytest.mark.skipif(

[xiaobochen-amd]

i think i see the issue, we need to unskip this test:

ao/test/prototype/moe_training/test_scaled_grouped_mm.py

Line 58 in 4e18d87

@pytest.mark.skipif(

When I tested locally, it was set to unskip and the test passed. I just noticed that #3788
hasn’t been fixed, so I left it as is.

[danielvegamyhre]

i think i see the issue, we need to unskip this test:

ao/test/prototype/moe_training/test_scaled_grouped_mm.py

Line 58 in 4e18d87

@pytest.mark.skipif(

When I tested locally, it was set to unskip and the test passed. I just noticed that #3788 hasn’t been fixed, so I left it as is.

I see, let me test on h100 if the issue still exists, i think it may have been a transient env/build issue we never root caused

[danielvegamyhre]

@xiaobochen-amd i tested and and the original cublas error is resolved, but there is a new error (numerics mismatch over the threshold). forward pass outputs are identical with torch.equal, however, gradients are slightly different, with most columns identical but some columns requiring atol/rtol=1 to pass.

I'll create an issue for this on the CUDA side. It doesn't block this, we can leave the test as skipped.