diff --git a/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py b/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py index 870b64d6ae..608032a434 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/training/callbacks.py @@ -4,24 +4,82 @@ """Training progress callback shared by the customization backends. Composes a :class:`nmp.customization_common.training.progress.JobsServiceProgressReporter` -and provides training-specific methods. Metric accumulation: ``train_loss`` and -``val_loss`` are accumulated as time-series lists and included in every -``status_details`` update under a ``metrics`` key, enabling loss-curve -reconstruction from job status. +and provides training-specific methods. Every numeric metric a backend reports is +accumulated as a time series and included in EVERY ``status_details`` update under +a ``metrics`` key, so any of them can be charted from job status alone. + +Every update matters because ``report_running`` REPLACES the task's +``status_details`` blob rather than merging into it. A report that omits +``metrics`` therefore erases the accumulated series from stored status until the +next train step resends it -- and if the job dies inside that window, the curve +is gone. Checkpoint and epoch-end reports fire mid-training, so they carry the +payload too. + +Series naming +------------- +Series are namespaced by the phase that produced them: ``train_`` and +``val_``, which is what the long-standing ``train_loss``/``val_loss`` pair +already did. The prefix is load-bearing rather than cosmetic -- GRPO reports +``truncation_rate`` in both its train and validation dicts, and DPO reports +``accuracy`` in both, so unprefixed names would interleave two different +quantities into one series. + +``train_loss`` and ``val_loss`` keep those exact names, so existing consumers +(the Studio loss chart) are unaffected. + +Payload size +------------ +Every series is resent in full on every update, so the stored blob grows as +``series x reports`` and total upload as the square of it. The driver of that +cost is the number of *reports*, not training steps -- backends throttle +reporting, so a 500-step GRPO run at ``log_interval=10`` accumulates 50 points +per series, not 500. + +Measured, for GRPO's ~22 series: + + 500 steps, log_interval 10 -> 42 KB final blob, 1.1 MB uploaded + 500 steps, log_interval 1 -> 413 KB final blob, 101.3 MB uploaded + +Deliberately accepted for batch training jobs. It does mean a backend that +reports every step of a long run pays quadratically, so if that becomes a real +configuration the transport should move to delta appends rather than the series +being trimmed here. Backends subclass this and set :attr:`_default_backend`: unsloth stamps a -``backend`` field on each report (``"unsloth"``); automodel leaves it ``None`` so -no ``backend`` key is added (preserving its status-detail shape). Callers may also -pass ``backend`` per call (e.g. unsloth's HF trainer callback). +``backend`` field on each report (``"unsloth"``); automodel and NeMo-RL leave it +``None`` so no ``backend`` key is added (preserving their status-detail shape). +Callers may also pass ``backend`` per call (e.g. unsloth's HF trainer callback). """ import logging -from typing import ClassVar +import math +import numbers +from typing import Any, ClassVar, cast from nmp.customization_common.training.progress import JobsServiceProgressReporter logger = logging.getLogger(__name__) +#: Series that keep their bare name instead of taking a phase prefix, because +#: they predate the prefixing scheme and are read by name downstream. +_UNPREFIXED = frozenset({"train_loss", "val_loss"}) + + +def is_chartable(value: Any) -> bool: + """Whether ``value`` is a finite scalar that can enter a metric series. + + Backends hand us whatever their framework produced, which is not always a + number: NeMo-RL's metric dicts interleave ``Histogram`` objects, tables and + nested dicts with the scalars, and ``math.isnan`` raises ``TypeError`` on all + of those rather than returning False. + + ``bool`` is rejected despite being an ``int`` subclass: no metric here is a + flag, and silently charting one as 0/1 is worse than dropping it. + """ + if isinstance(value, bool) or not isinstance(value, numbers.Real): + return False + return not math.isnan(float(value)) + class TrainingProgressCallback: """Report training progress to the Jobs service.""" @@ -33,30 +91,56 @@ class TrainingProgressCallback: def __init__(self, reporter: JobsServiceProgressReporter): self._reporter = reporter - prior = reporter.fetch_current_metrics() - self._train_metrics: list[dict[str, float | int]] = prior.get("train_loss", []) - self._val_metrics: list[dict[str, float | int]] = prior.get("val_loss", []) - if self._train_metrics or self._val_metrics: + #: series name -> [{step, epoch, value}], seeded from the server so a + #: resumed job continues its curves instead of restarting them. + self._series: dict[str, list[dict[str, float | int]]] = dict(reporter.fetch_current_metrics()) + if any(self._series.values()): logger.info( - "Seeded metrics from server: %d train_loss, %d val_loss entries", - len(self._train_metrics), - len(self._val_metrics), + "Seeded %d metric series from server (%d points): %s", + len(self._series), + sum(len(points) for points in self._series.values()), + ", ".join(sorted(self._series)), ) def _resolve_backend(self, backend: str | None) -> str | None: return backend if backend is not None else self._default_backend + def _record(self, phase: str, name: str, step: int, epoch: int, value: object) -> None: + """Append one point to the ``_`` series, if it is chartable. + + Silently drops non-numeric values rather than raising: a backend adding a + metric that turns out to be a histogram should lose that one series, not + fail the training run's progress reporting. + """ + if not is_chartable(value): + return + # Coerce to a built-in: numpy scalars satisfy numbers.Real but are not + # JSON-serializable. Counts stay ints rather than becoming 64.0. + real = cast(numbers.Real, value) + numeric: float | int = int(real) if isinstance(real, numbers.Integral) else float(real) + series = name if name in _UNPREFIXED else f"{phase}_{name}" + self._series.setdefault(series, []).append({"step": step, "epoch": epoch, "value": numeric}) + def _build_metrics_summary(self) -> dict[str, list[dict[str, float | int]]]: - """Build the accumulated metrics payload for inclusion in status_details.""" - return { - "train_loss": list(self._train_metrics), - "val_loss": list(self._val_metrics), - } + """Build the accumulated metrics payload for inclusion in status_details. + + ``train_loss``/``val_loss`` are always present, even when empty, so the + shape stays stable for consumers that index them directly. Lists are + copied: the payload must not mutate after it is handed over. + """ + summary: dict[str, list[dict[str, float | int]]] = {"train_loss": [], "val_loss": []} + summary.update({name: list(points) for name, points in self._series.items()}) + return summary def report_training_start(self, max_steps: int, num_epochs: int, *, backend: str | None = None) -> None: """Report that training has started with schedule information.""" self._reporter.configure_progress_tracking(max_steps, num_epochs) - details: dict[str, object] = {"step": 0, "max_steps": max_steps, "num_epochs": num_epochs} + details: dict[str, object] = { + "step": 0, + "max_steps": max_steps, + "num_epochs": num_epochs, + "metrics": self._build_metrics_summary(), + } resolved = self._resolve_backend(backend) if resolved is not None: details["backend"] = resolved @@ -71,10 +155,25 @@ def report_train_step( grad_norm: float | None = None, *, backend: str | None = None, + **additional_metrics: object, ) -> None: - """Report training step with metrics.""" - self._train_metrics.append({"step": step, "epoch": epoch, "value": loss}) + """Report training step with metrics. + + ``additional_metrics`` are backend-specific (DPO's ``preference_loss``, + GRPO's ``reward``/``kl_penalty``, ...). Each numeric one accumulates into + its own ``train_`` series *and* rides along as a current-step + scalar, so consumers can read either the curve or the latest value. + """ + self._record("train", "train_loss", step, epoch, loss) + self._record("train", "lr", step, epoch, lr) + self._record("train", "grad_norm", step, epoch, grad_norm) + for name, value in additional_metrics.items(): + self._record("train", name, step, epoch, value) + # `**additional_metrics` is splatted first, matching report_validation, so a + # backend metric cannot shadow the accumulated series or the step's own loss. + # `step`/`epoch`/`lr`/`grad_norm` are named parameters and so already safe. details: dict[str, object] = { + **additional_metrics, "step": step, "epoch": epoch, "train_loss": loss, @@ -87,15 +186,34 @@ def report_train_step( details["backend"] = resolved self._reporter.report_running(phase="training", **details) - def report_validation(self, step: int, epoch: int, val_loss: float, *, backend: str | None = None) -> None: - """Report validation results.""" - self._val_metrics.append({"step": step, "epoch": epoch, "value": val_loss}) + def report_validation( + self, + step: int, + epoch: int, + val_loss: float | None = None, + *, + backend: str | None = None, + **additional_metrics: object, + ) -> None: + """Report validation results. + + ``val_loss`` is optional because not every algorithm produces one: GRPO + validates on ``accuracy``/``avg_length`` and reports no loss at all. The + key is omitted rather than sent as null, which would chart as a real zero, + and the ``val_loss`` series simply stays empty for such runs. + """ details: dict[str, object] = { "step": step, "epoch": epoch, - "val_loss": val_loss, - "metrics": self._build_metrics_summary(), + **additional_metrics, } + for name, value in additional_metrics.items(): + self._record("val", name, step, epoch, value) + if val_loss is not None: + self._record("val", "val_loss", step, epoch, val_loss) + details["val_loss"] = val_loss + details["metrics"] = self._build_metrics_summary() + resolved = self._resolve_backend(backend) if resolved is not None: details["backend"] = resolved @@ -110,7 +228,12 @@ def report_checkpoint_saved( backend: str | None = None, ) -> None: """Report that a checkpoint was saved.""" - details: dict[str, object] = {"step": step, "epoch": epoch, "checkpoint_path": checkpoint_path} + details: dict[str, object] = { + "step": step, + "epoch": epoch, + "checkpoint_path": checkpoint_path, + "metrics": self._build_metrics_summary(), + } resolved = self._resolve_backend(backend) if resolved is not None: details["backend"] = resolved @@ -118,7 +241,11 @@ def report_checkpoint_saved( def report_epoch_end(self, step: int, epoch: int, *, backend: str | None = None) -> None: """Report that an epoch has completed.""" - details: dict[str, object] = {"step": step, "epoch": epoch} + details: dict[str, object] = { + "step": step, + "epoch": epoch, + "metrics": self._build_metrics_summary(), + } resolved = self._resolve_backend(backend) if resolved is not None: details["backend"] = resolved diff --git a/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py b/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py index f56221da11..03522d2bb1 100644 --- a/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py +++ b/packages/nmp_customization_common/src/nmp/customization_common/training/progress.py @@ -8,6 +8,10 @@ training runner; backends subclass it (or instantiate it directly) supplying their own ``service_name`` so the task SDK resolves the right credentials. +Every update REPLACES the task's ``status_details``, so a field is only as +durable as the next report that omits it. ``update_task`` carries a defined set +of fields across updates that don't restate them -- see :data:`_CARRY_FORWARD`. + For training-specific metrics (loss, validation, checkpoints) see the ``TrainingProgressCallback`` which composes this reporter. """ @@ -25,6 +29,43 @@ logger = logging.getLogger(__name__) +#: Fields restated on updates that don't supply their own. +#: +#: The rule is *what stays true after the update that stated it*: +#: +#: - ``metrics`` is cumulative -- the whole point is that it grows. +#: - ``max_steps``/``num_epochs`` are run constants, and are only ever stated +#: once, by ``report_training_start``. +#: - ``step``/``epoch`` are monotonic; a run does not un-reach step 30. +#: - ``checkpoint_path`` is a sticky latest-value, true until superseded. +#: +#: Deliberately excluded: ``phase`` (every report sets its own), and the +#: per-step observations (``train_loss``, ``lr``, ``grad_norm``, ``reward``, +#: ...). Those describe one instant, and a stale copy would misrepresent +#: "current" -- nothing is lost by letting them expire, because every one of +#: them is now recoverable from its series in ``metrics``. +#: +#: ``percentage_done`` is also excluded: it is derived from ``step`` and +#: ``max_steps``, both of which are carried, so a consumer can recompute it +#: rather than risk a copy that contradicts its own inputs. +_CARRY_FORWARD = frozenset({"metrics", "max_steps", "num_epochs", "step", "epoch", "checkpoint_path"}) + + +def _carries_information(value: Any) -> bool: + """Whether a stored value is worth restating on a later update. + + Empty containers are dropped so a task doesn't accumulate keys that say + nothing -- notably the all-empty ``metrics`` dict a job reports before its + first training step. + """ + if value is None: + return False + if isinstance(value, dict): + return any(_carries_information(item) for item in value.values()) + if isinstance(value, (list, str)): + return bool(value) + return True + class JobsServiceProgressReporter: """Reports high-level progress to the Jobs service.""" @@ -36,6 +77,10 @@ def __init__(self, job_ctx: NMPJobContext, service_name: str): self._max_steps = 0 self._num_epochs = 0 + #: Last-seen value of each :data:`_CARRY_FORWARD` field, populated as + #: updates pass through and from the stored blob when one is read back. + self._carried: dict[str, Any] = {} + # Gate on real job context, not bare truthiness: from_env() fills missing # identifiers with non-empty sentinel defaults, which would otherwise # enable reporting (and failing SDK calls) outside a real job run. @@ -53,6 +98,42 @@ def _calculate_percentage_done(self, step: int | None) -> int: # downstream progress consumers expect a bounded percentage. return min(100, int((step / self._max_steps) * 100)) + def _carry_forward(self, status_details: dict[str, Any] | None) -> dict[str, Any]: + """Restate the :data:`_CARRY_FORWARD` fields this update doesn't supply. + + ``status_details`` is REPLACED by the Jobs service, not merged, so a + field survives only as long as every subsequent report repeats it. Three + things were being lost to that: + + - the accumulated ``metrics``, on the runner's checkpoint/completion/ + failure reports -- so every job ended by erasing its own curves; + - ``max_steps``/``num_epochs``, stated once at training start and gone + from the first training step onward; + - ``checkpoint_path``, published by one report and wiped by the next. + + Values are remembered as they pass through (write-through), so a process + that has already stated a field can restate it for free. The stored blob + is read back only when the update omits ``metrics``, which is the tell + that it did not come from ``TrainingProgressCallback`` -- i.e. it is one + of the handful the runner makes, from a different process that holds no + state. Per-step training reports always carry ``metrics``, so the hot + path never pays for a round-trip. + """ + details = dict(status_details or {}) + self._remember(details) + + missing = [field for field in _CARRY_FORWARD if field not in details] + if not missing: + return details + + if "metrics" not in details: + self._fetch_status_details() + + for field in missing: + if field in self._carried: + details[field] = self._carried[field] + return details + def update_task( self, status: str = "active", @@ -65,6 +146,8 @@ def update_task( if not self._is_main_rank: return + details = self._carry_forward(status_details) + try: jobs = client_from_platform(self._sdk, JobsClient) jobs.update_job_step_task( @@ -74,16 +157,24 @@ def update_task( step=self._job_ctx.step, body=PlatformJobTaskUpdate( status=PlatformJobStatus(status), - status_details=status_details or {}, + status_details=details, error_details=error_details or {}, ), ) except Exception as e: logger.warning(f"Failed to update task progress: {e}") - def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: + def _fetch_status_details(self) -> dict[str, Any]: + """Read back the task's stored ``status_details`` blob. + + Refreshes the carry-forward cache as a side effect, so that the + resume-seeding fetch ``TrainingProgressCallback`` makes at construction + doubles as the seed for :meth:`_carry_forward`. Without that, a resumed + run would drop the previous run's ``checkpoint_path``: its first report + already carries ``metrics``, so it would never read the blob back. + """ if not self._enabled: - return {"train_loss": [], "val_loss": []} + return {} try: jobs = client_from_platform(self._sdk, JobsClient) @@ -93,14 +184,33 @@ def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: job=self._job_ctx.job_id, step=self._job_ctx.step, ).data() - metrics = cast(dict[str, Any], (task.status_details or {}).get("metrics", {}) or {}) - return { - "train_loss": metrics.get("train_loss", []), - "val_loss": metrics.get("val_loss", []), - } + stored = cast(dict[str, Any], task.status_details or {}) except Exception as e: - logger.info(f"No prior metrics to seed (expected on first run): {e}") - return {"train_loss": [], "val_loss": []} + # Expected on a first run, where the task has no stored details yet. + # Serves both resume seeding and update_task's carry-forward, so the + # message stays neutral about which caller hit it. + logger.info(f"No stored status details available: {e}") + return {} + + self._remember(stored) + return stored + + def _remember(self, source: dict[str, Any]) -> None: + """Cache the carry-forward fields present in ``source``.""" + self._carried.update( + {field: value for field, value in source.items() if field in _CARRY_FORWARD and _carries_information(value)} + ) + + def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: + """Read back every stored metric series, for resume seeding. + + Deliberately not restricted to a known set of names: backends decide what + they accumulate, and a resumed job that only seeded ``train_loss`` would + silently restart every other curve from empty. Non-list values are + dropped so a malformed blob cannot poison the accumulator. + """ + metrics = cast(dict[str, Any], self._fetch_status_details().get("metrics", {}) or {}) + return {name: points for name, points in metrics.items() if isinstance(points, list)} def report_running(self, phase: str, **details: Any) -> None: if "step" in details and "percentage_done" not in details and self._max_steps > 0: diff --git a/packages/nmp_customization_common/tests/training/test_callbacks.py b/packages/nmp_customization_common/tests/training/test_callbacks.py new file mode 100644 index 0000000000..dc0fd65fba --- /dev/null +++ b/packages/nmp_customization_common/tests/training/test_callbacks.py @@ -0,0 +1,241 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the shared TrainingProgressCallback. + +Focused on the contract that every backend depends on: ``report_running`` REPLACES +the task's ``status_details``, so the accumulated series must ride on every report +or it is erased from stored status. Basic accumulation and resume-seeding are +covered by the per-backend suites; this file covers the shared surface itself. +""" + +from __future__ import annotations + +from typing import Any, ClassVar, cast + +import pytest +from nmp.customization_common.training.callbacks import TrainingProgressCallback +from nmp.customization_common.training.progress import JobsServiceProgressReporter + + +class _RecordingReporter: + """Stands in for JobsServiceProgressReporter, capturing each report payload.""" + + def __init__(self, prior: dict[str, list[dict[str, Any]]] | None = None) -> None: + self._prior = prior or {"train_loss": [], "val_loss": []} + self.reports: list[dict[str, Any]] = [] + self.tracking: tuple[int, int] | None = None + self.closed = False + + def fetch_current_metrics(self) -> dict[str, list[dict[str, Any]]]: + return self._prior + + def configure_progress_tracking(self, max_steps: int, num_epochs: int) -> None: + self.tracking = (max_steps, num_epochs) + + def report_running(self, phase: str, **details: Any) -> None: + self.reports.append({"phase": phase, **details}) + + def close(self) -> None: + self.closed = True + + +@pytest.fixture +def reporter() -> _RecordingReporter: + return _RecordingReporter() + + +def _make_callback(reporter: _RecordingReporter) -> TrainingProgressCallback: + """Build the callback over a duck-typed reporter, narrowing the type once here.""" + return TrainingProgressCallback(cast(JobsServiceProgressReporter, reporter)) + + +# --------------------------------------------------------------------------- # +# Every report path carries the series +# --------------------------------------------------------------------------- # + + +def test_every_report_path_carries_the_series(reporter: _RecordingReporter) -> None: + """An omitted payload erases the curve from stored status_details.""" + callback = _make_callback(reporter) + callback.report_training_start(max_steps=10, num_epochs=1) + callback.report_train_step(step=1, epoch=1, loss=0.5) + callback.report_validation(step=1, epoch=1, val_loss=0.45) + callback.report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/ckpt") + callback.report_epoch_end(step=1, epoch=1) + + assert len(reporter.reports) == 5 + for report in reporter.reports: + assert "metrics" in report, report["phase"] + assert set(report["metrics"]) == {"train_loss", "val_loss"} + + +def test_training_start_does_not_erase_seeded_metrics() -> None: + """report_training_start fires before the first step; it must not blank the blob.""" + prior = {"train_loss": [{"step": 1, "epoch": 1, "value": 0.9}], "val_loss": []} + reporter = _RecordingReporter(prior) + _make_callback(reporter).report_training_start(max_steps=10, num_epochs=1) + + assert reporter.reports[0]["metrics"]["train_loss"] == prior["train_loss"] + + +def test_series_are_snapshots_not_live_references(reporter: _RecordingReporter) -> None: + """A shared list would retroactively mutate already-sent payloads.""" + callback = _make_callback(reporter) + callback.report_train_step(step=1, epoch=1, loss=0.5) + first_payload = reporter.reports[-1]["metrics"]["train_loss"] + callback.report_train_step(step=2, epoch=1, loss=0.4) + + assert len(first_payload) == 1 + + +# --------------------------------------------------------------------------- # +# Optional val_loss +# --------------------------------------------------------------------------- # + + +def test_validation_without_loss_omits_the_key(reporter: _RecordingReporter) -> None: + """GRPO validates on accuracy; a null val_loss would chart as a real zero.""" + _make_callback(reporter).report_validation(step=1, epoch=1, val_loss=None, accuracy=0.75) + + report = reporter.reports[-1] + assert "val_loss" not in report + assert report["accuracy"] == 0.75 + assert report["metrics"]["val_loss"] == [] + + +def test_validation_with_loss_records_both_key_and_series(reporter: _RecordingReporter) -> None: + _make_callback(reporter).report_validation(step=1, epoch=1, val_loss=0.25) + + report = reporter.reports[-1] + assert report["val_loss"] == 0.25 + assert report["metrics"]["val_loss"] == [{"step": 1, "epoch": 1, "value": 0.25}] + + +# --------------------------------------------------------------------------- # +# additional_metrics +# --------------------------------------------------------------------------- # + + +def test_additional_train_metrics_become_series_and_ride_along( + reporter: _RecordingReporter, +) -> None: + """Each backend metric is both a curve and a current-step scalar.""" + callback = _make_callback(reporter) + callback.report_train_step(step=1, epoch=1, loss=0.5, reward=0.62, kl_penalty=0.008) + callback.report_train_step(step=2, epoch=1, loss=0.4, reward=0.71, kl_penalty=0.009) + + report = reporter.reports[-1] + assert report["reward"] == 0.71, "latest value still rides along at the top level" + assert report["metrics"]["train_reward"] == [ + {"step": 1, "epoch": 1, "value": 0.62}, + {"step": 2, "epoch": 1, "value": 0.71}, + ] + assert report["metrics"]["train_kl_penalty"] == [ + {"step": 1, "epoch": 1, "value": 0.008}, + {"step": 2, "epoch": 1, "value": 0.009}, + ] + + +def test_lr_and_grad_norm_accumulate(reporter: _RecordingReporter) -> None: + """Both are curves people read; neither is an `additional_metric`.""" + _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, lr=5e-06, grad_norm=1.9) + + metrics = reporter.reports[-1]["metrics"] + assert metrics["train_lr"] == [{"step": 1, "epoch": 1, "value": 5e-06}] + assert metrics["train_grad_norm"] == [{"step": 1, "epoch": 1, "value": 1.9}] + + +def test_absent_lr_and_grad_norm_create_no_series(reporter: _RecordingReporter) -> None: + """A backend that reports neither should not get two empty keys.""" + _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5) + + metrics = reporter.reports[-1]["metrics"] + assert "train_lr" not in metrics + assert "train_grad_norm" not in metrics + + +def test_additional_validation_metrics_become_series_and_ride_along(reporter: _RecordingReporter) -> None: + _make_callback(reporter).report_validation(step=1, epoch=1, val_loss=0.25, accuracy=0.9) + + report = reporter.reports[-1] + assert report["accuracy"] == 0.9 + assert report["metrics"]["val_accuracy"] == [{"step": 1, "epoch": 1, "value": 0.9}] + + +def test_train_and_validation_metrics_of_the_same_name_stay_separate( + reporter: _RecordingReporter, +) -> None: + """GRPO reports `truncation_rate` in both dicts; one series would interleave them.""" + callback = _make_callback(reporter) + callback.report_train_step(step=1, epoch=1, loss=0.5, truncation_rate=0.18) + callback.report_validation(step=1, epoch=1, truncation_rate=0.04) + + metrics = reporter.reports[-1]["metrics"] + assert metrics["train_truncation_rate"] == [{"step": 1, "epoch": 1, "value": 0.18}] + assert metrics["val_truncation_rate"] == [{"step": 1, "epoch": 1, "value": 0.04}] + + +def test_non_numeric_metrics_are_dropped_from_the_series(reporter: _RecordingReporter) -> None: + """Histograms and tables ride in the same dict as the scalars upstream.""" + _make_callback(reporter).report_train_step( + step=1, epoch=1, loss=0.5, histogram=object(), nested={"a": 1}, flag=True, missing=float("nan") + ) + + metrics = reporter.reports[-1]["metrics"] + assert set(metrics) == {"train_loss", "val_loss"} + + +def test_series_survive_a_resume_beyond_the_loss_curves() -> None: + """A resumed job must continue every curve, not just train_loss.""" + prior = { + "train_loss": [{"step": 1, "epoch": 1, "value": 0.9}], + "train_reward": [{"step": 1, "epoch": 1, "value": 0.2}], + } + reporter = _RecordingReporter(prior) + _make_callback(reporter).report_train_step(step=2, epoch=1, loss=0.8, reward=0.3) + + assert reporter.reports[-1]["metrics"]["train_reward"] == [ + {"step": 1, "epoch": 1, "value": 0.2}, + {"step": 2, "epoch": 1, "value": 0.3}, + ] + + +def test_additional_metrics_cannot_shadow_the_series(reporter: _RecordingReporter) -> None: + """`metrics` is not a parameter, so only splat order stops a silent override. + + `step`/`epoch`/`lr`/`grad_norm`/`backend` are named parameters -- passing one + is a TypeError at the call site. `metrics` and `train_loss` would just win. + """ + _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, metrics="clobbered") + + report = reporter.reports[-1] + assert report["metrics"]["train_loss"] == [{"step": 1, "epoch": 1, "value": 0.5}] + + +def test_additional_metrics_cannot_shadow_the_step_loss(reporter: _RecordingReporter) -> None: + _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, train_loss="clobbered") + + assert reporter.reports[-1]["train_loss"] == 0.5 + + +def test_additional_metrics_do_not_collide_with_backend_stamping( + reporter: _RecordingReporter, +) -> None: + """`backend` is keyword-only, so **additional_metrics can never capture it.""" + + class _Stamped(TrainingProgressCallback): + _default_backend: ClassVar[str | None] = "test-backend" + + callback = _Stamped(cast(JobsServiceProgressReporter, reporter)) + callback.report_train_step(step=1, epoch=1, loss=0.5, reward=0.62) + + report = reporter.reports[-1] + assert report["backend"] == "test-backend" + assert report["reward"] == 0.62 + + +def test_close_delegates_to_the_reporter(reporter: _RecordingReporter) -> None: + _make_callback(reporter).close() + + assert reporter.closed diff --git a/packages/nmp_customization_common/tests/training/test_progress.py b/packages/nmp_customization_common/tests/training/test_progress.py new file mode 100644 index 0000000000..f34204794a --- /dev/null +++ b/packages/nmp_customization_common/tests/training/test_progress.py @@ -0,0 +1,311 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for JobsServiceProgressReporter's status_details handling. + +Focused on carry-forward: the Jobs service REPLACES ``status_details``, so a +field lives only as long as the next report repeats it. The runner reports +checkpoint processing, completion and failure from a different process than the +training driver, and the driver states the schedule once and the checkpoint path +once. Without the carry-forward set, each of those is erased by the next update. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from nmp.customization_common.training.progress import JobsServiceProgressReporter + +SERIES: dict[str, list[dict[str, Any]]] = { + "train_loss": [{"step": 10, "epoch": 1, "value": 0.5}], + "train_reward": [{"step": 10, "epoch": 1, "value": 0.62}], +} +#: A blob a mid-run job would have stored: series plus the sticky facts. +STORED: dict[str, Any] = { + "phase": "training", + "step": 10, + "epoch": 1, + "max_steps": 30, + "num_epochs": 3, + "train_loss": 0.5, + "lr": 5e-06, + "checkpoint_path": "/ckpt/step-10", + "metrics": SERIES, +} + + +class _JobCtx: + """The four identifiers update_task reads off the job context.""" + + normalized_task = "training" + workspace = "default" + job_id = "job-1" + step = "train" + + +class _Reporter(JobsServiceProgressReporter): + """Reporter with the SDK and job context stubbed out. + + Bypasses ``__init__`` rather than mocking the SDK factory: what is under test + is the status_details logic, and the real constructor calls ``get_task_sdk``, + which wants credentials. Every attribute ``update_task`` touches is set here. + + The SDK client itself is patched (see the ``jobs`` fixture) rather than + ``_fetch_status_details``, so the real fetch runs -- including its + carry-forward cache side effect. + """ + + def __init__(self) -> None: + self._job_ctx = _JobCtx() # type: ignore[assignment] - duck-typed stand-in + self._sdk = type("S", (), {"close": lambda self: None})() # type: ignore[assignment] + self._is_main_rank = True + self._enabled = True + self._max_steps = 0 + self._num_epochs = 0 + self._carried = {} + + +class _Task: + def __init__(self, status_details: dict[str, Any]) -> None: + self.status_details = status_details + + def data(self) -> "_Task": + return self + + +class _Jobs: + """A mini Jobs service: replace-on-write, readable back, counting fetches.""" + + def __init__(self) -> None: + self.sent: list[dict[str, Any]] = [] + self.stored: dict[str, Any] = {} + self.fetches = 0 + self.persist = False + + def client(self) -> Any: + harness = self + + class _Client: + def update_job_step_task(self, **kwargs: Any) -> None: + harness.sent.append(kwargs) + if harness.persist: + harness.stored = dict(kwargs["body"].status_details or {}) + + def get_job_step_task(self, **kwargs: Any) -> _Task: + harness.fetches += 1 + return _Task(harness.stored) + + return _Client() + + +@pytest.fixture +def jobs(monkeypatch: pytest.MonkeyPatch) -> _Jobs: + """Patch the SDK client seam so update_task and the fetch both run for real.""" + harness = _Jobs() + monkeypatch.setattr( + "nmp.customization_common.training.progress.client_from_platform", + lambda _sdk, _cls: harness.client(), + ) + return harness + + +def _reporter(jobs: _Jobs, stored: dict[str, Any] | None = None) -> _Reporter: + """A reporter over the harness, with the server pre-seeded if given.""" + if stored is not None: + jobs.stored = stored + return _Reporter() + + +def _details(jobs: _Jobs, index: int = -1) -> dict[str, Any]: + assert jobs.sent, "expected at least one task update" + return dict(jobs.sent[index]["body"].status_details or {}) + + +# --------------------------------------------------------------------------- # +# What carries forward +# --------------------------------------------------------------------------- # + + +def test_completion_carries_series_schedule_and_checkpoint(jobs: _Jobs) -> None: + """The last write of a successful job must not blank what it took to get there.""" + _reporter(jobs, STORED).report_completed("Training completed") + + details = _details(jobs) + assert details["metrics"] == SERIES + assert details["max_steps"] == 30 + assert details["num_epochs"] == 3 + assert details["step"] == 10 + assert details["checkpoint_path"] == "/ckpt/step-10" + assert details["phase"] == "completed", "the report's own phase still wins" + + +def test_failure_carries_the_same_set(jobs: _Jobs) -> None: + """A failed run is exactly when the partial curve and last checkpoint matter.""" + _reporter(jobs, STORED).report_error("boom") + + details = _details(jobs) + assert details["metrics"] == SERIES + assert details["step"] == 10 + assert details["checkpoint_path"] == "/ckpt/step-10" + + +def test_intermediate_phase_carries_forward(jobs: _Jobs) -> None: + """processing_checkpoint fires after the driver exits, before completion.""" + _reporter(jobs, STORED).report_running("processing_checkpoint") + + details = _details(jobs) + assert details["metrics"] == SERIES + assert details["max_steps"] == 30 + assert details["phase"] == "processing_checkpoint" + + +def test_per_step_observations_do_not_carry_forward(jobs: _Jobs) -> None: + """A completed task must not advertise a stale current loss or learning rate. + + Nothing is lost: each of these is recoverable from its series in `metrics`. + """ + _reporter(jobs, STORED).report_completed("Training completed") + + details = _details(jobs) + assert "train_loss" not in details + assert "lr" not in details + + +def test_caller_supplied_values_win(jobs: _Jobs) -> None: + fresher = {"train_loss": [{"step": 20, "epoch": 2, "value": 0.1}]} + _reporter(jobs, STORED).report_running("training", step=20, metrics=fresher, max_steps=99) + + details = _details(jobs) + assert details["metrics"] == fresher + assert details["step"] == 20 + assert details["max_steps"] == 99 + + +def test_empty_stored_values_add_no_keys(jobs: _Jobs) -> None: + """Before training starts there is nothing to carry; don't invent keys.""" + stored = {"metrics": {"train_loss": [], "val_loss": []}, "checkpoint_path": ""} + _reporter(jobs, stored).report_running("compiling_config") + + details = _details(jobs) + assert "metrics" not in details + assert "checkpoint_path" not in details + + +# --------------------------------------------------------------------------- # +# Write-through cache: the per-step hot path must not pay for a round-trip +# --------------------------------------------------------------------------- # + + +def test_reports_carrying_metrics_never_fetch(jobs: _Jobs) -> None: + """`metrics` marks an update as coming from the accumulating callback. + + Those are the per-step reports. They omit max_steps and checkpoint_path, so a + naive implementation would read the blob back on every single training step. + """ + reporter = _reporter(jobs, STORED) + for step in range(1, 11): + reporter.report_running("training", step=step, metrics=SERIES) + + assert jobs.fetches == 0 + + +def test_a_stated_value_is_restated_without_a_fetch(jobs: _Jobs) -> None: + """report_training_start states the schedule once; every later step needs it.""" + reporter = _reporter(jobs) + reporter.report_running("training", step=0, max_steps=30, num_epochs=3, metrics=SERIES) + reporter.report_running("training", step=1, metrics=SERIES) + + details = _details(jobs) + assert details["max_steps"] == 30 + assert details["num_epochs"] == 3 + assert jobs.fetches == 0 + + +def test_checkpoint_path_survives_the_next_training_step(jobs: _Jobs) -> None: + """It was published by one report and wiped by the very next one.""" + reporter = _reporter(jobs) + reporter.report_running("checkpoint_saved", step=10, checkpoint_path="/ckpt/step-10", metrics=SERIES) + reporter.report_running("training", step=11, metrics=SERIES) + + assert _details(jobs)["checkpoint_path"] == "/ckpt/step-10" + + +def test_a_newer_checkpoint_supersedes_the_carried_one(jobs: _Jobs) -> None: + reporter = _reporter(jobs) + reporter.report_running("checkpoint_saved", step=10, checkpoint_path="/ckpt/step-10", metrics=SERIES) + reporter.report_running("checkpoint_saved", step=20, checkpoint_path="/ckpt/step-20", metrics=SERIES) + reporter.report_running("training", step=21, metrics=SERIES) + + assert _details(jobs)["checkpoint_path"] == "/ckpt/step-20" + + +def test_updates_without_metrics_read_the_blob_back(jobs: _Jobs) -> None: + """The runner's reports come from a process that holds no state at all.""" + reporter = _reporter(jobs, STORED) + reporter.report_running("processing_checkpoint") + + assert jobs.fetches == 1 + + +def test_error_details_still_ride_along(jobs: _Jobs) -> None: + """Carry-forward must not displace the error payload.""" + _reporter(jobs, STORED).report_error({"message": "oom", "code": "OOM"}) + + assert jobs.sent[0]["body"].error_details == {"message": "oom", "code": "OOM"} + + +# --------------------------------------------------------------------------- # +# Resume seeding +# --------------------------------------------------------------------------- # + + +def test_fetch_current_metrics_returns_every_series(jobs: _Jobs) -> None: + """A resumed job that only seeded train_loss would restart the other curves.""" + assert _reporter(jobs, STORED).fetch_current_metrics() == SERIES + + +def test_fetch_current_metrics_drops_non_list_values(jobs: _Jobs) -> None: + """A malformed blob must not poison the accumulator.""" + stored = {"metrics": {"train_loss": [{"step": 1, "epoch": 1, "value": 1.0}], "junk": 3}} + + assert set(_reporter(jobs, stored).fetch_current_metrics()) == {"train_loss"} + + +def test_resume_seeding_also_seeds_the_carry_forward_cache(jobs: _Jobs) -> None: + """The callback's construction-time fetch must double as the carry-forward seed. + + Otherwise a resumed run drops the previous run's checkpoint_path: its very + first report already carries `metrics`, so it never reads the blob back. + """ + reporter = _reporter(jobs, STORED) + reporter.fetch_current_metrics() # what TrainingProgressCallback.__init__ does + fetches_after_seeding = jobs.fetches + + reporter.report_running("training", step=1, metrics=SERIES) + + assert _details(jobs)["checkpoint_path"] == "/ckpt/step-10" + assert jobs.fetches == fetches_after_seeding, "no second round-trip" + + +# --------------------------------------------------------------------------- # +# Gating +# --------------------------------------------------------------------------- # + + +def test_disabled_reporter_sends_nothing(jobs: _Jobs) -> None: + reporter = _reporter(jobs, STORED) + reporter._enabled = False + + reporter.report_completed("Training completed") + + assert jobs.sent == [] + + +def test_non_main_rank_sends_nothing(jobs: _Jobs) -> None: + reporter = _reporter(jobs, STORED) + reporter._is_main_rank = False + + reporter.report_completed("Training completed") + + assert jobs.sent == [] diff --git a/services/automodel/tests/tasks/training/backends/test_callbacks.py b/services/automodel/tests/tasks/training/backends/test_callbacks.py index 71dcb278ca..03276934c2 100644 --- a/services/automodel/tests/tasks/training/backends/test_callbacks.py +++ b/services/automodel/tests/tasks/training/backends/test_callbacks.py @@ -104,8 +104,8 @@ def test_seeds_from_server_on_init(self): } callback, reporter = self._make_callback(prior_metrics=prior) - assert len(callback._train_metrics) == 2 - assert len(callback._val_metrics) == 1 + assert len(callback._series["train_loss"]) == 2 + assert len(callback._series["val_loss"]) == 1 reporter.fetch_current_metrics.assert_called_once() def test_seeded_metrics_included_in_first_report(self): @@ -142,7 +142,13 @@ def test_report_training_start_delegates(self): callback.report_training_start(max_steps=500, num_epochs=2) reporter.configure_progress_tracking.assert_called_once_with(500, 2) - reporter.report_running.assert_called_once_with(phase="training", step=0, max_steps=500, num_epochs=2) + reporter.report_running.assert_called_once_with( + phase="training", + step=0, + max_steps=500, + num_epochs=2, + metrics={"train_loss": [], "val_loss": []}, + ) def test_report_checkpoint_saved_delegates(self): callback, reporter = self._make_callback() @@ -150,9 +156,37 @@ def test_report_checkpoint_saved_delegates(self): callback.report_checkpoint_saved(step=100, epoch=1, checkpoint_path="/tmp/ckpt") reporter.report_running.assert_called_once_with( - phase="checkpoint_saved", step=100, epoch=1, checkpoint_path="/tmp/ckpt" + phase="checkpoint_saved", + step=100, + epoch=1, + checkpoint_path="/tmp/ckpt", + metrics={"train_loss": [], "val_loss": []}, ) + def test_checkpoint_report_preserves_accumulated_series(self): + """report_running REPLACES status_details, so an omitted payload erases the curve. + + Checkpoint saves fire mid-training (finetune.py calls this from the save + hook), so a report without `metrics` would drop the series from stored + status until the next train step -- and lose it entirely if the job then died. + """ + callback, reporter = self._make_callback() + + callback.report_train_step(step=1, epoch=1, loss=3.21) + callback.report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/tmp/ckpt") + + kwargs = self._last_report_kwargs(reporter) + assert kwargs["metrics"]["train_loss"] == [{"step": 1, "epoch": 1, "value": 3.21}] + + def test_epoch_end_report_preserves_accumulated_series(self): + callback, reporter = self._make_callback() + + callback.report_train_step(step=1, epoch=1, loss=3.21) + callback.report_epoch_end(step=1, epoch=1) + + kwargs = self._last_report_kwargs(reporter) + assert kwargs["metrics"]["train_loss"] == [{"step": 1, "epoch": 1, "value": 3.21}] + def test_close_delegates(self): callback, reporter = self._make_callback() callback.close() diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py index 8ac0760844..f86444ae8d 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py @@ -8,88 +8,20 @@ # without an express license agreement from NVIDIA CORPORATION or # its affiliates is strictly prohibited. -import logging -from typing import Any +"""Training progress callbacks for NeMo-RL Jobs-service reporting. -from nmp.customization_common.training.progress import JobsServiceProgressReporter +Thin subclass of the shared +:class:`nmp.customization_common.training.callbacks.TrainingProgressCallback`, +matching the unsloth and automodel pattern. ``_default_backend`` stays ``None`` +so no ``backend`` key is added and the existing status-detail shape is preserved. +""" -logger = logging.getLogger(__name__) +from nmp.customization_common.training.callbacks import ( + TrainingProgressCallback as _BaseTrainingProgressCallback, +) +__all__ = ["TrainingProgressCallback"] -class TrainingProgressCallback: - """ - Callback for reporting NeMo RL training progress to the Jobs service. - This class composes JobsServiceProgressReporter and provides training-specific - methods for reporting detailed metrics during training. - """ - - def __init__(self, reporter: JobsServiceProgressReporter): - self._reporter = reporter - - def report_training_start(self, max_steps: int, num_epochs: int) -> None: - """Report that training has started with schedule information.""" - self._reporter.configure_progress_tracking(max_steps, num_epochs) - self._reporter.report_running(phase="training", step=0, max_steps=max_steps, num_epochs=num_epochs) - - def report_train_step( - self, - step: int, - epoch: int, - loss: float, - lr: float | None = None, - grad_norm: float | None = None, - **additional_metrics: Any, - ) -> None: - """Report training step with metrics. - - Args: - step: Training step number - epoch: Current epoch number - loss: Training loss value - lr: Learning rate (optional) - grad_norm: Gradient norm (optional) - **additional_metrics: Additional training metrics to report (e.g., num_valid_samples, - preference_loss, rewards_rejected_mean, global_valid_seqs, global_valid_toks) - """ - self._reporter.report_running( - phase="training", - step=step, - epoch=epoch, - train_loss=loss, - lr=lr, - grad_norm=grad_norm, - **additional_metrics, - ) - - def report_validation( - self, - step: int, - epoch: int, - val_loss: float, - **additional_metrics: Any, - ) -> None: - """Report validation results. - - Args: - step: Training step number - epoch: Current epoch number - val_loss: Validation loss value - **additional_metrics: Additional validation metrics to report (e.g., accuracy, - num_valid_samples, or any other validation-specific metrics) - """ - self._reporter.report_running( - phase="validation", - step=step, - epoch=epoch, - val_loss=val_loss, - **additional_metrics, - ) - - def report_checkpoint_saved(self, step: int, epoch: int, checkpoint_path: str | None = None) -> None: - """Report that a checkpoint was saved.""" - self._reporter.report_running(phase="checkpoint_saved", step=step, epoch=epoch, checkpoint_path=checkpoint_path) - - def close(self) -> None: - """Clean up resources.""" - self._reporter.close() +class TrainingProgressCallback(_BaseTrainingProgressCallback): + """Report NeMo-RL training progress to the Jobs service.""" diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py index 897a14a10b..53e093ad56 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/dpo_driver.py @@ -108,19 +108,14 @@ def main(): # Log only the non-sensitive job id; the full context carries service URLs # and identifiers that should not be dumped to stdout. print(f"Job context loaded (job_id={job_ctx.job_id})") + customizer_logger: NemoRLLogger | None = None if job_ctx.jobs_url: - # Extract training parameters for progress reporting - max_steps = config.dpo.max_num_steps - num_epochs = config.dpo.max_num_epochs - steps_per_epoch = config.dpo.steps_per_epoch # type: ignore[attr-defined] - extra (undeclared) DPOConfig field, allowed via extra="allow" - log_interval = (config.dpo.val_period // 10) + 1 - - customizer_logger = NemoRLLogger( - steps_per_epoch=steps_per_epoch, + customizer_logger = NemoRLLogger.for_schedule( job_ctx=job_ctx, - log_interval=log_interval, - max_steps=max_steps, - num_epochs=num_epochs, + max_steps=config.dpo.max_num_steps, + num_epochs=config.dpo.max_num_epochs, + val_period=config.dpo.val_period, + steps_per_epoch=config.dpo.steps_per_epoch, # type: ignore[attr-defined] - extra (undeclared) DPOConfig field, allowed via extra="allow" ) # The setup() logger is a composite with a `.loggers` list; guard in case # that internal shape changes. @@ -131,17 +126,24 @@ def main(): logger.log_hyperparams(config.model_dump()) - dpo_train( - policy, - train_dataloader, - val_dataloader, - tokenizer, - loss_fn, - master_config, - logger, - checkpointer, - dpo_save_state, - ) + try: + dpo_train( + policy, + train_dataloader, + val_dataloader, + tokenizer, + loss_fn, + master_config, + logger, + checkpointer, + dpo_save_state, + ) + finally: + # Flushes the final training step. NeMo-RL never closes the loggers it is + # handed, so without this the only fallback is NemoRLLogger.__del__ at + # interpreter shutdown, which does not run at all on an abnormal exit. + if customizer_logger is not None: + customizer_logger.close() if __name__ == "__main__": diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py index 0e3ca33ccb..5fdc16199a 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/grpo_driver.py @@ -139,17 +139,13 @@ def main() -> None: job_ctx = NMPJobContext.from_env() print(f"Job context loaded (job_id={job_ctx.job_id})") + customizer_logger: NemoRLLogger | None = None if job_ctx.jobs_url: - max_steps = config.grpo.max_num_steps - num_epochs = config.grpo.max_num_epochs - val_period = config.grpo.val_period or 1 - log_interval = max(val_period // 10, 1) - customizer_logger = NemoRLLogger( - steps_per_epoch=max(max_steps // max(num_epochs, 1), 1), + customizer_logger = NemoRLLogger.for_schedule( job_ctx=job_ctx, - log_interval=log_interval, - max_steps=max_steps, - num_epochs=num_epochs, + max_steps=config.grpo.max_num_steps, + num_epochs=config.grpo.max_num_epochs, + val_period=config.grpo.val_period, ) if hasattr(logger_inst, "loggers"): logger_inst.loggers.append(customizer_logger) @@ -172,6 +168,14 @@ def main() -> None: master_config, ) finally: + # Before the env teardown below, which can be slow and can itself raise: + # this flushes the final training step, and it is the only deterministic + # chance to do so. NeMo-RL never closes the loggers it is handed, so the + # fallback is NemoRLLogger.__del__ at interpreter shutdown, which does not + # run at all on an abnormal exit. + if customizer_logger is not None: + customizer_logger.close() + for task_name, env in task_to_env.items(): try: import ray diff --git a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py index f7534bc2c7..9aa5463288 100644 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py +++ b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py @@ -9,23 +9,100 @@ # its affiliates is strictly prohibited. import logging -import math from typing import Any, Mapping, Optional from nemo_rl.utils.logger import LoggerInterface from nmp.customization_common.service.context import NMPJobContext -from nmp.customization_common.training.progress import JobsServiceProgressReporter -from nmp.rl.app.constants import SERVICE_NAME +from nmp.customization_common.training.callbacks import is_chartable from nmp.rl.tasks.training.backends.nemo_rl.callbacks import TrainingProgressCallback +from nmp.rl.tasks.training.progress import JobsServiceProgressReporter _logger = logging.getLogger(__name__) +# How many progress reports to aim for across one validation period. NeMo-RL has no +# notion of a reporting cadence, so it is derived from val_period. +_REPORTS_PER_VAL_PERIOD = 10 + +# Metric keys forwarded to Jobs Service, in addition to loss/lr/grad_norm. +# +# A union across algorithms: selection is by presence, so a DPO run simply has no +# GRPO keys and vice versa. GRPO's `train` dict is the merge of the policy-loss +# metrics, the reward/advantage block and NeMo-Gym's rollout metrics +# (nemo_rl/algorithms/grpo.py builds it, nemo_rl/experience/rollouts.py supplies +# the rollout half), so all three families are represented here. +_TRAIN_METRIC_KEYS = ( + # Shared + "num_valid_samples", + "global_valid_seqs", + "global_valid_toks", + # DPO — both halves of the reward pair; either alone is hard to read. + "preference_loss", + "rewards_chosen_mean", + "rewards_rejected_mean", + # GRPO: reward and advantages — the signal that says whether RL is working. + # The total_reward spread is what a reward-distribution band is drawn from; + # `calculate_single_metric` emits the whole set (a /histogram rides along too, + # but it is a Histogram object and is dropped as non-scalar). + "reward", + "total_reward/mean", + "total_reward/median", + "total_reward/min", + "total_reward/max", + "total_reward/stddev", + "advantages/mean", + "advantages/min", + "advantages/max", + # GRPO: policy-optimization health. These are the diagnostics that predict a + # failing run -- entropy collapse, and the logprob-error family that says the + # trainer and the generation engine have drifted apart. All are returned + # unconditionally by ClippedPGLossFn and normalized via its + # `metric_normalizations` map, so the per-microbatch sum NeMo-RL applies is + # already the correct global value. + "kl_penalty", + "approx_entropy", + "token_mult_prob_error", + "gen_kl_error", + "policy_kl_error", + "js_divergence_error", + "sampling_importance_ratio", + # GRPO: rollout shape + "truncation_rate", + "natural_termination_rate", + "turns_per_sample/mean", + "mean_gen_tokens_per_sample", +) + +# GRPO reports `accuracy`/`avg_length` here and no loss at all; DPO reports loss. +_VALIDATION_METRIC_KEYS = _TRAIN_METRIC_KEYS + ("accuracy", "avg_length") + def has_metric_value(metric: Any) -> bool: - """Check if a metric has a valid value.""" - if metric is not None and not math.isnan(metric): - return True - return False + """Whether ``metric`` is a finite-enough scalar to forward to Jobs Service. + + The type check is load-bearing, not defensive. NeMo-RL's metric dicts carry + non-scalars alongside the numbers: ``calculate_single_metric`` emits a + ``/histogram`` holding a ``Histogram`` object, NeMo-Gym adds a + per-agent ``full_result`` ``Table``, and ``generation_logger_metrics`` is a + nested dict. ``math.isnan`` raises ``TypeError`` on all of those, so a bare + None-check would turn a widened key list into a crash mid-training. + + Delegates to the shared predicate so the wire filter and the series filter + cannot drift apart -- a metric this forwards must be one the callback can + chart. + """ + return is_chartable(metric) + + +def resolve_log_interval(val_period: int | None) -> int: + """Steps between progress reports, targeting ~10 reports per validation period.""" + return max((val_period or 0) // _REPORTS_PER_VAL_PERIOD, 1) + + +def resolve_steps_per_epoch(max_steps: int, num_epochs: int | None, explicit: int | None = None) -> int: + """Steps per epoch, preferring an explicit value from the algorithm config.""" + if explicit is not None and explicit >= 1: + return explicit + return max(max_steps // max(num_epochs or 1, 1), 1) class NemoRLLogger(LoggerInterface): @@ -74,7 +151,7 @@ def __init__( self._steps_per_epoch = steps_per_epoch # Create the callback for progress reporting - self._reporter = JobsServiceProgressReporter(self._job_ctx, SERVICE_NAME) + self._reporter = JobsServiceProgressReporter(self._job_ctx) self._callback = TrainingProgressCallback(self._reporter) # Track best metrics for monitoring @@ -82,12 +159,46 @@ def __init__( self._best_epoch: int | None = None self._closed = False + # Last train step built but withheld by the log_interval throttle. Flushed on + # close() so the final step is reported even when max_steps is not a multiple + # of log_interval -- otherwise the run's last recorded loss is stale. + self._pending_train_report: dict[str, Any] | None = None + _logger.info( f"Initialized NemoRLLogger with jobs_url={self._job_ctx.jobs_url}, " f"log_interval={log_interval}, max_steps={max_steps}, num_epochs={num_epochs}, " f"steps_per_epoch={steps_per_epoch}" ) + @classmethod + def for_schedule( + cls, + *, + max_steps: int, + num_epochs: int | None, + val_period: int | None, + steps_per_epoch: int | None = None, + job_ctx: NMPJobContext | None = None, + ) -> "NemoRLLogger": + """Build a logger from a NeMo-RL training schedule. + + The DPO and GRPO drivers previously derived ``log_interval`` and + ``steps_per_epoch`` with different formulas for the same intent -- at + ``val_period=100`` one produced 11 and the other 10 -- so the arithmetic + lives here instead of being restated per algorithm. + + Args: + steps_per_epoch: Authoritative value when the algorithm config carries + one (DPO does); otherwise derived from max_steps and num_epochs. + """ + return cls( + steps_per_epoch=resolve_steps_per_epoch(max_steps, num_epochs, steps_per_epoch), + job_ctx=job_ctx, + log_interval=resolve_log_interval(val_period), + max_steps=max_steps, + num_epochs=num_epochs, + ) + def log_metrics( self, metrics: dict[str, Any], @@ -105,58 +216,50 @@ def log_metrics( step_metric: Optional step metric name (ignored in this implementation) step_finished: Whether the step is finished (part of NeMo-RL's LoggerInterface; ignored here) """ - step = step + 1 # Increment step since we start counting from 1 - - # Calculate epoch from step (epochs start from 1) - epoch = ((step - 1) // self._steps_per_epoch) + 1 - - # Handle training loss + # `step` arrives 1-indexed and is used as-is. Both callers pass + # `total_steps + 1`, where total_steps is 0-based and incremented *after* + # logging (nemo_rl/algorithms/grpo.py, .../dpo.py), so it is already the + # 1-indexed step number. Incrementing again put the last step of an + # N-step run at N+1 and shifted the whole series one to the right of the + # axis Studio draws it against. + # + # Step 0 does arrive, from the validate-at-start path only; it belongs to + # epoch 1, hence the clamp rather than a bare `step - 1`. + epoch = (max(step - 1, 0) // self._steps_per_epoch) + 1 + + # Handle training loss. + # + # The loss gate also de-duplicates: GRPO logs under `train` twice per step — + # once mid-step with the rollout metrics alone (no loss), then again with the + # full merged dict. Requiring loss keeps the second, complete one. if prefix == "train" and has_metric_value(metrics.get("loss")): - # Only report at log_interval to reduce output + report = { + "step": step, + "epoch": epoch, + "loss": metrics["loss"], + "lr": metrics.get("lr"), + "grad_norm": metrics.get("grad_norm"), + **self._select_metrics(metrics, _TRAIN_METRIC_KEYS), + } + # Throttled to log_interval to reduce output. A withheld step is held as + # pending rather than dropped, so close() can flush the last one. if step % self._log_interval == 0: - # Extract core metrics - loss = metrics["loss"] - lr = metrics.get("lr") - grad_norm = metrics.get("grad_norm") - - # Extract additional training metrics (whitelisted only) - additional_metrics = {} - for key in [ - "num_valid_samples", - "preference_loss", - "rewards_rejected_mean", - "global_valid_seqs", - "global_valid_toks", - ]: - if has_metric_value(metrics.get(key)): - additional_metrics[key] = metrics[key] - - self._callback.report_train_step( - step=step, - epoch=epoch, - loss=loss, - lr=lr, - grad_norm=grad_norm, - **additional_metrics, - ) - - # Handle validation metrics + self._callback.report_train_step(**report) + self._pending_train_report = None + else: + self._pending_train_report = report + + # Handle validation metrics. + # + # `val_loss` is optional because GRPO has none: its validation dict is + # {accuracy, avg_length}. Gating on loss here dropped every GRPO validation + # report silently, so instead report whenever anything usable came through. elif prefix and prefix.startswith("validation"): - if has_metric_value(metrics.get("loss")): - val_loss = metrics["loss"] - - # Extract additional validation metrics (whitelisted only) - additional_metrics = {} - for key in [ - "num_valid_samples", - "preference_loss", - "rewards_rejected_mean", - "global_valid_seqs", - "global_valid_toks", - ]: - if has_metric_value(metrics.get(key)): - additional_metrics[key] = metrics[key] + raw_val_loss = metrics.get("loss") + val_loss = raw_val_loss if has_metric_value(raw_val_loss) else None + additional_metrics = self._select_metrics(metrics, _VALIDATION_METRIC_KEYS) + if val_loss is not None or additional_metrics: self._callback.report_validation( step=step, epoch=epoch, @@ -164,12 +267,17 @@ def log_metrics( **additional_metrics, ) # Track best validation loss - if val_loss < self._best_metric_value: + if val_loss is not None and val_loss < self._best_metric_value: self._best_metric_value = val_loss self._best_epoch = epoch _logger.debug(f"log_metrics: step={step}, prefix={prefix}, metrics={metrics}") + @staticmethod + def _select_metrics(metrics: dict[str, Any], keys: tuple[str, ...]) -> dict[str, Any]: + """Pick the whitelisted keys that carry a forwardable scalar.""" + return {key: metrics[key] for key in keys if has_metric_value(metrics.get(key))} + def log_hyperparams(self, params: Mapping[str, Any]) -> None: """Log hyperparameters and report training start. @@ -206,14 +314,42 @@ def log_plot(self, figure: Any, step: int, name: str) -> None: """ return None + def finish(self) -> None: + """Alias for :meth:`close` under the name NeMo-RL's composite fans out. + + ``nemo_rl.utils.logger.Logger`` has no ``close()`` at all; its only + teardown hook is ``finish()``, dispatched via + ``getattr(logger, "finish", None)``. Without this method the composite + silently skips us and the withheld final step is never flushed. The + drivers also call ``close()`` directly, because ``grpo_train``/ + ``dpo_train`` never invoke ``finish()`` either -- only the + single-controller path does. + """ + self.close() + def close(self) -> None: - """Clean up resources.""" + """Flush any withheld final step, then clean up resources.""" if self._closed: return self._closed = True + self._flush_pending_train_report() _logger.info("NemoRLLogger closing") self._callback.close() + def _flush_pending_train_report(self) -> None: + """Report the last step if the log_interval throttle withheld it. + + Reachable from ``__del__``, so failures must not propagate; the reporter + already swallows and logs transport errors, and this guards the rest. + """ + if self._pending_train_report is None: + return + report, self._pending_train_report = self._pending_train_report, None + try: + self._callback.report_train_step(**report) + except Exception as exc: # pragma: no cover - defensive, shutdown path + _logger.warning(f"Failed to flush final train step: {exc}") + def __del__(self): """Cleanup when the logger is destroyed.""" try: diff --git a/services/rl/src/nmp/rl/tasks/training/progress.py b/services/rl/src/nmp/rl/tasks/training/progress.py new file mode 100644 index 0000000000..f998b1adfa --- /dev/null +++ b/services/rl/src/nmp/rl/tasks/training/progress.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +"""Progress reporting for RL training tasks. + +Thin subclass of the shared +:class:`nmp.customization_common.training.progress.JobsServiceProgressReporter` +that bakes in the RL ``SERVICE_NAME`` so callers keep the +``JobsServiceProgressReporter(job_ctx)`` constructor. Mirrors the equivalent +modules in the unsloth and automodel services. +""" + +from nmp.customization_common.service.context import NMPJobContext +from nmp.customization_common.training.progress import ( + JobsServiceProgressReporter as _BaseJobsServiceProgressReporter, +) +from nmp.rl.app.constants import SERVICE_NAME + +__all__ = ["JobsServiceProgressReporter"] + + +class JobsServiceProgressReporter(_BaseJobsServiceProgressReporter): + """RL training progress reporter (binds the RL service name).""" + + def __init__(self, job_ctx: NMPJobContext): + super().__init__(job_ctx, service_name=SERVICE_NAME) diff --git a/services/rl/src/nmp/rl/tasks/training/runner.py b/services/rl/src/nmp/rl/tasks/training/runner.py index e9573dc68d..fe51f41150 100644 --- a/services/rl/src/nmp/rl/tasks/training/runner.py +++ b/services/rl/src/nmp/rl/tasks/training/runner.py @@ -19,8 +19,7 @@ import yaml from nmp.customization_common.service.context import NMPJobContext -from nmp.customization_common.training.progress import JobsServiceProgressReporter -from nmp.rl.app.constants import DEFAULT_TRAINING_RESULT_FILE_NAME, SERVICE_NAME +from nmp.rl.app.constants import DEFAULT_TRAINING_RESULT_FILE_NAME from nmp.rl.app.jobs.training.schemas import ( GPUInfo, TrainingMetrics, @@ -28,6 +27,7 @@ TrainingStepConfig, ) from nmp.rl.app.jobs.training.schemas import TrainingBackend as TrainingBackendEnum +from nmp.rl.tasks.training.progress import JobsServiceProgressReporter from .distributed import DistributedContext from .errors.converter import create_error_details @@ -72,7 +72,7 @@ def __init__(self, backend: TrainingBackend | None = None) -> None: self._job_ctx = NMPJobContext.from_env() self._config = self._load_config(self._job_ctx.config_path) - self._progress = JobsServiceProgressReporter(self._job_ctx, SERVICE_NAME) + self._progress = JobsServiceProgressReporter(self._job_ctx) self._dist_ctx = DistributedContext.from_env(self._get_barrier_dir()) self._backend = backend or self._load_backend(self._config.backend) # workspace_path and output_path are absolute paths from the config diff --git a/services/rl/tests/test_grpo_config.py b/services/rl/tests/test_grpo_config.py index 666d0eb11a..2b91440fb3 100644 --- a/services/rl/tests/test_grpo_config.py +++ b/services/rl/tests/test_grpo_config.py @@ -261,8 +261,15 @@ def test_tokenizer_omits_chat_template_when_none( one. Qwen3 has one, so a single-model GPU run never sees this. """ monkeypatch.setenv("NMP_JOB_STORAGE_PVC_CLAIM", "nmp-job-storage") + # Patched rather than relying on the fixture model dir having no tokenizer: + # resolution falls through to transformers.AutoTokenizer, and another suite in + # the same session installs a module-scope `transformers` MagicMock into + # sys.modules, whose truthy `.chat_template` silently invalidates the premise. + monkeypatch.setattr( + "nmp.rl.tasks.training.backends.nemo_rl.grpo_config.resolve_chat_template", + lambda *_, **__: None, + ) step, _ = _prepared_step(tmp_path) - # The fixture model dir has no tokenizer, so resolution falls through to None. tokenizer = compile_grpo_config(step, job_ctx)["policy"]["tokenizer"] assert "chat_template" not in tokenizer diff --git a/services/rl/tests/test_nemo_rl_callbacks.py b/services/rl/tests/test_nemo_rl_callbacks.py new file mode 100644 index 0000000000..df13896eb1 --- /dev/null +++ b/services/rl/tests/test_nemo_rl_callbacks.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Wiring tests for the RL TrainingProgressCallback subclass. + +The behaviour lives in the shared base +(``packages/nmp_customization_common/tests/training/test_callbacks.py``); what is +RL-specific is only which base it inherits and that it stays unstamped. +""" + +from __future__ import annotations + +from nmp.customization_common.training.callbacks import ( + TrainingProgressCallback as SharedTrainingProgressCallback, +) +from nmp.rl.tasks.training.backends.nemo_rl.callbacks import TrainingProgressCallback + + +def test_rl_callback_subclasses_the_shared_one() -> None: + """RL used to carry a standalone copy; accumulation fixes must land once.""" + assert issubclass(TrainingProgressCallback, SharedTrainingProgressCallback) + + +def test_rl_callback_adds_no_backend_field() -> None: + """Stamping `backend` would change RL's status-detail shape on the wire. + + unsloth opts in; automodel and RL deliberately do not. + """ + assert TrainingProgressCallback._default_backend is None diff --git a/services/rl/tests/test_nemo_rl_drivers.py b/services/rl/tests/test_nemo_rl_drivers.py new file mode 100644 index 0000000000..4fc2d28882 --- /dev/null +++ b/services/rl/tests/test_nemo_rl_drivers.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Source-level checks that the drivers tear the progress logger down. + +These are tripwires, not behaviour tests. The drivers cannot be imported outside +the training image -- they pull in nemo_rl, ray and omegaconf at module scope -- +so the wiring is asserted against the AST instead. + +It is worth asserting at all because the failure is silent: NeMo-RL never closes +the loggers it is handed, so if these calls are dropped the final training step +stops being reported and every unit test still passes. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +DRIVERS = Path(__file__).resolve().parents[1] / "src/nmp/rl/tasks/training/backends/nemo_rl" + + +def _closes_logger_in_finally(source: str) -> bool: + """Whether some `try/finally` closes `customizer_logger` in its finally body.""" + for node in ast.walk(ast.parse(source)): + if not isinstance(node, ast.Try): + continue + for stmt in node.finalbody: + for inner in ast.walk(stmt): + if ( + isinstance(inner, ast.Call) + and isinstance(inner.func, ast.Attribute) + and inner.func.attr == "close" + and isinstance(inner.func.value, ast.Name) + and inner.func.value.id == "customizer_logger" + ): + return True + return False + + +@pytest.mark.parametrize( + "source,expected", + [ + ("try:\n train()\nfinally:\n customizer_logger.close()\n", True), + # Guarded call — how the drivers actually write it. + ("try:\n train()\nfinally:\n if customizer_logger:\n customizer_logger.close()\n", True), + # Present, but not on the abnormal-exit path. + ("try:\n train()\nfinally:\n pass\ncustomizer_logger.close()\n", False), + ("try:\n train()\nfinally:\n other_logger.close()\n", False), + ("try:\n train()\nfinally:\n customizer_logger.flush()\n", False), + ], +) +def test_detector_discriminates(source: str, expected: bool) -> None: + """The tripwire is only worth having if it can actually trip.""" + assert _closes_logger_in_finally(source) is expected + + +@pytest.mark.parametrize("driver", ["grpo_driver.py", "dpo_driver.py"]) +def test_driver_closes_the_progress_logger_in_a_finally(driver: str) -> None: + """`finally`, not the happy path: an aborted run is when the flush matters.""" + source = (DRIVERS / driver).read_text() + + assert _closes_logger_in_finally(source), ( + f"{driver} must close customizer_logger from a finally block; " + "NeMo-RL does not close loggers, and __del__ does not run on abnormal exit" + ) diff --git a/services/rl/tests/test_nemo_rl_logger.py b/services/rl/tests/test_nemo_rl_logger.py new file mode 100644 index 0000000000..62113418c4 --- /dev/null +++ b/services/rl/tests/test_nemo_rl_logger.py @@ -0,0 +1,531 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for NemoRLLogger's translation of NeMo-RL metrics into Jobs Service reports. + +The metric dicts here mirror what NeMo-RL actually hands the logger, including the +non-scalar entries (``Histogram``, ``Table``, nested dicts) that share the dict with +the numbers. Those are the reason ``has_metric_value`` type-checks rather than just +None-checks, so they are exercised rather than sanitised away. +""" + +from __future__ import annotations + +import importlib.machinery +import importlib.util +import math +import sys +import types +from typing import Any + +import pytest + +# NeMo-RL is only installed inside the training image, so the LoggerInterface import +# at nemo_rl_logger module scope fails in a plain repo checkout. Stub just enough to +# import the module under test; when the real package IS present (the in-image smoke +# run) this is skipped and the genuine base class is used. +if importlib.util.find_spec("nemo_rl") is None: # pragma: no cover - env dependent + + class LoggerInterface: # minimal stand-in for the abstract base + pass + + def _stub(name: str) -> types.ModuleType: + """Build a stub module that survives a later importlib.util.find_spec. + + A bare ModuleType has ``__spec__ = None``, and find_spec consults + sys.modules first -- so leaving it unset makes a later + ``find_spec("nemo_rl")`` raise ValueError rather than return None. The + stub outlives this module (nothing tears it down), so it must not booby + trap whatever runs next in the session. + """ + module = types.ModuleType(name) + module.__spec__ = importlib.machinery.ModuleSpec(name, loader=None) + return module + + _logger_mod = _stub("nemo_rl.utils.logger") + setattr(_logger_mod, "LoggerInterface", LoggerInterface) + sys.modules.setdefault("nemo_rl", _stub("nemo_rl")) + sys.modules.setdefault("nemo_rl.utils", _stub("nemo_rl.utils")) + sys.modules.setdefault("nemo_rl.utils.logger", _logger_mod) + +from nmp.rl.tasks.training.backends.nemo_rl import nemo_rl_logger # noqa: E402 +from nmp.rl.tasks.training.backends.nemo_rl.nemo_rl_logger import ( # noqa: E402 + NemoRLLogger, + has_metric_value, + resolve_log_interval, + resolve_steps_per_epoch, +) + + +class _RecordingCallback: + """Stands in for TrainingProgressCallback, capturing what the logger forwards.""" + + def __init__(self) -> None: + self.train_steps: list[dict[str, Any]] = [] + self.validations: list[dict[str, Any]] = [] + self.training_starts: list[dict[str, Any]] = [] + self.closed = False + + def report_training_start(self, max_steps: int, num_epochs: int) -> None: + self.training_starts.append({"max_steps": max_steps, "num_epochs": num_epochs}) + + def report_train_step(self, step, epoch, loss, lr=None, grad_norm=None, **additional): + self.train_steps.append( + {"step": step, "epoch": epoch, "loss": loss, "lr": lr, "grad_norm": grad_norm, **additional} + ) + + def report_validation(self, step, epoch, val_loss=None, **additional): + self.validations.append({"step": step, "epoch": epoch, "val_loss": val_loss, **additional}) + + def close(self) -> None: + self.closed = True + + +@pytest.fixture +def callback(monkeypatch: pytest.MonkeyPatch) -> _RecordingCallback: + """Build a NemoRLLogger whose reporter/callback are inert local objects.""" + recorder = _RecordingCallback() + monkeypatch.setattr(nemo_rl_logger, "JobsServiceProgressReporter", lambda *a, **k: object()) + monkeypatch.setattr(nemo_rl_logger, "TrainingProgressCallback", lambda _reporter: recorder) + return recorder + + +def _make_logger(**kwargs: Any) -> NemoRLLogger: + params: dict[str, Any] = {"steps_per_epoch": 10, "log_interval": 1} + params.update(kwargs) + return NemoRLLogger(**params) + + +def _driver_steps(max_steps: int) -> range: + """The step sequence an N-step run actually produces. + + grpo.py and dpo.py both log `total_steps + 1` with total_steps 0-based and + incremented after the log, so an N-step run emits 1..N -- not 0..N-1. Tests + that use range(N) directly would validate the throttle against a convention + no caller uses. + """ + return range(1, max_steps + 1) + + +class _Histogram: + """Stand-in for nemo_rl's wandb Histogram — non-numeric, and NaN-hostile.""" + + +# Trimmed but faithful shape of GRPO's `train` dict: policy-loss metrics, the +# reward/advantage block, NeMo-Gym rollout metrics, and the non-scalars that ride +# along with them. +GRPO_TRAIN_METRICS: dict[str, Any] = { + "loss": 0.31, + "lr": 5e-6, + "grad_norm": 1.7, + "reward": 0.62, + "total_reward/mean": 0.62, + "total_reward/histogram": _Histogram(), + "advantages/mean": 0.04, + "advantages/min": -1.2, + "advantages/max": 1.4, + "kl_penalty": 0.008, + "approx_entropy": 0.55, + "token_mult_prob_error": 1.02, + "truncation_rate": 0.125, + "natural_termination_rate": 0.875, + "turns_per_sample/mean": 2.5, + "mean_gen_tokens_per_sample": 148.0, + "num_valid_samples": 64, + "global_valid_seqs": 64.0, + "global_valid_toks": 9472.0, + "generation_logger_metrics": {"inflight": [1, 2, 3]}, + "per_worker_token_counts": [{0: 100, 1: 120}], + "ascii_tree_agent/full_result": object(), +} + +# GRPO validation, verbatim in shape: no `loss` key anywhere. +GRPO_VALIDATION_METRICS: dict[str, Any] = {"accuracy": 0.75, "avg_length": 143.2} + +DPO_TRAIN_METRICS: dict[str, Any] = { + "loss": 0.5, + "lr": 1e-5, + "grad_norm": 0.9, + "preference_loss": 0.42, + "rewards_rejected_mean": -0.3, + "num_valid_samples": 8, + "global_valid_seqs": 8.0, + "global_valid_toks": 1024.0, +} + + +# --------------------------------------------------------------------------- # +# has_metric_value +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + "value,expected", + [ + (0.5, True), + (0, True), + (-1.5, True), + (float("nan"), False), + (None, False), + # Non-scalars that genuinely appear in NeMo-RL metric dicts. Each of these + # raises TypeError under a bare math.isnan, which is the regression guarded here. + (_Histogram(), False), + ({"inflight": [1, 2]}, False), + ([1, 2, 3], False), + ("0.5", False), + # bool is an int subclass; charting a flag as 0/1 is not wanted. + (True, False), + (False, False), + ], +) +def test_has_metric_value(value: Any, expected: bool) -> None: + assert has_metric_value(value) is expected + + +def test_has_metric_value_does_not_raise_on_any_grpo_metric() -> None: + """Every value in a real GRPO dict must be classifiable without raising.""" + for key, value in GRPO_TRAIN_METRICS.items(): + assert isinstance(has_metric_value(value), bool), key + + +def test_module_stub_does_not_break_find_spec() -> None: + """The stub installed at import time outlives this module; it must be inert. + + find_spec consults sys.modules first and raises on a `__spec__` of None, so a + bare ModuleType here would turn an unrelated later `find_spec("nemo_rl")` + into a ValueError -- the same kind of cross-suite leak this file's sibling + test_grpo_config had to be rewritten around. + """ + assert importlib.util.find_spec("nemo_rl") is not None + + +def test_has_metric_value_accepts_numpy_scalars() -> None: + np = pytest.importorskip("numpy") + assert has_metric_value(np.float32(0.5)) is True + assert has_metric_value(np.float64(0.5)) is True + assert has_metric_value(np.int64(3)) is True + assert has_metric_value(np.float32("nan")) is False + + +# --------------------------------------------------------------------------- # +# GRPO train metrics +# --------------------------------------------------------------------------- # + + +def test_grpo_train_step_forwards_reward_and_rollout_metrics(callback: _RecordingCallback) -> None: + """The reward signal is the point of GRPO; it must reach Jobs Service.""" + _make_logger().log_metrics(GRPO_TRAIN_METRICS, step=0, prefix="train") + + assert len(callback.train_steps) == 1 + reported = callback.train_steps[0] + assert reported["loss"] == 0.31 + assert reported["reward"] == 0.62 + assert reported["total_reward/mean"] == 0.62 + assert reported["advantages/mean"] == 0.04 + assert reported["kl_penalty"] == 0.008 + assert reported["approx_entropy"] == 0.55 + assert reported["truncation_rate"] == 0.125 + assert reported["turns_per_sample/mean"] == 2.5 + + +def test_grpo_train_step_drops_non_scalar_metrics(callback: _RecordingCallback) -> None: + """Histograms/Tables/nested dicts must not be forwarded, and must not raise.""" + _make_logger().log_metrics(GRPO_TRAIN_METRICS, step=0, prefix="train") + + reported = callback.train_steps[0] + for key in ( + "total_reward/histogram", + "generation_logger_metrics", + "per_worker_token_counts", + "ascii_tree_agent/full_result", + ): + assert key not in reported + + +def test_rollout_only_train_call_is_ignored(callback: _RecordingCallback) -> None: + """GRPO logs `train` twice per step; the mid-step call has no loss and is a partial.""" + rollout_only = {k: v for k, v in GRPO_TRAIN_METRICS.items() if k != "loss"} + _make_logger().log_metrics(rollout_only, step=0, prefix="train") + + assert callback.train_steps == [] + + +def test_dpo_train_metrics_still_forwarded(callback: _RecordingCallback) -> None: + """Widening the key list for GRPO must not drop DPO's existing metrics.""" + _make_logger().log_metrics(DPO_TRAIN_METRICS, step=0, prefix="train") + + reported = callback.train_steps[0] + assert reported["loss"] == 0.5 + assert reported["preference_loss"] == 0.42 + assert reported["rewards_rejected_mean"] == -0.3 + assert reported["global_valid_toks"] == 1024.0 + + +def test_log_interval_throttles_train_reports(callback: _RecordingCallback) -> None: + logger = _make_logger(log_interval=5) + for step in _driver_steps(10): + logger.log_metrics(GRPO_TRAIN_METRICS, step=step, prefix="train") + + assert [r["step"] for r in callback.train_steps] == [5, 10] + + +# --------------------------------------------------------------------------- # +# Final-step flush +# --------------------------------------------------------------------------- # + + +def test_close_flushes_the_withheld_final_step(callback: _RecordingCallback) -> None: + """When max_steps is not a multiple of log_interval the last step is throttled out. + + Without a flush the run's last recorded loss is stale — for 23 steps at an + interval of 10 it would be step 20's, and steps 21-23 would never be seen. + """ + logger = _make_logger(log_interval=10) + for step in _driver_steps(23): + logger.log_metrics(GRPO_TRAIN_METRICS, step=step, prefix="train") + + assert [r["step"] for r in callback.train_steps] == [10, 20] + + logger.close() + + assert [r["step"] for r in callback.train_steps] == [10, 20, 23] + + +def test_close_does_not_duplicate_an_already_reported_step(callback: _RecordingCallback) -> None: + logger = _make_logger(log_interval=10) + for step in _driver_steps(20): + logger.log_metrics(GRPO_TRAIN_METRICS, step=step, prefix="train") + + logger.close() + + assert [r["step"] for r in callback.train_steps] == [10, 20] + + +def test_flushed_step_carries_the_full_metric_payload(callback: _RecordingCallback) -> None: + logger = _make_logger(log_interval=10) + logger.log_metrics(GRPO_TRAIN_METRICS, step=1, prefix="train") + logger.close() + + flushed = callback.train_steps[-1] + assert flushed["step"] == 1 + assert flushed["reward"] == 0.62 + assert flushed["loss"] == 0.31 + + +def test_double_close_flushes_once(callback: _RecordingCallback) -> None: + logger = _make_logger(log_interval=10) + logger.log_metrics(GRPO_TRAIN_METRICS, step=1, prefix="train") + + logger.close() + logger.close() + + assert len(callback.train_steps) == 1 + + +def test_finish_flushes_like_close(callback: _RecordingCallback) -> None: + """`finish` is the name NeMo-RL's composite Logger actually dispatches. + + nemo_rl.utils.logger.Logger has no close(); its teardown fan-out is + `getattr(logger, "finish", None)`. Without this alias the composite skips us + entirely and the withheld final step is never flushed. + """ + logger = _make_logger(log_interval=10) + logger.log_metrics(GRPO_TRAIN_METRICS, step=1, prefix="train") + + logger.finish() + + assert [r["step"] for r in callback.train_steps] == [1] + assert callback.closed + + +def test_finish_is_reachable_through_the_composite_dispatch(callback: _RecordingCallback) -> None: + """Mirrors Logger.finish()'s exact lookup, so a rename here fails loudly.""" + logger = _make_logger(log_interval=10) + logger.log_metrics(GRPO_TRAIN_METRICS, step=1, prefix="train") + + finish = getattr(logger, "finish", None) + assert callable(finish) + finish() + + assert [r["step"] for r in callback.train_steps] == [1] + + +def test_finish_then_close_flushes_once(callback: _RecordingCallback) -> None: + """Both the composite and the driver may call in; the step reports once.""" + logger = _make_logger(log_interval=10) + logger.log_metrics(GRPO_TRAIN_METRICS, step=1, prefix="train") + + logger.finish() + logger.close() + + assert len(callback.train_steps) == 1 + + +def test_close_with_nothing_pending_reports_nothing(callback: _RecordingCallback) -> None: + _make_logger().close() + + assert callback.train_steps == [] + + +# --------------------------------------------------------------------------- # +# Schedule resolution — one formula for both drivers +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + "val_period,expected", + [ + (100, 10), + (10, 1), + (5, 1), # floors to 0 -> clamped + (1, 1), + (0, 1), + (None, 1), # GRPO's val_period is Optional + ], +) +def test_resolve_log_interval(val_period: int | None, expected: int) -> None: + assert resolve_log_interval(val_period) == expected + + +@pytest.mark.parametrize( + "max_steps,num_epochs,explicit,expected", + [ + (100, 4, None, 25), + (100, None, None, 100), + (100, 0, None, 100), # guard against a zero divisor + (3, 10, None, 1), # floors to 0 -> clamped + (100, 4, 40, 40), # explicit wins (DPO carries steps_per_epoch) + (100, 4, 0, 25), # ...unless it is unusable + ], +) +def test_resolve_steps_per_epoch(max_steps: int, num_epochs: int | None, explicit: int | None, expected: int) -> None: + assert resolve_steps_per_epoch(max_steps, num_epochs, explicit) == expected + + +def test_for_schedule_builds_a_consistent_logger(callback: _RecordingCallback) -> None: + """Both drivers now share this path; DPO used to derive a different interval.""" + logger = NemoRLLogger.for_schedule(max_steps=100, num_epochs=4, val_period=100) + + assert logger._log_interval == 10 + assert logger._steps_per_epoch == 25 + assert logger._max_steps == 100 + + +# --------------------------------------------------------------------------- # +# Step and epoch arithmetic +# --------------------------------------------------------------------------- # + + +def test_step_is_reported_as_the_caller_numbered_it(callback: _RecordingCallback) -> None: + """The caller's step is already 1-indexed; re-incrementing shifted the curve.""" + logger = _make_logger(max_steps=23) + for step in _driver_steps(23): + logger.log_metrics(GRPO_TRAIN_METRICS, step=step, prefix="train") + + reported = [r["step"] for r in callback.train_steps] + assert reported[0] == 1, "an N-step run starts at 1" + assert reported[-1] == 23, "...and ends at N, not N+1" + + +@pytest.mark.parametrize( + "step,expected_epoch", + [ + (0, 1), # validate-at-start, before any training + (1, 1), + (10, 1), # last step of epoch 1 at steps_per_epoch=10 + (11, 2), # first of epoch 2 + (20, 2), + (21, 3), + ], +) +def test_epoch_boundaries(callback: _RecordingCallback, step: int, expected_epoch: int) -> None: + """Epoch flips on the step after a full epoch, not the last step of one.""" + _make_logger().log_metrics({"loss": 0.5}, step=step, prefix="train") + + assert callback.train_steps[0]["epoch"] == expected_epoch + + +def test_validate_at_start_reports_step_zero(callback: _RecordingCallback) -> None: + """Both algorithms run an optional validation pass at step 0 before training.""" + _make_logger().log_metrics(GRPO_VALIDATION_METRICS, step=0, prefix="validation") + + reported = callback.validations[0] + assert reported["step"] == 0 + assert reported["epoch"] == 1 + + +# --------------------------------------------------------------------------- # +# Validation — the branch GRPO never reached +# --------------------------------------------------------------------------- # + + +def test_grpo_validation_is_reported_without_a_loss(callback: _RecordingCallback) -> None: + """GRPO validates on accuracy/avg_length and reports no loss. + + Gating this branch on `loss` silently dropped every GRPO validation report. + """ + _make_logger().log_metrics(GRPO_VALIDATION_METRICS, step=10, prefix="validation") + + assert len(callback.validations) == 1 + reported = callback.validations[0] + assert reported["val_loss"] is None + assert reported["accuracy"] == 0.75 + assert reported["avg_length"] == 143.2 + assert reported["step"] == 10 + assert reported["epoch"] == 1 + + +def test_dpo_validation_still_reports_loss(callback: _RecordingCallback) -> None: + _make_logger().log_metrics({"loss": 0.25, "num_valid_samples": 8}, step=10, prefix="validation") + + reported = callback.validations[0] + assert reported["val_loss"] == 0.25 + assert reported["num_valid_samples"] == 8 + + +def test_validation_with_nothing_usable_is_not_reported(callback: _RecordingCallback) -> None: + """An empty or all-non-scalar dict must not produce a hollow report.""" + _make_logger().log_metrics({}, step=9, prefix="validation") + _make_logger().log_metrics({"histogram/x": _Histogram()}, step=9, prefix="validation") + + assert callback.validations == [] + + +def test_best_validation_loss_tracks_minimum(callback: _RecordingCallback) -> None: + logger = _make_logger() + logger.log_metrics({"loss": 0.5}, step=10, prefix="validation") + logger.log_metrics({"loss": 0.2}, step=20, prefix="validation") + logger.log_metrics({"loss": 0.7}, step=30, prefix="validation") + + assert logger._best_metric_value == 0.2 + assert logger._best_epoch == 2 + + +def test_grpo_validation_leaves_best_loss_untouched(callback: _RecordingCallback) -> None: + """No loss means no best-loss update — and no crash comparing None.""" + logger = _make_logger() + logger.log_metrics(GRPO_VALIDATION_METRICS, step=10, prefix="validation") + + assert math.isinf(logger._best_metric_value) + assert logger._best_epoch is None + + +@pytest.mark.parametrize("prefix", ["validation", "validation-0", "validation/nemo_gym"]) +def test_all_validation_prefixes_are_handled(callback: _RecordingCallback, prefix: str) -> None: + """NeMo-RL suffixes the prefix per dataloader; all must route to validation.""" + _make_logger().log_metrics(GRPO_VALIDATION_METRICS, step=9, prefix=prefix) + + assert len(callback.validations) == 1 + + +# --------------------------------------------------------------------------- # +# Prefixes we intentionally ignore +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("prefix", ["timing/train", "timing/validation", "timing/setup", "performance", "refit", ""]) +def test_unhandled_prefixes_produce_no_reports(callback: _RecordingCallback, prefix: str) -> None: + _make_logger().log_metrics({"loss": 0.1, "total_step_time": 12.0}, step=0, prefix=prefix) + + assert callback.train_steps == [] + assert callback.validations == [] diff --git a/services/unsloth/tests/test_callbacks.py b/services/unsloth/tests/test_callbacks.py index bdd09282c7..15a83a6bd0 100644 --- a/services/unsloth/tests/test_callbacks.py +++ b/services/unsloth/tests/test_callbacks.py @@ -65,6 +65,7 @@ def test_report_training_start_delegates(self): step=0, max_steps=500, num_epochs=2, + metrics={"train_loss": [], "val_loss": []}, backend="unsloth", )