Repository navigation
Honor fp16 hysteresis in the ZeRO stage 0 fp16 optimizers - #8647
Open
vineethsaivs wants to merge 4 commits into
Open
vineethsaivs wants to merge 4 commits into
vineethsaivs wants to merge 4 commits into
Conversation
LossScaleConfig copied the initial scale, window and minimum out of the dynamic loss scale args and dropped delayed_shift and consecutive_hysteresis, and FP16_Optimizer and FP16_UnfusedOptimizer halved the scale on every overflow. So at stage 0, fp16.hysteresis and fp16.consecutive_hysteresis did nothing, while the same config at ZeRO 1, 2 and 3 goes through DynamicLossScaler and honors both. Carry both settings into LossScaleConfig and apply them in _update_scale the way DynamicLossScaler does. Save the remaining hysteresis in the optimizer state dict, falling back to the configured value for older checkpoints, and stop the overflow log claiming a reduction when hysteresis absorbed it. The four stage 0 overflow tests that relied on the ignored default now pin hysteresis to 1, as the ZeRO overflow tests already do. Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
vineethsaivs
requested review from
loadams,
tjruwase and
tohtana
as code owners
September 23, 2026 22:30
Contributor
|
I think this is a legitimate fix for stage 0, @tjruwase, what's your thoughts? |
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The tests bypass the public stage-0 engine and training-step path, leaving the external behavior without integration coverage.
Review effort: Balanced
Findings: 1
What changed in this PR
Adds fp16 hysteresis support to ZeRO stage-0 optimizers, aligning their loss-scaling behavior with later ZeRO stages.
Changes:
- Propagates hysteresis settings into stage-0 loss scaling.
- Persists hysteresis state across checkpoints.
- Adds focused tests and updates existing expectations.
| File | Description |
|---|---|
deepspeed/runtime/fp16/loss_scaler.py |
Stores hysteresis configuration and state. |
deepspeed/runtime/fp16/fused_optimizer.py |
Applies and checkpoints hysteresis for fused optimization. |
deepspeed/runtime/fp16/unfused_optimizer.py |
Applies and checkpoints hysteresis for unfused optimization. |
tests/unit/runtime/half_precision/test_dynamic_loss_scale.py |
Adds hysteresis tests and adjusts existing overflow tests. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+306
to
+315
| def _optimizer(self, optimizer_class, **fp16_config): | ||
| fp16 = DeepSpeedFP16Config(enabled=True, loss_scale=0, initial_scale_power=8, **fp16_config) | ||
| profile = LossScaleProfile.FUSED if optimizer_class is FP16_Optimizer else LossScaleProfile.UNFUSED | ||
| loss_scale_config = LossScaleConfig(low_precision_dtype=torch.float16, | ||
| dynamic_loss_scale=True, | ||
| static_loss_scale=0, | ||
| dynamic_loss_args=fp16.dynamic_loss_scale_args(), | ||
| profile=profile) | ||
| param = torch.nn.Parameter(torch.zeros(4, dtype=torch.float16)) | ||
| return optimizer_class(torch.optim.SGD([param], lr=0.1), loss_scale_config=loss_scale_config) |
Contributor
Author
There was a problem hiding this comment.
Added in 7f0482c: an engine test that sets fp16.hysteresis: 2, overflows through model.step() and resumes from the optimizer state dict, for both stage 0 optimizers.
Configure fp16.hysteresis in the engine config, overflow through model.step(), and resume from the optimizer state dict, for both the fused and unfused stage 0 optimizers. Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
…unfused_optimizer.py Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
pengdurice
approved these changes
Oct 25, 2026
pengdurice
enabled auto-merge
October 25, 2026 16:29
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Oct 25, 2026
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
At ZeRO stage 0 (fp16 with no
zero_optimizationblock, or"stage": 0),fp16.hysteresisandfp16.consecutive_hysteresisare silently ignored. The loss scale halves on every overflow. The same config at ZeRO 1, 2 or 3 goes throughDynamicLossScalerand honors both.LossScaleConfigcopies the initial scale, window and minimum out ofdynamic_loss_scale_args()and dropsdelayed_shiftandconsecutive_hysteresis.FP16_Optimizer._update_scaleandFP16_UnfusedOptimizer._update_scalethen have no hysteresis logic at all. This carries both settings through and applies them the wayDynamicLossScaler.update_scaledoes. 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:This changes the default behavior at stage 0: the default
hysteresis=2now 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
TestZeroStage0Hysteresisintest_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.