Add mixed-precision (FP32) LM output logits - #6252
Conversation
Signed-off-by: mkhona <mkhona@nvidia.com>
Signed-off-by: mkhona <mkhona@nvidia.com>
Signed-off-by: mkhona <mkhona@nvidia.com>
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
/claude review |
|
Light review — one substantive concern, otherwise LGTM. Default
The fast path in
Since the stated goal is opt-in FP32 logits, a dtype-neutral default (e.g. Could you confirm the fp16/fp32 training paths were considered, or add a guard so the default does not alter those runs? |
Signed-off-by: mkhona <mkhona@nvidia.com>
|
/ok to test 2800628 |
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
Explicitly set inference logits dtype
|
|
||
| # Explicitly cast logits to expected dtype | ||
| logits = logits.to(self.config.params_dtype) | ||
| logits = logits.to(self.logit_dtype) |
There was a problem hiding this comment.
Would this be a no-op with fp32 logit dtype?
|
/ok to test c2ca628 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31563547355 |
11 participants
What does this PR do?
Allow the language model to compute
output_layerwith logits dtype leaving the accumulator in FP32, rather than casting down to BF16 as default. We use TE'sgeneral_gemmAPI to do this and modify autograd to cast the gradients back down to BF16, following Megatron-LM's contract with TE for GEMMs.The legacy path performs FP32 accumulation internally but rounds every vocabulary logit to BF16 before cross entropy. Casting those rounded values back to FP32 cannot recover the discarded bits. Requesting an FP32 GEMM output preserves the accumulator result and materially reduces error at every tested reduction width.
The explicit backward cast does not introduce a new gradient contract: it reproduces the cast that autograd already performs through the legacy post-GEMM
logits.float()operation.Tests
[-8, 8]withK=512. SinceK * 8^2 = 32,768 < 2^24, every product and partial sum is exactly representable in FP32; the test compares raw FP32 words against a PyTorch FP32 reference with no tolerance.transformer_engine.pytorch.cpp_extensions.general_gemmcall configured with the same FP32 output contract, requiring bitwise-identical outputs.2825327).2825327).torch.matmulagainst Transformer Enginegeneral_gemmwith an FP32 output over 15 shapes (Min 16, 128, 1,024;Kin 1,024–16,384;N=32,768). The median overhead was 5.4% overall, 2.2% forK >= 4,096, and 1.9% forK >= 8,192; 11/15 shapes were within 10% and the full range was 0.3%–15.9% (NVIDIA GB300, Slurm job2825411).--output-logit-dtypedefaults tobf16, accepts onlybf16andfp32, and rejectsfp16.GEMM kernel benchmark
The benchmark uses preallocated output buffers, 10 warmup pairs, 30 alternating-order CUDA-event samples per kernel, and reports medians. The vocabulary dimension is 32,768 for every shape. The larger-width results show that preserving FP32 GEMM output has low kernel-level overhead for representative LM heads.
Integer bitwise exactness results
Both operands are BF16 tensors containing integers in
[-8, 8]. For every width,K * 8^2 < 2^24, so the products and possible partial sums are exactly representable in FP32. The raw FP32 output words matched the FP32 PyTorch reference with no tolerance at every width (NVIDIA GB300).Method
general_gemmwithout_dtype=torch.float32.16 x Kactivations,32,768 x Kvocabulary weights, K in512,1024,2048,4096,8192,16384.1234; commit `27ef2b465.Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.