Log the skipped optimizer step on gradient overflow (#8417)

ZeRO returns early from `step()` on non-finite gradients.
`DynamicLossScaler` reports that from inside `update_scale`, but
`CreateLossScaler` only returns it for `fp16 + dynamic_scaling` — bf16
and fp32 always get the static `LossScaler`, whose `update_scale` is an
inherited no-op.

The skip is counted in `engine.skipped_steps`, but that is only surfaced
by `_report_progress`, gated on `steps_per_print`, and
`STEPS_PER_PRINT_DEFAULT` is `None`. So on a default bf16 config the
loop advances, the LR scheduler steps and metrics are emitted while no
weights are updated, with no output at all.

Fixes #8415

---------

Signed-off-by: Abhishek Kulkarni <abhishek.kulkarni@snowflake.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>

Author: sfc-gh-abkulkarni

deepspeed/runtime/fp16/loss_scaler.py

@@ -21,13 +21,15 @@
 Commit: 93ab4bea59dc5cbf97c079d313741866af4deac9
 """
 
+import logging
+
 import torch
 from dataclasses import dataclass
 from typing import Optional
 from enum import Enum
 from deepspeed.runtime.config_utils import DeepSpeedConfigObject
 from deepspeed import comm as dist
-from deepspeed.utils import logger
+from deepspeed.utils import log_dist, logger
 from deepspeed.runtime.utils import has_inf_or_nan
 
 INITIAL_LOSS_SCALE = 'init_scale'
@@ -147,7 +149,10 @@ def scale_gradient(self, module, grad_in, grad_out):
         return tuple(self.loss_scale * g for g in grad_in)
 
     def update_scale(self, overflow):
-        pass
+        if overflow:
+            log_dist("[deepspeed] OVERFLOW! Skipping step. Gradients are non-finite, weights were not updated.",
+                     ranks=[0],
+                     level=logging.WARNING)
 
     def scale_loss(self, loss):
         """ Scales the loss by the current loss scale.

tests/unit/runtime/half_precision/test_static_loss_scale_overflow.py

@@ -0,0 +1,55 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import logging
+
+import pytest
+import torch
+
+from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler, DynamicLossScaler, LossScaler
+
+
+@pytest.fixture
+def warnings(monkeypatch):
+    recorded = []
+    monkeypatch.setattr("deepspeed.utils.logging.logger.log", lambda level, msg, *a, **kw: recorded.append(
+        (level, msg)))
+    return recorded
+
+
+def test_static_scaler_reports_the_skipped_step(warnings):
+    LossScaler(scale=1.0).update_scale(overflow=True)
+
+    assert len(warnings) == 1
+    level, msg = warnings[0]
+    assert level == logging.WARNING
+    assert "OVERFLOW" in msg and "weights were not updated" in msg
+
+
+def test_static_scaler_is_quiet_without_overflow(warnings):
+    LossScaler(scale=1.0).update_scale(overflow=False)
+
+    assert warnings == []
+
+
+def test_dynamic_scaler_is_not_double_reported(warnings, monkeypatch):
+    """DynamicLossScaler overrides update_scale and logs the skip itself."""
+    monkeypatch.setattr("deepspeed.comm.get_rank", lambda: 0)
+    scaler = DynamicLossScaler(init_scale=2**16,
+                               scale_window=1000,
+                               min_scale=1,
+                               delayed_shift=1,
+                               consecutive_hysteresis=False)
+    scaler.update_scale(overflow=True)
+
+    assert not [msg for _, msg in warnings if "weights were not updated" in msg]
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32, torch.half])
+def test_non_dynamic_configs_get_the_reporting_scaler(dtype):
+    """bf16 and fp32 never get a DynamicLossScaler, so the static path is the only one."""
+    scaler = CreateLossScaler(dtype=dtype, static_loss_scale=1.0, dynamic_scaling=False, dynamic_loss_args=None)
+
+    assert not isinstance(scaler, DynamicLossScaler)