Add Autograd Backward Support for FP4 GEMM Custom Ops - #393
Conversation
|
I put the FP4 autograd logic in its own file rather than inlining into I put all of the dequantization and gradient logic in the same file, but this felt a bit awkward. Please let me know if I should organize this differently. Thanks! |
jwfromm
left a comment
jwfromm
left a comment
There was a problem hiding this comment.
Overall looks great, thank you for the contribution! I think youll just need to rebase and do a little bit of merge cleanup and we can land.
@jwfromm Thanks! I made the necessary changes. If there is anything else, just lmk. |
|
@colalb1 could you rebase this PR onto latest main? |
|
@q10 has imported this pull request. If you are a Meta employee, you can view this in D113691571. |
Summary: Pull Request resolved: meta-pytorch#473 Adds backward pass support for the three FP4 GEMM custom ops (`mslk::f4f4bf16`, `mslk::f4f4bf16_grouped_mm`, `mslk::f4f4bf16_ultra_grouped_mm`) so that `.backward()` works through them during FP4 training. Previously these operations were forward-only. **Implementation Method:** Dequantize-then-compute in BF16. FP4 (E2M1, only 8 representable magnitudes) cannot carry gradient signal, so the backward dequantizes saved packed FP4 inputs to BF16 and computes standard matmul gradients: - `grad_X = grad_output @ W_bf16` - `grad_W = grad_output.T @ X_bf16` This matches the FP8‑forward / BF16‑backward pattern in the repository. ## Files changed - `mslk/gemm/fp4_autograd.py` (new): Dequantization helpers + `register_autograd` for all 3 ops - `mslk/gemm/__init__.py`: Import `fp4_autograd` to trigger registration - `test/gemm/gemm_test.py`: `MXFP4BackwardTests` and `NVFP4BackwardTests` classes Pull Request resolved: meta-pytorch#393 Test Plan: Ran on 2x NVIDIA GB200 (SM 10.0, CUDA 12.8), `@//mode/opt`: ``` buck2 build @//mode/opt fbcode//mslk/test/gemm:gemm_test gemm_test.par -r '(MXFP4BackwardTests|NVFP4BackwardTests)' -> Ran 19 tests ... OK (skipped=1) ``` The single skip is `test_ultra_grouped_mm_wrapper_forward_matches_op` ("requires CUDA toolkit 13+"); this host is on CUDA 12.8, so the ultra grouped op itself is not exercised here. Full `gemm_test` suite: 237 tests, 8 failures + 7 errors. All 15 are pre-existing on this host and reproduce identically at the parent commit (37e09ecf731d: 218 tests, same 8 failures + 7 errors, same test names) — they are in FP8Tests / FP8GroupwiseTests / MXFP4Tests / NVFP4Tests and are unrelated to this change. Mutation-tested the backward math to confirm the tests are not vacuous. Injecting a deliberate error into `_F4F4BF16.backward` and `_F4F4BF16GroupedMM.backward` fails 10 of 19 tests; before this revision the same mutation failed only 2 of 17, because `test_backward` and `test_gradient_flow_nonzero` never called the code under test. Reviewed By: jwfromm Differential Revision: D113691571 Pulled By: q10
3 participants
Adds backward pass support for the three FP4 GEMM custom ops (
mslk::f4f4bf16,mslk::f4f4bf16_grouped_mm,mslk::f4f4bf16_ultra_grouped_mm) so that.backward()works through them during FP4 training. Previously these operations were forward-only.Implementation Method:
Dequantize-then-compute in BF16. FP4 (E2M1, only 8 representable magnitudes) cannot carry gradient signal, so the backward dequantizes saved packed FP4 inputs to BF16 and computes standard matmul gradients:
grad_X = grad_output @ W_bf16grad_W = grad_output.T @ X_bf16This matches the FP8‑forward / BF16‑backward pattern in the repository.
Files changed
mslk/gemm/fp4_autograd.py(new): Dequantization helpers +register_autogradfor all 3 opsmslk/gemm/__init__.py: Importfp4_autogradto trigger registrationtest/gemm/gemm_test.py:MXFP4BackwardTestsandNVFP4BackwardTestsclasses