[bug fix] Compute fp16 batch norm in fp32 instead of downcasting the params - #54
Merged
gokulkrishna98 merged 1 commit intoAug 7, 2026
Conversation
`nn.BatchNorm2d` keeps gamma/beta/running_mean/running_var in fp32 even when the activation is fp16. ATen handles that mixed-precision case by promoting the input to fp32, using the params and eps at full precision, and narrowing only the final result (`native_batch_norm_helper` in torch/_decomp/decompositions.py). The composite body did the opposite: it cast the fp32 params and eps down to the input element type and ran every op in fp16. That loses mantissa bits on the params (eps 1e-5 becomes 1.00136e-05), and any channel whose running_var or running_mean exceeds the fp16 max of 65504 becomes inf -- sqrt(inf) then divides the activation to a silent zero. Use the existing `prepare_compute_type_for_norm` helper, as layer_norm and group_norm already do, so the fp16 input is promoted, the arithmetic runs in fp32, and there is a single narrowing cast before the output. fp32 inputs are unaffected apart from dropping now-redundant f32->f32 casts, which the optimizer already folded. Tests: a numerical case with large and tiny running_var (max abs error 4.945 and 0.5 against eager before the fix) and an IR case pinning the promote/compute/narrow shape of the fp16 composite body.
gokulkrishna98
force-pushed
the
dev/gokul/batchnorm-fp32-compute
branch
from
August 7, 2026 19:38
9b11e80 to
b3345b3
Compare
gokulkrishna98
marked this pull request as ready for review
August 7, 2026 19:39
gokulkrishna98
requested review from
Lewis300,
TobyRoseman,
YifanShenSZ,
cymbalrush and
jakesabathia2
August 7, 2026 19:39
TobyRoseman
approved these changes
Aug 7, 2026
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.
Description:
BatchNorm keeps its params in fp32 even when the activation is fp16. ATen promotes the input to fp32, uses the params and eps at full precision, and narrows only the result (
native_batch_norm_helper). The composite body did the opposite: it cast the params down to fp16, losing mantissa bits and turning anyrunning_varabove the fp16 max of 65504 intoinf, which silently zeroes the output.Fix: use
prepare_compute_type_for_norm, aslayer_normandgroup_normalready do. fp32 inputs are unchanged apart from dropping redundantf32 -> f32casts.Tests: numerical case with large/tiny
running_var(fails on main with max abs error 4.945 and 0.5), plus an IR case pinning the promote/compute/narrow body.