DeepSpeed

f0253c86 - Enable bf16 check_grad_overflow by default (matching fp16) (#8035)

Enable bf16 check_grad_overflow by default (matching fp16) (#8035) ## 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 #5242 (open) - bf16+ZeRO-2 NaN on real training runs (Baichuan2-7B + 8x A800) - PR #6976 - introduced `check_grad_overflow` option and underlying skip logic ## Test plan - [x] Reproducer (private repo): with default `False`, run dies at step ~22; with `True`, training survives via DeepSpeed's existing skip-step path. - [x] Existing CI should pass unchanged; this PR only changes a default value in `precision_config.py`. Signed-off-by: Yongzhe Wang <yzwang2020@gmail.com>