Repository navigation
Respect BF16 overflow-check opt-out in ZeRO-1/2 - #8818
Conversation
Gate invalid gradient-norm overflow detection on the existing setting and cover clipping enabled and disabled in the BF16 integration tests. Assisted-by: OpenAI Codex Signed-off-by: zhaoye <772971548@qq.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 74e4d5d7d3
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # Group-norm helpers report inf/NaN as -1; combining that sentinel must | ||
| # not produce a finite 1.0 and then skip overflow / apply an unclipped step. | ||
| if is_invalid_grad_norm(scaled_global_grad_norm): | ||
| if self.check_grad_overflow and is_invalid_grad_norm(scaled_global_grad_norm): |
There was a problem hiding this comment.
When ZeRO-1/2 runs in ordinary FP32 (neither fp16 nor bf16 is enabled), _configure_zero_optimizer sets check_grad_overflow=False at engine.py:2642-2647. This new conjunction therefore disables the only invalid-global-norm guard for that mode, so a NaN/Inf gradient now reaches Adam/SGD and poisons parameters instead of being skipped as it was before this commit; gate only the BF16 opt-out or separate it from FP32 norm validation.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Addressed in follow-up d89a883.
The norm guard now permits bypass only for BF16 with an explicit opt-out. Ordinary FP32 keeps invalid-norm rejection even though its engine-provided check_grad_overflow flag is false; FP16 protection is retained.
The real two-rank FP32 engine regression covers ZeRO-1/2, clipping 0/1 and finite/Inf/NaN gradients. The published 74e4d5d reproduced eight failed invalid-gradient checks per rank and nonfinite parameters; the follow-up passes all 12 checks per rank and preserves finite parameters. New native distributed regression cases are in tests/unit/runtime/half_precision/test_zero_grad_norm_overflow.py.
The two-H200 FP32/BF16 matrix passes all 36 step checks per rank, including exact finite-update comparisons against an independent Torch SGD reference. The existing overflow suite rerun passes 36 tests, and scoped hooks pass. The follow-up was human-reviewed before pushing.
| # Group-norm helpers report inf/NaN as -1; combining that sentinel must | ||
| # not produce a finite 1.0 and then skip overflow / apply an unclipped step. | ||
| if is_invalid_grad_norm(scaled_global_grad_norm): | ||
| if self.check_grad_overflow and is_invalid_grad_norm(scaled_global_grad_norm): |
There was a problem hiding this comment.
When zero_optimization.zenflow.overlap_step=true, the engine selects ZenFlowZeroOptimizerParallel, whose step() override still unconditionally marks an invalid norm as overflow at zenflow_stage_1_and_2.py:793-797; this new base-class condition is never reached. Consequently BF16 check_grad_overflow=false still skips the step in that ZeRO-1/2 configuration, so apply the same guard in the override and cover the overlap path.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Addressed in follow-up d89a883.
The same BF16-only opt-out rule is now applied directly in ZenFlowZeroOptimizerParallel.step(), so the override no longer bypasses the fix. The new regression calls the factory-selected overlap step and checks dispatch/cleanup for valid norms, the invalid sentinel, Inf and NaN across FP32, FP16 and both BF16 check settings.
Separately, a real two-rank ZenFlow CPU-offload run used the native optimizer process with ZeRO-1/2, clipping 0/1 and both BF16 check settings. The original overlap guard reproduced eight failed opt-out checks per rank; the follow-up passes 24/24 step checks per rank.
This native run covers the full-update warmup path (full_warm_up_rounds=100), not selective/top-k updates, GPU ZenFlow or performance. The unit tests' mocked dispatch is not presented as native integration. The follow-up was human-reviewed, and scoped hooks plus the 36-test existing overflow suite pass.
Keep FP32 and FP16 invalid-norm rejection while honoring the explicit BF16 opt-out in both ordinary ZeRO and the ZenFlow overlap override. Add native distributed FP32 coverage and overlap admission regressions. Assisted-by: OpenAI Codex Signed-off-by: zhaoye <772971548@qq.com>
2 participants
Fixes #8805.
Why
The post-norm overflow guard ignored the explicit BF16
check_grad_overflow=falsesetting. Gating it on that flag alone also disabled ordinary FP32 invalid-norm rejection, because the engine sets the flag to false in FP32 mode.Scope
The follow-up
d89a883addresses both P1 findings on the original submitted head.Testing
All runs used isolated GPFS environments, Python 3.12 and PyTorch 2.11.0+cu130.
74e4d5d: 4 passed / 8 failed step checks per rank, with nonfinite parametersThe CPU and H200 engine matrices cover ZeRO-1/2, clipping 0/1 and finite/Inf/NaN gradients. The H200 matrix includes FP32, BF16 default checking and explicit BF16 opt-out; finite updates match an independent FP32 Torch SGD reference converted to the target dtype exactly (
rtol=0, atol=0).ZenFlow integration uses actual CPU offload and its native optimizer process with both BF16 check settings. It covers the full-update warmup path (
full_warm_up_rounds=100), not the selective/top-k phase. The overlap unit cases mock dispatch/cleanup and are not presented as native-worker validation.Not claimed: GPU ZenFlow/offload, ZeRO-3, other accelerator vendors, full convergence, performance or full CI. Keep-alive was restored and verified on all four H200s after device tests.
AI assistance: OpenAI Codex assisted with implementation and validation. The submitting human has reviewed and approved the original changes and this follow-up line by line; the reported validation was run with Codex assistance.