3733e0b274* fix(studio): surface live step with null loss through the SSE progress stream The metric histories skip non-finite steps, so during a NaN stretch the SSE live loop and final complete event replayed the last finite step/loss pair. Follow the live progress step when it is ahead of the history tail and report its loss honestly (null until recovery). Completes the NaN honesty fix for the SSE consumer flagged in review. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Apply live-step handling to inactive streams and clear the UI loss on null for PR #6206 Fresh /progress connections after a finished run took the inactive branch which still replayed the last finite step and loss pair; apply the same live-step correction there. On the frontend, applyProgress kept the stale currentLoss when a payload advanced the step with a null loss; clear it so the display shows -- until the loss recovers. Widen the runtime state type to number | null, which the view layer already handles. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
110 lines
4.2 KiB
Python
110 lines
4.2 KiB
Python
# 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
|