From d5790a54979edf5f3a2805c33beddbcc58081693 Mon Sep 17 00:00:00 2001 From: Bardia Koopah Date: Thu, 4 Jun 2026 14:45:23 -0700 Subject: [PATCH 1/2] fix(studio): surface NaN loss honestly instead of laundering to last finite value 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 0af3349c6f..fc5e0b00df 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -561,8 +561,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): @@ -574,6 +587,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( 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 From f8854772ea874857a7403cd827d153a8c03a0394 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 11 Jun 2026 12:34:23 +0000 Subject: [PATCH 2/2] Fix NaN loss surfacing in /metrics, SSE stream and frontend store for PR #6016 --- studio/backend/core/training/training.py | 13 +- studio/backend/routes/training.py | 98 +++++++------ .../tests/test_training_nan_loss_handling.py | 49 ++++--- .../tests/test_training_nan_loss_routes.py | 132 ++++++++++++++++++ .../training/stores/training-runtime-store.ts | 17 ++- .../src/features/training/types/runtime.ts | 7 +- 6 files changed, 244 insertions(+), 72 deletions(-) create mode 100644 studio/backend/tests/test_training_nan_loss_routes.py diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index fc5e0b00df..0fc7ff7be5 100644 --- a/studio/backend/core/training/training.py +++ b/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..." ) @@ -565,12 +567,10 @@ class TrainingBackend: _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. + # Report None instead of the stale finite loss; run continues _safe_loss = None - if not getattr(self._progress, "_nonfinite_loss_warned", False): - self._progress._nonfinite_loss_warned = True + 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.", @@ -588,8 +588,7 @@ class TrainingBackend: 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. + # 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 diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 41a9e15562..c27546a753 100644 --- a/studio/backend/routes/training.py +++ b/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, diff --git a/studio/backend/tests/test_training_nan_loss_handling.py b/studio/backend/tests/test_training_nan_loss_handling.py index 0ab444f9bb..3ea9e1e5f9 100644 --- a/studio/backend/tests/test_training_nan_loss_handling.py +++ b/studio/backend/tests/test_training_nan_loss_handling.py @@ -3,11 +3,10 @@ """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. +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 @@ -47,7 +46,7 @@ class TestNonfiniteLossSoftHandling: 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 + assert b._nonfinite_loss_warned is False def test_nan_loss_clears_progress_loss(self): b = _make_backend() @@ -60,7 +59,7 @@ class TestNonfiniteLossSoftHandling: 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 + assert b._nonfinite_loss_warned is True def test_inf_loss_clears_progress_loss(self): b = _make_backend() @@ -68,7 +67,7 @@ class TestNonfiniteLossSoftHandling: 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 + assert b._nonfinite_loss_warned is True def test_negative_inf_loss_clears_progress_loss(self): b = _make_backend() @@ -76,19 +75,37 @@ class TestNonfiniteLossSoftHandling: 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 + assert b._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.""" + 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._progress._nonfinite_loss_warned is True - # Further NaN steps don't change anything we care about + 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 b._progress._nonfinite_loss_warned is True + 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 @@ -103,4 +120,4 @@ class TestNonfiniteLossSoftHandling: 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 + assert b._nonfinite_loss_warned is True diff --git a/studio/backend/tests/test_training_nan_loss_routes.py b/studio/backend/tests/test_training_nan_loss_routes.py new file mode 100644 index 0000000000..2252923613 --- /dev/null +++ b/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) diff --git a/studio/frontend/src/features/training/stores/training-runtime-store.ts b/studio/frontend/src/features/training/stores/training-runtime-store.ts index f1966f03fb..a74d69a040 100644 --- a/studio/frontend/src/features/training/stores/training-runtime-store.ts +++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts @@ -212,8 +212,13 @@ export const useTrainingRuntimeStore = create()((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()((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()((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, diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts index d30c3f75c2..688a336a6c 100644 --- a/studio/frontend/src/features/training/types/runtime.ts +++ b/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;