Fix NaN loss surfacing in /metrics, SSE stream and frontend store for PR #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

@@ -135,6 +135,7 @@ class TrainingBackend:
         self._progress = TrainingProgress()
         self._should_stop = False
         self._cancel_requested = False  # True only for stop(save=False)
+        self._nonfinite_loss_warned = False  # one-shot warning per run
 
         # Training Metrics (consumed by routes for SSE and /metrics)
         self.loss_history: list = []
@@ -305,6 +306,7 @@ class TrainingBackend:
         self.current_job_id = job_id
         self._should_stop = False
         self._cancel_requested = False
+        self._nonfinite_loss_warned = False
         self._progress = TrainingProgress(
             is_training = True, status_message = "Initializing training..."
         )
@@ -561,8 +563,19 @@ 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:
+                    # Report None instead of the stale finite loss; run continues
                     _safe_loss = None
+                    if not self._nonfinite_loss_warned:
+                        self._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):
@@ -574,6 +587,9 @@ class TrainingBackend:
                     _safe_lr = None
                 if _safe_loss is not None:
                     self._progress.loss = _safe_loss
+                elif _loss_is_nonfinite:
+                    # Clear the stale finite loss while loss is non-finite
+                    self._progress.loss = None
                 if _safe_lr is not None:
                     self._progress.learning_rate = _safe_lr
                 self._progress.total_steps = event.get(

studio/backend/routes/training.py

@@ -528,10 +528,17 @@ async def get_training_metrics(
         grad_norm_history = getattr(backend, "grad_norm_history", [])
         grad_norm_step_history = getattr(backend, "grad_norm_step_history", [])
 
-        # Get current values
-        current_loss = loss_history[-1] if loss_history else None
-        current_lr = lr_history[-1] if lr_history else None
-        current_step = step_history[-1] if step_history else None
+        # Current values come from live progress: histories are finite-only
+        # and would replay the last finite loss after a NaN/Inf step
+        progress = getattr(backend, "_progress", None)
+        if progress is not None and getattr(progress, "step", 0) > 0:
+            current_loss = progress.loss
+            current_lr = progress.learning_rate
+            current_step = progress.step
+        else:
+            current_loss = loss_history[-1] if loss_history else None
+            current_lr = lr_history[-1] if lr_history else None
+            current_step = step_history[-1] if step_history else None
 
         return TrainingMetricsResponse(
             loss_history = loss_history,
@@ -642,6 +649,25 @@ async def stream_training_progress(
             lines.append("")  # double newline terminates the event
             return "\n".join(lines)
 
+        def current_progress_values(fallback_step: int):
+            """Current sample from live progress. Histories are finite-only,
+            so they would replay the last finite loss during NaN/Inf steps."""
+            tp = getattr(backend, "_progress", None)
+            if tp is not None and getattr(tp, "step", 0) > 0:
+                total = tp.total_steps or tp.step
+                return tp.step, tp.loss, tp.learning_rate, total, tp.epoch, tp
+            if backend.step_history:
+                step = backend.step_history[-1]
+                return (
+                    step,
+                    backend.loss_history[-1] if backend.loss_history else None,
+                    backend.lr_history[-1] if backend.lr_history else None,
+                    step,
+                    None,
+                    tp,
+                )
+            return fallback_step, None, None, 0, None, tp
+
         # ── Retry directive ──────────────────────────────────────
         # Tell the browser to reconnect after 3 seconds if the connection drops
         yield "retry: 3000\n\n"
@@ -714,23 +740,22 @@ async def stream_training_progress(
 
             # If not active, send final state and exit
             if not is_active:
-                if backend.step_history:
-                    final_step = backend.step_history[-1]
-                    final_loss = (
-                        backend.loss_history[-1] if backend.loss_history else None
-                    )
-                    final_lr = backend.lr_history[-1] if backend.lr_history else None
-                    final_total_steps = (
-                        getattr(tp, "total_steps", final_step) if tp else final_step
-                    )
-                    final_epoch = getattr(tp, "epoch", None) if tp else None
+                (
+                    final_step,
+                    final_loss,
+                    final_lr,
+                    final_total_steps,
+                    final_epoch,
+                    tp_final,
+                ) = current_progress_values(-1)
+                if final_step > 0:
                     payload = build_progress(
                         final_step,
                         final_loss,
                         final_lr,
                         final_total_steps,
                         final_epoch,
-                        progress = tp,
+                        progress = tp_final,
                     )
                     yield format_sse(
                         payload.model_dump_json(), event = "complete", event_id = final_step
@@ -754,24 +779,15 @@ async def stream_training_progress(
 
         while backend.is_training_active():
             try:
-                if backend.step_history:
-                    current_step = backend.step_history[-1]
-                    current_loss = (
-                        backend.loss_history[-1] if backend.loss_history else None
-                    )
-                    current_lr = backend.lr_history[-1] if backend.lr_history else None
-                    tp_inner = getattr(
-                        getattr(backend, "trainer", None), "training_progress", None
-                    )
-                    current_total_steps = (
-                        getattr(tp_inner, "total_steps", current_step)
-                        if tp_inner
-                        else current_step
-                    )
-                    current_epoch = (
-                        getattr(tp_inner, "epoch", None) if tp_inner else None
-                    )
-
+                (
+                    current_step,
+                    current_loss,
+                    current_lr,
+                    current_total_steps,
+                    current_epoch,
+                    tp_inner,
+                ) = current_progress_values(0)
+                if current_step > 0:
                     # Only send if step changed
                     if current_step != last_step:
                         progress_payload = build_progress(
@@ -865,14 +881,14 @@ async def stream_training_progress(
                 break
 
         # ── Final "complete" event ───────────────────────────────
-        final_step = backend.step_history[-1] if backend.step_history else last_step
-        final_loss = backend.loss_history[-1] if backend.loss_history else None
-        final_lr = backend.lr_history[-1] if backend.lr_history else None
-        final_tp = getattr(getattr(backend, "trainer", None), "training_progress", None)
-        final_total_steps = (
-            getattr(final_tp, "total_steps", final_step) if final_tp else final_step
-        )
-        final_epoch = getattr(final_tp, "epoch", None) if final_tp else None
+        (
+            final_step,
+            final_loss,
+            final_lr,
+            final_total_steps,
+            final_epoch,
+            final_tp,
+        ) = current_progress_values(last_step)
         final_payload = build_progress(
             final_step,
             final_loss,

studio/backend/tests/test_training_nan_loss_handling.py

@@ -0,0 +1,123 @@
+# 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 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. We now clear it: 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 b._nonfinite_loss_warned 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._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._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._nonfinite_loss_warned is True
+
+    def test_repeated_nan_only_warns_once(self, monkeypatch):
+        """Only the first NaN logs a warning; later NaNs stay quiet."""
+        import core.training.training as training_module
+
+        warnings = []
+        monkeypatch.setattr(
+            training_module,
+            "logger",
+            type(
+                "LoggerStub",
+                (),
+                {
+                    "warning": lambda self, *a, **k: warnings.append(a),
+                    "info": lambda self, *a, **k: None,
+                    "debug": lambda self, *a, **k: None,
+                    "error": lambda self, *a, **k: 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._nonfinite_loss_warned is True
+        assert len(warnings) == 1
+        # Further NaN steps stay quiet
+        b._handle_event(_progress_event(step=3, loss=float("nan")))
+        b._handle_event(_progress_event(step=4, loss=float("nan")))
+        assert len(warnings) == 1
+        assert b._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._nonfinite_loss_warned is True

studio/backend/tests/test_training_nan_loss_routes.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+"""Route-level regression tests for non-finite loss reporting.
+
+/metrics and the SSE stream used to derive "current" values from the
+finite-only history arrays, which replayed the last finite loss during
+NaN/Inf steps. They must read live progress instead.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+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
+import routes.training as training_routes
+
+
+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,
+    }
+
+
+def _finish_run(b: TrainingBackend) -> TrainingBackend:
+    """Mark the run as finished so SSE takes the final-state path."""
+    b._progress.is_training = False
+    b._progress.is_completed = True
+    return b
+
+
+def _backend_after_nan() -> TrainingBackend:
+    b = TrainingBackend()
+    b._handle_event(_progress_event(step=1, loss=0.97))
+    b._handle_event(_progress_event(step=2, loss=float("nan")))
+    return b
+
+
+class _FakeRequest:
+    headers: dict = {}
+
+
+def _collect_sse_events(response) -> list[tuple[str, dict]]:
+    """Drain a StreamingResponse of SSE messages into (event, payload) pairs."""
+
+    async def drain():
+        chunks = []
+        async for chunk in response.body_iterator:
+            chunks.append(chunk)
+        return "".join(chunks)
+
+    raw = asyncio.run(drain())
+    events = []
+    for block in raw.split("\n\n"):
+        event_name, data = None, None
+        for line in block.splitlines():
+            if line.startswith("event: "):
+                event_name = line[len("event: "):]
+            elif line.startswith("data: "):
+                data = json.loads(line[len("data: "):])
+        if event_name is not None and data is not None:
+            events.append((event_name, data))
+    return events
+
+
+class TestNonfiniteLossRoutes:
+    def test_metrics_reports_null_loss_and_current_step_after_nan(self, monkeypatch):
+        b = _backend_after_nan()
+        monkeypatch.setattr(training_routes, "get_training_backend", lambda: b)
+        resp = asyncio.run(
+            training_routes.get_training_metrics(current_subject="test")
+        )
+        # Live progress, not the stale finite history point
+        assert resp.current_step == 2
+        assert resp.current_loss is None
+        # Chart history stays finite-only
+        assert resp.loss_history == [0.97]
+        assert resp.step_history == [1]
+
+    def test_metrics_falls_back_to_history_when_no_progress(self, monkeypatch):
+        b = TrainingBackend()
+        monkeypatch.setattr(training_routes, "get_training_backend", lambda: b)
+        resp = asyncio.run(
+            training_routes.get_training_metrics(current_subject="test")
+        )
+        assert resp.current_step is None
+        assert resp.current_loss is None
+
+    def test_sse_complete_event_reports_nan_step_with_null_loss(self, monkeypatch):
+        b = _finish_run(_backend_after_nan())
+        monkeypatch.setattr(training_routes, "get_training_backend", lambda: b)
+        resp = asyncio.run(
+            training_routes.stream_training_progress(
+                _FakeRequest(), current_subject="test"
+            )
+        )
+        events = _collect_sse_events(resp)
+        completes = [payload for name, payload in events if name == "complete"]
+        assert len(completes) == 1
+        # The NaN step is surfaced, not the last finite one
+        assert completes[0]["step"] == 2
+        assert completes[0]["loss"] is None
+
+    def test_sse_complete_event_reports_finite_loss_normally(self, monkeypatch):
+        b = TrainingBackend()
+        b._handle_event(_progress_event(step=1, loss=0.97))
+        _finish_run(b)
+        monkeypatch.setattr(training_routes, "get_training_backend", lambda: b)
+        resp = asyncio.run(
+            training_routes.stream_training_progress(
+                _FakeRequest(), current_subject="test"
+            )
+        )
+        events = _collect_sse_events(resp)
+        completes = [payload for name, payload in events if name == "complete"]
+        assert len(completes) == 1
+        assert completes[0]["step"] == 1
+        assert completes[0]["loss"] == pytest.approx(0.97)

studio/frontend/src/features/training/stores/training-runtime-store.ts

@@ -212,8 +212,13 @@ export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => (
           typeof detailTotal === "number"
             ? Math.max(detailTotal, 0)
             : state.totalSteps,
+        // Explicit null means a non-finite loss step; do not keep the stale value
         currentLoss:
-          typeof detailLoss === "number" ? detailLoss : state.currentLoss,
+          detailLoss === null
+            ? null
+            : typeof detailLoss === "number"
+              ? detailLoss
+              : state.currentLoss,
         currentLearningRate:
           typeof detailLr === "number" ? detailLr : state.currentLearningRate,
         currentEpoch:
@@ -251,9 +256,11 @@ export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => (
             ? Math.max(latestStep, state.currentStep)
             : state.currentStep,
         currentLoss:
-          typeof payload.current_loss === "number"
-            ? payload.current_loss
-            : state.currentLoss,
+          payload.current_loss === null
+            ? null
+            : typeof payload.current_loss === "number"
+              ? payload.current_loss
+              : state.currentLoss,
         currentLearningRate:
           typeof payload.current_lr === "number"
             ? payload.current_lr
@@ -274,7 +281,7 @@ export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => (
         jobId: payload.job_id || state.jobId,
         currentStep: step,
         totalSteps: Math.max(payload.total_steps, state.totalSteps),
-        currentLoss: currentLoss ?? state.currentLoss,
+        currentLoss: payload.loss === null ? null : (currentLoss ?? state.currentLoss),
         currentLearningRate: currentLearningRate ?? state.currentLearningRate,
         progressPercent: payload.progress_percent,
         currentEpoch: payload.epoch ?? state.currentEpoch,

studio/frontend/src/features/training/types/runtime.ts

@@ -24,8 +24,8 @@ export interface TrainingStatusResponse {
     epoch?: number;
     step?: number;
     total_steps?: number;
-    loss?: number;
-    learning_rate?: number;
+    loss?: number | null;
+    learning_rate?: number | null;
     output_dir?: string;
   } | null;
   metric_history?: {
@@ -90,7 +90,8 @@ export interface TrainingRuntimeState {
   currentStep: number;
   totalSteps: number;
   currentEpoch: number;
-  currentLoss: number;
+  // null means the latest step reported a non-finite (NaN/Inf) loss
+  currentLoss: number | null;
   currentLearningRate: number;
   progressPercent: number;
   elapsedSeconds: number | null;