fix(studio): surface NaN loss honestly instead of laundering to last finite value (#6016)

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

@@ -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)

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