diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 5bd9bc5338..2764101d93 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -534,9 +534,7 @@ class TrainingBackend: except (TypeError, ValueError): logger.debug("Could not convert loss to float: %s", _raw_loss) _safe_loss = None - _loss_is_nonfinite = ( - _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 diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 427687451c..18b25cb4fe 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1858,7 +1858,7 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 11. Run training ── gc.collect() mx.synchronize() - trainer.train(resume_from_checkpoint=resume_from_checkpoint) + trainer.train(resume_from_checkpoint = resume_from_checkpoint) # ── 12. Save and finalize ── if trainer.stop_requested and not _stop_save[0]: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 96403e21a5..d7687ffdee 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -665,10 +665,17 @@ 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] + _live = (getattr(tp, "step", 0) or 0) if tp else 0 + if backend.step_history or _live > 0: + final_step = backend.step_history[-1] if backend.step_history else 0 final_loss = backend.loss_history[-1] if backend.loss_history else None final_lr = backend.lr_history[-1] if backend.lr_history else None + # Histories skip non-finite steps; report the live step with + # loss=None instead of the last finite pair. + if _live > final_step: + final_step = _live + final_loss = getattr(tp, "loss", None) + final_lr = getattr(tp, "learning_rate", final_lr) final_total_steps = getattr(tp, "total_steps", final_step) if tp else final_step final_epoch = getattr(tp, "epoch", None) if tp else None payload = build_progress( @@ -697,11 +704,18 @@ async def stream_training_progress( while backend.is_training_active(): try: - if backend.step_history: - current_step = backend.step_history[-1] + tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None) + live_step = (getattr(tp_inner, "step", 0) or 0) if tp_inner else 0 + if backend.step_history or live_step > 0: + current_step = backend.step_history[-1] if backend.step_history else 0 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) + # Histories skip non-finite steps; follow the live progress + # step and report its loss (None until it recovers). + if live_step > current_step: + current_step = live_step + current_loss = getattr(tp_inner, "loss", None) + current_lr = getattr(tp_inner, "learning_rate", current_lr) current_total_steps = ( getattr(tp_inner, "total_steps", current_step) if tp_inner else current_step ) @@ -798,6 +812,13 @@ async def stream_training_progress( 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) + # If the run ended on a non-finite stretch, report the live step with + # loss=None instead of rolling back to the last finite pair. + _final_live_step = (getattr(final_tp, "step", 0) or 0) if final_tp else 0 + if _final_live_step > (final_step if final_step is not None else -1): + final_step = _final_live_step + final_loss = getattr(final_tp, "loss", None) + final_lr = getattr(final_tp, "learning_rate", final_lr) 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_payload = build_progress( diff --git a/studio/backend/tests/test_training_nan_loss_handling.py b/studio/backend/tests/test_training_nan_loss_handling.py index 0ab444f9bb..a2dc78bee2 100644 --- a/studio/backend/tests/test_training_nan_loss_handling.py +++ b/studio/backend/tests/test_training_nan_loss_handling.py @@ -29,7 +29,11 @@ def _make_backend() -> TrainingBackend: return TrainingBackend() -def _progress_event(step: int, loss: float, lr: float = 1e-4) -> dict: +def _progress_event( + step: int, + loss: float, + lr: float = 1e-4, +) -> dict: return { "type": "progress", "step": step, @@ -43,7 +47,7 @@ def _progress_event(step: int, loss: float, lr: float = 1e-4) -> dict: class TestNonfiniteLossSoftHandling: def test_finite_loss_updates_progress_normally(self): b = _make_backend() - b._handle_event(_progress_event(step=1, loss=0.97)) + 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 @@ -51,9 +55,9 @@ class TestNonfiniteLossSoftHandling: def test_nan_loss_clears_progress_loss(self): b = _make_backend() - b._handle_event(_progress_event(step=1, loss=0.97)) + 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"))) + 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 @@ -64,7 +68,7 @@ class TestNonfiniteLossSoftHandling: def test_inf_loss_clears_progress_loss(self): b = _make_backend() - b._handle_event(_progress_event(step=1, loss=float("inf"))) + 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 @@ -72,7 +76,7 @@ class TestNonfiniteLossSoftHandling: def test_negative_inf_loss_clears_progress_loss(self): b = _make_backend() - b._handle_event(_progress_event(step=1, loss=float("-inf"))) + 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 @@ -82,12 +86,12 @@ class TestNonfiniteLossSoftHandling: """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"))) + 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"))) + 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 @@ -97,10 +101,10 @@ class TestNonfiniteLossSoftHandling: """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"))) + 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)) + 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 diff --git a/studio/backend/tests/test_training_progress_stream_nan.py b/studio/backend/tests/test_training_progress_stream_nan.py new file mode 100644 index 0000000000..899527a04d --- /dev/null +++ b/studio/backend/tests/test_training_progress_stream_nan.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The SSE progress stream must follow the live progress step during +non-finite-loss stretches (loss reported as null) instead of replaying the +last finite step/loss pair from the metric histories, which skip NaN steps.""" + +import asyncio +import json +import sys +import types + +import pytest + +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ) + +import routes.training as rt + + +class _Progress: + def __init__(self): + self.step = 5 + self.total_steps = 10 + self.loss = None # cleared by the NaN honesty fix in core training + self.learning_rate = 8e-5 + self.epoch = 0.1 + self.grad_norm = None + self.num_tokens = None + self.eval_loss = None + self.elapsed_seconds = None + self.eta_seconds = None + + +class _FakeBackend: + """Finite history stops at step 2; live progress is at step 5 with NaN + (loss=None). Active for a few polls, then done.""" + + def __init__(self, active_polls = 2): + self.current_job_id = "job-1" + self.step_history = [1, 2] + self.loss_history = [2.0, 1.5] + self.lr_history = [1e-4, 9e-5] + self.eval_enabled = False + self._active_calls = 0 + self._active_polls = active_polls + self.trainer = types.SimpleNamespace(training_progress = _Progress()) + + def is_training_active(self): + self._active_calls += 1 + return self._active_calls <= self._active_polls + + +class _FakeRequest: + headers = {} + + +def _collect_events(response, timeout = 15): + async def _drain(): + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + return "".join(c.decode() if isinstance(c, bytes) else c for c in chunks) + + return asyncio.run(asyncio.wait_for(_drain(), timeout)) + + +def _progress_payloads(raw): + payloads = [] + for block in raw.split("\n\n"): + lines = block.strip().splitlines() + data = next((l[6:] for l in lines if l.startswith("data: ")), None) + if data: + payloads.append(json.loads(data)) + return payloads + + +def test_stream_reports_live_step_with_null_loss_during_nan(monkeypatch): + backend = _FakeBackend(active_polls = 2) + monkeypatch.setattr(rt, "get_training_backend", lambda: backend) + + response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester")) + raw = _collect_events(response) + payloads = _progress_payloads(raw) + assert payloads, f"no SSE payloads parsed from: {raw!r}" + + live = [p for p in payloads if p.get("step") == 5] + assert live, ( + "stream never advanced to the live progress step during the NaN " + f"stretch; steps seen: {[p.get('step') for p in payloads]}" + ) + assert live[0]["loss"] is None + # The stale finite pair must not be re-emitted as the latest progress. + stale = [p for p in payloads if p.get("step") == 2 and p.get("loss") == 1.5] + assert not stale + + +def test_inactive_stream_completes_with_live_step_and_null_loss(monkeypatch): + # Fresh connection after the run already ended during a NaN stretch: the + # immediate complete event must not replay the stale finite pair either. + backend = _FakeBackend(active_polls = 0) + monkeypatch.setattr(rt, "get_training_backend", lambda: backend) + + response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester")) + payloads = _progress_payloads(_collect_events(response)) + final = payloads[-1] + assert final["step"] == 5 + assert final["loss"] is None + + +def test_stream_uses_finite_history_when_progress_in_sync(monkeypatch): + backend = _FakeBackend(active_polls = 2) + # Live progress agrees with the history tail: normal finite behavior. + backend.trainer.training_progress.step = 2 + backend.trainer.training_progress.loss = 1.5 + monkeypatch.setattr(rt, "get_training_backend", lambda: backend) + + response = asyncio.run(rt.stream_training_progress(_FakeRequest(), current_subject = "tester")) + payloads = _progress_payloads(_collect_events(response)) + finite = [p for p in payloads if p.get("step") == 2] + assert finite and finite[0]["loss"] == 1.5 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..97fbd32d57 100644 --- a/studio/frontend/src/features/training/stores/training-runtime-store.ts +++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts @@ -274,7 +274,10 @@ 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, + // A null loss at a new step means the backend reported a non-finite + // loss; clear the display instead of keeping the stale value. + currentLoss: + currentLoss ?? (step > state.currentStep ? null : 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..5fd8286d22 100644 --- a/studio/frontend/src/features/training/types/runtime.ts +++ b/studio/frontend/src/features/training/types/runtime.ts @@ -90,7 +90,8 @@ export interface TrainingRuntimeState { currentStep: number; totalSteps: number; currentEpoch: number; - currentLoss: number; + // null while the latest reported loss is non-finite + currentLoss: number | null; currentLearningRate: number; progressPercent: number; elapsedSeconds: number | null;