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.