Enable bf16 check_grad_overflow by default (matching fp16)

[yongzhe-wang]

Summary

Flip DeepSpeedBF16Config.check_grad_overflow default from False to True, so bf16 users get the same gradient-overflow protection that fp16 users already get by default.

Motivation

The bf16 documentation states bf16 "does not require loss scaling" (deepspeed.ai/docs/config-json/), but this overstates the safety guarantee for the bf16 + ZeRO-2 (non-offload) partition-flat gradient accumulation path. We reproduced a deterministic catastrophic NaN under a small set of training conditions:

• ZeRO-2 (non-offload) + bf16
• Mixture-of-Transformers (modality-specific transformer branches)
• Heterogeneous per-sample loss masks (e.g. 50% action-invalid samples in robotics VLA training)

Under these conditions, a single bf16 element in engine.optimizer.averaged_gradients[i] overflows to +inf. The downstream Adam.step then computes inf / sqrt(inf) = NaN in a fused kernel, which simultaneously corrupts the partition slice's exp_avg, exp_avg_sq, and fp32 master weights. The next forward pass propagates NaN through every layer; the training run is dead with no useful diagnostic. Reproduced consistently in DeepSpeed 0.16.9 - 0.17.1 at step ~22 with our internal repro.

The infrastructure to detect and skip such steps was correctly added by #6976 (check_grad_overflow option, DeepSpeedZeroOptimizer.check_overflow method, and step-skip logic at stage_1_and_2.step() lines ~2128-2143). However the default was set to False for bf16, so users hitting this condition do not receive the protection.

Change

Single line: check_grad_overflow: bool = False -> check_grad_overflow: bool = True in DeepSpeedBF16Config. Updated docstring + bf16 example block accordingly.

Backward compatibility

Users who have benchmarked the check as too expensive AND have separately confirmed their bf16 path cannot overflow can opt out by setting:

```json
"bf16": {
"enabled": true,
"check_grad_overflow": false
}
```

The runtime cost is one isfinite-style scan over the gradient partition per optimizer step (already implemented in DeepSpeedZeroOptimizer.check_overflow); typically under 1% of step wallclock.

Related

• Issue [BUG] grad_norm and loss is nan when deepspeed==0.13.5 but ok with deepspeed==0.10.2 #5242 (open) - bf16+ZeRO-2 NaN on real training runs (Baichuan2-7B + 8x A800)
• PR Improve overflow handling in ZeRO #6976 - introduced check_grad_overflow option and underlying skip logic

Test plan

• Reproducer (private repo): with default False, run dies at step ~22; with True, training survives via DeepSpeed's existing skip-step path.
• Existing CI should pass unchanged; this PR only changes a default value in precision_config.py.

[tohtana]

Thank you @yongzhe-wang,
This change overall looks good to me, but I'm still not sure about the performance impact. Adding a synchronization point might bring a noticeable difference.

What is your thought? @sfc-gh-truwase

[yongzhe-wang]

Thanks @tohtana — good question. A few things that I think bound the sync-cost concern:

1. This isn't a new synchronization point for DeepSpeed — fp16 already does exactly this, unconditionally, every step. In engine.py the flag is hard-wired on for fp16:

if self.bfloat16_enabled():
    check_grad_overflow =
self._config.bfloat16_config.check_grad_overflow
elif self.fp16_enabled():
    check_grad_overflow = True      # fp16 always pays this
else:
    check_grad_overflow = False

The work it gates — has_overflow(): an isfinite scan over the gradient partition, a scalar all_reduce(MAX), and one .item() — is the same code path fp16 has run by default for years. So this is a well-characterized production cost, not a new one; the PR just gives bf16 the same protection.

2. ZeRO-1/2 already incurs a per-step device sync regardless. In stage_1_and_2.step(), every non-overflow step calls scaled_global_norm() for gradient clipping, which does an all_reduce + .item(). The engine therefore isn't running fully async across the optimizer step anyway — the overflow check's .item() lands just before a sync that was already going to happen.

[sfc-gh-truwase]

Make this terse, and move the details to the PR. Okay to leave the issue and PR links here.

[sfc-gh-truwase]

@yongzhe-wang, apologies for the delayed response. Thanks for this PR.

Similar to @tohtana's point, performance concern was what prevented making this default. But, I think your arguments are solid: (1) parity with known production-cost of fp16, and (2) potentially reusing the current grad norm scans to reduce cost.

I left one minor code cleanup comment.

[tohtana]

@yongzhe-wang Let me know when you think this PR is ready to merge.

[yongzhe-wang]

@sfc-gh-truwase Rebased onto latest master — branch is up to date, approved, and DCO green. The remaining required checks are waiting on maintainer workflow approval (fork PR). Could you approve the workflow runs and merge when green? Thanks!