Conversation
The forward projected hidden states with a low-precision GEMM and then cast the rounded result to FP32, discarding the FP32 accumulator cuBLAS already computed. JSD gradients difference near-equal distributions, so the rounding cancels catastrophically (4-23% grad error in bf16). Use torch.mm(..., out_dtype=torch.float32) where supported (torch>=2.8, CUDA sm_80+, fp16/bf16) to keep the accumulator at no cost, with an explicit FP32 GEMM fallback elsewhere. Apply the same FP32 projection to the Ascend and cuTile forwards. Fix the test oracle, which shared the same cast-after-matmul defect, and switch its degenerate rand inputs to randn so the bug is no longer masked. Signed-off-by: fei <204683769+feiiiiii5@users.noreply.github.com>
Author
|
Closing as duplicate of #1433: same issue (#1432), same fix direction, and the reporter's own PR has priority. One thing #1433 may want as a follow-up: the same round-then-cast pattern exists in the ascend and cutile backends of this kernel, which my branch had covered. Happy to help verify those if useful. |
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
1 participant
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.
Fixes #1432.
Root cause: forward projected hidden states with low-precision GEMM then cast rounded result to FP32, discarding FP32 accumulator. JSD gradients difference near-equal distributions so rounding cancels catastrophically (4-23% grad error in bf16). Test oracle shared same defect with degenerate rand inputs masking it.
Fix: use torch.mm(..., out_dtype=torch.float32) where supported with explicit FP32 GEMM fallback. Same fix to Ascend and cuTile forwards. Fix oracle + rand to randn.
Tests: py_compile clean, needs GPU CI for numerical verification.