[Pallas] Run kl_div at fp32 on the TPU bridge to match run_tpu.py - #3208
Open
norx1991 wants to merge 1 commit into
Open
[Pallas] Run kl_div at fp32 on the TPU bridge to match run_tpu.py#3208norx1991 wants to merge 1 commit into
norx1991 wants to merge 1 commit into
Conversation
The bridge forced --precision bf16 for every bridged kernel, but run_tpu.py (and GPU) run kl_div at fp32 — its loss reduction stays fp32. Forcing bf16 made the bridge diverge from run_tpu for kl_div and produced an fp32-impl vs bf16-baseline accuracy mismatch (torch.testing.assert_close fails on dtype). Set precision per kernel: fp32 for kl_div, bf16 for the rest (TPU-native, and matching run_tpu.py). Fixes kl_div accuracy on the bridge at the source, with no change to the accuracy comparison.
meta-cla
Bot
added
the
CLA Signed
norx1991
marked this pull request as ready for review
July 30, 2026 21:41
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.
The TPU bridge forced
--precision bf16for every bridged kernel, butrun_tpu.py(and GPU) runkl_divat fp32 — its loss reduction stays fp32. Forcing bf16 made the bridge diverge from run_tpu for kl_div and produced an fp32-impl vs bf16-baseline accuracy mismatch (torch.testing.assert_closefails on dtype), so kl_div reported a spurious accuracy failure on the bridge.Set
--precisionper kernel: fp32 forkl_div, bf16 for the rest (TPU-native and matching run_tpu.py). This fixes kl_div accuracy at the source and keeps its dashboard number consistent with both run_tpu.py and GPU — no change to the accuracy comparison needed.