ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix NaN loss surfacing in /metrics, SSE stream and frontend store for PR #6016

Browse source
This commit is contained in:
Daniel Han 2026-06-11 12:34:23 +00:00
commit f8854772ea
6 changed files with 245 additions and 73 deletions
Download patch file Download diff file Expand all files Collapse all files

13
studio/backend/core/training/training.py
View file

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

98
studio/backend/routes/training.py
View file

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

49
studio/backend/tests/test_training_nan_loss_handling.py
View file

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

132
studio/backend/tests/test_training_nan_loss_routes.py Normal file
View file

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

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

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

7
studio/frontend/src/features/training/types/runtime.ts
View file

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