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.