Keep invalid group-norm sentinels from becoming a finite 1.0 step (#8638)

## Summary
- DeepSpeed encodes inf/NaN group norms as `-1`. Combining those with
`vector_norm` squares the sentinel to `1.0`, so clipping and Adam treat
a failed group as healthy and apply an unclipped step.
- Treat non-finite or negative group norms as overflow (`inf`), skip the
optimizer step, reset overflow each step, and advance the loss scaler
once (including ZeRO-3 / SuperOffload / ZenFlow / FP16 fused+unfused /
BF16).
- Direct `unscale_and_clip_grads` callers map the sentinel to `inf` so
they still zero gradients instead of using clip_coef=1.

This is split out of the ZeRO-2 oversized-gradient / stream-lifetime
work so the sentinel contract can land independently.

## Test plan
- [x] CPU, rebased onto current master (`75eb961d1`, includes #8730): 17
passed in `tests/unit/v1/zero/test_invalid_grad_norm.py`
(`combine_grad_norm_groups` / `get_global_norm` / `CheckOverflow`,
ZeRO-1 clip, CPU-offload combine, ZeRO-3 / BF16 / SuperOffload overflow
paths)
- [x] GPU (2× H200, prior run):
`TestZeROBFloat16Stage3InvalidNormRecovery`. Skipped on the rebase
machine (2 GPUs required, 1 available)
- [ ] CI unit tests for the touched runtime files, started by the rebase
push

---------

Signed-off-by: tunji-ruwase_snow <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: sfc-gh-truwase

deepspeed/runtime/bf16_optimizer.py

@@ -14,7 +14,8 @@
 from deepspeed.git_version_info import version
 from deepspeed.runtime.utils import (get_global_norm_of_tensors, clip_tensors_by_global_norm, DummyOptim,
                                      align_dense_tensors, all_gather_dp_groups, is_model_parallel_parameter,
-                                     see_memory_usage, get_norm_with_moe_layers, is_optimized_parameter)
+                                     see_memory_usage, get_norm_with_moe_layers, is_invalid_grad_norm,
+                                     is_optimized_parameter)
 from deepspeed.utils import link_hp_params, lazy_init_hp_params_optimizer_state, fragment_address, groups
 from deepspeed.moe.utils import is_moe_param, is_moe_param_group
 from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank
@@ -74,6 +75,7 @@ def __init__(self,
         self.custom_loss_scaler = False
         self.external_loss_scale = None
         self.torch_autocast_gradscaler = None
+        self.overflow = False
 
         self.immediate_grad_update = bfloat16_config.immediate_grad_update
 
@@ -379,6 +381,7 @@ def step(self, closure=None):
         if closure is not None:
             raise NotImplementedError(f'{self.__class__} does not support closure.')
 
+        self.overflow = False
         non_expert_grads_for_norm, expert_grads_for_norm = self.get_grads_for_norm()
         non_expert_groups_norm = get_global_norm_of_tensors(input_tensors=non_expert_grads_for_norm,
                                                             mpu=self.mpu,
@@ -392,7 +395,12 @@ def step(self, closure=None):
 
         self._global_grad_norm = all_groups_norm
 
-        assert all_groups_norm > 0.
+        if is_invalid_grad_norm(all_groups_norm):
+            self.overflow = True
+            self._global_grad_norm = float("inf")
+            self.clear_hp_grads()
+            self.clear_lp_grads()
+            return
         if self.clip_grad > 0.:
             clip_tensors_by_global_norm(input_tensors=self.get_grads_for_norm(for_clipping=True),
                                         max_norm=self.clip_grad,

deepspeed/runtime/fp16/fused_optimizer.py

@@ -10,9 +10,11 @@
 import torch
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer
-from deepspeed.runtime.utils import get_global_norm, get_flattened_grad_norm, CheckOverflow, get_weight_norm, get_norm_with_moe_layers, is_model_parallel_parameter, is_optimized_parameter
+from deepspeed.runtime.utils import (CheckOverflow, get_flattened_grad_norm, get_global_norm, get_norm_with_moe_layers,
+                                     get_weight_norm, is_invalid_grad_norm, is_model_parallel_parameter,
+                                     is_optimized_parameter)
 from deepspeed.runtime.fp16.loss_scaler import LossScaleConfig, LossScaleProfile
-from deepspeed.utils import logger, log_dist
+from deepspeed.utils import logger
 from deepspeed.utils.torch import required_torch_version
 from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, CLIP_GRAD
 from deepspeed.accelerator import get_accelerator
@@ -146,6 +148,18 @@ def zero_grad(self, set_to_none=True):
                         p.grad.detach_()
                         p.grad.zero_()
 
+    def _skip_invalid_grad_norm(self):
+        self.overflow = True
+        self._global_grad_norm = float("inf")
+        if self.flatten_grad_norm_mask_list:
+            self.has_executed_step = True
+        if self.loss_scale_config.use_grad_scaling:
+            self._update_scale(self.overflow)
+        self.zero_grad(set_to_none=True)
+        for group in self.fp32_groups_flat:
+            group.grad = None
+        return self.overflow
+
     def step_fused_adam(self, closure=None):
         """
         Not supporting closure.
@@ -162,17 +176,14 @@ def step_fused_adam(self, closure=None):
             norm_groups.append(get_weight_norm(grads_groups_flat[i], mpu=self.mpu))
 
         self.overflow = self.overflow_checker.check_using_norm(norm_groups)
-        if self.loss_scale_config.use_grad_scaling:
-            prev_scale = self.loss_scale_config.cur_scale
-            self._update_scale(self.overflow)
-
-            if self.overflow:
-                if self.verbose:
-                    logger.info("[deepspeed] fp16 dynamic loss scale overflow! Skipping step. Attempted loss "
-                                "scale: {}, reducing to {}".format(prev_scale, self.loss_scale_config.cur_scale))
-                return self.overflow
+        if self.overflow:
+            return self._skip_invalid_grad_norm()
 
         scaled_grad_norm = get_global_norm(norm_list=norm_groups)
+        if is_invalid_grad_norm(scaled_grad_norm):
+            return self._skip_invalid_grad_norm()
+        if self.loss_scale_config.use_grad_scaling:
+            self._update_scale(False)
 
         combined_scale = self.unscale_and_clip_grads(grads_groups_flat, scaled_grad_norm, apply_scale=False)
 
@@ -260,22 +271,8 @@ def step(self, closure=None):
         if self.timers:
             self.timers(OVERFLOW_CHECK_TIMER).stop()
 
-        if self.loss_scale_config.use_grad_scaling:
-            prev_scale = self.loss_scale_config.cur_scale
-            self._update_scale(self.overflow)
-            if self.overflow:
-                if self.verbose:
-                    log_dist(
-                        "Overflow detected. Skipping step. Attempted loss "
-                        f"scale: {prev_scale}, reducing to {self.loss_scale_config.cur_scale}",
-                        ranks=[0])
-                # Clear gradients
-                for i, group in enumerate(self.fp16_groups):
-                    for p in group:
-                        p.grad = None
-                if self.timers:
-                    self.timers.log(OVERFLOW_TIMERS)
-                return self.overflow
+        if self.overflow:
+            return self._skip_invalid_grad_norm()
 
         grads_groups_flat = []
         non_experts_grads_for_norm = []
@@ -330,8 +327,14 @@ def step(self, closure=None):
                                                        norm_type=self.norm_type)
 
         scaled_global_grad_norm = get_global_norm(norm_list=[all_groups_norm])
+        if is_invalid_grad_norm(scaled_global_grad_norm):
+            if self.timers:
+                self.timers(COMPUTE_NORM_TIMER).stop()
+            return self._skip_invalid_grad_norm()
         if self.timers:
             self.timers(COMPUTE_NORM_TIMER).stop()
+        if self.loss_scale_config.use_grad_scaling:
+            self._update_scale(False)
 
         # Stash unscaled gradient norm
         self._global_grad_norm = scaled_global_grad_norm / self.loss_scale_config.cur_scale

deepspeed/runtime/fp16/unfused_optimizer.py

@@ -12,7 +12,8 @@
 from torch._utils import _flatten_dense_tensors
 
 from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer
-from deepspeed.runtime.utils import get_global_norm, CheckOverflow, get_weight_norm, is_optimized_parameter
+from deepspeed.runtime.utils import (CheckOverflow, get_global_norm, get_weight_norm, is_invalid_grad_norm,
+                                     is_optimized_parameter)
 from deepspeed.runtime.fp16.loss_scaler import LossScaleConfig, LossScaleProfile
 from deepspeed.utils import logger
 from deepspeed.utils.torch import required_torch_version
@@ -121,6 +122,16 @@ def zero_grad(self, set_to_none=True):
                         p.grad.detach_()
                         p.grad.zero_()
 
+    def _skip_invalid_grad_norm(self):
+        self.overflow = True
+        self._global_grad_norm = float("inf")
+        self._update_scale(self.overflow)
+        self.zero_grad(set_to_none=True)
+        for group in self.fp32_groups:
+            for param in group:
+                param.grad = None
+        return self.overflow
+
     def step_fused_lamb(self, closure=None):
         """
         Not supporting closure.
@@ -147,16 +158,13 @@ def step_fused_lamb(self, closure=None):
             expert_norm_groups.append(expert_norm_group_value)
 
         self.overflow = self.overflow_checker.check_using_norm(norm_groups + expert_norm_groups)
-        prev_scale = self.loss_scale_config.cur_scale
-
-        self._update_scale(self.overflow)
         if self.overflow:
-            if self.verbose:
-                logger.info("[deepspeed] fp16 dynamic loss scale overflow! Skipping step. Attempted loss "
-                            "scale: {}, reducing to {}".format(prev_scale, self.loss_scale_config.cur_scale))
-            return self.overflow
+            return self._skip_invalid_grad_norm()
 
         self._global_grad_norm = get_global_norm(norm_list=norm_groups)
+        if is_invalid_grad_norm(self._global_grad_norm):
+            return self._skip_invalid_grad_norm()
+        self._update_scale(False)
         combined_scale = self.unscale_and_clip_grads(self._global_grad_norm, apply_scale=False)
         self.optimizer.step(grads=grads_groups, output_params=self.fp16_groups, scale=combined_scale)
 
@@ -195,14 +203,8 @@ def step(self, closure=None):
             return self.step_fused_lamb()
 
         self.overflow = self.overflow_checker.check()
-        prev_scale = self.loss_scale_config.cur_scale
-
-        self._update_scale(self.overflow)
         if self.overflow:
-            if self.verbose:
-                logger.info("[deepspeed] fp16 dynamic loss scale overflow! Skipping step. Attempted loss "
-                            "scale: {}, reducing to {}".format(prev_scale, self.loss_scale_config.cur_scale))
-            return self.overflow
+            return self._skip_invalid_grad_norm()
 
         norm_groups = []
         for i, group in enumerate(self.fp16_groups):
@@ -220,6 +222,9 @@ def step(self, closure=None):
                     fp32_param.grad = fp16_param.grad.to(fp32_param.dtype)
 
         self._global_grad_norm = get_global_norm(norm_list=norm_groups)
+        if is_invalid_grad_norm(self._global_grad_norm):
+            return self._skip_invalid_grad_norm()
+        self._update_scale(False)
         self.unscale_and_clip_grads(self._global_grad_norm)
 
         self.optimizer.step()

deepspeed/runtime/superoffload/superoffload_stage3.py

@@ -4,12 +4,12 @@
 # DeepSpeed Team
 
 import time
-import torch
 from typing import List
 
 from deepspeed.runtime.superoffload.superoffload_utils import SuperOffloadCPUOptimizer, TaskKeys, ResultKeys, EventTypes
 from deepspeed.runtime.zero.partition_parameters import Parameter, Tensor
 from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3
+from deepspeed.runtime.utils import combine_grad_norm_groups, is_invalid_grad_norm
 from deepspeed.utils.nvtx import instrument_w_nvtx
 from deepspeed.utils import logger
 from deepspeed.accelerator import get_accelerator
@@ -39,6 +39,7 @@ def __init__(
 
         self.sub_group_to_param_num = {}
         self.sub_group_grad_partition_counts = {}
+        self._submitted_cpu_sub_groups = set()
         self.async_cpuadam_num = 0
         self.max_grad_numel = 0
 
@@ -176,6 +177,7 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L
                                                                fp32_param.data,
                                                                fp32_param.grad.data,
                                                                lr=current_lr)
+                    self._submitted_cpu_sub_groups.add(i)
                     self.async_cpuadam_num += 1
 
                     result = self.superoffload_cpu_optimizer.get_result()
@@ -200,13 +202,22 @@ def step(self, closure=None):
         self._pre_step()
         self._partition_all_parameters()
 
-        if self._overflow_check_and_loss_scale_update():
+        self._overflow_check_and_loss_scale_update(update_scale=False)
+        if self.overflow:
             if not self.clip_grad:
                 self._handle_overflow_rollback()
+            self._loss_scale_update_and_overflow_cleanup()
             return
 
         norm_groups = self._get_norm_groups()
-        scaled_global_grad_norm = torch.linalg.vector_norm(torch.stack(norm_groups))
+        scaled_global_grad_norm = combine_grad_norm_groups(norm_groups)
+        if is_invalid_grad_norm(scaled_global_grad_norm):
+            self.overflow = True
+            if not self.clip_grad:
+                self._handle_overflow_rollback()
+            self._loss_scale_update_and_overflow_cleanup()
+            return
+        self._loss_scale_update_and_overflow_cleanup()
         self._global_grad_norm = scaled_global_grad_norm / self.loss_scale
 
         timer_names = set()
@@ -221,6 +232,10 @@ def step(self, closure=None):
         self.timers(OPTIMIZER_STEP_TIMER).stop()
         self._post_step(timer_names)
 
+    def _post_step(self, timer_names):
+        super()._post_step(timer_names)
+        self._submitted_cpu_sub_groups.clear()
+
     def _step_without_clipping(self, scaled_global_grad_norm, timer_names):
         """Fast path: async CPU steps already completed during backward."""
         for sub_group_id, group in enumerate(self.fp16_groups):
@@ -332,17 +347,15 @@ def _sync_cpu_optimizer_step(self,
 
     def _handle_overflow_rollback(self):
         """Handle gradient overflow by rolling back CPU optimizer states."""
-        for sub_group_id, _ in enumerate(self.fp16_groups):
-            if self.subgroup_to_device[sub_group_id] == 'cpu':
-                param_group_id = self.sub_group_to_group_id[sub_group_id]
-                fp32_param = self.fp32_partitioned_groups_flat[sub_group_id]
-
-                # Trigger rollback
-                self._sync_cpu_optimizer_step(param_group_id,
-                                              sub_group_id,
-                                              fp32_param.data,
-                                              fp32_param.grad.data,
-                                              rollback=True)
+        for sub_group_id in self._submitted_cpu_sub_groups:
+            param_group_id = self.sub_group_to_group_id[sub_group_id]
+            fp32_param = self.fp32_partitioned_groups_flat[sub_group_id]
+            self._sync_cpu_optimizer_step(param_group_id,
+                                          sub_group_id,
+                                          fp32_param.data,
+                                          fp32_param.grad.data,
+                                          rollback=True)
+        self._submitted_cpu_sub_groups.clear()
 
     def _handle_gradient_clipping(self, scaled_global_grad_norm):
         """Handle gradient clipping with CPU optimizer rollback and re-optimization."""

deepspeed/runtime/utils.py

@@ -12,7 +12,7 @@
 import os
 import psutil
 import gc
-from math import sqrt
+from math import isfinite, sqrt
 
 from numpy import prod
 
@@ -227,6 +227,12 @@ def has_inf_or_nan(x):
     return torch.isfinite(x.amax()).logical_and(torch.isfinite(x.amin())).logical_not()
 
 
+def is_invalid_grad_norm(norm):
+    if torch.is_tensor(norm):
+        return bool((~torch.isfinite(norm) | (norm < 0)).any())
+    return not isfinite(float(norm)) or norm < 0
+
+
 class CheckOverflow(object):
     '''Checks for overflow in gradient across parallel process'''
 
@@ -245,7 +251,7 @@ def __init__(self, param_groups=None, mpu=None, zero_reduce_scatter=False, deeps
 
     def check_using_norm(self, norm_group, reduce_overflow=True):
         # TODO: I don't think reduce_overflow is needed if mpu is None
-        overflow = -1 in norm_group
+        overflow = any(is_invalid_grad_norm(norm) for norm in norm_group)
         overflow_gpu = get_accelerator().FloatTensor([overflow])
         if self.has_moe_params:
             # In this case, we need to do an all_reduce across
@@ -337,6 +343,8 @@ def _handle_overflow(cpu_sum, x, i):
 def get_global_norm(norm_list):
     """ Compute total from a list of norms
     """
+    if any(is_invalid_grad_norm(norm) for norm in norm_list):
+        return float("inf")
     total_norm = 0.0
     for norm in norm_list:
         total_norm += norm**2.0
@@ -879,6 +887,24 @@ def mask_nan_or_inf_with_val_inplace(input, device=None, val=-1.):
     input.masked_fill_(inf_or_nan, err)
 
 
+def combine_grad_norm_groups(norm_groups, norm_type=2):
+    """L2-combine per-group gradient norms without laundering the -1 invalid sentinel.
+
+    Several DeepSpeed helpers replace inf/NaN group norms with -1. Combining those
+    with ``vector_norm`` squares the sentinel to 1.0, so clipping treats a failed
+    group as a healthy unit-norm group. Non-finite or negative group norms stay
+    invalid (positive infinity) instead.
+    """
+    norm_groups = list(norm_groups)
+    device = next((norm.device for norm in norm_groups if torch.is_tensor(norm)), None)
+    stacked = torch.stack([
+        norm if torch.is_tensor(norm) else torch.tensor(norm, dtype=torch.float, device=device) for norm in norm_groups
+    ])
+    invalid = (~torch.isfinite(stacked) | (stacked < 0)).any()
+    combined = torch.linalg.vector_norm(stacked, ord=norm_type)
+    return torch.where(invalid, torch.full_like(combined, float("inf")), combined)
+
+
 def get_global_norm_of_tensors(input_tensors, norm_type=2, mpu=None, moe_ep_group=None):
     """Get norm of an iterable of tensors.
 
@@ -1109,7 +1135,7 @@ def to_tensor(v):
 
     # check if all norms are valid
     group_norms = torch.stack([to_tensor(norm) for norm in group_norms])
-    if group_norms.eq(-1).any():
+    if is_invalid_grad_norm(group_norms):
         return -1
 
     # combine norms