Repository navigation
Keep RunningNorm estimate stable when a slice is all zeros - #2142
Open
rootkiller6788 wants to merge 4 commits into
Open
rootkiller6788 wants to merge 4 commits into
rootkiller6788 wants to merge 4 commits into
Conversation
An all-zero reduction slice comes out of the masked nanmean as NaN. It was then replaced by 0.0 and fed through log(), producing -inf which poisoned ln_running_mean permanently (and made subsequent outputs inf). Fall back to the current running estimate for those slices so the EWMA update is a no-op, which is what the docstring already promises. Slices with any non-zero entry are unaffected.
Cover the all-zero slice path (first batch and after seeding) plus the zero-ignoring and calibration behavior. The zero-slice tests fail against the old log(0) = -inf handling.
Contributor
CODEOWNERS review mapCurrent for commit ⏳ @loliverhennigh — 1 file(s)
No CODEOWNER
Comment |
Contributor
|
[Low impact] The change affects one edge case in RunningNorm’s training updates. The half-precision overflow should be fixed before merging. Findings
SummaryChanges
Reviews (1) · Last reviewed commit: "Merge branch 'main' into fix-running-nor..." · Reviewed by Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
The all-zero-slice fallback rebuilt the linear running mean with exp(ln_running_mean) and then took log() of it again. In half precision that exp overflows to inf once the finite estimate approaches the dtype maximum (e.g. ln(65504) ~= 11.09), so an all-zero batch still poisoned the EWMA and every later finite input normalized to zero. Select the fallback directly in log space instead and add an fp16 regression test.
Author
|
Fixed — the all-zero fallback now stays in log space instead of round-tripping through exp, so it can't overflow to inf in half precision, and I added an fp16 regression test covering the 65504 seed case. |
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.
RunningNorm keeps a running mean of |x| in log space. If a channel is all zeros for a whole batch, the slice comes out of the masked reduction as NaN, and the old code replaced that with 0.0 before taking the log. The estimate then became -inf and stayed there, so every later batch for that channel came out as inf.
The docstring already promises that an all-zero slice leaves the running estimate unchanged, so this makes it true: those slices now fall back to the current estimate (or the neutral value 1.0 before the module has seen any data), and the EWMA update is a no-op for them. Slices with any non-zero entry behave exactly as before.
Added regression tests covering the all-zero first batch, an all-zero batch after the estimate was seeded, and the existing zero-ignoring/calibration behavior. Verified on the CPU path: before the change a seeded estimate went to -inf after one all-zero batch; after the change it stays put and outputs stay finite.