Skip to content

Say why a bf16 gradient norm cannot be clipped from - #8614

Open
vineethsaivs wants to merge 1 commit into
deepspeedai:masterfrom
vineethsaivs:fix-bf16-zero-norm-assert-20260920
Open

vineethsaivs wants to merge 1 commit into
deepspeedai:masterfrom
vineethsaivs:fix-bf16-zero-norm-assert-20260920

Conversation

@vineethsaivs

Copy link
Copy Markdown
Contributor

BF16_Optimizer.step guards clipping with a bare assert all_groups_norm > 0.. It stops a step whose gradients are all exactly zero, which a fully masked micro-batch produces, and it explains nothing when it fires.

The norm helpers do not raise on a non-finite norm; mask_nan_or_inf_with_val_inplace writes -1 and get_norm_with_moe_layers returns -1. Handed that, clip_tensors_by_global_norm computes max_norm / (-1 + 1e-6), a negative clip coefficient, and negates every gradient rather than shrinking it: measured here, 2.0 becomes -2.000002. So the assert has real work to do, but it conflates that with a legitimate zero and disappears under python -O. Reject only a negative norm, with a message naming the sentinel, and let a zero norm clip to the no-op it is.

Test: TORCHDYNAMO_DISABLE=1 DS_ACCELERATOR=cpu PYTHONPATH=.:tests python -m pytest tests/unit/runtime/half_precision/ -q. Two new cases fail before: an all-zero-gradient step raises, and a non-finite norm gives a bare AssertionError instead of a reason. Three more pin the premises, that the helpers report -1 and that clipping from -1 negates the gradients, and that an ordinary norm still clips 2.0 to 0.5. Nine pass and 32 skip after. yapf 0.43.0 and flake8 pass on the changed files. Executed on Apple M2 Pro CPU; no GPU or distributed run. The guard and the clip arithmetic are local to one rank.

Reject only a negative norm, with a message naming the sentinel, and let a zero norm clip to the no-op it is.

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant