DeepSpeed
62656fe7 - Honor fp16 hysteresis in the ZeRO stage 0 fp16 optimizers (#8647)

Commit
5 days ago
Honor fp16 hysteresis in the ZeRO stage 0 fp16 optimizers (#8647) At ZeRO stage 0 (fp16 with no `zero_optimization` block, or `"stage": 0`), `fp16.hysteresis` and `fp16.consecutive_hysteresis` are silently ignored. The loss scale halves on every overflow. The same config at ZeRO 1, 2 or 3 goes through `DynamicLossScaler` and honors both. `LossScaleConfig` copies the initial scale, window and minimum out of `dynamic_loss_scale_args()` and drops `delayed_shift` and `consecutive_hysteresis`. `FP16_Optimizer._update_scale` and `FP16_UnfusedOptimizer._update_scale` then have no hysteresis logic at all. This carries both settings through and applies them the way `DynamicLossScaler.update_scale` does. It also saves the remaining hysteresis in the stage 0 state dict (with a fallback for older checkpoints), and stops the overflow log from saying the scale was reduced when hysteresis absorbed it. Four overflows with `hysteresis=3`, `initial_scale_power=8`: ``` before [128, 64, 32, 16] after [256, 256, 128, 64] (same as ZeRO 1/2/3) ``` This changes the default behavior at stage 0: the default `hysteresis=2` now takes effect, which matches ZeRO and the docs. The four existing stage 0 overflow tests relied on the default being ignored, so they now pin `"hysteresis": 1`, the way the ZeRO overflow tests already do. The separate open #8393 (stage 0 fp16 backward not scaling the loss) is on the same path but is a different bug. Tests: new `TestZeroStage0Hysteresis` in `test_dynamic_loss_scale.py`, over both optimizers, with three cases: hysteresis, consecutive hysteresis, and a checkpoint round trip. 6 pass here. On master, a single-process run of the same checks gives the "before" row above. yapf 0.43.0 and flake8 5.0.4 are clean. Executed on Apple M2 Pro CPU. The GPU-gated stage 0 overflow tests could not run here. --------- Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com> Co-authored-by: pengdurice <pengduhit@gmail.com>
Author
Parents
Loading