Repository navigation
MxFP8 gradient cross entropy option - #3758
Open
sraman-rgb wants to merge 9 commits into
Open
sraman-rgb wants to merge 9 commits into
sraman-rgb wants to merge 9 commits into
Conversation
… launch Add a grad_quantizer argument to parallel_cross_entropy. With an MXFP8Quantizer (E4M3 or E5M2, compact scales), a BF16 input on a single rank, and row count and vocabulary that are multiples of 128, one Triton kernel writes the derivative as an MXFP8 tensor with row-wise and column-wise data, bit-identical to quantizing the BF16 derivative with the quantizer. A following MXFP8 linear layer (an LM head) then skips its own grad-output quantization pass. Other cases compute the derivative as before and quantize it. Launch the single-rank forward with 4096-element blocks and 4 warps and the backward with 4096-element blocks and 8 warps instead of one 32-warp program with up to 32K elements per row, and find the argmax with a min-reduction over the block maximum instead of tl.max(return_indices=True). Only the accumulation order of the softmax statistics changes. Signed-off-by: Siddhartha Raman S <sraman@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… scale The NaN/Inf selects in the E8M0 exponent and reciprocal helpers slowed the fused MXFP8 backward from about 340 us to 420 us at 4096 x 129280 on a Rubin GPU. Finite gradients, the only ones a training step uses, still get the bytes of TE's MXFP8 cast; this is now stated in the docstring. Signed-off-by: Siddhartha Raman S <sraman@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Contributor
 [Medium impact] The PR appears safe to merge; no new actionable issues remain. SummaryAdds optional quantized gradients to
No new actionable issues found. Tests were not run. Diagram%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[Backward] --> B{grad_quantizer supplied?}
B -->|No| C[Existing gradient kernel]
B -->|Yes| D{Supported MXFP8 configuration?}
D -->|Yes| E[Fused gradient and MXFP8 kernel]
D -->|No| F[Existing gradient kernel]
F --> G[Quantize gradient]
Reviews (7) · Last reviewed commit: "Small testing refactors" · Reviewed by Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
…sum into NaN When every logit seen so far in a row is -inf, the online softmax shifted the exponentials by a running maximum of -inf, so exp(-inf - (-inf)) made the sum of exponentials NaN for the rest of the row. With 4096-wide forward blocks this happens for any row whose first block is masked. Shift by 0 until a finite maximum is seen, in the single-rank and tensor-parallel forward kernels and in the cross-rank merge. Rows with a finite maximum compute exactly as before; rows whose logits are all -inf still give NaN, as in PyTorch. Add tests with masked logits against PyTorch, including the MXFP8 gradient, CUDA graph replay and tensor parallelism. Signed-off-by: Siddhartha Raman S <sraman@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Resolve conflicts with the upstream NaN fix (e41b2be): keep its _safe_shift kernels and masked-logits tests, which supersede the equivalent changes on ksivamani/nan_fix. Keep from ksivamani/nan_fix the tensor-parallel fused MXFP8 backward, the removal of the dead contiguity check, the fused-kernel launch checks, the 128-column tile coverage and the tensor-parallel MXFP8 test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Member
|
/te-ci pytorch |
This branch has not been deployed
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
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.
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: