[0z5a]

With offload_optimizer, ZeRO-1/2 checks each partition gradient for inf/NaN in update_offload_overflow_tracker(). The check had two problems.

1. A host sync per gradient. It put the result of has_inf_or_nan() straight into a Python if. has_inf_or_nan() returns a device tensor so that callers can accumulate it without synchronizing (its docstring says so), and the if forces a sync for every gradient anyway.
2. It runs when the check is off. The tracker also runs with check_grad_overflow disabled: bf16 with the check turned off, and fp32 training, where the engine always passes False. Only has_overflow() reads the flag, and step() calls it only when check_grad_overflow is set.

Measured

Setup: one GH200, Qwen3-4B (398 parameter tensors), ZeRO-1 with CPU offload, bf16, fp32 gradient accumulation.

The backward of the accumulation-boundary micro-step showed 399 aten::item and 404 cudaStreamSynchronize, about one per parameter tensor. The absolute cost is small here, a few tenths of a second in a ~62 s training step. It grows with the number of parameter tensors, and it serializes the host with the GPU right where the boundary reduction and offload copies run.

Change

• update_offload_overflow_tracker() returns early when check_grad_overflow is off.
• Otherwise it ORs has_inf_or_nan() into a device flag. local_overflow starts as False and becomes a device bool tensor.
• has_overflow() turns that flag into the byte tensor it all-reduces. The host reads it once, in the existing .item() after the all-reduce.
• The method signatures are unchanged, so the safe_set_full_grad path (tensor_fragment.py) needs no change.

Tests

New file tests/unit/runtime/zero/test_offload_overflow_flag.py: world size 2, ZeRO stages 1 and 2 with offload, torch Adam, runs on the CPU accelerator. Three tests:

• the flag is never truth-tested while it accumulates (a tensor subclass that raises on __bool__); it stays a tensor, catches an inf and resets;
• a NaN gradient still skips the step;
• with check_grad_overflow=False, no gradient is checked.

On the unchanged code, 4 of the 6 cases fail: the first and third tests, at both stages.

test_zero2_gradient_safety.py builds the optimizer with __new__ and sets the __init__ attributes by hand. It now sets check_grad_overflow as well.

The following pass on the CPU accelerator (LOCAL_SIZE=2), 84 passed and 2 skipped:

• test_zero_optim_overflow.py
• test_zero1_optimizer_fastpath.py
• test_zero2_gradient_safety.py
• the new file

yapf, flake8, codespell and the license check pass.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-11-05T10:01:41.782981Z	c4570ac	Draft marked ready

ℹ️ 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[Copilot]

🟡 Changes recommended

The six regression cases are outside the configured CI test scope and cannot run in the selected GPU or multirank CPU jobs, leaving the overflow and disabled-check behavior unverified by those jobs.

1 open issue

• Move multirank regression tests into the configured CI test scope · New

What changed in this PR

This PR reduces host synchronization in ZeRO-1/2 optimizer offload by keeping overflow tracking on the accelerator until the final collective result is read.

Changes:

• Skip gradient overflow tracking when check_grad_overflow is disabled.
• Accumulate overflow results in a device tensor and convert it for the existing all-reduce.
• Add regression tests for flag accumulation, overflow step skipping, and disabled checks; update the existing optimizer fixture.

File	Description
tests/​unit/​v1/​zero/​test_zero2_gradient_safety.py	Initialize check_grad_overflow in the optimizer fixture.
tests/​unit/​runtime/​zero/​test_offload_overflow_flag.py	Add two-rank regression tests for offload overflow tracking.
deepspeed/​runtime/​zero/​stage_1_and_2.py	Avoid per-gradient host truth tests and bypass disabled overflow checks.

_{Review effort: Balanced}

[0z5a]

Moved to tests/unit/v1/zero in 5257815; the selector now picks it in both workflows.