fix(utils): accumulate param L2 norm in float32 for bf16/fp16 models - #3625
Open
ralovets wants to merge 2 commits into
Open
fix(utils): accumulate param L2 norm in float32 for bf16/fp16 models#3625ralovets wants to merge 2 commits into
ralovets wants to merge 2 commits into
Conversation
_get_model_param_stats accumulated the squared parameter norm in the parameter dtype, so the reported "Param L2 norm" drifted for bf16 models and overflowed to inf for fp16 models. Pass an accumulation dtype to norm() instead. promote_types rather than a literal float32, because vector_norm rejects a narrowing dtype: float64 parameters would raise and be silently dropped by the surrounding except. The dtype= kwarg folds the upcast into the reduction, so this keeps the device-side accumulation introduced in NVIDIA-NeMo#1463. Signed-off-by: Roman Ralovets <roman@ralovets.com>
ralovets
marked this pull request as ready for review
August 22, 2026 05:41
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
Contributor
|
/ok to test 0c5f12c |
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
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
3 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.
What does this PR do ?
Fixes the
Param L2 normlogged byprint_trainable_parameters._get_model_param_statsaccumulateslocal_sq_norm += p.detach().norm(2) ** 2, which takes the parameter dtype, so the value drifts for bf16 models and isinffor fp16 models — a real 3.62B-parameter bf16 checkpoint logs3007.3191against a true value of3167.8195, 5.1% low.The fix passes an accumulation dtype to
norm(), usingtorch.promote_types(p.dtype, torch.float32)rather than a literaltorch.float32to avoid narrowing float64 parameters, whichlinalg.vector_normrejects. Regression from #1463, which dropped the.float()upcast along with a per-parameter host sync. Metric only, no effect on training numerics.Performance
Benchmarked on GPU against the real 1120-tensor / 3.62B-param shape distribution, bf16, variants interleaved within each round:
.float()upcast copy, plus a.item()host sync per parameternorm(2), accumulating in the parameter dtype — the bugnorm(2, dtype=...), upcast folded into the reduction+1.3% (0.15 ms) against current main, from the per-parameter
promote_typescall. One-time cost, since all call sites run once after sharding.Changelog
_get_model_param_statsaccumulates the squared parameter norm in at least float32 instead of the parameter dtype, fixing the reportedParam L2 normfor bf16 (drifts >1%) and fp16 (reportsinf) models.Before your PR is "Ready for review"
Pre checks: