Spaces:
Running on L40S
Running on L40S
File size: 2,945 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
from __future__ import annotations
import torch
import torch.distributed as dist
from cosmos_framework.model._base import ImaginaireModel
from cosmos_framework.utils import log
from cosmos_framework.utils.callback import Callback
class SkipNaNStep(Callback):
"""Skip the optimizer step only when ALL ranks produce NaN/Inf loss.
When only some ranks produce NaN, the existing GradClip callback's
nan_to_num handling is sufficient (NaN gradients become 0, valid
gradients from clean ranks are still used). This callback only
intervenes when every rank has NaN, meaning no useful gradient
signal exists.
The all-reduce ensures all ranks agree on skip/no-skip, preventing
NCCL desync.
Args:
max_consecutive_nan: Abort training after this many consecutive
all-rank-NaN optimizer steps. Set to 0 to disable the limit.
"""
def __init__(self, max_consecutive_nan: int = 100) -> None:
super().__init__()
self.max_consecutive_nan = max_consecutive_nan
self._nan_detected = False
self._consecutive_nan_count = 0
def on_before_backward(
self,
model: ImaginaireModel,
loss: torch.Tensor,
iteration: int = 0,
) -> None:
if torch.isnan(loss).any() or torch.isinf(loss).any():
self._nan_detected = True
def on_before_optimizer_step(
self,
model: ImaginaireModel,
optimizer: torch.optim.Optimizer,
scheduler: torch.optim.lr_scheduler.LRScheduler,
grad_scaler: torch.amp.GradScaler,
iteration: int = 0,
) -> None:
nan_flag = torch.tensor([1.0 if self._nan_detected else 0.0], device="cuda")
dist.all_reduce(nan_flag, op=dist.ReduceOp.SUM)
nan_rank_count = int(nan_flag.item())
world_size = dist.get_world_size()
if nan_rank_count > 0 and nan_rank_count < world_size:
self._consecutive_nan_count = 0
elif nan_rank_count == world_size:
for param in model.parameters():
if param.grad is not None:
param.grad.zero_()
self._consecutive_nan_count += 1
log.warning(
f"ALL ranks NaN/Inf at iteration {iteration}, skipping optimizer step "
f"(consecutive: {self._consecutive_nan_count})",
)
if self.max_consecutive_nan > 0 and self._consecutive_nan_count >= self.max_consecutive_nan:
raise RuntimeError(
f"Training unstable: all-rank NaN/Inf loss for {self._consecutive_nan_count} "
f"consecutive optimizer steps at iteration {iteration}. Aborting.",
)
else:
self._consecutive_nan_count = 0
self._nan_detected = False
|