diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 0af3349c6f..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..." ) @@ -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( 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 new file mode 100644 index 0000000000..3ea9e1e5f9 --- /dev/null +++ b/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 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;