From 14f679da8acd3f083a9f2a5e2fe84d2e38f06351 Mon Sep 17 00:00:00 2001 From: BardiaKoopah Date: Thu, 11 Jun 2026 07:25:24 -0700 Subject: [PATCH] fix(studio): surface NaN loss honestly instead of laundering to last finite value (#6016) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When training produced a NaN or Inf loss event, the handler filtered the value to None but never updated progress.loss — clients kept seeing the last finite value as if everything were fine. Now: on non-finite loss, clear progress.loss to None and log a one-shot warning. Training continues (no phase=error, no _should_stop), matching the expected behavior for a non-fatal numerical event. Test: tests/test_training_nan_loss_handling.py with 6 cases covering finite, NaN, +/-Inf, idempotency of the one-shot warning, and recovery when a finite step follows a non-finite one. --- studio/backend/core/training/training.py | 19 +++- .../tests/test_training_nan_loss_handling.py | 106 ++++++++++++++++++ 2 files changed, 124 insertions(+), 1 deletion(-) create mode 100644 studio/backend/tests/test_training_nan_loss_handling.py diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index c5bc3dbad5..5bd9bc5338 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -534,8 +534,21 @@ class TrainingBackend: except (TypeError, ValueError): logger.debug("Could not convert loss to float: %s", _raw_loss) _safe_loss = None - if _safe_loss is not None and not math.isfinite(_safe_loss): + _loss_is_nonfinite = ( + _safe_loss is not None and not math.isfinite(_safe_loss) + ) + if _loss_is_nonfinite: + # Drop the value rather than laundering it back to the last + # finite loss; clients see loss=None at this step so the NaN + # is not hidden behind a stale value. Training continues. _safe_loss = None + if not getattr(self._progress, "_nonfinite_loss_warned", False): + self._progress._nonfinite_loss_warned = True + logger.warning( + "Training produced non-finite loss at step %s; " + "loss field will report null until it recovers.", + event.get("step", "?"), + ) try: _safe_lr = float(_raw_lr) if _raw_lr is not None else None except (TypeError, ValueError): @@ -545,6 +558,10 @@ class TrainingBackend: _safe_lr = None if _safe_loss is not None: self._progress.loss = _safe_loss + elif _loss_is_nonfinite: + # Clear stale finite loss so the API doesn't keep + # reporting the last good value while NaN is happening. + self._progress.loss = None if _safe_lr is not None: self._progress.learning_rate = _safe_lr self._progress.total_steps = event.get("total_steps", self._progress.total_steps) diff --git a/studio/backend/tests/test_training_nan_loss_handling.py b/studio/backend/tests/test_training_nan_loss_handling.py new file mode 100644 index 0000000000..0ab444f9bb --- /dev/null +++ b/studio/backend/tests/test_training_nan_loss_handling.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Pin Studio's behavior when a training event reports non-finite (NaN/Inf) loss. + +The training event handler used to filter NaN/Inf to None silently while +leaving the previous finite loss in progress.loss — so the API kept reporting +the stale value as if everything were fine. We now drop the stale value: +clients see loss=None at the affected step and a one-shot warning is logged. +Training continues; the run is not marked failed. +""" + +from __future__ import annotations + +import math +import os +import sys + +import pytest + +_BACKEND = os.path.join(os.path.dirname(__file__), "..") +if _BACKEND not in sys.path: + sys.path.insert(0, _BACKEND) + +from core.training.training import TrainingBackend + + +def _make_backend() -> TrainingBackend: + return TrainingBackend() + + +def _progress_event(step: int, loss: float, lr: float = 1e-4) -> dict: + return { + "type": "progress", + "step": step, + "loss": loss, + "learning_rate": lr, + "epoch": 0.0, + "total_steps": 100, + } + + +class TestNonfiniteLossSoftHandling: + def test_finite_loss_updates_progress_normally(self): + b = _make_backend() + b._handle_event(_progress_event(step=1, loss=0.97)) + assert b._progress.loss == pytest.approx(0.97) + assert b._progress.error is None + assert b._should_stop is False + assert getattr(b._progress, "_nonfinite_loss_warned", False) is False + + def test_nan_loss_clears_progress_loss(self): + b = _make_backend() + b._handle_event(_progress_event(step=1, loss=0.97)) + assert b._progress.loss == pytest.approx(0.97) + b._handle_event(_progress_event(step=2, loss=float("nan"))) + # Stale finite loss must NOT leak through + assert b._progress.loss is None + # Run is not marked failed + assert b._progress.error is None + assert b._should_stop is False + # Warning flag is set so we don't re-log on every subsequent NaN step + assert b._progress._nonfinite_loss_warned is True + + def test_inf_loss_clears_progress_loss(self): + b = _make_backend() + b._handle_event(_progress_event(step=1, loss=float("inf"))) + assert b._progress.loss is None + assert b._progress.error is None + assert b._should_stop is False + assert b._progress._nonfinite_loss_warned is True + + def test_negative_inf_loss_clears_progress_loss(self): + b = _make_backend() + b._handle_event(_progress_event(step=1, loss=float("-inf"))) + assert b._progress.loss is None + assert b._progress.error is None + assert b._should_stop is False + assert b._progress._nonfinite_loss_warned is True + + def test_repeated_nan_only_warns_once(self): + """Subsequent NaN events must not re-fire the warning flag setter. + The flag should already be True after the first NaN.""" + b = _make_backend() + b._handle_event(_progress_event(step=1, loss=0.97)) + b._handle_event(_progress_event(step=2, loss=float("nan"))) + assert b._progress._nonfinite_loss_warned is True + # Further NaN steps don't change anything we care about + b._handle_event(_progress_event(step=3, loss=float("nan"))) + b._handle_event(_progress_event(step=4, loss=float("nan"))) + assert b._progress._nonfinite_loss_warned is True + assert b._progress.loss is None + assert b._progress.error is None + assert b._should_stop is False + + def test_recovery_updates_loss_when_finite_again(self): + """If a NaN step is followed by a finite step, progress.loss must + reflect the new finite value (not stay stuck at None).""" + b = _make_backend() + b._handle_event(_progress_event(step=1, loss=0.97)) + b._handle_event(_progress_event(step=2, loss=float("nan"))) + assert b._progress.loss is None + b._handle_event(_progress_event(step=3, loss=0.85)) + assert b._progress.loss == pytest.approx(0.85) + # Warning flag stays set (we don't reset it on recovery) + assert b._progress._nonfinite_loss_warned is True