Add _MXFP8GroupedMM variant with bf16 backward - #4649
Draft
andrewor14 wants to merge 1 commit into
Draft
andrewor14 wants to merge 1 commit into
andrewor14 wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4649
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 344526e with merge base f10e69b ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
Contributor
|
why do we need cutedsl 4.6.0? can we use 4.5.2? |
Contributor
Author
It's because vllm nightly (since 7/16/26) has 4.6.0 has a hard dependency, and I feel it's unlikely they'll downgrade. torchtitan RL is expected to work with vllm nightly |
vkuzo
reviewed
Aug 3, 2026
| out_dtype: Optional[torch.dtype] = torch.bfloat16 | ||
|
|
||
| # Whether to compute the gradient of the weights in high precision (True) or use MXFP8 (False). | ||
| wgrad_with_hp: bool = False |
Contributor
There was a problem hiding this comment.
@andrewor14 , i think instead of bf16_bwd you can set this flag to true and also add dgrad_with_hp to get the same effect but matching the existing code better.
Contributor
Author
|
Going to close this to reopen with ghstack |
**Summary:** Add `_MXFP8GroupedMMFwdBF16Bwd`, which does real MXFP8 forward but BF16 backward backward (straight-through), for QAT or RL trainer where a low-precision backward isn't needed. Selectable via a new `bf16_bwd` flag on `_to_mxfp8_then_scaled_grouped_mm` and `MXFP8TrainingOpConfig`. Also add KernelPreference.TRITON, which runs pure triton kernels for the forward quantization to work around an mxfp8 cutedsl issue (see #4647). Example usage: ``` from torchao.prototype.moe_training.config import MXFP8TrainingOpConfig from torchao.prototype.mx_formats.config import ScaleCalculationMode from torchao.quantization import quantize_ from torchao.quantization.quantize_.common import KernelPreference from torchtitan.models.common.moe import GroupedExperts config = MXFP8TrainingOpConfig( kernel_preference=KernelPreference.TRITON, scale_calculation_mode=ScaleCalculationMode.RCEIL, bf16_bwd=True, ) quantize_(model, config=config, filter_fn=lambda m, fqn: isinstance(m, GroupedExperts)) ``` **Test Plan:** ``` pytest test/prototype/moe_training/test_mxfp8_grouped_mm.py -k bf16_bwd_forward_matches_default pytest test/prototype/moe_training/test_mxfp8_grouped_mm.py -k bf16_bwd_triton_matches_emulated ``` Real e2e RL tests TBD.
andrewor14
force-pushed
the
mxfp8_fwd_bf16_bwd
branch
2 times, most recently
from
August 17, 2026 19:43
0814732 to
344526e
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
Add
_MXFP8GroupedMMFwdBF16Bwd, which does real MXFP8 forward but BF16 backward backward (straight-through), for QAT or RL trainer where a low-precision backward isn't needed. Selectable via a newbf16_bwdflag on_to_mxfp8_then_scaled_grouped_mmandMXFP8TrainingOpConfig. Also add KernelPreference.TRITON, which runs pure triton kernels for the forward quantization to work around an mxfp8 cutedsl issue (see #4647).Example usage:
Test Plan:
Real e2e RL tests TBD.