From 73750b1fb3f4614676f67c1707b96d856e3af1ca Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 14:40:04 -0400 Subject: [PATCH 01/23] refactor(rl): use the shared training progress callback directly RL carried a standalone TrainingProgressCallback that duplicated the shared one in packages/nmp_customization_common, minus its metric accumulation. Fixes to progress reporting therefore had to be made twice or -- in practice -- only once, in whichever copy the author happened to be looking at. The copy is deleted outright rather than replaced by a subclass, because a subclass would add nothing: `_default_backend` is already None on the base, and that default is what keeps RL's status-detail shape unchanged on the wire (no `backend` key is added). automodel imports the shared class directly for exactly this reason; unsloth is the only service that subclasses it, and only to stamp `backend="unsloth"`. Two additive changes to the shared class make it a drop-in for what RL's copy supported: **additional_metrics backend-specific scalars alongside loss/lr/grad_norm. Splatted first, so a backend metric cannot shadow the accumulated series or the step's own loss; every other colliding name is a real parameter and already errors at the call site. optional val_loss not every algorithm produces one. The key is omitted rather than sent as null, which would chart as a zero. Also adds the missing services/rl/.../training/progress.py, matching the unsloth and automodel modules that bind SERVICE_NAME, so the two RL construction sites stop passing it by hand. Signed-off-by: Albert Cui --- .../training/callbacks.py | 34 +++++-- .../training/backends/nemo_rl/callbacks.py | 95 ------------------- .../backends/nemo_rl/nemo_rl_logger.py | 7 +- .../rl/src/nmp/rl/tasks/training/progress.py | 33 +++++++ .../rl/src/nmp/rl/tasks/training/runner.py | 6 +- 5 files changed, 67 insertions(+), 108 deletions(-) delete mode 100644 services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py create mode 100644 services/rl/src/nmp/rl/tasks/training/progress.py 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..fd6ef87b61 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 @@ -71,10 +71,18 @@ def report_train_step( grad_norm: float | None = None, *, backend: str | None = None, + **additional_metrics: object, ) -> None: - """Report training step with metrics.""" + """Report training step with metrics. + + ``additional_metrics`` are backend-specific scalars (DPO's + ``preference_loss``, ...). Splatted first so a backend metric cannot + shadow the accumulated series or the step's own loss; every other + colliding name is a real parameter and so already errors at the call. + """ self._train_metrics.append({"step": step, "epoch": epoch, "value": loss}) details: dict[str, object] = { + **additional_metrics, "step": step, "epoch": epoch, "train_loss": loss, @@ -87,15 +95,29 @@ 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. The + key is omitted rather than sent as null, which would chart as a real zero. + """ details: dict[str, object] = { "step": step, "epoch": epoch, - "val_loss": val_loss, - "metrics": self._build_metrics_summary(), + **additional_metrics, } + if val_loss is not None: + self._val_metrics.append({"step": step, "epoch": epoch, "value": 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 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 deleted file mode 100644 index 8ac0760844..0000000000 --- a/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/callbacks.py +++ /dev/null @@ -1,95 +0,0 @@ -# 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. - -import logging -from typing import Any - -from nmp.customization_common.training.progress import JobsServiceProgressReporter - -logger = logging.getLogger(__name__) - - -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() 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..767a9599ce 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 @@ -14,9 +14,8 @@ 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.rl.tasks.training.backends.nemo_rl.callbacks import TrainingProgressCallback +from nmp.customization_common.training.callbacks import TrainingProgressCallback +from nmp.rl.tasks.training.progress import JobsServiceProgressReporter _logger = logging.getLogger(__name__) @@ -74,7 +73,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 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 From d94bc200728928f3b3c87d65692c77d7c15a420e Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 14:40:05 -0400 Subject: [PATCH 02/23] fix(customization): stop non-step reports from erasing the metric series `report_running` REPLACES the task's status_details blob rather than merging into it, so a report that omits `metrics` erases the accumulated series from stored status until the next train step resends it -- and loses it outright if the job dies in that window. report_training_start, report_epoch_end and report_checkpoint_saved all omitted it. automodel calls both report_epoch_end and report_checkpoint_saved mid-training, so this was reachable in practice, not theoretical. On report_training_start it also blanked a resumed job's seeded series before the first step could restate it. Two automodel tests and one unsloth test pinned the buggy payload with exact-kwargs assertions; they now assert the series survives instead. Signed-off-by: Albert Cui --- .../training/callbacks.py | 29 ++++++++++++-- .../tasks/training/backends/test_callbacks.py | 38 ++++++++++++++++++- services/unsloth/tests/test_callbacks.py | 1 + 3 files changed, 62 insertions(+), 6 deletions(-) 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 fd6ef87b61..09de0e9b2e 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 @@ -5,10 +5,17 @@ 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 +``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. +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. + 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 @@ -56,7 +63,12 @@ def _build_metrics_summary(self) -> dict[str, list[dict[str, float | int]]]: 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 @@ -132,7 +144,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 @@ -140,7 +157,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/services/automodel/tests/tasks/training/backends/test_callbacks.py b/services/automodel/tests/tasks/training/backends/test_callbacks.py index 71dcb278ca..46132c3903 100644 --- a/services/automodel/tests/tasks/training/backends/test_callbacks.py +++ b/services/automodel/tests/tasks/training/backends/test_callbacks.py @@ -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/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", ) From 3148e2324c8b3734cfdd7c84cd98278c61a7f104 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 14:40:05 -0400 Subject: [PATCH 03/23] feat(customization): accumulate a time series for every reported metric Only train_loss and val_loss were series; every other metric a backend reported rode as a current-step scalar that the next update overwrote. So the only thing a finished job could be charted on was its loss, no matter how much else the backend knew. Now every numeric metric accumulates into its own series in the same {step, epoch, value} shape Studio already renders. The current-step scalars stay on the blob alongside, so consumers can read either the curve or the latest value. Series are namespaced by phase: train_ / val_. The prefix is load-bearing, not cosmetic -- backends report the same metric name in both their train and validation dicts (NeMo-RL does this with truncation_rate, and DPO with accuracy), so unprefixed names would interleave two different quantities into one curve. train_loss and val_loss keep their bare names, so the existing Studio loss chart is unaffected. lr and grad_norm accumulate too; they are curves people read, and they were only excluded because they happen to be named parameters rather than **additional_metrics. fetch_current_metrics had to stop hardcoding the two names, or a resumed job would silently restart every other curve from empty. It now returns whatever list-valued series are stored. The numeric guard lands here as is_chartable(), and NemoRLLogger's has_metric_value delegates to it: a metric the logger forwards must be one the callback can chart, and letting those drift is how a histogram object ends up in a series. It also removes a latent crash -- math.isnan raises TypeError on the non-scalars a framework metric dict can carry. Size scales with the number of *reports*, not training steps, since backends throttle reporting. Measured for a 22-series RL run: 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 Accepted for batch training jobs. A backend that reports every step of a long run pays quadratically; if that becomes a real configuration the fix is delta appends in the transport, not trimming the series here. Signed-off-by: Albert Cui --- .../training/callbacks.py | 140 ++++++++-- .../customization_common/training/progress.py | 18 +- .../tests/training/test_callbacks.py | 258 ++++++++++++++++++ .../tasks/training/backends/test_callbacks.py | 4 +- .../backends/nemo_rl/nemo_rl_logger.py | 15 +- 5 files changed, 392 insertions(+), 43 deletions(-) create mode 100644 packages/nmp_customization_common/tests/training/test_callbacks.py 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 09de0e9b2e..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,10 +4,9 @@ """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 @@ -16,19 +15,71 @@ 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.""" @@ -40,25 +91,46 @@ 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.""" @@ -87,12 +159,19 @@ def report_train_step( ) -> None: """Report training step with metrics. - ``additional_metrics`` are backend-specific scalars (DPO's - ``preference_loss``, ...). Splatted first so a backend metric cannot - shadow the accumulated series or the step's own loss; every other - colliding name is a real parameter and so already errors at the call. + ``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._train_metrics.append({"step": step, "epoch": epoch, "value": loss}) + 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, @@ -118,18 +197,23 @@ def report_validation( ) -> None: """Report validation results. - ``val_loss`` is optional because not every algorithm produces one. The - key is omitted rather than sent as null, which would chart as a real zero. + ``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, **additional_metrics, } + for name, value in additional_metrics.items(): + self._record("val", name, step, epoch, value) if val_loss is not None: - self._val_metrics.append({"step": step, "epoch": epoch, "value": val_loss}) + 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 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..3495af9898 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 @@ -82,8 +82,15 @@ def update_task( logger.warning(f"Failed to update task progress: {e}") 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. + """ if not self._enabled: - return {"train_loss": [], "val_loss": []} + return {} try: jobs = client_from_platform(self._sdk, JobsClient) @@ -94,13 +101,10 @@ def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: 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", []), - } + return {name: points for name, points in metrics.items() if isinstance(points, list)} except Exception as e: - logger.info(f"No prior metrics to seed (expected on first run): {e}") - return {"train_loss": [], "val_loss": []} + logger.info(f"No stored metrics available: {e}") + return {} 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..efb71ed2ed --- /dev/null +++ b/packages/nmp_customization_common/tests/training/test_callbacks.py @@ -0,0 +1,258 @@ +# 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_no_backend_field_when_the_default_is_unset(reporter: _RecordingReporter) -> None: + """automodel and NeMo-RL both depend on this: their reports carry no `backend`. + + Neither subclasses this class, so the absence of the key is a property of the + default here rather than of anything on their side. unsloth opts in. + """ + 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.4) + callback.report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/ckpt") + callback.report_epoch_end(step=1, epoch=1) + + assert reporter.reports, "expected reports to assert against" + assert all("backend" not in report for report in reporter.reports) + + +def test_close_delegates_to_the_reporter(reporter: _RecordingReporter) -> None: + _make_callback(reporter).close() + + assert reporter.closed diff --git a/services/automodel/tests/tasks/training/backends/test_callbacks.py b/services/automodel/tests/tasks/training/backends/test_callbacks.py index 46132c3903..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): 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 767a9599ce..e7d412063c 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,22 +9,25 @@ # 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.callbacks import TrainingProgressCallback +from nmp.customization_common.training.callbacks import TrainingProgressCallback, is_chartable from nmp.rl.tasks.training.progress import JobsServiceProgressReporter _logger = logging.getLogger(__name__) 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. + + 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. It also stops being a crash risk: ``math.isnan`` raises ``TypeError`` + on the non-scalars a metric dict can carry, rather than returning False. + """ + return is_chartable(metric) class NemoRLLogger(LoggerInterface): From 87dd5916b5b385ffe7286e3278e72fef04efde4a Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 14:40:05 -0400 Subject: [PATCH 04/23] fix(customization): carry sticky status_details fields across updates status_details is REPLACED on every update, so a field survives only as long as the next report repeats it. Three kinds of field were being lost to that: metrics erased by the runner's checkpoint/completion/failure reports, which come from a different process than the training driver and hold no series to resend -- so every job ended by erasing its own curves, worst of all on the failure path where the partial curve is worth the most max_steps, num_epochs stated once by report_training_start, gone from the first training step onward checkpoint_path published by one report, wiped by the next Studio reads max_steps and checkpoint_path straight out of status_details, so "step / max steps" fell back to a bare step number for the whole run, and the latest-checkpoint row appeared and vanished. _CARRY_FORWARD names the rule: what stays true after the update that stated it. Cumulative (metrics), run constants (max_steps, num_epochs), monotonic progress (step, epoch), and sticky latest-values (checkpoint_path). Excluded deliberately: `phase`, which every report sets for itself, and the per-step observations (train_loss, lr, grad_norm, ...). Those describe one instant and a stale copy would misrepresent "current" -- and nothing is lost, because each is now recoverable from its series. percentage_done is excluded too: it is derived from step and max_steps, both carried, so a consumer can recompute it rather than risk a copy that contradicts its own inputs. Keeping the GET off the hot path is the design constraint. Values are remembered as they pass through, so a process that has already stated a field restates it for free; the stored blob is read back only when an update omits `metrics`, which is the tell that it did not come from TrainingProgressCallback. Per-step reports always carry `metrics` and never fetch. The runner's handful always do. One subtlety: on resume the driver's first report already carries `metrics`, so it would never read the blob back and would drop the previous run's checkpoint_path. _fetch_status_details therefore refreshes the cache as a side effect, which makes the resume-seeding fetch the callback already performs at construction double as the carry-forward seed -- no extra round-trip. Tests drive the SDK client seam rather than stubbing the fetch, so the real _fetch_status_details runs, cache side effect included. First test coverage for this module. Signed-off-by: Albert Cui --- .../customization_common/training/progress.py | 126 ++++++- .../tests/training/test_progress.py | 311 ++++++++++++++++++ 2 files changed, 427 insertions(+), 10 deletions(-) create mode 100644 packages/nmp_customization_common/tests/training/test_progress.py 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 3495af9898..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,20 +157,21 @@ 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]]]: - """Read back every stored metric series, for resume seeding. + def _fetch_status_details(self) -> dict[str, Any]: + """Read back the task's stored ``status_details`` blob. - 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. + 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 {} @@ -100,12 +184,34 @@ 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 {name: points for name, points in metrics.items() if isinstance(points, list)} + stored = cast(dict[str, Any], task.status_details or {}) except Exception as e: - logger.info(f"No stored metrics available: {e}") + # 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: details["percentage_done"] = self._calculate_percentage_done(details["step"]) 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 == [] From 5198b650623b6ab1baeb6bb8144649660d8bafe1 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 14:40:05 -0400 Subject: [PATCH 05/23] fix(rl): report the final training step, and stop double-counting it Two defects in NemoRLLogger, both in how it counts and reports steps. The final training step was never reported. The throttle is `step % log_interval == 0`, so when max_steps is not a multiple of log_interval the last steps are dropped -- at 23 steps and an interval of 10 the run's last recorded loss was step 20's. A withheld step is now held as pending and flushed by close(). Nothing called close(). The driver appends the logger to `logger_inst.loggers` and never tears it down; nemo_rl.utils.logger.Logger has no close() at all -- its only teardown hook is finish(), dispatched as `getattr(logger, "finish", None)`, which skipped us because NemoRLLogger did not define one. And dpo_train never calls finish() either; the only caller upstream is the single-controller path. So the flush would have run only from __del__, at GC or interpreter shutdown, where every failure is swallowed. Two hooks now, because neither alone is sufficient: finish() aliases close() under the name the composite dispatches, and the driver calls close() from a finally, which is the case that matters -- an abnormal exit is exactly when the last step is worth having. Steps were double-counted. `log_metrics` opened with `step = step + 1`, but the caller already counts from 1: dpo.py logs `total_steps + 1`, where total_steps is 0-based and incremented *after* the log. A 23-step run therefore recorded steps 2..24 against max_steps=23, and the log_interval throttle fired on true steps 9, 19, 29 -- withholding the last step even when max_steps *was* a multiple of the interval. Epoch derivation read the same inflated step and flipped an epoch early at the boundary; it now clamps at zero, because step 0 does arrive, from the validate-at-start path, and belongs to epoch 1. for_schedule owns the log_interval and steps_per_epoch arithmetic that the DPO driver used to derive inline. Its `(val_period // 10) + 1` had a `+1` that was a divide-by-zero guard and also skewed every value it produced, and it raised outright when val_period was None. DPO's reporting cadence changes slightly as a result. The driver teardown is asserted against the AST -- the drivers cannot be imported outside the training image -- with the detector's own negative cases pinned, since a tripwire that cannot trip is worse than none. Signed-off-by: Albert Cui --- .../training/backends/nemo_rl/dpo_driver.py | 46 +- .../backends/nemo_rl/nemo_rl_logger.py | 177 +++++-- services/rl/tests/test_nemo_rl_drivers.py | 68 +++ services/rl/tests/test_nemo_rl_logger.py | 460 ++++++++++++++++++ 4 files changed, 682 insertions(+), 69 deletions(-) create mode 100644 services/rl/tests/test_nemo_rl_drivers.py create mode 100644 services/rl/tests/test_nemo_rl_logger.py 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/nemo_rl_logger.py b/services/rl/src/nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py index e7d412063c..fd5fec3694 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 @@ -18,18 +18,53 @@ _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. +# Selection is by presence, so an algorithm that does not produce one of these +# simply omits it. +_TRAIN_METRIC_KEYS = ( + "num_valid_samples", + "preference_loss", + "rewards_rejected_mean", + "global_valid_seqs", + "global_valid_toks", +) + +_VALIDATION_METRIC_KEYS = _TRAIN_METRIC_KEYS + def has_metric_value(metric: Any) -> bool: """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. It also stops being a crash risk: ``math.isnan`` raises ``TypeError`` - on the non-scalars a metric dict can carry, rather than returning False. + 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): """ NemoRLLogger is a logger implementation that reports training updates to Jobs Service. @@ -84,12 +119,47 @@ 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 arithmetic lives here rather than in each driver. DPO's copy read + ``(val_period // 10) + 1``, where the ``+1`` was a divide-by-zero guard + that also skewed every value it produced, and raised outright when + ``val_period`` was None. Owning it here fixes both and gives any further + algorithm one place to call. + + 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], @@ -107,63 +177,44 @@ 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 + # `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 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, - ) + self._callback.report_train_step(**report) + self._pending_train_report = None + else: + self._pending_train_report = report # Handle validation metrics 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] - self._callback.report_validation( step=step, epoch=epoch, val_loss=val_loss, - **additional_metrics, + **self._select_metrics(metrics, _VALIDATION_METRIC_KEYS), ) # Track best validation loss if val_loss < self._best_metric_value: @@ -172,6 +223,11 @@ def log_metrics( _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. @@ -208,14 +264,41 @@ 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 + driver also calls ``close()`` directly, because ``dpo_train`` never + invokes ``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/tests/test_nemo_rl_drivers.py b/services/rl/tests/test_nemo_rl_drivers.py new file mode 100644 index 0000000000..6cdf8c2057 --- /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 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", ["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..4e58ce000f --- /dev/null +++ b/services/rl/tests/test_nemo_rl_logger.py @@ -0,0 +1,460 @@ +# 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`` objects, tables, 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 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 a non-numeric metric value — NaN-hostile, like the real thing.""" + + +# A DPO `train` dict: the whitelisted scalars, plus the non-scalars that ride +# along with them in a real NeMo-RL metric dict. +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, + "some/histogram": _Histogram(), + "generation_logger_metrics": {"inflight": [1, 2, 3]}, + "per_worker_token_counts": [{0: 100, 1: 120}], +} + +VALIDATION_METRICS: dict[str, Any] = {"loss": 0.25, "num_valid_samples": 8} + + +# --------------------------------------------------------------------------- # +# 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_real_metric() -> None: + """Every value in a real metric dict must be classifiable without raising.""" + for key, value in 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_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(TRAIN_METRICS, step=0, prefix="train") + + reported = callback.train_steps[0] + for key in ("some/histogram", "generation_logger_metrics", "per_worker_token_counts"): + assert key not in reported + + +def test_train_call_without_a_loss_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 TRAIN_METRICS.items() if k != "loss"} + _make_logger().log_metrics(rollout_only, step=0, prefix="train") + + assert callback.train_steps == [] + + +def test_whitelisted_train_metrics_are_forwarded(callback: _RecordingCallback) -> None: + """The whitelisted scalars ride along with loss/lr/grad_norm.""" + _make_logger().log_metrics(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(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(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(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(TRAIN_METRICS, step=1, prefix="train") + logger.close() + + flushed = callback.train_steps[-1] + assert flushed["step"] == 1 + assert flushed["loss"] == 0.5 + assert flushed["preference_loss"] == 0.42 + + +def test_double_close_flushes_once(callback: _RecordingCallback) -> None: + logger = _make_logger(log_interval=10) + logger.log_metrics(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(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(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(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: + """DPO's own formula produced 11 here, via a `+1` that skewed every value.""" + 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(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({"loss": 0.5}, step=0, prefix="validation") + + reported = callback.validations[0] + assert reported["step"] == 0 + assert reported["epoch"] == 1 + + +# --------------------------------------------------------------------------- # +# Validation — the branch GRPO never reached +# --------------------------------------------------------------------------- # + + +def test_validation_reports_loss_and_whitelisted_metrics(callback: _RecordingCallback) -> None: + _make_logger().log_metrics(VALIDATION_METRICS, 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 + + +@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(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 == [] From db49abd5e137a2c9f9a3e239e9fe464245896b23 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 15:34:38 -0400 Subject: [PATCH 06/23] thanks CodeRabbit Signed-off-by: Albert Cui --- .../training/callbacks.py | 24 ++++++++++--- .../tests/training/test_callbacks.py | 34 +++++++++++++++++++ .../backends/nemo_rl/nemo_rl_logger.py | 6 ++-- services/rl/tests/test_nemo_rl_logger.py | 5 ++- 4 files changed, 60 insertions(+), 9 deletions(-) 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 608032a434..baa884d86a 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 @@ -70,15 +70,20 @@ def is_chartable(value: Any) -> bool: 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. + nested dicts with the scalars, and ``math.isfinite`` raises ``TypeError`` on + all of those rather than returning False. + + NaN and both infinities are rejected. Neither is a chart value, and a + diverged loss that reaches the wire serializes as a bare ``NaN``/ + ``Infinity`` token, which is not valid JSON -- one such point would put the + whole ``status_details`` blob beyond a strict parser's reach. ``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)) + return math.isfinite(float(value)) class TrainingProgressCallback: @@ -227,13 +232,22 @@ def report_checkpoint_saved( *, backend: str | None = None, ) -> None: - """Report that a checkpoint was saved.""" + """Report that a checkpoint was saved. + + The ``checkpoint_path`` key is omitted when the backend has no path to + state -- both automodel and unsloth pass ``None`` when their framework + doesn't hand one back. Sending it as null would not merely say nothing: + the field is a sticky latest-value carried across updates, and an + explicit null counts as this report's own value, so it would overwrite + the last known checkpoint rather than let it carry forward. + """ details: dict[str, object] = { "step": step, "epoch": epoch, - "checkpoint_path": checkpoint_path, "metrics": self._build_metrics_summary(), } + if checkpoint_path: + details["checkpoint_path"] = checkpoint_path resolved = self._resolve_backend(backend) if resolved is not None: details["backend"] = resolved diff --git a/packages/nmp_customization_common/tests/training/test_callbacks.py b/packages/nmp_customization_common/tests/training/test_callbacks.py index efb71ed2ed..edfb9a4f83 100644 --- a/packages/nmp_customization_common/tests/training/test_callbacks.py +++ b/packages/nmp_customization_common/tests/training/test_callbacks.py @@ -112,6 +112,30 @@ def test_validation_with_loss_records_both_key_and_series(reporter: _RecordingRe assert report["metrics"]["val_loss"] == [{"step": 1, "epoch": 1, "value": 0.25}] +# --------------------------------------------------------------------------- # +# Optional checkpoint_path +# --------------------------------------------------------------------------- # + + +def test_checkpoint_report_without_a_path_omits_the_key(reporter: _RecordingReporter) -> None: + """A null would overwrite the last known checkpoint instead of carrying it. + + `checkpoint_path` is a sticky carry-forward field: the reporter restates it + only on updates that don't state one of their own, and an explicit null + counts as stating one. automodel and unsloth both pass None when their + framework hands back no path. + """ + _make_callback(reporter).report_checkpoint_saved(step=1, epoch=1) + + assert "checkpoint_path" not in reporter.reports[-1] + + +def test_checkpoint_report_with_a_path_states_it(reporter: _RecordingReporter) -> None: + _make_callback(reporter).report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/ckpt") + + assert reporter.reports[-1]["checkpoint_path"] == "/ckpt" + + # --------------------------------------------------------------------------- # # additional_metrics # --------------------------------------------------------------------------- # @@ -186,6 +210,16 @@ def test_non_numeric_metrics_are_dropped_from_the_series(reporter: _RecordingRep assert set(metrics) == {"train_loss", "val_loss"} +def test_infinite_metrics_are_dropped_from_the_series(reporter: _RecordingReporter) -> None: + """A diverged run's inf is not a chart value, and `Infinity` is not valid JSON.""" + _make_callback(reporter).report_train_step( + step=1, epoch=1, loss=0.5, diverged=float("inf"), collapsed=float("-inf") + ) + + 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 = { 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 fd5fec3694..745a94fb4f 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 @@ -37,14 +37,14 @@ def has_metric_value(metric: Any) -> bool: - """Whether ``metric`` is a finite-enough scalar to forward to Jobs Service. + """Whether ``metric`` is a finite 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. + nested dict. ``math.isfinite`` 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 diff --git a/services/rl/tests/test_nemo_rl_logger.py b/services/rl/tests/test_nemo_rl_logger.py index 4e58ce000f..5276ab3feb 100644 --- a/services/rl/tests/test_nemo_rl_logger.py +++ b/services/rl/tests/test_nemo_rl_logger.py @@ -141,9 +141,12 @@ class _Histogram: (0, True), (-1.5, True), (float("nan"), False), + # A diverged loss. Not a chart value, and `Infinity` is not valid JSON. + (float("inf"), False), + (float("-inf"), 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. + # raises TypeError under a bare math.isfinite, which is the regression guarded here. (_Histogram(), False), ({"inflight": [1, 2]}, False), ([1, 2, 3], False), From 8c1f86c8b36b11ee483c89d79d4ec5b311407076 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 16:15:57 -0400 Subject: [PATCH 07/23] refactor(customization): drop the carry-forward machinery The Jobs service MERGES task status_details key-wise rather than replacing the blob -- JobDispatcher._update_status_details_object, applied both to the task and to the copy propagated up to the job. A field therefore survives every later update that does not restate it. Verified end-to-end against a running platform: a full mid-training report followed by a bare {"phase": "processing_checkpoint"} leaves the series, the schedule and the checkpoint path stored intact. Nothing was ever erased. That makes _CARRY_FORWARD, the read-back GET and the metrics payload on the training-start / checkpoint / epoch-end reports redundant. Removing them also drops a network round-trip from every non-step report, including report_error, where it sat between the exception and the error being recorded. The merge is shallow, so a report that does send `metrics` still replaces the stored series wholesale -- the train and validation reports keep resending every series in full. Dropping the payload from report_training_start closes a real hole while it is at it: when the seeding fetch failed, that report wrote an empty accumulator over a resumed job's stored curves. Two smaller fixes in the blast radius: fetch_current_metrics copies each point list so the callback's accumulator no longer aliases the response, and is_chartable's docstring no longer claims NaN/Inf reach the wire as bare JSON tokens -- the SDK coerces both to null, so the cost of letting one through is a hole in the curve, not a malformed blob. Signed-off-by: Albert Cui --- .../training/callbacks.py | 43 ++-- .../customization_common/training/progress.py | 138 +++---------- .../tests/training/test_callbacks.py | 42 +++- .../tests/training/test_progress.py | 186 +++++------------- .../tasks/training/backends/test_callbacks.py | 19 +- services/unsloth/tests/test_callbacks.py | 1 - 6 files changed, 143 insertions(+), 286 deletions(-) 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 baa884d86a..8de76b73f5 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 @@ -5,15 +5,14 @@ Composes a :class:`nmp.customization_common.training.progress.JobsServiceProgressReporter` 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. +accumulated as a time series and sent under a ``metrics`` key on the train and +validation reports, 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. +The Jobs service merges ``status_details`` key-wise, so the stored series +survives every report that does not mention it: checkpoint, epoch-end and +training-start reports state only what they observed. The merge is shallow, +though -- a report that does send ``metrics`` replaces the stored value +wholesale, so every series goes out in full rather than as a delta. Series naming ------------- @@ -29,8 +28,9 @@ 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 +Every series is resent in full on every train and validation report, 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. @@ -73,10 +73,10 @@ def is_chartable(value: Any) -> bool: nested dicts with the scalars, and ``math.isfinite`` raises ``TypeError`` on all of those rather than returning False. - NaN and both infinities are rejected. Neither is a chart value, and a - diverged loss that reaches the wire serializes as a bare ``NaN``/ - ``Infinity`` token, which is not valid JSON -- one such point would put the - whole ``status_details`` blob beyond a strict parser's reach. + NaN and both infinities are rejected: neither is a value a chart can place on + an axis, and a single infinity flattens every real point in the series + against it. The SDK coerces both to ``null`` on the wire, so letting one + through costs a hole in the curve rather than a malformed blob. ``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. @@ -138,13 +138,17 @@ def _build_metrics_summary(self) -> dict[str, list[dict[str, float | int]]]: 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.""" + """Report that training has started with schedule information. + + Carries no ``metrics``: it fires before the first step, so it has nothing + to add to the curves, and sending the empty accumulator would replace a + resumed job's stored series with two empty lists. + """ self._reporter.configure_progress_tracking(max_steps, 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: @@ -237,14 +241,12 @@ def report_checkpoint_saved( The ``checkpoint_path`` key is omitted when the backend has no path to state -- both automodel and unsloth pass ``None`` when their framework doesn't hand one back. Sending it as null would not merely say nothing: - the field is a sticky latest-value carried across updates, and an - explicit null counts as this report's own value, so it would overwrite - the last known checkpoint rather than let it carry forward. + the server merges key-wise, so an explicit null overwrites the last known + checkpoint, while omitting the key leaves it standing. """ details: dict[str, object] = { "step": step, "epoch": epoch, - "metrics": self._build_metrics_summary(), } if checkpoint_path: details["checkpoint_path"] = checkpoint_path @@ -258,7 +260,6 @@ def report_epoch_end(self, step: int, epoch: int, *, backend: str | None = None) details: dict[str, object] = { "step": step, "epoch": epoch, - "metrics": self._build_metrics_summary(), } resolved = self._resolve_backend(backend) if resolved is not None: 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 03522d2bb1..e220e93ebb 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,9 +8,16 @@ 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`. +The Jobs service MERGES ``status_details`` key-wise rather than replacing the +blob -- ``JobDispatcher._update_status_details_object``, applied both to the task +and to the copy propagated up to the job. A field therefore survives every later +update that does not restate it, and each report sends only what it observed. + +The merge is shallow, which matters for exactly one key: a report that sends +``metrics`` replaces the stored series wholesale. That is why +``TrainingProgressCallback`` resends every series in full on the reports that +carry it, and why the reports with nothing new to say about the curves omit the +key rather than sending a partial copy. For training-specific metrics (loss, validation, checkpoints) see the ``TrainingProgressCallback`` which composes this reporter. @@ -29,43 +36,6 @@ 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.""" @@ -77,10 +47,6 @@ 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. @@ -98,42 +64,6 @@ 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", @@ -146,8 +76,6 @@ 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( @@ -157,21 +85,25 @@ def update_task( step=self._job_ctx.step, body=PlatformJobTaskUpdate( status=PlatformJobStatus(status), - status_details=details, + status_details=status_details or {}, error_details=error_details or {}, ), ) except Exception as e: logger.warning(f"Failed to update task progress: {e}") - def _fetch_status_details(self) -> dict[str, Any]: - """Read back the task's stored ``status_details`` blob. + def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: + """Read back every stored metric series, for resume seeding. + + The only read this reporter makes, and it happens once per process. + Because the server merges, no write path needs to know what is already + stored -- a report that omits a field leaves it standing. - 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. + 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, and each list + is copied so the caller's accumulator does not alias the response. """ if not self._enabled: return {} @@ -187,31 +119,15 @@ def _fetch_status_details(self) -> dict[str, Any]: stored = cast(dict[str, Any], task.status_details or {}) except Exception as e: # 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}") + logger.info(f"No stored status details to seed from: {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)} + metrics = cast(dict[str, Any], stored.get("metrics", {}) or {}) + return cast( + dict[str, list[dict[str, float | int]]], + {name: list(points) for name, points in metrics.items() if isinstance(points, list)}, ) - 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: details["percentage_done"] = self._calculate_percentage_done(details["step"]) diff --git a/packages/nmp_customization_common/tests/training/test_callbacks.py b/packages/nmp_customization_common/tests/training/test_callbacks.py index edfb9a4f83..92ec4821ac 100644 --- a/packages/nmp_customization_common/tests/training/test_callbacks.py +++ b/packages/nmp_customization_common/tests/training/test_callbacks.py @@ -55,28 +55,50 @@ def _make_callback(reporter: _RecordingReporter) -> TrainingProgressCallback: # --------------------------------------------------------------------------- # -def test_every_report_path_carries_the_series(reporter: _RecordingReporter) -> None: - """An omitted payload erases the curve from stored status_details.""" +def test_step_reports_carry_the_series(reporter: _RecordingReporter) -> None: + """The reports that move a curve send every curve, since the merge is shallow.""" 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 + assert len(reporter.reports) == 2 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.""" +def test_non_step_reports_omit_the_series(reporter: _RecordingReporter) -> None: + """The server merges, so a report with nothing to add leaves the stored curve alone. + + Sending it anyway is pure upload on every checkpoint, and at training start + it is worse than useless: the accumulator is empty there, and a shallow merge + would replace a resumed job's stored series with two empty lists. + """ + callback = _make_callback(reporter) + callback.report_training_start(max_steps=10, num_epochs=1) + callback.report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/ckpt") + callback.report_epoch_end(step=1, epoch=1) + + assert len(reporter.reports) == 3 + for report in reporter.reports: + assert "metrics" not in report, report["phase"] + + +def test_training_start_does_not_resend_seeded_metrics() -> None: + """It fires before the first step, so it has nothing to say about the curves.""" 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) + callback = _make_callback(reporter) + callback.report_training_start(max_steps=10, num_epochs=1) - assert reporter.reports[0]["metrics"]["train_loss"] == prior["train_loss"] + assert "metrics" not in reporter.reports[0] + + # ...and the seeded series is still there for the first step that does report. + callback.report_train_step(step=2, epoch=1, loss=0.8) + assert reporter.reports[-1]["metrics"]["train_loss"] == [ + {"step": 1, "epoch": 1, "value": 0.9}, + {"step": 2, "epoch": 1, "value": 0.8}, + ] def test_series_are_snapshots_not_live_references(reporter: _RecordingReporter) -> None: diff --git a/packages/nmp_customization_common/tests/training/test_progress.py b/packages/nmp_customization_common/tests/training/test_progress.py index f34204794a..7aedfe366f 100644 --- a/packages/nmp_customization_common/tests/training/test_progress.py +++ b/packages/nmp_customization_common/tests/training/test_progress.py @@ -3,11 +3,11 @@ """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. +The Jobs service merges ``status_details`` key-wise rather than replacing the +blob (``JobDispatcher._update_status_details_object``, verified end-to-end +against a running platform), so a report only ever has to state what it +observed. These tests pin that: each report sends its own fields and nothing +else, and the only read the reporter makes is the one-shot resume seeding. """ from __future__ import annotations @@ -21,7 +21,7 @@ "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. +#: A blob a mid-run job would have stored, as read back on resume. STORED: dict[str, Any] = { "phase": "training", "step": 10, @@ -52,8 +52,7 @@ class _Reporter(JobsServiceProgressReporter): 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. + ``fetch_current_metrics``, so the real fetch runs. """ def __init__(self) -> None: @@ -63,7 +62,6 @@ def __init__(self) -> None: self._enabled = True self._max_steps = 0 self._num_epochs = 0 - self._carried = {} class _Task: @@ -75,13 +73,12 @@ def data(self) -> "_Task": class _Jobs: - """A mini Jobs service: replace-on-write, readable back, counting fetches.""" + """A mini Jobs service: records writes, serves reads, counts 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 @@ -89,8 +86,6 @@ def client(self) -> Any: 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 @@ -123,143 +118,65 @@ def _details(jobs: _Jobs, index: int = -1) -> dict[str, Any]: # --------------------------------------------------------------------------- # -# What carries forward +# A report states what it observed, and nothing else # --------------------------------------------------------------------------- # -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.""" +def test_a_report_sends_only_its_own_fields(jobs: _Jobs) -> None: + """The server merges, so restating untouched fields would be pure upload.""" _reporter(jobs, STORED).report_running("processing_checkpoint") - details = _details(jobs) - assert details["metrics"] == SERIES - assert details["max_steps"] == 30 - assert details["phase"] == "processing_checkpoint" + assert _details(jobs) == {"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`. - """ +def test_completion_sends_only_its_own_fields(jobs: _Jobs) -> None: + """The stored series, schedule and checkpoint path survive on the server.""" _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 + assert _details(jobs) == {"message": "Training completed", "phase": "completed"} + assert jobs.sent[-1]["body"].status == "completed" -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.""" +def test_percentage_done_is_derived_from_a_stated_step(jobs: _Jobs) -> None: 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) + reporter.configure_progress_tracking(max_steps=40, num_epochs=1) + reporter.report_running("training", step=10, metrics=SERIES) - details = _details(jobs) - assert details["max_steps"] == 30 - assert details["num_epochs"] == 3 - assert jobs.fetches == 0 + assert _details(jobs)["percentage_done"] == 25 -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.""" +def test_percentage_done_is_clamped(jobs: _Jobs) -> None: + """A resumed or over-run job can report past max_steps.""" 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" - + reporter.configure_progress_tracking(max_steps=10, num_epochs=1) + reporter.report_running("training", step=99, metrics=SERIES) -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 + assert _details(jobs)["percentage_done"] == 100 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 +# Reads: exactly one, for resume seeding # --------------------------------------------------------------------------- # +def test_no_report_reads_the_blob_back(jobs: _Jobs) -> None: + """Writes are fire-and-forget. A read per report was the old carry-forward tax.""" + reporter = _reporter(jobs, STORED) + for step in range(1, 11): + reporter.report_running("training", step=step, metrics=SERIES) + reporter.report_running("processing_checkpoint") + reporter.report_completed("Training completed") + reporter.report_error("boom") + + assert jobs.fetches == 0 + + 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 @@ -272,20 +189,17 @@ def test_fetch_current_metrics_drops_non_list_values(jobs: _Jobs) -> None: 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. +def test_fetch_current_metrics_copies_the_point_lists(jobs: _Jobs) -> None: + """The caller appends to what it gets back; it must not alias the response.""" + seeded = _reporter(jobs, STORED).fetch_current_metrics() + seeded["train_loss"].append({"step": 20, "epoch": 2, "value": 0.4}) - 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 + assert len(SERIES["train_loss"]) == 1 - 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" +def test_fetch_current_metrics_survives_an_unseeded_task(jobs: _Jobs) -> None: + """First run: the task has no stored details, which is not an error.""" + assert _reporter(jobs, {}).fetch_current_metrics() == {} # --------------------------------------------------------------------------- # @@ -302,6 +216,14 @@ def test_disabled_reporter_sends_nothing(jobs: _Jobs) -> None: assert jobs.sent == [] +def test_disabled_reporter_does_not_fetch(jobs: _Jobs) -> None: + reporter = _reporter(jobs, STORED) + reporter._enabled = False + + assert reporter.fetch_current_metrics() == {} + assert jobs.fetches == 0 + + def test_non_main_rank_sends_nothing(jobs: _Jobs) -> None: reporter = _reporter(jobs, STORED) reporter._is_main_rank = False diff --git a/services/automodel/tests/tasks/training/backends/test_callbacks.py b/services/automodel/tests/tasks/training/backends/test_callbacks.py index 03276934c2..98778b3c65 100644 --- a/services/automodel/tests/tasks/training/backends/test_callbacks.py +++ b/services/automodel/tests/tasks/training/backends/test_callbacks.py @@ -147,7 +147,6 @@ def test_report_training_start_delegates(self): step=0, max_steps=500, num_epochs=2, - metrics={"train_loss": [], "val_loss": []}, ) def test_report_checkpoint_saved_delegates(self): @@ -160,32 +159,30 @@ def test_report_checkpoint_saved_delegates(self): 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. + def test_checkpoint_report_leaves_the_accumulated_series_alone(self): + """The Jobs service merges status_details, so an omitted key is not an erasure. 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. + hook) and have nothing to add to the curves. Resending every series on + each one would be pure upload; omitting the key leaves the stored series + exactly as the last train step left it. """ 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}] + assert "metrics" not in self._last_report_kwargs(reporter) - def test_epoch_end_report_preserves_accumulated_series(self): + def test_epoch_end_report_leaves_the_accumulated_series_alone(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}] + assert "metrics" not in self._last_report_kwargs(reporter) def test_close_delegates(self): callback, reporter = self._make_callback() diff --git a/services/unsloth/tests/test_callbacks.py b/services/unsloth/tests/test_callbacks.py index 15a83a6bd0..bdd09282c7 100644 --- a/services/unsloth/tests/test_callbacks.py +++ b/services/unsloth/tests/test_callbacks.py @@ -65,7 +65,6 @@ def test_report_training_start_delegates(self): step=0, max_steps=500, num_epochs=2, - metrics={"train_loss": [], "val_loss": []}, backend="unsloth", ) From a38f1031e0c8cf05f48c820552b89ec009b11941 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 16:19:38 -0400 Subject: [PATCH 08/23] fix(customization): correct three defects in the reported payload All three reproduced against a live platform before fixing. The unprefixed-name exemption was keyed on the metric name alone, so a backend reporting `val_loss` among a *train* step's metrics appended it to the validation loss curve -- the exact cross-phase interleaving the prefix exists to prevent. It is keyed on (phase, name) now, so such a metric lands in `train_val_loss` and the curve Studio draws as the validation loss stays clean. A non-chartable metric was dropped from its series but still splatted into status_details, where a Histogram makes the whole update fail to serialize. update_task swallows that error, so every metric in the report was lost while the job went on looking healthy -- the opposite of what _record's docstring promises. additional_metrics are now filtered once and the filtered set feeds both the series and the payload, so a metric rides along as a current-step scalar exactly when it entered a series. A metric named `phase` goes out through the same filter: it collides with report_running's own parameter and raised TypeError into the training loop rather than being shadowed by splat order. train_loss, lr, grad_norm and val_loss are now stated only when observed, which is what val_loss already did alone. An absent lr or a NaN grad_norm -- routine on a skipped step -- otherwise reached the server as a null, and a chart reads null as a real zero. Also hardens is_chartable against the OverflowError float() raises on an unbounded int: it was the one input that could still raise out of a predicate whose two call sites both rely on it never raising. Signed-off-by: Albert Cui --- .../training/callbacks.py | 68 ++++++++--- .../tests/training/test_callbacks.py | 114 ++++++++++++++++-- 2 files changed, 157 insertions(+), 25 deletions(-) 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 8de76b73f5..11bb24f33e 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 @@ -60,9 +60,20 @@ 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"}) +#: ``(phase, name)`` pairs that keep the bare name instead of taking a phase +#: prefix, because they predate the prefixing scheme and are read by name +#: downstream. +#: +#: Keyed on the phase as well as the name, not the name alone: a backend that +#: reports ``val_loss`` among a *train* step's metrics would otherwise append it +#: to the validation loss curve, which is exactly the cross-phase interleaving +#: the prefix exists to prevent. Such a metric becomes ``train_val_loss``. +_UNPREFIXED = frozenset({("train", "train_loss"), ("val", "val_loss")}) + +#: Names a backend metric may not use, because ``report_running`` takes ``phase`` +#: as its own parameter -- a collision is a TypeError out of the training loop +#: rather than the silent shadowing that splat order gives the other names. +_RESERVED = frozenset({"phase"}) def is_chartable(value: Any) -> bool: @@ -83,7 +94,30 @@ def is_chartable(value: Any) -> bool: """ if isinstance(value, bool) or not isinstance(value, numbers.Real): return False - return math.isfinite(float(value)) + try: + return math.isfinite(float(value)) + except (OverflowError, ValueError): + # An unbounded Python int overflows float(). The contract here is that + # any value a backend hands us gets classified without raising, so an + # absurd counter loses its series rather than the run losing reporting. + return False + + +def _forwardable(additional_metrics: dict[str, object]) -> dict[str, object]: + """The subset of ``additional_metrics`` that may ride along in status_details. + + Keeps one invariant: a metric appears as a current-step scalar exactly when + it also entered a series. Both drops are silent, for the same reason -- one + bad metric should cost its own curve, never the whole report: + + - Values :func:`is_chartable` rejects. ``_record`` already skipped these, but + they were still splatted into the payload, where a ``Histogram`` makes the + whole update fail to serialize. ``update_task`` swallows that error, so + every metric in the report is lost while the job goes on looking healthy. + - :data:`_RESERVED` names, which collide with ``report_running``'s own + parameters and raise ``TypeError`` into the training loop. + """ + return {name: value for name, value in additional_metrics.items() if name not in _RESERVED and is_chartable(value)} class TrainingProgressCallback: @@ -123,7 +157,7 @@ def _record(self, phase: str, name: str, step: int, epoch: int, value: object) - # 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}" + series = name if (phase, 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]]]: @@ -172,24 +206,29 @@ def report_train_step( 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. + + Every scalar is stated only when it was actually observed, matching + ``val_loss``: an absent ``lr`` or a NaN ``grad_norm`` reaches the server + as a null, which a chart reads as a real zero. """ + forwardable = _forwardable(additional_metrics) 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(): + for name, value in forwardable.items(): self._record("train", name, step, epoch, value) - # `**additional_metrics` is splatted first, matching report_validation, so a + # `**forwardable` 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, + **forwardable, "step": step, "epoch": epoch, - "train_loss": loss, - "lr": lr, - "grad_norm": grad_norm, "metrics": self._build_metrics_summary(), } + for name, value in (("train_loss", loss), ("lr", lr), ("grad_norm", grad_norm)): + if is_chartable(value): + details[name] = value resolved = self._resolve_backend(backend) if resolved is not None: details["backend"] = resolved @@ -211,14 +250,15 @@ def report_validation( 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. """ + forwardable = _forwardable(additional_metrics) details: dict[str, object] = { "step": step, "epoch": epoch, - **additional_metrics, + **forwardable, } - for name, value in additional_metrics.items(): + for name, value in forwardable.items(): self._record("val", name, step, epoch, value) - if val_loss is not None: + if is_chartable(val_loss): self._record("val", "val_loss", step, epoch, val_loss) details["val_loss"] = val_loss details["metrics"] = self._build_metrics_summary() diff --git a/packages/nmp_customization_common/tests/training/test_callbacks.py b/packages/nmp_customization_common/tests/training/test_callbacks.py index 92ec4821ac..ac73edc365 100644 --- a/packages/nmp_customization_common/tests/training/test_callbacks.py +++ b/packages/nmp_customization_common/tests/training/test_callbacks.py @@ -3,10 +3,12 @@ """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. +Focused on the contract every backend depends on. The Jobs service merges +``status_details`` key-wise but does so shallowly, so a report either resends a +series in full or leaves the key out entirely -- and a report states a scalar +only when it actually observed one, because a null charts as a real zero. Basic +accumulation and resume-seeding are covered by the per-backend suites; this file +covers the shared surface itself. """ from __future__ import annotations @@ -14,7 +16,7 @@ from typing import Any, ClassVar, cast import pytest -from nmp.customization_common.training.callbacks import TrainingProgressCallback +from nmp.customization_common.training.callbacks import TrainingProgressCallback, is_chartable from nmp.customization_common.training.progress import JobsServiceProgressReporter @@ -51,7 +53,7 @@ def _make_callback(reporter: _RecordingReporter) -> TrainingProgressCallback: # --------------------------------------------------------------------------- # -# Every report path carries the series +# Which reports carry the series # --------------------------------------------------------------------------- # @@ -140,12 +142,11 @@ def test_validation_with_loss_records_both_key_and_series(reporter: _RecordingRe def test_checkpoint_report_without_a_path_omits_the_key(reporter: _RecordingReporter) -> None: - """A null would overwrite the last known checkpoint instead of carrying it. + """A null would overwrite the last known checkpoint instead of leaving it. - `checkpoint_path` is a sticky carry-forward field: the reporter restates it - only on updates that don't state one of their own, and an explicit null - counts as stating one. automodel and unsloth both pass None when their - framework hands back no path. + The server merges key-wise, so omitting the key leaves the stored path + standing while an explicit null replaces it. automodel and unsloth both pass + None when their framework hands back no path. """ _make_callback(reporter).report_checkpoint_saved(step=1, epoch=1) @@ -232,6 +233,97 @@ def test_non_numeric_metrics_are_dropped_from_the_series(reporter: _RecordingRep assert set(metrics) == {"train_loss", "val_loss"} +def test_a_dropped_metric_costs_only_itself(reporter: _RecordingReporter) -> None: + """It must not ride along in the payload either, or the whole report dies. + + A `Histogram` in status_details makes the SDK update fail to serialize, and + `update_task` swallows that error -- so every metric in the report is lost + while the job goes on looking healthy. Verified against a live platform + before this filter existed: the step never landed. + """ + _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, hist=object(), reward=0.62) + + report = reporter.reports[-1] + assert "hist" not in report + assert report["reward"] == 0.62, "a well-behaved metric in the same report still lands" + + +def test_a_metric_named_phase_does_not_break_the_report(reporter: _RecordingReporter) -> None: + """`phase` is report_running's own parameter; a collision is a TypeError. + + Splat order silently shadows the other reserved names. This one raised + straight out of the training loop, so it is filtered rather than shadowed. + """ + callback = _make_callback(reporter) + callback.report_train_step(step=1, epoch=1, loss=0.5, phase=1.0) + callback.report_validation(step=1, epoch=1, val_loss=0.4, phase=1.0) + + assert [r["phase"] for r in reporter.reports] == ["training", "validation"] + + +def test_a_train_step_metric_named_val_loss_stays_out_of_the_val_curve( + reporter: _RecordingReporter, +) -> None: + """The legacy bare names are exempt from prefixing per phase, not per name. + + Matching on the name alone put a training-side number in the series Studio + draws as the validation loss. + """ + _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, val_loss=99.0) + + metrics = reporter.reports[-1]["metrics"] + assert metrics["val_loss"] == [] + assert metrics["train_val_loss"] == [{"step": 1, "epoch": 1, "value": 99.0}] + + +def test_the_legacy_names_still_go_unprefixed_in_their_own_phase(reporter: _RecordingReporter) -> None: + """The Studio loss chart reads `train_loss`/`val_loss` by those exact names.""" + callback = _make_callback(reporter) + callback.report_train_step(step=1, epoch=1, loss=0.5) + callback.report_validation(step=1, epoch=1, val_loss=0.4) + + metrics = reporter.reports[-1]["metrics"] + assert metrics["train_loss"] == [{"step": 1, "epoch": 1, "value": 0.5}] + assert metrics["val_loss"] == [{"step": 1, "epoch": 1, "value": 0.4}] + + +# --------------------------------------------------------------------------- # +# A scalar is stated only when it was observed +# --------------------------------------------------------------------------- # + + +def test_absent_lr_and_grad_norm_are_omitted_not_nulled(reporter: _RecordingReporter) -> None: + """A null charts as a real zero -- the same reason val_loss is omitted.""" + _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5) + + report = reporter.reports[-1] + assert "lr" not in report + assert "grad_norm" not in report + + +def test_non_finite_scalars_are_omitted(reporter: _RecordingReporter) -> None: + """A NaN grad_norm is routine on a skipped step; the SDK sends it as null.""" + _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, lr=float("inf"), grad_norm=float("nan")) + + report = reporter.reports[-1] + assert "lr" not in report + assert "grad_norm" not in report + assert report["train_loss"] == 0.5 + + +def test_a_non_finite_val_loss_is_omitted(reporter: _RecordingReporter) -> None: + _make_callback(reporter).report_validation(step=1, epoch=1, val_loss=float("nan")) + + report = reporter.reports[-1] + assert "val_loss" not in report + assert report["metrics"]["val_loss"] == [] + + +def test_is_chartable_classifies_an_unbounded_int_without_raising() -> None: + """float() overflows on a big enough int; the predicate must not propagate it.""" + assert is_chartable(10**400) is False + + def test_infinite_metrics_are_dropped_from_the_series(reporter: _RecordingReporter) -> None: """A diverged run's inf is not a chart value, and `Infinity` is not valid JSON.""" _make_callback(reporter).report_train_step( From 6cb5a1814361888812ef124ec81a78d3aa9dac69 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 16:27:40 -0400 Subject: [PATCH 09/23] chore(rl): give the new progress module the header its siblings use The file was added with an Apache-2.0 SPDX identifier followed by an NVIDIA proprietary "any use ... is strictly prohibited" clause -- two mutually exclusive licenses on one file -- plus a 2026-only copyright year. The block came from the deleted backends/nemo_rl/callbacks.py. Its own docstring says it mirrors the equivalent modules in the unsloth and automodel services; both of those, and its directory neighbour runner.py, use the plain two-line 2025-2026 Apache-2.0 header. Match them. Scoped to the file this branch adds. The same block sits on 17 other files under services/rl and services/automodel, which is a pre-existing repo-wide question rather than this PR's to answer. Note that check-copyright-headers cannot catch any of it: the fixer only adds a header where one is missing and never inspects an existing one, so all 6210 files currently report as correct. Signed-off-by: Albert Cui --- services/rl/src/nmp/rl/tasks/training/progress.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/services/rl/src/nmp/rl/tasks/training/progress.py b/services/rl/src/nmp/rl/tasks/training/progress.py index f998b1adfa..b1ff8d7d0f 100644 --- a/services/rl/src/nmp/rl/tasks/training/progress.py +++ b/services/rl/src/nmp/rl/tasks/training/progress.py @@ -1,12 +1,5 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-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. From 9e66baa3fe921f917b71ae4879267430d1c92cea Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 16:36:10 -0400 Subject: [PATCH 10/23] docs(customization): stop justifying the design with an unwired GRPO path grpo_driver.py is a 108-line stub on this branch: it never constructs a NemoRLLogger, so no GRPO run reports anything. The docstrings added here nonetheless leaned on GRPO behaviour to explain three decisions -- the phase prefix (`truncation_rate` in both metric dicts), the optional val_loss (validating on accuracy/avg_length with no loss), and the payload measurements ("GRPO's ~22 series"). A reader on this branch cannot check any of it. Restated against what is actually wired. DPO's `accuracy` already appears in both its train and validation dicts, so it carries the prefix argument on its own; the optional val_loss is explained by the general case rather than one algorithm; and the payload numbers are real measurements, now attributed to "a backend reporting ~22 series" instead of to a path that does not run. Two stale references fixed while in here: the step-indexing comment cited nemo_rl/algorithms/grpo.py as a second caller when only dpo.py calls in, and two test docstrings pointed at a sibling named test_grpo_config -- the file is test_dpo_config.py. The test that pins cross-phase series separation now uses `accuracy`, the collision that actually occurs, rather than GRPO's `truncation_rate`. Signed-off-by: Albert Cui --- .../training/callbacks.py | 26 +++++++++---------- .../tests/training/test_callbacks.py | 12 ++++----- .../backends/nemo_rl/nemo_rl_logger.py | 4 +-- services/rl/tests/test_nemo_rl_logger.py | 17 ++++++------ 4 files changed, 28 insertions(+), 31 deletions(-) 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 11bb24f33e..1fba2fc180 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 @@ -18,10 +18,9 @@ ------------- 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. +already did. The prefix is load-bearing rather than cosmetic -- DPO reports +``accuracy`` in both its train and validation dicts, 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. @@ -30,12 +29,11 @@ ------------ Every series is resent in full on every train and validation report, 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. +The driver of that cost is the number of *reports*, not training steps -- +backends throttle reporting, so a 500-step run at ``log_interval=10`` +accumulates 50 points per series, not 500. -Measured, for GRPO's ~22 series: +Measured for a backend reporting ~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 @@ -203,7 +201,7 @@ def report_train_step( """Report training step with metrics. ``additional_metrics`` are backend-specific (DPO's ``preference_loss``, - GRPO's ``reward``/``kl_penalty``, ...). Each numeric one accumulates into + ``rewards_rejected_mean``, ...). 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. @@ -245,10 +243,10 @@ def report_validation( ) -> 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. + ``val_loss`` is optional because not every algorithm produces one -- an + algorithm may validate purely on task metrics and report 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. """ forwardable = _forwardable(additional_metrics) details: dict[str, object] = { diff --git a/packages/nmp_customization_common/tests/training/test_callbacks.py b/packages/nmp_customization_common/tests/training/test_callbacks.py index ac73edc365..d061d36d6b 100644 --- a/packages/nmp_customization_common/tests/training/test_callbacks.py +++ b/packages/nmp_customization_common/tests/training/test_callbacks.py @@ -119,7 +119,7 @@ def test_series_are_snapshots_not_live_references(reporter: _RecordingReporter) 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.""" + """An algorithm may validate on task metrics alone; a null val_loss charts as zero.""" _make_callback(reporter).report_validation(step=1, epoch=1, val_loss=None, accuracy=0.75) report = reporter.reports[-1] @@ -213,14 +213,14 @@ def test_additional_validation_metrics_become_series_and_ride_along(reporter: _R 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.""" + """DPO reports `accuracy` 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) + callback.report_train_step(step=1, epoch=1, loss=0.5, accuracy=0.18) + callback.report_validation(step=1, epoch=1, accuracy=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}] + assert metrics["train_accuracy"] == [{"step": 1, "epoch": 1, "value": 0.18}] + assert metrics["val_accuracy"] == [{"step": 1, "epoch": 1, "value": 0.04}] def test_non_numeric_metrics_are_dropped_from_the_series(reporter: _RecordingReporter) -> None: 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 745a94fb4f..1fc488e3a6 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 @@ -177,9 +177,9 @@ 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` arrives 1-indexed and is used as-is. Both callers pass + # `step` arrives 1-indexed and is used as-is. The caller passes # `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 + # logging (nemo_rl/algorithms/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. diff --git a/services/rl/tests/test_nemo_rl_logger.py b/services/rl/tests/test_nemo_rl_logger.py index 5276ab3feb..2f49e83488 100644 --- a/services/rl/tests/test_nemo_rl_logger.py +++ b/services/rl/tests/test_nemo_rl_logger.py @@ -98,10 +98,9 @@ def _make_logger(**kwargs: Any) -> NemoRLLogger: 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. + dpo.py logs `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) @@ -172,7 +171,7 @@ def test_module_stub_does_not_break_find_spec() -> None: 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. + test_dpo_config had to be rewritten around. """ assert importlib.util.find_spec("nemo_rl") is not None @@ -186,7 +185,7 @@ def test_has_metric_value_accepts_numpy_scalars() -> None: # --------------------------------------------------------------------------- # -# GRPO train metrics +# Train metrics # --------------------------------------------------------------------------- # @@ -200,7 +199,7 @@ def test_train_step_drops_non_scalar_metrics(callback: _RecordingCallback) -> No def test_train_call_without_a_loss_is_ignored(callback: _RecordingCallback) -> None: - """GRPO logs `train` twice per step; the mid-step call has no loss and is a partial.""" + """A `train` log without a loss is a partial, mid-step call rather than a step.""" rollout_only = {k: v for k, v in TRAIN_METRICS.items() if k != "loss"} _make_logger().log_metrics(rollout_only, step=0, prefix="train") @@ -337,7 +336,7 @@ def test_close_with_nothing_pending_reports_nothing(callback: _RecordingCallback (5, 1), # floors to 0 -> clamped (1, 1), (0, 1), - (None, 1), # GRPO's val_period is Optional + (None, 1), # val_period is Optional ], ) def test_resolve_log_interval(val_period: int | None, expected: int) -> None: @@ -412,7 +411,7 @@ def test_validate_at_start_reports_step_zero(callback: _RecordingCallback) -> No # --------------------------------------------------------------------------- # -# Validation — the branch GRPO never reached +# Validation # --------------------------------------------------------------------------- # From 8fdf28f371ad198c5fbcc344b079a752861e0bb6 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 16:38:18 -0400 Subject: [PATCH 11/23] fix(rl): make the steps_per_epoch fallback reachable, and drop a string hint for_schedule takes `steps_per_epoch: int | None = None` and derives the value from max_steps and num_epochs when it is missing. That fallback could never run: dpo_driver read `config.dpo.steps_per_epoch` as a plain attribute, and the field is an undeclared extra that exists only because DPOConfig allows extras. pydantic raises AttributeError for a missing extra, so a config compiled anywhere other than dpo_config.py crashed at driver startup -- the exact failure the fallback was written to absorb. Read with a defaulted getattr instead. Guarded with an AST tripwire alongside the existing close()-in-finally one, for the same reason that file gives: the drivers pull in nemo_rl and omegaconf at module scope, so they cannot be imported in a unit test and a regression here would be silent. Separately, for_schedule's return annotation was the string "NemoRLLogger". AGENTS.md asks for concrete hints over string-based ones; typing.Self is the concrete form for a classmethod constructor, and matches NMPJobContext.from_env. Signed-off-by: Albert Cui --- .../training/backends/nemo_rl/dpo_driver.py | 7 ++- .../backends/nemo_rl/nemo_rl_logger.py | 4 +- services/rl/tests/test_nemo_rl_drivers.py | 51 +++++++++++++++++++ 3 files changed, 59 insertions(+), 3 deletions(-) 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 53e093ad56..cd07b80474 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 @@ -115,7 +115,12 @@ def main(): 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" + # An extra (undeclared) DPOConfig field, present only because the model + # allows extras -- dpo_config.py puts it there. Read with getattr, not + # attribute access: a config compiled elsewhere simply omits it, and + # pydantic raises AttributeError for a missing extra. None is what + # for_schedule's derive-from-the-schedule fallback takes. + steps_per_epoch=getattr(config.dpo, "steps_per_epoch", None), ) # The setup() logger is a composite with a `.loggers` list; guard in case # that internal shape changes. 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 1fc488e3a6..f97128d2bd 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,7 +9,7 @@ # its affiliates is strictly prohibited. import logging -from typing import Any, Mapping, Optional +from typing import Any, Mapping, Optional, Self from nemo_rl.utils.logger import LoggerInterface from nmp.customization_common.service.context import NMPJobContext @@ -139,7 +139,7 @@ def for_schedule( val_period: int | None, steps_per_epoch: int | None = None, job_ctx: NMPJobContext | None = None, - ) -> "NemoRLLogger": + ) -> Self: """Build a logger from a NeMo-RL training schedule. The arithmetic lives here rather than in each driver. DPO's copy read diff --git a/services/rl/tests/test_nemo_rl_drivers.py b/services/rl/tests/test_nemo_rl_drivers.py index 6cdf8c2057..215a031a4a 100644 --- a/services/rl/tests/test_nemo_rl_drivers.py +++ b/services/rl/tests/test_nemo_rl_drivers.py @@ -66,3 +66,54 @@ def test_driver_closes_the_progress_logger_in_a_finally(driver: str) -> None: f"{driver} must close customizer_logger from a finally block; " "NeMo-RL does not close loggers, and __del__ does not run on abnormal exit" ) + + +def _reads_via_defaulted_getattr(source: str, field: str) -> bool: + """Whether `field` is read with a three-argument getattr and never as an attribute.""" + via_getattr = False + for node in ast.walk(ast.parse(source)): + if isinstance(node, ast.Attribute) and node.attr == field: + return False + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "getattr" + and len(node.args) == 3 + and isinstance(node.args[1], ast.Constant) + and node.args[1].value == field + ): + via_getattr = True + return via_getattr + + +@pytest.mark.parametrize( + "source,expected", + [ + ("x = getattr(config.dpo, 'steps_per_epoch', None)\n", True), + # Attribute access on an undeclared extra: AttributeError at startup. + ("x = config.dpo.steps_per_epoch\n", False), + # No default, so it raises exactly as attribute access would. + ("x = getattr(config.dpo, 'steps_per_epoch')\n", False), + # Never read at all. + ("x = 1\n", False), + ], +) +def test_optional_field_detector_discriminates(source: str, expected: bool) -> None: + assert _reads_via_defaulted_getattr(source, "steps_per_epoch") is expected + + +@pytest.mark.parametrize("driver", ["dpo_driver.py"]) +def test_driver_reads_steps_per_epoch_defensively(driver: str) -> None: + """It is an undeclared extra, so a config compiled elsewhere simply omits it. + + pydantic raises AttributeError for a missing extra, and this read happens at + driver startup -- before `for_schedule` can apply the fallback that derives + steps_per_epoch from max_steps and num_epochs. Hard attribute access turns + that fallback into dead code and the missing key into a crash. + """ + source = (DRIVERS / driver).read_text() + + assert _reads_via_defaulted_getattr(source, "steps_per_epoch"), ( + f"{driver} must read steps_per_epoch via getattr with a None default; " + "it is an extra=allow field that a config need not carry" + ) From 366840714a04cf883b6d18b5016fed3920bf1250 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 17:07:47 -0400 Subject: [PATCH 12/23] docs(customization): keep the callback's rationale out of the reporter progress.py was explaining why TrainingProgressCallback resends whole series and why some reports omit the metrics key. The dependency runs the other way -- the callback composes the reporter, not the reverse -- so the reporter should state the transport property and stop there. progress.py now says only what it owns: the service merges key-wise and the merge is shallow. The consequence for the accumulator moves into callbacks.py, next to the code that acts on it. No behaviour change. Signed-off-by: Albert Cui --- .../src/nmp/customization_common/training/callbacks.py | 6 ++++-- .../src/nmp/customization_common/training/progress.py | 7 +------ 2 files changed, 5 insertions(+), 8 deletions(-) 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 1fba2fc180..26e06b0b76 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 @@ -11,8 +11,10 @@ The Jobs service merges ``status_details`` key-wise, so the stored series survives every report that does not mention it: checkpoint, epoch-end and training-start reports state only what they observed. The merge is shallow, -though -- a report that does send ``metrics`` replaces the stored value -wholesale, so every series goes out in full rather than as a delta. +though -- a key that is sent replaces the stored value wholesale -- so a report +carrying ``metrics`` sends every series in full rather than a delta, and a report +with nothing new to say about the curves omits the key rather than sending a +partial copy. Series naming ------------- 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 e220e93ebb..917179ebf4 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 @@ -12,12 +12,7 @@ blob -- ``JobDispatcher._update_status_details_object``, applied both to the task and to the copy propagated up to the job. A field therefore survives every later update that does not restate it, and each report sends only what it observed. - -The merge is shallow, which matters for exactly one key: a report that sends -``metrics`` replaces the stored series wholesale. That is why -``TrainingProgressCallback`` resends every series in full on the reports that -carry it, and why the reports with nothing new to say about the curves omit the -key rather than sending a partial copy. +The merge is shallow: a key that is sent replaces the stored value wholesale. For training-specific metrics (loss, validation, checkpoints) see the ``TrainingProgressCallback`` which composes this reporter. From e72b354be8e56b64ed5421bdbce08cdc8f5d8460 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 17:09:33 -0400 Subject: [PATCH 13/23] refactor(rl): stop holding a reference the logger never reads self._reporter was assigned and never touched again. The reporter exists only to be composed into TrainingProgressCallback, which owns it from that point: close() reaches it through self._callback.close(), not through the logger. Signed-off-by: Albert Cui --- .../nmp/rl/tasks/training/backends/nemo_rl/nemo_rl_logger.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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 f97128d2bd..66dd191cf1 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 @@ -110,9 +110,7 @@ def __init__( self._num_epochs = num_epochs self._steps_per_epoch = steps_per_epoch - # Create the callback for progress reporting - self._reporter = JobsServiceProgressReporter(self._job_ctx) - self._callback = TrainingProgressCallback(self._reporter) + self._callback = TrainingProgressCallback(JobsServiceProgressReporter(self._job_ctx)) # Track best metrics for monitoring self._best_metric_value = float("inf") From 73fb71edf6ad9ae5dd76e3e180e1a405deeb164f Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 17:49:14 -0400 Subject: [PATCH 14/23] refactor(customization)!: one naming rule for every training metric train_loss and val_loss were special: named parameters on the callback, exempt from the phase prefix, recorded and forwarded by a different code path than the `**additional_metrics` bag. lr and grad_norm were a third case -- named parameters, but prefixed like ordinary metrics in the series and unprefixed at the top level. Four treatments for four kinds of the same thing. There is now one. A backend hands over its framework's metric dict under its own names, and `_` is the stored series name AND the current-value key. `train_loss` and `val_loss` are what that rule produces for a metric called `loss`, which is why the two series Studio charts did not have to move -- the special case existed only because callers passed them pre-prefixed. Prefixing also retires a whole bug class rather than filtering it. A metric named `phase` used to raise TypeError into the training loop, and `step`/`epoch`/`metrics` needed splat ordering to avoid being shadowed; none of them is reachable from a `_` name, so _RESERVED and the ordering comments are gone. The metric bag is now a Mapping parameter rather than **kwargs. Backends forward whatever their framework emits and a framework is free to call something `step`, which as **kwargs was a hard TypeError. BREAKING: top-level `lr` and `grad_norm` in status_details are now `train_lr` and `train_grad_norm`; Studio is updated to match. The series payload and the top-level `train_loss`/`val_loss` are unchanged. Also drops NeMo-RL's metric allow-list. The callback already keeps the finite scalars and drops the rest, so the list was a second gate doing a weaker version of the same check -- and it silently dropped DPO's accuracy, sft_loss and rewards_chosen_mean for never having been added to it. NeMo-RL's dict is forwarded whole, so a metric it adds charts without a change here. has_metric_value, _select_metrics and the _VALIDATION_METRIC_KEYS alias go with it. Verified against a running platform: train_loss/val_loss keep their names, train and val `accuracy` stay separate, the DPO scalars the allow-list dropped now chart, a Histogram and a nested dict are dropped without costing the report, and metrics named phase/step/metrics land as train_phase/train_step/train_metrics with the real fields intact. Signed-off-by: Albert Cui --- .../training/callbacks.py | 207 ++++---- .../tests/training/test_callbacks.py | 451 +++++++++--------- .../tasks/training/backends/finetune.py | 10 +- .../tasks/training/backends/test_callbacks.py | 38 +- .../backends/nemo_rl/nemo_rl_logger.py | 64 +-- services/rl/tests/test_nemo_rl_logger.py | 103 ++-- .../training/backends/hf_trainer_callback.py | 13 +- services/unsloth/tests/test_callbacks.py | 14 +- .../unsloth/tests/test_hf_trainer_callback.py | 4 +- .../mocks/customizer/customization-jobs.ts | 4 +- .../studio/src/util/customizations.test.ts | 8 +- .../studio/src/util/customizations.tsx | 4 +- 12 files changed, 430 insertions(+), 490 deletions(-) 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 26e06b0b76..500899526d 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 @@ -18,14 +18,21 @@ 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 -- DPO reports -``accuracy`` in both its train and validation dicts, 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. +One rule, no exceptions: a metric is stored and reported as ``_``, +where the backend supplies its framework's own ```` (``loss``, ``lr``, +``accuracy``) and the phase that produced it supplies the prefix. The same +namespaced name is used for the accumulated series and for the current-step +value in ``status_details``. + +The prefix is load-bearing rather than cosmetic -- DPO reports ``accuracy`` in +both its train and validation dicts, so bare names would interleave two +different quantities into one series. It also makes the payload collision-proof: +nothing a ``_`` name can spell reaches ``phase``, ``step``, ``epoch`` or +``metrics``. + +``train_loss`` and ``val_loss``, the two series Studio charts, are what this rule +produces for a metric named ``loss``. They are not special cases, and there are +none. Payload size ------------ @@ -54,27 +61,13 @@ import logging import math import numbers +from collections.abc import Mapping from typing import Any, ClassVar, cast from nmp.customization_common.training.progress import JobsServiceProgressReporter logger = logging.getLogger(__name__) -#: ``(phase, name)`` pairs that keep the bare name instead of taking a phase -#: prefix, because they predate the prefixing scheme and are read by name -#: downstream. -#: -#: Keyed on the phase as well as the name, not the name alone: a backend that -#: reports ``val_loss`` among a *train* step's metrics would otherwise append it -#: to the validation loss curve, which is exactly the cross-phase interleaving -#: the prefix exists to prevent. Such a metric becomes ``train_val_loss``. -_UNPREFIXED = frozenset({("train", "train_loss"), ("val", "val_loss")}) - -#: Names a backend metric may not use, because ``report_running`` takes ``phase`` -#: as its own parameter -- a collision is a TypeError out of the training loop -#: rather than the silent shadowing that splat order gives the other names. -_RESERVED = frozenset({"phase"}) - def is_chartable(value: Any) -> bool: """Whether ``value`` is a finite scalar that can enter a metric series. @@ -103,21 +96,36 @@ def is_chartable(value: Any) -> bool: return False -def _forwardable(additional_metrics: dict[str, object]) -> dict[str, object]: - """The subset of ``additional_metrics`` that may ride along in status_details. +def _namespace(phase: str, metrics: Mapping[str, object]) -> dict[str, float | int]: + """The chartable subset of ``metrics``, keyed by ``_``. + + The single naming rule. A backend passes its framework's own metric names -- + ``loss``, ``lr``, ``accuracy`` -- and the phase that produced them supplies + the prefix, so the train and validation copies of one name cannot collide. + ``train_loss`` and ``val_loss`` are what the rule produces for ``loss``, not + exceptions carved out of it. - Keeps one invariant: a metric appears as a current-step scalar exactly when - it also entered a series. Both drops are silent, for the same reason -- one - bad metric should cost its own curve, never the whole report: + Prefixing is also what makes the payload collision-proof: ``report_running`` + owns ``phase``, and this callback owns ``step``, ``epoch`` and ``metrics``, + none of which a ``_`` name can reach. - - Values :func:`is_chartable` rejects. ``_record`` already skipped these, but - they were still splatted into the payload, where a ``Histogram`` makes the - whole update fail to serialize. ``update_task`` swallows that error, so - every metric in the report is lost while the job goes on looking healthy. - - :data:`_RESERVED` names, which collide with ``report_running``'s own - parameters and raise ``TypeError`` into the training loop. + Values :func:`is_chartable` rejects are dropped silently, and dropped from + the report as well as the series -- one bad metric should cost its own curve, + never the whole report. A ``Histogram`` left in ``status_details`` makes the + update fail to serialize, and ``update_task`` swallows that error, so every + metric in the report is lost while the job goes on looking healthy. """ - return {name: value for name, value in additional_metrics.items() if name not in _RESERVED and is_chartable(value)} + return {f"{phase}_{name}": _coerce(value) for name, value in metrics.items() if is_chartable(value)} + + +def _coerce(value: object) -> float | int: + """Narrow a chartable value to a JSON-serializable builtin. + + numpy scalars satisfy ``numbers.Real`` but are not serializable. Counts stay + ``int`` rather than becoming ``64.0``. + """ + real = cast(numbers.Real, value) + return int(real) if isinstance(real, numbers.Integral) else float(real) class TrainingProgressCallback: @@ -144,21 +152,35 @@ def __init__(self, reporter: JobsServiceProgressReporter): 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. + def _report_metrics( + self, + phase: str, + report_phase: str, + step: int, + epoch: int, + metrics: Mapping[str, object], + backend: str | None, + ) -> None: + """Record every metric as a point and report them as current values. - 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. + The one path both ``report_train_step`` and ``report_validation`` take. + ``phase`` namespaces the series (``train``/``val``); ``report_phase`` is + what the Jobs service records as the task's phase. """ - 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 (phase, name) in _UNPREFIXED else f"{phase}_{name}" - self._series.setdefault(series, []).append({"step": step, "epoch": epoch, "value": numeric}) + namespaced = _namespace(phase, metrics) + for name, value in namespaced.items(): + self._series.setdefault(name, []).append({"step": step, "epoch": epoch, "value": value}) + + details: dict[str, object] = { + "step": step, + "epoch": epoch, + **namespaced, + "metrics": self._build_metrics_summary(), + } + resolved = self._resolve_backend(backend) + if resolved is not None: + details["backend"] = resolved + self._reporter.report_running(phase=report_phase, **details) def _build_metrics_summary(self) -> dict[str, list[dict[str, float | int]]]: """Build the accumulated metrics payload for inclusion in status_details. @@ -190,83 +212,36 @@ def report_training_start(self, max_steps: int, num_epochs: int, *, backend: str self._reporter.report_running(phase="training", **details) def report_train_step( - self, - step: int, - epoch: int, - loss: float, - lr: float | None = None, - grad_norm: float | None = None, - *, - backend: str | None = None, - **additional_metrics: object, + self, step: int, epoch: int, metrics: Mapping[str, object], *, backend: str | None = None ) -> None: - """Report training step with metrics. + """Report one training step. + + Hand over the framework's metric dict as-is, under its own names -- + ``loss``, ``lr``, ``grad_norm``, ``preference_loss``, whatever it + produces. Each chartable entry becomes a ``train_`` series *and* a + ``train_`` current value, so a consumer can read either the curve or + the latest point. - ``additional_metrics`` are backend-specific (DPO's ``preference_loss``, - ``rewards_rejected_mean``, ...). 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. + No metric is required and none is privileged: a step that produces no + loss reports no ``train_loss``, and a name is stated only when it was + observed, because a null charts as a real zero. - Every scalar is stated only when it was actually observed, matching - ``val_loss``: an absent ``lr`` or a NaN ``grad_norm`` reaches the server - as a null, which a chart reads as a real zero. + Taken as a dict rather than ``**kwargs`` so that the metric namespace and + this method's own parameters cannot collide. Backends forward whatever + their framework emits, and a framework is free to call something ``step``. """ - forwardable = _forwardable(additional_metrics) - 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 forwardable.items(): - self._record("train", name, step, epoch, value) - # `**forwardable` 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] = { - **forwardable, - "step": step, - "epoch": epoch, - "metrics": self._build_metrics_summary(), - } - for name, value in (("train_loss", loss), ("lr", lr), ("grad_norm", grad_norm)): - if is_chartable(value): - details[name] = value - resolved = self._resolve_backend(backend) - if resolved is not None: - details["backend"] = resolved - self._reporter.report_running(phase="training", **details) + self._report_metrics("train", "training", step, epoch, metrics, backend) def report_validation( - self, - step: int, - epoch: int, - val_loss: float | None = None, - *, - backend: str | None = None, - **additional_metrics: object, + self, step: int, epoch: int, metrics: Mapping[str, object], *, backend: str | None = None ) -> None: - """Report validation results. + """Report one validation pass. - ``val_loss`` is optional because not every algorithm produces one -- an - algorithm may validate purely on task metrics and report 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. + The same rule under the ``val_`` prefix. An algorithm that validates on + task metrics alone and reports no ``loss`` simply leaves ``val_loss`` + empty, which is why nothing here is required either. """ - forwardable = _forwardable(additional_metrics) - details: dict[str, object] = { - "step": step, - "epoch": epoch, - **forwardable, - } - for name, value in forwardable.items(): - self._record("val", name, step, epoch, value) - if is_chartable(val_loss): - 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 - self._reporter.report_running(phase="validation", **details) + self._report_metrics("val", "validation", step, epoch, metrics, backend) def report_checkpoint_saved( self, diff --git a/packages/nmp_customization_common/tests/training/test_callbacks.py b/packages/nmp_customization_common/tests/training/test_callbacks.py index d061d36d6b..fffdcf8331 100644 --- a/packages/nmp_customization_common/tests/training/test_callbacks.py +++ b/packages/nmp_customization_common/tests/training/test_callbacks.py @@ -3,12 +3,12 @@ """Unit tests for the shared TrainingProgressCallback. -Focused on the contract every backend depends on. The Jobs service merges -``status_details`` key-wise but does so shallowly, so a report either resends a -series in full or leaves the key out entirely -- and a report states a scalar -only when it actually observed one, because a null charts as a real zero. Basic -accumulation and resume-seeding are covered by the per-backend suites; this file -covers the shared surface itself. +Focused on the contract every backend depends on. There is one naming rule -- +a metric is stored and reported as ``_`` -- and no metric is +privileged, so most of this file is about proving the rule holds with no +exceptions hiding in it. The rest covers the transport: the Jobs service merges +``status_details`` key-wise but shallowly, so a report either resends a series in +full or leaves the key out, and states a scalar only when it observed one. """ from __future__ import annotations @@ -53,180 +53,141 @@ def _make_callback(reporter: _RecordingReporter) -> TrainingProgressCallback: # --------------------------------------------------------------------------- # -# Which reports carry the series +# One naming rule, no exceptions # --------------------------------------------------------------------------- # -def test_step_reports_carry_the_series(reporter: _RecordingReporter) -> None: - """The reports that move a curve send every curve, since the merge is shallow.""" +def test_the_phase_supplies_the_prefix(reporter: _RecordingReporter) -> None: + """Backends pass their framework's own names; the phase namespaces them.""" callback = _make_callback(reporter) - callback.report_train_step(step=1, epoch=1, loss=0.5) - callback.report_validation(step=1, epoch=1, val_loss=0.45) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 0.5, "lr": 5e-06, "grad_norm": 1.9}) + callback.report_validation(step=1, epoch=1, metrics={"loss": 0.4, "accuracy": 0.9}) - assert len(reporter.reports) == 2 - for report in reporter.reports: - assert "metrics" in report, report["phase"] - assert set(report["metrics"]) == {"train_loss", "val_loss"} + metrics = reporter.reports[-1]["metrics"] + assert metrics["train_loss"] == [{"step": 1, "epoch": 1, "value": 0.5}] + assert metrics["train_lr"] == [{"step": 1, "epoch": 1, "value": 5e-06}] + assert metrics["train_grad_norm"] == [{"step": 1, "epoch": 1, "value": 1.9}] + assert metrics["val_loss"] == [{"step": 1, "epoch": 1, "value": 0.4}] + assert metrics["val_accuracy"] == [{"step": 1, "epoch": 1, "value": 0.9}] -def test_non_step_reports_omit_the_series(reporter: _RecordingReporter) -> None: - """The server merges, so a report with nothing to add leaves the stored curve alone. +def test_train_loss_and_val_loss_fall_out_of_the_rule(reporter: _RecordingReporter) -> None: + """The two names Studio charts are what `loss` produces, not special cases. - Sending it anyway is pure upload on every checkpoint, and at training start - it is worse than useless: the accumulator is empty there, and a shallow merge - would replace a resumed job's stored series with two empty lists. + They used to be carved out of the prefixing scheme because backends passed + them pre-prefixed. Passing the framework's own `loss` regenerates both names + from the ordinary rule, which is why nothing downstream had to change. """ callback = _make_callback(reporter) - callback.report_training_start(max_steps=10, num_epochs=1) - callback.report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/ckpt") - callback.report_epoch_end(step=1, epoch=1) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 0.5}) + callback.report_validation(step=1, epoch=1, metrics={"loss": 0.4}) - assert len(reporter.reports) == 3 - for report in reporter.reports: - assert "metrics" not in report, report["phase"] + metrics = reporter.reports[-1]["metrics"] + assert metrics["train_loss"] == [{"step": 1, "epoch": 1, "value": 0.5}] + assert metrics["val_loss"] == [{"step": 1, "epoch": 1, "value": 0.4}] -def test_training_start_does_not_resend_seeded_metrics() -> None: - """It fires before the first step, so it has nothing to say about the curves.""" - prior = {"train_loss": [{"step": 1, "epoch": 1, "value": 0.9}], "val_loss": []} - reporter = _RecordingReporter(prior) +def test_the_same_name_in_both_phases_stays_separate(reporter: _RecordingReporter) -> None: + """DPO reports `accuracy` in both dicts; one series would interleave them.""" callback = _make_callback(reporter) - callback.report_training_start(max_steps=10, num_epochs=1) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 0.5, "accuracy": 0.18}) + callback.report_validation(step=1, epoch=1, metrics={"accuracy": 0.04}) - assert "metrics" not in reporter.reports[0] - - # ...and the seeded series is still there for the first step that does report. - callback.report_train_step(step=2, epoch=1, loss=0.8) - assert reporter.reports[-1]["metrics"]["train_loss"] == [ - {"step": 1, "epoch": 1, "value": 0.9}, - {"step": 2, "epoch": 1, "value": 0.8}, - ] - - -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) + metrics = reporter.reports[-1]["metrics"] + assert metrics["train_accuracy"] == [{"step": 1, "epoch": 1, "value": 0.18}] + assert metrics["val_accuracy"] == [{"step": 1, "epoch": 1, "value": 0.04}] - assert len(first_payload) == 1 +def test_a_pre_prefixed_name_is_not_re_interpreted(reporter: _RecordingReporter) -> None: + """A backend passing `val_loss` on a train step gets `train_val_loss`. -# --------------------------------------------------------------------------- # -# Optional val_loss -# --------------------------------------------------------------------------- # - - -def test_validation_without_loss_omits_the_key(reporter: _RecordingReporter) -> None: - """An algorithm may validate on task metrics alone; a null val_loss charts as zero.""" - _make_callback(reporter).report_validation(step=1, epoch=1, val_loss=None, accuracy=0.75) + The rule is mechanical: it never inspects the name for meaning, so a train + step cannot reach the validation curve however its metric is spelled. + """ + _make_callback(reporter).report_train_step(step=1, epoch=1, metrics={"loss": 0.5, "val_loss": 99.0}) - report = reporter.reports[-1] - assert "val_loss" not in report - assert report["accuracy"] == 0.75 - assert report["metrics"]["val_loss"] == [] + metrics = reporter.reports[-1]["metrics"] + assert metrics["val_loss"] == [] + assert metrics["train_val_loss"] == [{"step": 1, "epoch": 1, "value": 99.0}] -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) +def test_the_current_value_uses_the_series_name(reporter: _RecordingReporter) -> None: + """One name per metric, whether you read the curve or the latest point.""" + _make_callback(reporter).report_train_step(step=1, epoch=1, metrics={"loss": 0.5, "lr": 5e-06, "reward": 0.62}) report = reporter.reports[-1] - assert report["val_loss"] == 0.25 - assert report["metrics"]["val_loss"] == [{"step": 1, "epoch": 1, "value": 0.25}] - - -# --------------------------------------------------------------------------- # -# Optional checkpoint_path -# --------------------------------------------------------------------------- # + assert report["train_loss"] == 0.5 + assert report["train_lr"] == 5e-06 + assert report["train_reward"] == 0.62 + assert "loss" not in report and "lr" not in report and "reward" not in report -def test_checkpoint_report_without_a_path_omits_the_key(reporter: _RecordingReporter) -> None: - """A null would overwrite the last known checkpoint instead of leaving it. +def test_reserved_names_cannot_be_reached_by_a_metric(reporter: _RecordingReporter) -> None: + """Prefixing removes the collision class outright rather than filtering it. - The server merges key-wise, so omitting the key leaves the stored path - standing while an explicit null replaces it. automodel and unsloth both pass - None when their framework hands back no path. + `phase` is report_running's own parameter and used to raise TypeError into + the training loop; `step`, `epoch` and `metrics` are this callback's own keys + and used to need splat ordering to avoid being shadowed. """ - _make_callback(reporter).report_checkpoint_saved(step=1, epoch=1) - - assert "checkpoint_path" not in reporter.reports[-1] - - -def test_checkpoint_report_with_a_path_states_it(reporter: _RecordingReporter) -> None: - _make_callback(reporter).report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/ckpt") + callback = _make_callback(reporter) + collide = {"phase": 1.0, "step": 2.0, "epoch": 3.0, "metrics": 4.0, "loss": 0.5} + callback.report_train_step(step=10, epoch=1, metrics=collide) + callback.report_validation(step=10, epoch=1, metrics=collide) - assert reporter.reports[-1]["checkpoint_path"] == "/ckpt" + for report in reporter.reports: + assert report["step"] == 10, "the real step survives" + assert report["epoch"] == 1 + assert isinstance(report["metrics"], dict), "the series payload survives" + assert [r["phase"] for r in reporter.reports] == ["training", "validation"] + assert reporter.reports[0]["train_phase"] == 1.0, "the metric is kept, under its namespaced name" + assert reporter.reports[0]["train_step"] == 2.0 # --------------------------------------------------------------------------- # -# additional_metrics +# Nothing is required # --------------------------------------------------------------------------- # -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) +def test_a_step_without_a_loss_reports_the_rest(reporter: _RecordingReporter) -> None: + """No metric is privileged, so none of them is mandatory either.""" + _make_callback(reporter).report_train_step(step=1, epoch=1, metrics={"reward": 0.62}) 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}, - ] - + assert report["train_reward"] == 0.62 + assert report["metrics"]["train_loss"] == [] -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_validation_without_a_loss_leaves_the_curve_empty(reporter: _RecordingReporter) -> None: + """An algorithm may validate on task metrics alone; a null charts as zero.""" + _make_callback(reporter).report_validation(step=1, epoch=1, metrics={"accuracy": 0.75}) - -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 + report = reporter.reports[-1] + assert "val_loss" not in report + assert report["val_accuracy"] == 0.75 + assert report["metrics"]["val_loss"] == [] -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) +def test_an_empty_metric_dict_still_reports_progress(reporter: _RecordingReporter) -> None: + """step/epoch are progress, not metrics; they land with or without a curve.""" + _make_callback(reporter).report_train_step(step=7, epoch=2, metrics={}) report = reporter.reports[-1] - assert report["accuracy"] == 0.9 - assert report["metrics"]["val_accuracy"] == [{"step": 1, "epoch": 1, "value": 0.9}] + assert report["step"] == 7 + assert report["epoch"] == 2 + assert report["metrics"] == {"train_loss": [], "val_loss": []} -def test_train_and_validation_metrics_of_the_same_name_stay_separate( - reporter: _RecordingReporter, -) -> None: - """DPO reports `accuracy` in both dicts; one series would interleave them.""" - callback = _make_callback(reporter) - callback.report_train_step(step=1, epoch=1, loss=0.5, accuracy=0.18) - callback.report_validation(step=1, epoch=1, accuracy=0.04) - - metrics = reporter.reports[-1]["metrics"] - assert metrics["train_accuracy"] == [{"step": 1, "epoch": 1, "value": 0.18}] - assert metrics["val_accuracy"] == [{"step": 1, "epoch": 1, "value": 0.04}] +# --------------------------------------------------------------------------- # +# What cannot be charted is dropped from both places +# --------------------------------------------------------------------------- # -def test_non_numeric_metrics_are_dropped_from_the_series(reporter: _RecordingReporter) -> None: +def test_unchartable_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") + step=1, + epoch=1, + metrics={"loss": 0.5, "histogram": object(), "nested": {"a": 1}, "flag": True, "missing": float("nan")}, ) metrics = reporter.reports[-1]["metrics"] @@ -241,97 +202,134 @@ def test_a_dropped_metric_costs_only_itself(reporter: _RecordingReporter) -> Non while the job goes on looking healthy. Verified against a live platform before this filter existed: the step never landed. """ - _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, hist=object(), reward=0.62) + _make_callback(reporter).report_train_step(step=1, epoch=1, metrics={"hist": object(), "reward": 0.62}) report = reporter.reports[-1] - assert "hist" not in report - assert report["reward"] == 0.62, "a well-behaved metric in the same report still lands" + assert "hist" not in report and "train_hist" not in report + assert report["train_reward"] == 0.62, "a well-behaved metric in the same report still lands" -def test_a_metric_named_phase_does_not_break_the_report(reporter: _RecordingReporter) -> None: - """`phase` is report_running's own parameter; a collision is a TypeError. +def test_non_finite_scalars_are_omitted(reporter: _RecordingReporter) -> None: + """A NaN grad_norm is routine on a skipped step; the SDK sends it as null.""" + _make_callback(reporter).report_train_step( + step=1, epoch=1, metrics={"loss": 0.5, "lr": float("inf"), "grad_norm": float("nan"), "absent": None} + ) - Splat order silently shadows the other reserved names. This one raised - straight out of the training loop, so it is filtered rather than shadowed. - """ - callback = _make_callback(reporter) - callback.report_train_step(step=1, epoch=1, loss=0.5, phase=1.0) - callback.report_validation(step=1, epoch=1, val_loss=0.4, phase=1.0) + report = reporter.reports[-1] + assert "train_lr" not in report + assert "train_grad_norm" not in report + assert "train_absent" not in report + assert report["train_loss"] == 0.5 - assert [r["phase"] for r in reporter.reports] == ["training", "validation"] +def test_counts_stay_integers(reporter: _RecordingReporter) -> None: + """`num_valid_samples: 8` should not chart as 8.0.""" + _make_callback(reporter).report_train_step(step=1, epoch=1, metrics={"num_valid_samples": 8}) -def test_a_train_step_metric_named_val_loss_stays_out_of_the_val_curve( - reporter: _RecordingReporter, -) -> None: - """The legacy bare names are exempt from prefixing per phase, not per name. + assert reporter.reports[-1]["train_num_valid_samples"] == 8 - Matching on the name alone put a training-side number in the series Studio - draws as the validation loss. - """ - _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, val_loss=99.0) - - metrics = reporter.reports[-1]["metrics"] - assert metrics["val_loss"] == [] - assert metrics["train_val_loss"] == [{"step": 1, "epoch": 1, "value": 99.0}] +def test_numpy_scalars_are_coerced(reporter: _RecordingReporter) -> None: + """numpy satisfies numbers.Real but is not JSON-serializable.""" + np = pytest.importorskip("numpy") + _make_callback(reporter).report_train_step(step=1, epoch=1, metrics={"loss": np.float32(0.5), "count": np.int64(3)}) -def test_the_legacy_names_still_go_unprefixed_in_their_own_phase(reporter: _RecordingReporter) -> None: - """The Studio loss chart reads `train_loss`/`val_loss` by those exact names.""" - callback = _make_callback(reporter) - callback.report_train_step(step=1, epoch=1, loss=0.5) - callback.report_validation(step=1, epoch=1, val_loss=0.4) - - metrics = reporter.reports[-1]["metrics"] - assert metrics["train_loss"] == [{"step": 1, "epoch": 1, "value": 0.5}] - assert metrics["val_loss"] == [{"step": 1, "epoch": 1, "value": 0.4}] + report = reporter.reports[-1] + assert type(report["train_loss"]) is float + assert type(report["train_count"]) is int + + +class _Histogram: + """Stand-in for a non-numeric metric value -- NaN-hostile, like the real thing.""" + + def __float__(self) -> float: + raise TypeError("Histogram is not a scalar") + + +@pytest.mark.parametrize( + "value,expected", + [ + (0.5, True), + (0, True), + (-1.5, True), + (float("nan"), False), + (float("inf"), False), + (float("-inf"), False), + (None, False), + # Non-scalars that genuinely appear in NeMo-RL metric dicts. Each raises + # TypeError under a bare math.isfinite, 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), + # float() overflows here; the predicate must classify, never propagate. + (10**400, False), + ], +) +def test_is_chartable(value: Any, expected: bool) -> None: + assert is_chartable(value) is expected + + +def test_is_chartable_accepts_numpy_scalars() -> None: + np = pytest.importorskip("numpy") + assert is_chartable(np.float32(0.5)) is True + assert is_chartable(np.float64(0.5)) is True + assert is_chartable(np.int64(3)) is True + assert is_chartable(np.float32("nan")) is False # --------------------------------------------------------------------------- # -# A scalar is stated only when it was observed +# Which reports carry the series # --------------------------------------------------------------------------- # -def test_absent_lr_and_grad_norm_are_omitted_not_nulled(reporter: _RecordingReporter) -> None: - """A null charts as a real zero -- the same reason val_loss is omitted.""" - _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5) - - report = reporter.reports[-1] - assert "lr" not in report - assert "grad_norm" not in report - - -def test_non_finite_scalars_are_omitted(reporter: _RecordingReporter) -> None: - """A NaN grad_norm is routine on a skipped step; the SDK sends it as null.""" - _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, lr=float("inf"), grad_norm=float("nan")) +def test_step_reports_carry_the_series(reporter: _RecordingReporter) -> None: + """The reports that move a curve send every curve, since the merge is shallow.""" + callback = _make_callback(reporter) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 0.5}) + callback.report_validation(step=1, epoch=1, metrics={"loss": 0.45}) - report = reporter.reports[-1] - assert "lr" not in report - assert "grad_norm" not in report - assert report["train_loss"] == 0.5 + assert len(reporter.reports) == 2 + for report in reporter.reports: + assert "metrics" in report, report["phase"] + assert set(report["metrics"]) == {"train_loss", "val_loss"} -def test_a_non_finite_val_loss_is_omitted(reporter: _RecordingReporter) -> None: - _make_callback(reporter).report_validation(step=1, epoch=1, val_loss=float("nan")) +def test_non_step_reports_omit_the_series(reporter: _RecordingReporter) -> None: + """The server merges, so a report with nothing to add leaves the stored curve alone. - report = reporter.reports[-1] - assert "val_loss" not in report - assert report["metrics"]["val_loss"] == [] + Sending it anyway is pure upload on every checkpoint, and at training start + it is worse than useless: the accumulator is empty there, and a shallow merge + would replace a resumed job's stored series with two empty lists. + """ + callback = _make_callback(reporter) + callback.report_training_start(max_steps=10, num_epochs=1) + callback.report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/ckpt") + callback.report_epoch_end(step=1, epoch=1) + assert len(reporter.reports) == 3 + for report in reporter.reports: + assert "metrics" not in report, report["phase"] -def test_is_chartable_classifies_an_unbounded_int_without_raising() -> None: - """float() overflows on a big enough int; the predicate must not propagate it.""" - assert is_chartable(10**400) is False +def test_training_start_does_not_resend_seeded_metrics() -> None: + """It fires before the first step, so it has nothing to say about the curves.""" + prior = {"train_loss": [{"step": 1, "epoch": 1, "value": 0.9}], "val_loss": []} + reporter = _RecordingReporter(prior) + callback = _make_callback(reporter) + callback.report_training_start(max_steps=10, num_epochs=1) -def test_infinite_metrics_are_dropped_from_the_series(reporter: _RecordingReporter) -> None: - """A diverged run's inf is not a chart value, and `Infinity` is not valid JSON.""" - _make_callback(reporter).report_train_step( - step=1, epoch=1, loss=0.5, diverged=float("inf"), collapsed=float("-inf") - ) + assert "metrics" not in reporter.reports[0] - metrics = reporter.reports[-1]["metrics"] - assert set(metrics) == {"train_loss", "val_loss"} + # ...and the seeded series is still there for the first step that does report. + callback.report_train_step(step=2, epoch=1, metrics={"loss": 0.8}) + assert reporter.reports[-1]["metrics"]["train_loss"] == [ + {"step": 1, "epoch": 1, "value": 0.9}, + {"step": 2, "epoch": 1, "value": 0.8}, + ] def test_series_survive_a_resume_beyond_the_loss_curves() -> None: @@ -341,7 +339,7 @@ def test_series_survive_a_resume_beyond_the_loss_curves() -> None: "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) + _make_callback(reporter).report_train_step(step=2, epoch=1, metrics={"loss": 0.8, "reward": 0.3}) assert reporter.reports[-1]["metrics"]["train_reward"] == [ {"step": 1, "epoch": 1, "value": 0.2}, @@ -349,38 +347,54 @@ def test_series_survive_a_resume_beyond_the_loss_curves() -> None: ] -def test_additional_metrics_cannot_shadow_the_series(reporter: _RecordingReporter) -> None: - """`metrics` is not a parameter, so only splat order stops a silent override. +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, metrics={"loss": 0.5}) + first_payload = reporter.reports[-1]["metrics"]["train_loss"] + callback.report_train_step(step=2, epoch=1, metrics={"loss": 0.4}) + + assert len(first_payload) == 1 + + +# --------------------------------------------------------------------------- # +# Optional checkpoint_path +# --------------------------------------------------------------------------- # + - `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. +def test_checkpoint_report_without_a_path_omits_the_key(reporter: _RecordingReporter) -> None: + """A null would overwrite the last known checkpoint instead of leaving it. + + The server merges key-wise, so omitting the key leaves the stored path + standing while an explicit null replaces it. automodel and unsloth both pass + None when their framework hands back no path. """ - _make_callback(reporter).report_train_step(step=1, epoch=1, loss=0.5, metrics="clobbered") + _make_callback(reporter).report_checkpoint_saved(step=1, epoch=1) - report = reporter.reports[-1] - assert report["metrics"]["train_loss"] == [{"step": 1, "epoch": 1, "value": 0.5}] + assert "checkpoint_path" not in reporter.reports[-1] -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") +def test_checkpoint_report_with_a_path_states_it(reporter: _RecordingReporter) -> None: + _make_callback(reporter).report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/ckpt") + + assert reporter.reports[-1]["checkpoint_path"] == "/ckpt" - assert reporter.reports[-1]["train_loss"] == 0.5 +# --------------------------------------------------------------------------- # +# Backend stamping +# --------------------------------------------------------------------------- # -def test_additional_metrics_do_not_collide_with_backend_stamping( - reporter: _RecordingReporter, -) -> None: - """`backend` is keyword-only, so **additional_metrics can never capture it.""" +def test_the_default_backend_is_stamped_on_every_report(reporter: _RecordingReporter) -> None: 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) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 0.5, "reward": 0.62}) report = reporter.reports[-1] assert report["backend"] == "test-backend" - assert report["reward"] == 0.62 + assert report["train_reward"] == 0.62 def test_no_backend_field_when_the_default_is_unset(reporter: _RecordingReporter) -> None: @@ -391,8 +405,8 @@ def test_no_backend_field_when_the_default_is_unset(reporter: _RecordingReporter """ 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.4) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 0.5}) + callback.report_validation(step=1, epoch=1, metrics={"loss": 0.4}) callback.report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/ckpt") callback.report_epoch_end(step=1, epoch=1) @@ -400,6 +414,13 @@ def test_no_backend_field_when_the_default_is_unset(reporter: _RecordingReporter assert all("backend" not in report for report in reporter.reports) +def test_a_per_call_backend_overrides_the_default(reporter: _RecordingReporter) -> None: + """unsloth's HF trainer callback passes it per call.""" + _make_callback(reporter).report_train_step(step=1, epoch=1, metrics={"loss": 0.5}, backend="unsloth") + + assert reporter.reports[-1]["backend"] == "unsloth" + + def test_close_delegates_to_the_reporter(reporter: _RecordingReporter) -> None: _make_callback(reporter).close() diff --git a/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py b/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py index 609089a0d0..9073e87ab6 100644 --- a/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py +++ b/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py @@ -130,9 +130,7 @@ def _log_train_metrics(self, log_data: Any) -> None: self.callback.report_train_step( step=getattr(log_data, "step", 0) + 1, # Convert to 1-based epoch=getattr(log_data, "epoch", 0) + 1, # Convert to 1-based - loss=metrics.get("loss", 0.0), - lr=metrics.get("lr"), - grad_norm=metrics.get("grad_norm"), + metrics=dict(metrics), ) except Exception as e: logger.warning(f"Failed to report training progress: {e}") @@ -173,7 +171,11 @@ def _log_val_metrics(self, *args: Any, **kwargs: Any) -> None: self.callback.report_validation( step=getattr(log_data, "step", 0) + 1, # Convert to 1-based epoch=getattr(log_data, "epoch", 0) + 1, # Convert to 1-based - val_loss=metrics.get("val_loss", 0.0), + # Automodel names it `val_loss`; renaming to `loss` lets the + # phase prefix produce `val_loss` the way it does elsewhere. + # An absent one now yields no point, where it used to chart + # a fabricated 0.0. + metrics={("loss" if name == "val_loss" else name): value for name, value in metrics.items()}, ) except Exception as e: logger.warning(f"Failed to report validation progress: {e}") diff --git a/services/automodel/tests/tasks/training/backends/test_callbacks.py b/services/automodel/tests/tasks/training/backends/test_callbacks.py index 98778b3c65..a106865cc8 100644 --- a/services/automodel/tests/tasks/training/backends/test_callbacks.py +++ b/services/automodel/tests/tasks/training/backends/test_callbacks.py @@ -26,9 +26,9 @@ def _last_report_kwargs(self, mock_reporter: MagicMock) -> dict: def test_train_step_accumulates_metrics(self): callback, reporter = self._make_callback() - callback.report_train_step(step=1, epoch=1, loss=3.21) - callback.report_train_step(step=2, epoch=1, loss=2.89) - callback.report_train_step(step=3, epoch=1, loss=2.56) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 3.21}) + callback.report_train_step(step=2, epoch=1, metrics={"loss": 2.89}) + callback.report_train_step(step=3, epoch=1, metrics={"loss": 2.56}) kwargs = self._last_report_kwargs(reporter) assert kwargs["metrics"]["train_loss"] == [ @@ -40,8 +40,8 @@ def test_train_step_accumulates_metrics(self): def test_validation_accumulates_metrics(self): callback, reporter = self._make_callback() - callback.report_validation(step=250, epoch=1, val_loss=3.19) - callback.report_validation(step=500, epoch=2, val_loss=2.09) + callback.report_validation(step=250, epoch=1, metrics={"loss": 3.19}) + callback.report_validation(step=500, epoch=2, metrics={"loss": 2.09}) kwargs = self._last_report_kwargs(reporter) assert kwargs["metrics"]["val_loss"] == [ @@ -53,10 +53,10 @@ def test_validation_accumulates_metrics(self): def test_mixed_train_and_val_both_present(self): callback, reporter = self._make_callback() - callback.report_train_step(step=1, epoch=1, loss=3.21) - callback.report_train_step(step=2, epoch=1, loss=2.89) - callback.report_validation(step=2, epoch=1, val_loss=3.19) - callback.report_train_step(step=3, epoch=1, loss=2.56) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 3.21}) + callback.report_train_step(step=2, epoch=1, metrics={"loss": 2.89}) + callback.report_validation(step=2, epoch=1, metrics={"loss": 3.19}) + callback.report_train_step(step=3, epoch=1, metrics={"loss": 2.56}) kwargs = self._last_report_kwargs(reporter) assert len(kwargs["metrics"]["train_loss"]) == 3 @@ -65,19 +65,19 @@ def test_mixed_train_and_val_both_present(self): def test_metrics_included_in_every_update(self): callback, reporter = self._make_callback() - callback.report_train_step(step=1, epoch=1, loss=3.21) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 3.21}) first_call_kwargs = reporter.report_running.call_args_list[0].kwargs assert len(first_call_kwargs["metrics"]["train_loss"]) == 1 assert first_call_kwargs["metrics"]["val_loss"] == [] - callback.report_train_step(step=2, epoch=1, loss=2.89) + callback.report_train_step(step=2, epoch=1, metrics={"loss": 2.89}) second_call_kwargs = reporter.report_running.call_args_list[1].kwargs assert len(second_call_kwargs["metrics"]["train_loss"]) == 2 def test_train_step_uses_train_loss_flat_field(self): callback, reporter = self._make_callback() - callback.report_train_step(step=1, epoch=1, loss=3.21) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 3.21}) kwargs = self._last_report_kwargs(reporter) assert kwargs["train_loss"] == 3.21 @@ -86,11 +86,11 @@ def test_train_step_uses_train_loss_flat_field(self): def test_train_step_passes_optional_fields(self): callback, reporter = self._make_callback() - callback.report_train_step(step=1, epoch=1, loss=3.21, lr=0.0002, grad_norm=1.5) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 3.21, "lr": 0.0002, "grad_norm": 1.5}) kwargs = self._last_report_kwargs(reporter) - assert kwargs["lr"] == 0.0002 - assert kwargs["grad_norm"] == 1.5 + assert kwargs["train_lr"] == 0.0002 + assert kwargs["train_grad_norm"] == 1.5 def test_seeds_from_server_on_init(self): prior = { @@ -115,7 +115,7 @@ def test_seeded_metrics_included_in_first_report(self): } callback, reporter = self._make_callback(prior_metrics=prior) - callback.report_train_step(step=2, epoch=1, loss=2.89) + callback.report_train_step(step=2, epoch=1, metrics={"loss": 2.89}) kwargs = self._last_report_kwargs(reporter) assert kwargs["metrics"]["train_loss"] == [ @@ -130,7 +130,7 @@ def test_seeded_val_metrics_preserved_across_train_steps(self): } callback, reporter = self._make_callback(prior_metrics=prior) - callback.report_train_step(step=2, epoch=1, loss=2.89) + callback.report_train_step(step=2, epoch=1, metrics={"loss": 2.89}) kwargs = self._last_report_kwargs(reporter) assert len(kwargs["metrics"]["val_loss"]) == 1 @@ -171,7 +171,7 @@ def test_checkpoint_report_leaves_the_accumulated_series_alone(self): """ callback, reporter = self._make_callback() - callback.report_train_step(step=1, epoch=1, loss=3.21) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 3.21}) callback.report_checkpoint_saved(step=1, epoch=1, checkpoint_path="/tmp/ckpt") assert "metrics" not in self._last_report_kwargs(reporter) @@ -179,7 +179,7 @@ def test_checkpoint_report_leaves_the_accumulated_series_alone(self): def test_epoch_end_report_leaves_the_accumulated_series_alone(self): callback, reporter = self._make_callback() - callback.report_train_step(step=1, epoch=1, loss=3.21) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 3.21}) callback.report_epoch_end(step=1, epoch=1) assert "metrics" not in self._last_report_kwargs(reporter) 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 66dd191cf1..031703f1a1 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 @@ -22,35 +22,16 @@ # 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. -# Selection is by presence, so an algorithm that does not produce one of these -# simply omits it. -_TRAIN_METRIC_KEYS = ( - "num_valid_samples", - "preference_loss", - "rewards_rejected_mean", - "global_valid_seqs", - "global_valid_toks", -) - -_VALIDATION_METRIC_KEYS = _TRAIN_METRIC_KEYS - - -def has_metric_value(metric: Any) -> bool: - """Whether ``metric`` is a finite 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.isfinite`` 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) +# NeMo-RL's metric dicts are forwarded whole. There is no allow-list: the +# callback keeps the finite scalars and drops everything else, so a metric +# NeMo-RL adds charts itself instead of waiting on a change here. That is what +# the old list cost -- DPO's `accuracy`, `sft_loss` and `rewards_chosen_mean` +# were dropped silently for never having been added to it. +# +# The callback's filter is load-bearing, not defensive: `calculate_single_metric` +# emits a `/histogram` holding a `Histogram`, NeMo-Gym adds a per-agent +# `full_result` `Table`, and `generation_logger_metrics` is a nested dict. Each +# one rides in the same dict as the scalars. def resolve_log_interval(val_period: int | None) -> int: @@ -187,15 +168,8 @@ def log_metrics( epoch = (max(step - 1, 0) // self._steps_per_epoch) + 1 # Handle training loss - if prefix == "train" and has_metric_value(metrics.get("loss")): - 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), - } + if prefix == "train" and is_chartable(metrics.get("loss")): + report = {"step": step, "epoch": epoch, "metrics": dict(metrics)} # 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: @@ -206,14 +180,9 @@ def log_metrics( # Handle validation metrics elif prefix and prefix.startswith("validation"): - if has_metric_value(metrics.get("loss")): + if is_chartable(metrics.get("loss")): val_loss = metrics["loss"] - self._callback.report_validation( - step=step, - epoch=epoch, - val_loss=val_loss, - **self._select_metrics(metrics, _VALIDATION_METRIC_KEYS), - ) + self._callback.report_validation(step=step, epoch=epoch, metrics=dict(metrics)) # Track best validation loss if val_loss < self._best_metric_value: self._best_metric_value = val_loss @@ -221,11 +190,6 @@ def log_metrics( _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. diff --git a/services/rl/tests/test_nemo_rl_logger.py b/services/rl/tests/test_nemo_rl_logger.py index 2f49e83488..070a15d6db 100644 --- a/services/rl/tests/test_nemo_rl_logger.py +++ b/services/rl/tests/test_nemo_rl_logger.py @@ -5,8 +5,10 @@ The metric dicts here mirror what NeMo-RL actually hands the logger, including the non-scalar entries (``Histogram`` objects, tables, 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. +with the numbers, so they are exercised rather than sanitised away. + +The logger forwards the dict whole and decides only *whether* and *when* to report: +which entries survive is the shared callback's business, covered by its own suite. """ from __future__ import annotations @@ -50,7 +52,6 @@ def _stub(name: str) -> types.ModuleType: 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, ) @@ -68,13 +69,11 @@ def __init__(self) -> None: 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_train_step(self, step, epoch, metrics, *, backend=None): + self.train_steps.append({"step": step, "epoch": epoch, "metrics": metrics}) - def report_validation(self, step, epoch, val_loss=None, **additional): - self.validations.append({"step": step, "epoch": epoch, "val_loss": val_loss, **additional}) + def report_validation(self, step, epoch, metrics, *, backend=None): + self.validations.append({"step": step, "epoch": epoch, "metrics": metrics}) def close(self) -> None: self.closed = True @@ -109,8 +108,8 @@ class _Histogram: """Stand-in for a non-numeric metric value — NaN-hostile, like the real thing.""" -# A DPO `train` dict: the whitelisted scalars, plus the non-scalars that ride -# along with them in a real NeMo-RL metric dict. +# A DPO `train` dict: the scalars, plus the non-scalars that ride along with them +# in a real NeMo-RL metric dict. TRAIN_METRICS: dict[str, Any] = { "loss": 0.5, "lr": 1e-5, @@ -129,42 +128,10 @@ class _Histogram: # --------------------------------------------------------------------------- # -# has_metric_value +# Module stub hygiene # --------------------------------------------------------------------------- # -@pytest.mark.parametrize( - "value,expected", - [ - (0.5, True), - (0, True), - (-1.5, True), - (float("nan"), False), - # A diverged loss. Not a chart value, and `Infinity` is not valid JSON. - (float("inf"), False), - (float("-inf"), False), - (None, False), - # Non-scalars that genuinely appear in NeMo-RL metric dicts. Each of these - # raises TypeError under a bare math.isfinite, 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_real_metric() -> None: - """Every value in a real metric dict must be classifiable without raising.""" - for key, value in 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. @@ -176,26 +143,31 @@ def test_module_stub_does_not_break_find_spec() -> None: 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 - - # --------------------------------------------------------------------------- # # Train metrics # --------------------------------------------------------------------------- # -def test_train_step_drops_non_scalar_metrics(callback: _RecordingCallback) -> None: - """Histograms/Tables/nested dicts must not be forwarded, and must not raise.""" +def test_the_metric_dict_is_forwarded_whole(callback: _RecordingCallback) -> None: + """No allow-list: a metric NeMo-RL adds charts without a change here. + + The old list silently dropped anything never added to it -- DPO's `accuracy`, + `sft_loss` and `rewards_chosen_mean` among them. Deciding which entries are + chartable is the callback's job, not a second gate doing a weaker version of + the same check. + """ _make_logger().log_metrics(TRAIN_METRICS, step=0, prefix="train") - reported = callback.train_steps[0] - for key in ("some/histogram", "generation_logger_metrics", "per_worker_token_counts"): - assert key not in reported + assert callback.train_steps[0]["metrics"] == TRAIN_METRICS + + +def test_the_forwarded_dict_is_a_copy(callback: _RecordingCallback) -> None: + """NeMo-RL reuses its metric dict across steps; a reference would alias it.""" + metrics = dict(TRAIN_METRICS) + _make_logger().log_metrics(metrics, step=0, prefix="train") + metrics["loss"] = 99.0 + + assert callback.train_steps[0]["metrics"]["loss"] == 0.5 def test_train_call_without_a_loss_is_ignored(callback: _RecordingCallback) -> None: @@ -206,11 +178,11 @@ def test_train_call_without_a_loss_is_ignored(callback: _RecordingCallback) -> N assert callback.train_steps == [] -def test_whitelisted_train_metrics_are_forwarded(callback: _RecordingCallback) -> None: - """The whitelisted scalars ride along with loss/lr/grad_norm.""" +def test_every_train_scalar_is_forwarded(callback: _RecordingCallback) -> None: + """Including the ones the old allow-list had no entry for.""" _make_logger().log_metrics(TRAIN_METRICS, step=0, prefix="train") - reported = callback.train_steps[0] + reported = callback.train_steps[0]["metrics"] assert reported["loss"] == 0.5 assert reported["preference_loss"] == 0.42 assert reported["rewards_rejected_mean"] == -0.3 @@ -264,8 +236,8 @@ def test_flushed_step_carries_the_full_metric_payload(callback: _RecordingCallba flushed = callback.train_steps[-1] assert flushed["step"] == 1 - assert flushed["loss"] == 0.5 - assert flushed["preference_loss"] == 0.42 + assert flushed["metrics"]["loss"] == 0.5 + assert flushed["metrics"]["preference_loss"] == 0.42 def test_double_close_flushes_once(callback: _RecordingCallback) -> None: @@ -415,11 +387,12 @@ def test_validate_at_start_reports_step_zero(callback: _RecordingCallback) -> No # --------------------------------------------------------------------------- # -def test_validation_reports_loss_and_whitelisted_metrics(callback: _RecordingCallback) -> None: +def test_validation_forwards_the_dict_whole(callback: _RecordingCallback) -> None: + """`loss` is forwarded under its own name; the phase prefix makes it val_loss.""" _make_logger().log_metrics(VALIDATION_METRICS, step=10, prefix="validation") - reported = callback.validations[0] - assert reported["val_loss"] == 0.25 + reported = callback.validations[0]["metrics"] + assert reported["loss"] == 0.25 assert reported["num_valid_samples"] == 8 diff --git a/services/unsloth/src/nmp/unsloth/tasks/training/backends/hf_trainer_callback.py b/services/unsloth/src/nmp/unsloth/tasks/training/backends/hf_trainer_callback.py index c8c498b871..016fd6b487 100644 --- a/services/unsloth/src/nmp/unsloth/tasks/training/backends/hf_trainer_callback.py +++ b/services/unsloth/src/nmp/unsloth/tasks/training/backends/hf_trainer_callback.py @@ -50,9 +50,14 @@ def on_log( self._progress.report_train_step( step=int(state.global_step), epoch=_epoch_from_value(epoch_raw, self._num_epochs), - loss=float(logs["loss"]), - lr=float(logs["learning_rate"]) if logs.get("learning_rate") is not None else None, - grad_norm=float(logs["grad_norm"]) if logs.get("grad_norm") is not None else None, + # `learning_rate` is renamed to the `lr` every other backend uses, + # so the series is `train_lr` regardless of who reported it. The + # callback drops whichever of these the trainer did not produce. + metrics={ + "loss": float(logs["loss"]), + "lr": logs.get("learning_rate"), + "grad_norm": logs.get("grad_norm"), + }, backend=self._backend, ) @@ -71,7 +76,7 @@ def on_evaluate( self._progress.report_validation( step=int(state.global_step), epoch=_epoch_from_value(epoch_raw, self._num_epochs), - val_loss=float(metrics["eval_loss"]), + metrics={"loss": float(metrics["eval_loss"])}, backend=self._backend, ) diff --git a/services/unsloth/tests/test_callbacks.py b/services/unsloth/tests/test_callbacks.py index bdd09282c7..3adf700447 100644 --- a/services/unsloth/tests/test_callbacks.py +++ b/services/unsloth/tests/test_callbacks.py @@ -24,9 +24,9 @@ def _last_report_kwargs(self, mock_reporter: MagicMock) -> dict: def test_train_step_accumulates_metrics(self): callback, reporter = self._make_callback() - callback.report_train_step(step=1, epoch=1, loss=3.21) - callback.report_train_step(step=2, epoch=1, loss=2.89) - callback.report_train_step(step=3, epoch=1, loss=2.56) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 3.21}) + callback.report_train_step(step=2, epoch=1, metrics={"loss": 2.89}) + callback.report_train_step(step=3, epoch=1, metrics={"loss": 2.56}) kwargs = self._last_report_kwargs(reporter) assert kwargs["metrics"]["train_loss"] == [ @@ -38,7 +38,7 @@ def test_train_step_accumulates_metrics(self): def test_train_step_uses_train_loss_flat_field(self): callback, reporter = self._make_callback() - callback.report_train_step(step=1, epoch=1, loss=3.21) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 3.21}) kwargs = self._last_report_kwargs(reporter) assert kwargs["train_loss"] == 3.21 @@ -48,11 +48,11 @@ def test_train_step_uses_train_loss_flat_field(self): def test_train_step_passes_optional_fields(self): callback, reporter = self._make_callback() - callback.report_train_step(step=1, epoch=1, loss=3.21, lr=0.0002, grad_norm=1.5) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 3.21, "lr": 0.0002, "grad_norm": 1.5}) kwargs = self._last_report_kwargs(reporter) - assert kwargs["lr"] == 0.0002 - assert kwargs["grad_norm"] == 1.5 + assert kwargs["train_lr"] == 0.0002 + assert kwargs["train_grad_norm"] == 1.5 def test_report_training_start_delegates(self): callback, reporter = self._make_callback() diff --git a/services/unsloth/tests/test_hf_trainer_callback.py b/services/unsloth/tests/test_hf_trainer_callback.py index 800fbbe186..6f6c4d39f6 100644 --- a/services/unsloth/tests/test_hf_trainer_callback.py +++ b/services/unsloth/tests/test_hf_trainer_callback.py @@ -46,8 +46,8 @@ def test_on_log_reports_train_step(self, progress: tuple[TrainingProgressCallbac assert kwargs["phase"] == "training" assert kwargs["step"] == 8 assert kwargs["train_loss"] == 2.89 - assert kwargs["lr"] == 5e-5 - assert kwargs["grad_norm"] == 10.6 + assert kwargs["train_lr"] == 5e-5 + assert kwargs["train_grad_norm"] == 10.6 assert kwargs["backend"] == "unsloth" assert kwargs["metrics"]["train_loss"][-1]["value"] == 2.89 diff --git a/web/packages/studio/src/mocks/customizer/customization-jobs.ts b/web/packages/studio/src/mocks/customizer/customization-jobs.ts index 990fccda61..48d9109208 100644 --- a/web/packages/studio/src/mocks/customizer/customization-jobs.ts +++ b/web/packages/studio/src/mocks/customizer/customization-jobs.ts @@ -16,8 +16,8 @@ const completedStatusDetails = { percentage_done: 100, train_loss: 0.9, val_loss: 0.9, - lr: 0.000005, - grad_norm: 1.2345, + train_lr: 0.000005, + train_grad_norm: 1.2345, checkpoint_path: 'default/output-fileset/checkpoints/step-10', metrics: { train_loss: [ diff --git a/web/packages/studio/src/util/customizations.test.ts b/web/packages/studio/src/util/customizations.test.ts index d20cada6f0..5751554f53 100644 --- a/web/packages/studio/src/util/customizations.test.ts +++ b/web/packages/studio/src/util/customizations.test.ts @@ -234,8 +234,8 @@ describe('getTrainingTelemetry', () => { epoch: 1, train_loss: 0.42, val_loss: 0.55, - lr: 0.000005, - grad_norm: 1.25, + train_lr: 0.000005, + train_grad_norm: 1.25, checkpoint_path: 'ws/fileset/checkpoints/step-4', }) ) @@ -259,8 +259,8 @@ describe('getTrainingTelemetry', () => { jobWithDetails({ phase: '', step: Number.NaN, - lr: null, - grad_norm: 'oops', + train_lr: null, + train_grad_norm: 'oops', checkpoint_path: '', }) ) diff --git a/web/packages/studio/src/util/customizations.tsx b/web/packages/studio/src/util/customizations.tsx index 2c16c02cb1..8eb53a29d8 100644 --- a/web/packages/studio/src/util/customizations.tsx +++ b/web/packages/studio/src/util/customizations.tsx @@ -182,8 +182,8 @@ export const getTrainingTelemetry = ( epoch: asFiniteNumber(details.epoch), trainLoss: asFiniteNumber(details.train_loss), valLoss: asFiniteNumber(details.val_loss), - learningRate: asFiniteNumber(details.lr), - gradNorm: asFiniteNumber(details.grad_norm), + learningRate: asFiniteNumber(details.train_lr), + gradNorm: asFiniteNumber(details.train_grad_norm), checkpointPath: asNonEmptyString(details.checkpoint_path), }; }; From 316a08221fb83fc26054feb27faadb093783649c Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 18:01:12 -0400 Subject: [PATCH 15/23] docs(customizer): document the metric naming rule and the series payload Both pages described a fixed pair of metrics and a flat set of status_details fields, which stopped being true when every reported metric started accumulating as a series, and stopped being accurate at all when `lr` and `grad_norm` became `train_lr` and `train_grad_norm`. get-job-status now lists the progress fields, the per-metric latest values and the `metrics` history separately, and both example responses carry a real `metrics` payload rather than only the flat scalars. The metrics tutorial gains the naming rule -- `_`, with train_loss and val_loss as what it produces for `loss` -- plus where each metric appears and why a missing field is not a zero. Its API sample reads the renamed fields and gains a loop over `metrics`, which is how a caller picks up backend-specific curves without naming them in advance. optimize-throughput.mdx needed no change: it reads train_loss and val_loss, and neither name moved. Signed-off-by: Albert Cui --- .../get-job-status.mdx | 51 ++++++++++++++++--- docs/customizer/tutorials/metrics.mdx | 49 +++++++++++++++--- 2 files changed, 88 insertions(+), 12 deletions(-) diff --git a/docs/customizer/manage-customization-jobs/get-job-status.mdx b/docs/customizer/manage-customization-jobs/get-job-status.mdx index b7f852f887..66a3b622b2 100644 --- a/docs/customizer/manage-customization-jobs/get-job-status.mdx +++ b/docs/customizer/manage-customization-jobs/get-job-status.mdx @@ -10,9 +10,13 @@ Get detailed execution status for a customization job, including step-by-step pr This endpoint provides granular execution details including: - **Step-level status**: `model-and-dataset-download` → `training` → `model-upload` → `model-entity-creation` -- **Training metrics**: `step`, `epoch`, `train_loss`, `lr` (learning rate), `grad_norm`, `val_loss` +- **Training progress**: `step`, `epoch`, `max_steps`, `num_epochs`, `percentage_done` +- **Latest metric values**: one field per metric, named `_` — `train_loss`, `train_lr`, `train_grad_norm`, `val_loss`, and whatever else your backend reports +- **Metric history**: `metrics`, holding each metric as a series of `{step, epoch, value}` points - **Progress tracking**: `downloaded_files`, `uploaded_bytes`, `progress_pct` +Which metrics appear depends on the backend and the algorithm. See [Checking Your Customization Job Metrics](/documentation/customizer-reference/tutorials/metrics) for how the names are formed. + To list jobs or get job definitions (model entity, hyperparameters, spec), use [List Active Jobs](/documentation/customizer-reference/manage-customization-jobs/list-active-jobs) instead. @@ -146,9 +150,25 @@ curl -X GET \ "num_epochs": 2, "step": 8, "epoch": 1, + "percentage_done": 8, "train_loss": 2.8918895721435547, - "lr": 4.9101714686276044e-05, - "grad_norm": 26.0 + "train_lr": 4.9101714686276044e-05, + "train_grad_norm": 26.0, + "metrics": { + "train_loss": [ + { "step": 4, "epoch": 1, "value": 3.2087905406951904 }, + { "step": 8, "epoch": 1, "value": 2.8918895721435547 } + ], + "val_loss": [], + "train_lr": [ + { "step": 4, "epoch": 1, "value": 4.9550857343138022e-05 }, + { "step": 8, "epoch": 1, "value": 4.9101714686276044e-05 } + ], + "train_grad_norm": [ + { "step": 4, "epoch": 1, "value": 31.5 }, + { "step": 8, "epoch": 1, "value": 26.0 } + ] + } } } ] @@ -222,11 +242,30 @@ curl -X GET \ "num_epochs": 2, "step": 94, "epoch": 2, + "percentage_done": 100, "train_loss": 0.3437718152999878, - "lr": 5.000000000000001e-07, - "grad_norm": 20.125, + "train_lr": 5.000000000000001e-07, + "train_grad_norm": 20.125, "val_loss": 0.5527229905128479, - "checkpoint_path": "/var/run/scratch/job/training/checkpoints" + "checkpoint_path": "/var/run/scratch/job/training/checkpoints", + "metrics": { + "train_loss": [ + { "step": 47, "epoch": 1, "value": 1.1204545497894287 }, + { "step": 94, "epoch": 2, "value": 0.3437718152999878 } + ], + "val_loss": [ + { "step": 47, "epoch": 1, "value": 0.9182837605476379 }, + { "step": 94, "epoch": 2, "value": 0.5527229905128479 } + ], + "train_lr": [ + { "step": 47, "epoch": 1, "value": 2.5e-05 }, + { "step": 94, "epoch": 2, "value": 5.000000000000001e-07 } + ], + "train_grad_norm": [ + { "step": 47, "epoch": 1, "value": 24.75 }, + { "step": 94, "epoch": 2, "value": 20.125 } + ] + } } } ] diff --git a/docs/customizer/tutorials/metrics.mdx b/docs/customizer/tutorials/metrics.mdx index a7794d62ca..46fd0d7c89 100644 --- a/docs/customizer/tutorials/metrics.mdx +++ b/docs/customizer/tutorials/metrics.mdx @@ -27,10 +27,38 @@ The time to complete this tutorial is approximately 10 minutes. ## Available Metrics -Each customization job tracks two key metrics: +A customization job tracks every numeric metric its backend reports, not a fixed +list. Training loss and validation loss are always among them: -- **Training Loss**: Calculated during training, logged every 10 steps (default, configurable via hyperparameters) -- **Validation Loss**: Calculated during validation, logged at each validation interval +- **Training Loss** (`train_loss`): Calculated during training, logged every 10 steps (default, configurable via hyperparameters) +- **Validation Loss** (`val_loss`): Calculated during validation, logged at each validation interval + +Alongside those you will typically see `train_lr` (learning rate) and +`train_grad_norm`, plus whatever else the algorithm produces — a DPO job also +reports `train_preference_loss` and `val_accuracy`, for example. + +### How Metrics Are Named + +Each metric is named `_`, where the phase is `train` or `val` and +the metric keeps whatever name the training framework gave it. A metric reported +during both training and validation therefore stays separate: DPO's `accuracy` +becomes `train_accuracy` and `val_accuracy` rather than one interleaved series. + +`train_loss` and `val_loss` are simply what this rule produces for a metric named +`loss`. + +### Where Metrics Appear + +Each metric shows up in two places in a training task's `status_details`: + +- **The latest value**, as a top-level field under its full name (`train_loss`, + `train_lr`, ...). Present only when the metric was actually reported, so a + missing field means no value rather than a zero. +- **The full history**, under `metrics`, as a list of `{step, epoch, value}` + points per metric. This is what the loss curves in the UI are drawn from. + +Non-numeric values a framework emits alongside the scalars — histograms, tables, +nested dictionaries — are not charted and do not appear in either place. ## Viewing Your Metrics @@ -64,11 +92,20 @@ for step in status.steps or []: print(f"Epoch: {details.get('epoch')}/{details.get('num_epochs')}") print(f"Training Loss: {details.get('train_loss')}") print(f"Validation Loss: {details.get('val_loss')}") - print(f"Learning Rate: {details.get('lr')}") - print(f"Gradient Norm: {details.get('grad_norm')}") + print(f"Learning Rate: {details.get('train_lr')}") + print(f"Gradient Norm: {details.get('train_grad_norm')}") ``` -The response includes training progress and metrics including loss, learning rate, and validation loss. +To read the curves rather than the latest values, use the `metrics` payload. It +carries every metric the job reported, so iterating it picks up backend-specific +ones without naming them in advance: + +```python +for name, points in (details.get("metrics") or {}).items(): + if not points: + continue + print(f"{name}: {len(points)} points, latest {points[-1]['value']}") +``` ### Using MLflow From 58307cc0ba09b23ca6da7737061b4930127db6d4 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 21:12:15 -0400 Subject: [PATCH 16/23] fix(rl): bound progress reports by run length, not just val_period resolve_log_interval derived the reporting cadence from val_period alone, targeting ~10 reports per validation period. val_period is the user's val_check_interval, so any value below 10 -- "validate every 5 steps" is an ordinary request -- floored the term to zero, clamped to 1, and reported every step of a run of any length. That is not a linear cost. Every train report resends every accumulated series in full, and the Jobs service persists each one twice (the task, then the copy propagated up to the job), so upload and stored-blob writes both grow as the square of the report count. A 20k-step run at ~22 series was reporting 20,000 times. There is now a second floor at _MAX_REPORTS_PER_RUN reports for the whole run, and the coarser of the two wins. Ceiling division, so the bound is a real <=200 rather than up to twice that. Nothing in the existing regime moves: val_period=100 over 100 steps still gives an interval of 10. The payload note in callbacks.py is corrected while it is being cited. Its measured figures are client-side upload only, and it claimed that a backend reporting every step of a long run was a hypothetical the transport would have to grow delta appends for -- it was reachable from a documented hyperparameter, and a backend can simply bound itself. Signed-off-by: Albert Cui --- .../training/callbacks.py | 13 ++++-- .../backends/nemo_rl/nemo_rl_logger.py | 28 +++++++++-- services/rl/tests/test_nemo_rl_logger.py | 46 +++++++++++++++---- 3 files changed, 70 insertions(+), 17 deletions(-) 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 500899526d..bebae038f4 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 @@ -47,10 +47,15 @@ 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. +Those are client-side upload figures. The server pays twice over: ``JobDispatcher`` +persists each report to the task and then again to the copy propagated up to the +job, so the write volume is double the numbers above. + +Deliberately accepted for batch training jobs, on the understanding that a +backend bounds its own report count -- NeMo-RL's ``resolve_log_interval`` caps a +run for exactly this reason. A backend that instead reports every step of a long +run pays quadratically, and the answer to that is delta appends in the transport +rather than trimming the series here. Backends subclass this and set :attr:`_default_backend`: unsloth stamps a ``backend`` field on each report (``"unsloth"``); automodel and NeMo-RL leave it 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 031703f1a1..de9ecb156f 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 @@ -22,6 +22,12 @@ # notion of a reporting cadence, so it is derived from val_period. _REPORTS_PER_VAL_PERIOD = 10 +# Upper bound on progress reports for one run, whatever the validation cadence +# asks for. Every train report resends every accumulated series in full and the +# Jobs service stores the blob twice, so the cost of a report is not constant -- +# see the payload note in nmp.customization_common.training.callbacks. +_MAX_REPORTS_PER_RUN = 200 + # NeMo-RL's metric dicts are forwarded whole. There is no allow-list: the # callback keeps the finite scalars and drops everything else, so a metric # NeMo-RL adds charts itself instead of waiting on a change here. That is what @@ -34,9 +40,23 @@ # one rides in the same dict as the scalars. -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_log_interval(val_period: int | None, max_steps: int) -> int: + """Steps between progress reports: the coarser of two floors. + + The first targets ~10 reports per validation period, which is the cadence + someone watching the job expects. The second bounds the whole run at + ``_MAX_REPORTS_PER_RUN`` reports, because each report resends every series in + full -- so upload and stored-blob writes both grow as the square of the + report count, and it is the report count, not the step count, that drives it. + + The second floor is not a corner case guard: ``val_period`` is the user's + ``val_check_interval``, and any value below ``_REPORTS_PER_VAL_PERIOD`` -- + "validate every 5 steps" is an ordinary request -- floors the first one to + zero, which clamps to 1 and reports every step of an arbitrarily long run. + """ + per_val_period = (val_period or 0) // _REPORTS_PER_VAL_PERIOD + per_run = (max(max_steps, 0) + _MAX_REPORTS_PER_RUN - 1) // _MAX_REPORTS_PER_RUN + return max(per_val_period, per_run, 1) def resolve_steps_per_epoch(max_steps: int, num_epochs: int | None, explicit: int | None = None) -> int: @@ -134,7 +154,7 @@ def for_schedule( 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), + log_interval=resolve_log_interval(val_period, max_steps), max_steps=max_steps, num_epochs=num_epochs, ) diff --git a/services/rl/tests/test_nemo_rl_logger.py b/services/rl/tests/test_nemo_rl_logger.py index 070a15d6db..b9701a17be 100644 --- a/services/rl/tests/test_nemo_rl_logger.py +++ b/services/rl/tests/test_nemo_rl_logger.py @@ -51,6 +51,7 @@ def _stub(name: str) -> types.ModuleType: 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 + _MAX_REPORTS_PER_RUN, NemoRLLogger, resolve_log_interval, resolve_steps_per_epoch, @@ -301,18 +302,45 @@ def test_close_with_nothing_pending_reports_nothing(callback: _RecordingCallback @pytest.mark.parametrize( - "val_period,expected", + "val_period,max_steps,expected", [ - (100, 10), - (10, 1), - (5, 1), # floors to 0 -> clamped - (1, 1), - (0, 1), - (None, 1), # val_period is Optional + # Driven by val_period: ~10 reports across one validation period. + (100, 200, 10), + (10, 200, 1), + (5, 200, 1), # floors to 0 -> clamped + (1, 200, 1), + (0, 200, 1), + (None, 200, 1), # val_period is Optional + # Driven by the run-length cap, once val_period has stopped bounding + # anything. This is the regime a small val_check_interval lands in. + (5, 2_000, 10), + (5, 20_000, 100), + (None, 20_000, 100), + # Whichever floor is coarser wins; here it is val_period's. + (10_000, 20_000, 1_000), + # A run shorter than the cap is never throttled past its own length. + (0, 1, 1), ], ) -def test_resolve_log_interval(val_period: int | None, expected: int) -> None: - assert resolve_log_interval(val_period) == expected +def test_resolve_log_interval(val_period: int | None, max_steps: int, expected: int) -> None: + assert resolve_log_interval(val_period, max_steps) == expected + + +def test_the_run_length_cap_bounds_the_report_count(callback: _RecordingCallback) -> None: + """val_period alone does not bound reporting, and the report is what costs. + + `val_check_interval=5` is an ordinary request, and it floors the + reports-per-validation-period term to zero -- so before the cap this run + reported all 20,000 steps. Each report resends every series in full, so that + is quadratic in upload and in stored-blob writes, not linear. + """ + max_steps = 20_000 + logger = NemoRLLogger.for_schedule(max_steps=max_steps, num_epochs=1, val_period=5) + for step in _driver_steps(max_steps): + logger.log_metrics(TRAIN_METRICS, step=step, prefix="train") + + assert len(callback.train_steps) <= _MAX_REPORTS_PER_RUN + assert len(callback.train_steps) >= _MAX_REPORTS_PER_RUN // 2, "still a usable curve, not a throttle to nothing" @pytest.mark.parametrize( From 2889bf609557ed5883c8c685334eca42f4a27c5d Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 21:12:24 -0400 Subject: [PATCH 17/23] fix(automodel): strip the val_ prefix the recipes already applied The shared callback's naming rule is that a backend passes its framework's own metric name and the phase supplies the prefix. Automodel is the one backend whose framework prefixes some of its validation metrics itself, so a name arriving pre-prefixed came back doubled. That was handled by renaming exactly `val_loss` to `loss`, which fixed the one curve Studio charts and left every other prefixed name alone: train_bi_encoder reports val_acc1 and val_mrr, which landed as val_val_acc1 and val_val_mrr. strip_val_prefix takes the prefix off wherever the recipe happened to put one. It has to be unconditional rather than a list, because the recipes are inconsistent about which metrics carry it -- train_ft pairs `val_loss` with a bare `lr`, `num_label_tokens` and `mem`, all of which still pick up the phase prefix normally. removeprefix, not a replace, so an interior `val_` stays part of the name. finetune.py imports the recipes at module scope and nemo_automodel exists only inside the training image, so the test stubs the six leaf modules to import it. Through monkeypatch.setitem rather than the module-scope sys.modules assignment its neighbours use: those outlive the file that installed them and leak into whatever shares the xdist worker, which is already why two unsloth tests fail whenever this directory runs first. Signed-off-by: Albert Cui --- .../tasks/training/backends/finetune.py | 25 ++++-- .../tasks/training/backends/test_finetune.py | 83 +++++++++++++++++++ 2 files changed, 103 insertions(+), 5 deletions(-) create mode 100644 services/automodel/tests/tasks/training/backends/test_finetune.py diff --git a/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py b/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py index 9073e87ab6..db16cd29e6 100644 --- a/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py +++ b/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py @@ -10,6 +10,7 @@ from __future__ import annotations import logging +from collections.abc import Mapping from typing import Any, Protocol, runtime_checkable from nemo_automodel.components.checkpoint.checkpointing import Checkpointer @@ -71,6 +72,24 @@ def save_checkpoint( ... +def strip_val_prefix(metrics: Mapping[str, Any]) -> dict[str, Any]: + """Drop the ``val_`` the Automodel recipes already put on a metric name. + + The shared callback's naming rule is that the backend supplies its + framework's own name and the phase supplies the prefix, so a name that + arrives pre-prefixed comes back doubled. Stripping is what makes ``val_loss`` + land as ``val_loss`` rather than ``val_val_loss``. + + It applies to every name, not just ``val_loss``, because the recipes are + inconsistent about which metrics they prefix: ``train_ft`` reports + ``val_loss`` alongside a bare ``lr`` and ``num_label_tokens``, and + ``train_bi_encoder`` adds ``val_acc1`` and ``val_mrr``. Only the callback + should be deciding the phase, so the prefix comes off wherever the recipe + happened to put one. + """ + return {name.removeprefix("val_"): value for name, value in metrics.items()} + + class AutomodelRecipeWrapper: """Wraps an Automodel recipe with Jobs-service progress reporting.""" @@ -171,11 +190,7 @@ def _log_val_metrics(self, *args: Any, **kwargs: Any) -> None: self.callback.report_validation( step=getattr(log_data, "step", 0) + 1, # Convert to 1-based epoch=getattr(log_data, "epoch", 0) + 1, # Convert to 1-based - # Automodel names it `val_loss`; renaming to `loss` lets the - # phase prefix produce `val_loss` the way it does elsewhere. - # An absent one now yields no point, where it used to chart - # a fabricated 0.0. - metrics={("loss" if name == "val_loss" else name): value for name, value in metrics.items()}, + metrics=strip_val_prefix(metrics), ) except Exception as e: logger.warning(f"Failed to report validation progress: {e}") diff --git a/services/automodel/tests/tasks/training/backends/test_finetune.py b/services/automodel/tests/tasks/training/backends/test_finetune.py new file mode 100644 index 0000000000..27efd33726 --- /dev/null +++ b/services/automodel/tests/tasks/training/backends/test_finetune.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the Automodel training entry point's metric naming. + +The shared callback's rule is that the backend passes its framework's own metric +name and the phase supplies the prefix. Automodel is the one backend whose +framework already prefixes some of its validation metrics, so it has to undo +that first -- otherwise the prefix arrives twice. +""" + +from __future__ import annotations + +import importlib +import sys +from collections.abc import Iterator +from types import ModuleType +from unittest.mock import MagicMock + +import pytest + +#: The leaf modules finetune.py imports from. nemo_automodel exists only inside +#: the training image, so the module cannot be imported in a plain checkout. +#: Only the leaves are needed: `from a.b.c import D` consults sys.modules for +#: `a.b.c` before it ever reaches for the parents. +_STUBBED_MODULES = ( + "nemo_automodel.components.checkpoint.checkpointing", + "nemo_automodel.components.config._arg_parser", + "nemo_automodel.components.training.step_scheduler", + "nemo_automodel.recipes.llm.kd", + "nemo_automodel.recipes.llm.train_ft", + "nemo_automodel.recipes.retrieval.train_bi_encoder", +) +_FINETUNE = "nmp.automodel.tasks.training.backends.finetune" + + +@pytest.fixture +def finetune(monkeypatch: pytest.MonkeyPatch) -> Iterator[ModuleType]: + """Import finetune.py under throwaway stubs. + + The stubs go in via monkeypatch rather than a bare ``sys.modules`` + assignment so they are torn down with the test: a module-scope stub outlives + the file that installed it and leaks into whatever else shares the xdist + worker. finetune itself is popped for the same reason -- left behind, it is + a cached module holding references to mocks. + """ + for name in _STUBBED_MODULES: + monkeypatch.setitem(sys.modules, name, MagicMock()) + monkeypatch.delitem(sys.modules, _FINETUNE, raising=False) + yield importlib.import_module(_FINETUNE) + sys.modules.pop(_FINETUNE, None) + + +def test_the_recipes_own_val_prefix_comes_off(finetune: ModuleType) -> None: + """Otherwise the phase prefix arrives twice: `val_loss` as `val_val_loss`.""" + assert finetune.strip_val_prefix({"val_loss": 0.5}) == {"loss": 0.5} + + +def test_every_prefixed_name_is_stripped_not_just_the_loss(finetune: ModuleType) -> None: + """train_bi_encoder reports val_acc1 and val_mrr alongside val_loss. + + Special-casing `val_loss` fixes the one curve Studio charts and leaves the + rest as `val_val_acc1` / `val_val_mrr`. + """ + stripped = finetune.strip_val_prefix({"val_loss": 0.5, "val_acc1": 0.8, "val_mrr": 0.7}) + + assert stripped == {"loss": 0.5, "acc1": 0.8, "mrr": 0.7} + + +def test_an_unprefixed_name_is_left_alone(finetune: ModuleType) -> None: + """The recipes are inconsistent: train_ft pairs `val_loss` with a bare `lr`.""" + stripped = finetune.strip_val_prefix({"val_loss": 0.5, "lr": 5e-06, "num_label_tokens": 128, "mem": 4.2}) + + assert stripped == {"loss": 0.5, "lr": 5e-06, "num_label_tokens": 128, "mem": 4.2} + + +def test_only_a_leading_occurrence_is_removed(finetune: ModuleType) -> None: + """`removeprefix`, not a replace: an interior `val_` is part of the name.""" + assert finetune.strip_val_prefix({"interval_val_x": 1.0}) == {"interval_val_x": 1.0} + + +def test_an_empty_metric_dict_survives(finetune: ModuleType) -> None: + assert finetune.strip_val_prefix({}) == {} From 934d8b1cdbe6d432ebf1f8b9c278d34768e5e531 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 21:17:16 -0400 Subject: [PATCH 18/23] fix(customization): discard the steps a resumed run replays Reporting runs ahead of checkpointing, so the two do not line up: a job that checkpoints at step 100 and is interrupted at 150 has already reported 110 through 150. Resuming rolls training back to the checkpoint and replays those steps, and the accumulator -- seeded from the server precisely so a resumed job continues its curves -- appended the replayed points after the ones they superseded. The curve then doubled back on itself, with two values at each replayed step. A report below the high-water mark is what identifies that: training only ever moves forward within a run, so a step behind the furthest one recorded means a rewind. Every point from that step on is dropped before the new one lands. The replayed values are the real ones; what goes is work that was rolled back, and the gap that leaves is honest where the doubled-back curve was not. Across every series, not only the one being written. Validation runs on its own cadence, so pruning per-series would leave a val curve carrying rolled-back points until the next validation pass -- hundreds of steps later, or never on a short run. The comparison is strict. NeMo-RL validates at step N before logging train N, so a report that merely fails to advance the mark is not a rewind; treating it as one would have the train report delete the validation point that legitimately shares its step. A rewind to 0 is a task rerunning without a checkpoint, and clearing the curves is right. Stored points are read back from a blob this process did not write, so seeding now drops any that record no step: unplaceable on a curve and unplaceable against a rewind. That also keeps the new comparison total -- report_train_step is called straight from NeMo-RL's log_metrics with nothing catching underneath, so it must not raise on a malformed point. Signed-off-by: Albert Cui --- .../training/callbacks.py | 58 ++++++++- .../tests/training/test_callbacks.py | 112 ++++++++++++++++++ 2 files changed, 168 insertions(+), 2 deletions(-) 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 bebae038f4..ebdb65cc1b 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 @@ -133,6 +133,25 @@ def _coerce(value: object) -> float | int: return int(real) if isinstance(real, numbers.Integral) else float(real) +def _point_step(point: object) -> float | int | None: + """The step a stored series point records, or ``None`` if it records none. + + Defensive because the accumulator is seeded from whatever the server had + stored: a malformed blob must cost the points it corrupted, never raise out + of a report and into the training loop. + """ + if not isinstance(point, dict): + return None + step = point.get("step") + return step if isinstance(step, (int, float)) else None + + +def _highest_step(series: Mapping[str, list[dict[str, float | int]]]) -> float | int: + """The furthest step any series in ``series`` has recorded, or 0 for none.""" + steps = [step for points in series.values() for point in points if (step := _point_step(point)) is not None] + return max(steps, default=0) + + class TrainingProgressCallback: """Report training progress to the Jobs service.""" @@ -144,8 +163,19 @@ def __init__(self, reporter: JobsServiceProgressReporter): self._reporter = reporter #: 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()) + #: resumed job continues its curves instead of restarting them. A seeded + #: point that records no step is dropped here: it can be placed neither + #: on a curve nor against a rewind, so carrying it only defers the + #: problem to whoever reads it. + self._series: dict[str, list[dict[str, float | int]]] = { + name: [point for point in points if _point_step(point) is not None] + for name, points in reporter.fetch_current_metrics().items() + } + + #: Furthest step any series has reached, seeded points included. A report + #: below it means training resumed from a checkpoint and is replaying + #: steps that were already recorded; see :meth:`_discard_from`. + self._high_water_step: float | int = _highest_step(self._series) if any(self._series.values()): logger.info( "Seeded %d metric series from server (%d points): %s", @@ -173,6 +203,9 @@ def _report_metrics( what the Jobs service records as the task's phase. """ namespaced = _namespace(phase, metrics) + if step < self._high_water_step: + self._discard_from(step) + self._high_water_step = max(self._high_water_step, step) for name, value in namespaced.items(): self._series.setdefault(name, []).append({"step": step, "epoch": epoch, "value": value}) @@ -187,6 +220,27 @@ def _report_metrics( details["backend"] = resolved self._reporter.report_running(phase=report_phase, **details) + def _discard_from(self, step: int) -> None: + """Drop every recorded point at or after ``step``, across all series. + + A report below the high-water mark means training resumed from a + checkpoint and is replaying steps it already reported. The replayed + values are the real ones -- what is discarded belongs to work that was + rolled back -- so keeping both makes the curve double back on itself, + which is worse than the gap that removing them leaves. + + Every series, not just the one being appended to: validation runs on its + own cadence, so a val curve would otherwise carry its rolled-back points + until the next validation pass, potentially hundreds of steps later. + Points whose step cannot be read go too, there being no way to place + them relative to the rewind. + + Lists are truncated in place; :meth:`_build_metrics_summary` copies each + one before handing it over, so payloads already sent are unaffected. + """ + for points in self._series.values(): + points[:] = [point for point in points if (recorded := _point_step(point)) is not None and recorded < step] + def _build_metrics_summary(self) -> dict[str, list[dict[str, float | int]]]: """Build the accumulated metrics payload for inclusion in status_details. diff --git a/packages/nmp_customization_common/tests/training/test_callbacks.py b/packages/nmp_customization_common/tests/training/test_callbacks.py index fffdcf8331..6f48adb671 100644 --- a/packages/nmp_customization_common/tests/training/test_callbacks.py +++ b/packages/nmp_customization_common/tests/training/test_callbacks.py @@ -347,6 +347,118 @@ def test_series_survive_a_resume_beyond_the_loss_curves() -> None: ] +# --------------------------------------------------------------------------- # +# Resuming from a checkpoint rewinds the curves +# --------------------------------------------------------------------------- # + + +def test_a_resumed_run_discards_the_steps_it_replays() -> None: + """Reporting runs ahead of checkpointing, so a resume replays reported steps. + + Checkpoint at 100, reports as far as 150, interruption. Training resumes + from the checkpoint and reports 110 again with a different value. Both + points kept, the curve doubles back on itself; the replayed one is the real + one, because 110-150 belong to work that was rolled back. + """ + prior = { + "train_loss": [{"step": step, "epoch": 1, "value": 1.0} for step in (100, 110, 120, 130, 140, 150)], + } + reporter = _RecordingReporter(prior) + _make_callback(reporter).report_train_step(step=110, epoch=1, metrics={"loss": 0.7}) + + assert reporter.reports[-1]["metrics"]["train_loss"] == [ + {"step": 100, "epoch": 1, "value": 1.0}, + {"step": 110, "epoch": 1, "value": 0.7}, + ] + + +def test_a_rewind_prunes_every_series_not_only_the_one_being_written() -> None: + """Validation has its own cadence, so its curve cannot wait to self-heal. + + Pruning only the series being appended to would leave val_loss holding its + rolled-back points until the next validation pass, which may be hundreds of + steps away -- and on a short run may never come. + """ + prior = { + "train_loss": [{"step": 100, "epoch": 1, "value": 1.0}, {"step": 150, "epoch": 2, "value": 0.9}], + "val_loss": [{"step": 100, "epoch": 1, "value": 1.1}, {"step": 150, "epoch": 2, "value": 1.2}], + "train_reward": [{"step": 150, "epoch": 2, "value": 0.3}], + } + reporter = _RecordingReporter(prior) + _make_callback(reporter).report_train_step(step=110, epoch=1, metrics={"loss": 0.7}) + + metrics = reporter.reports[-1]["metrics"] + assert [point["step"] for point in metrics["train_loss"]] == [100, 110] + assert [point["step"] for point in metrics["val_loss"]] == [100] + assert metrics["train_reward"] == [] + + +def test_forward_progress_never_prunes(reporter: _RecordingReporter) -> None: + """The ordinary path: every step is past the high-water mark, nothing moves.""" + callback = _make_callback(reporter) + for step in (1, 2, 3, 4): + callback.report_train_step(step=step, epoch=1, metrics={"loss": 1.0 / step}) + + assert [point["step"] for point in reporter.reports[-1]["metrics"]["train_loss"]] == [1, 2, 3, 4] + + +def test_validation_and_training_at_the_same_step_both_land(reporter: _RecordingReporter) -> None: + """Not a rewind: NeMo-RL validates at step N before logging train N. + + The high-water comparison is strict for this reason -- treating "not ahead" + as a rewind would have the train report delete the validation point that + legitimately shares its step. + """ + callback = _make_callback(reporter) + callback.report_validation(step=10, epoch=1, metrics={"loss": 0.9}) + callback.report_train_step(step=10, epoch=1, metrics={"loss": 0.8}) + + metrics = reporter.reports[-1]["metrics"] + assert metrics["val_loss"] == [{"step": 10, "epoch": 1, "value": 0.9}] + assert metrics["train_loss"] == [{"step": 10, "epoch": 1, "value": 0.8}] + + +def test_a_restart_from_scratch_clears_the_curves() -> None: + """A task that reruns without a checkpoint rewinds all the way to zero.""" + prior = {"train_loss": [{"step": 50, "epoch": 1, "value": 1.0}]} + reporter = _RecordingReporter(prior) + _make_callback(reporter).report_train_step(step=1, epoch=1, metrics={"loss": 2.0}) + + assert reporter.reports[-1]["metrics"]["train_loss"] == [{"step": 1, "epoch": 1, "value": 2.0}] + + +def test_a_malformed_stored_point_cannot_raise_into_the_training_loop() -> None: + """The blob is read back from the server, so its shape is not guaranteed. + + report_train_step is called straight from NeMo-RL's log_metrics with nothing + catching underneath, so a stored point that is not a {step, ...} dict has to + be dropped rather than raise. + """ + # Cast because the shape is the point: this is what a corrupted blob reads + # back as, and no annotation describes it honestly. + prior = cast( + dict[str, list[dict[str, Any]]], + {"train_loss": [{"step": 100, "epoch": 1, "value": 1.0}, "junk", {"epoch": 1}, {"step": "eight"}]}, + ) + reporter = _RecordingReporter(prior) + _make_callback(reporter).report_train_step(step=110, epoch=1, metrics={"loss": 0.7}) + + assert reporter.reports[-1]["metrics"]["train_loss"] == [ + {"step": 100, "epoch": 1, "value": 1.0}, + {"step": 110, "epoch": 1, "value": 0.7}, + ] + + +def test_a_malformed_stored_point_does_not_skew_the_high_water_mark() -> None: + """It is unplaceable, so it cannot be what a later report is compared against.""" + prior = {"train_loss": [{"step": "eight"}, {"step": 10, "epoch": 1, "value": 1.0}]} + reporter = _RecordingReporter(prior) + callback = _make_callback(reporter) + callback.report_train_step(step=11, epoch=1, metrics={"loss": 0.7}) + + assert [point["step"] for point in reporter.reports[-1]["metrics"]["train_loss"]] == [10, 11] + + def test_series_are_snapshots_not_live_references(reporter: _RecordingReporter) -> None: """A shared list would retroactively mutate already-sent payloads.""" callback = _make_callback(reporter) From 4095302b6be15eef09d2a634bdc8413241c0f8c5 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Thu, 13 Aug 2026 21:17:23 -0400 Subject: [PATCH 19/23] fix(rl): report a validation pass that scores on something other than loss The validation branch required a chartable `loss` before it reported anything, which is the last place the old "train_loss and val_loss are special" assumption survived. GRPO scores validation on rewards: its validate() returns `accuracy` and `avg_length` and no loss at all, so the gate dropped not its loss curve but every validation pass it ran. It now reports whatever the pass produced, gated only against a hollow report -- a pass whose metrics are all histograms still says nothing. Best-so-far still tracks the validation loss and so is updated only when there is one. The train branch keeps its `loss` check, which is doing different work and is worth saying so explicitly. GRPO and PPO log twice under `prefix="train"` at a single step -- rollout stats first, then the training metrics, both at `total_steps + 1` -- and only the second carries a loss. There it is a discriminator, not a requirement: without it one step reports twice, each report resending every series, and a throttled step ends up pending as the rollout half with the loss lost. Signed-off-by: Albert Cui --- .../backends/nemo_rl/nemo_rl_logger.py | 27 +++++++--- services/rl/tests/test_nemo_rl_logger.py | 51 +++++++++++++++++++ 2 files changed, 70 insertions(+), 8 deletions(-) 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 de9ecb156f..4b62ee1f64 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 @@ -187,7 +187,12 @@ def log_metrics( # epoch 1, hence the clamp rather than a bare `step - 1`. epoch = (max(step - 1, 0) // self._steps_per_epoch) + 1 - # Handle training loss + # `loss` gates the train branch as a discriminator, not as a requirement. + # GRPO and PPO log twice under `prefix="train"` at one step -- rollout + # stats first, then the training metrics, both at `total_steps + 1` + # (grpo.py) -- and only the second carries a loss. Keying on it is what + # keeps one step from producing two reports, and what keeps the rollout + # log from displacing a pending report that has the loss in it. if prefix == "train" and is_chartable(metrics.get("loss")): report = {"step": step, "epoch": epoch, "metrics": dict(metrics)} # Throttled to log_interval to reduce output. A withheld step is held as @@ -198,15 +203,21 @@ def log_metrics( else: self._pending_train_report = report - # Handle validation metrics + # Validation reports whatever the pass produced. There is one validation + # log per pass and no rollout twin to tell apart, so requiring a `loss` + # here bought nothing and cost whole passes: GRPO validates on + # `accuracy` and `avg_length` and reports no loss at all, so every + # validation it ran went unrecorded. The gate is only against a hollow + # report -- a pass whose metrics are all histograms says nothing. elif prefix and prefix.startswith("validation"): - if is_chartable(metrics.get("loss")): - val_loss = metrics["loss"] + if any(is_chartable(value) for value in metrics.values()): self._callback.report_validation(step=step, epoch=epoch, metrics=dict(metrics)) - # Track best validation loss - if val_loss < self._best_metric_value: - self._best_metric_value = val_loss - self._best_epoch = epoch + # Track best validation loss, when the algorithm reports one. + if is_chartable(metrics.get("loss")): + val_loss = float(metrics["loss"]) + if 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}") diff --git a/services/rl/tests/test_nemo_rl_logger.py b/services/rl/tests/test_nemo_rl_logger.py index b9701a17be..f0c948cf0a 100644 --- a/services/rl/tests/test_nemo_rl_logger.py +++ b/services/rl/tests/test_nemo_rl_logger.py @@ -127,6 +127,18 @@ class _Histogram: VALIDATION_METRICS: dict[str, Any] = {"loss": 0.25, "num_valid_samples": 8} +# GRPO/PPO log this under `prefix="train"` at the same step as TRAIN_METRICS, +# from the rollout rather than the training step. Generation stats only: no loss. +ROLLOUT_METRICS: dict[str, Any] = { + "total_turns": 64, + "avg_turns_per_sample": 1.0, + "mean_gen_tokens_per_sample": 412.5, + "truncation_rate": 0.02, +} + +# What GRPO's validate() returns. It scores on rewards, so there is no loss in it. +REWARD_VALIDATION_METRICS: dict[str, Any] = {"accuracy": 0.61, "avg_length": 412.5} + # --------------------------------------------------------------------------- # # Module stub hygiene @@ -179,6 +191,22 @@ def test_train_call_without_a_loss_is_ignored(callback: _RecordingCallback) -> N assert callback.train_steps == [] +def test_the_rollout_log_does_not_produce_a_second_train_report(callback: _RecordingCallback) -> None: + """GRPO and PPO log twice under `prefix="train"` at one step. + + grpo.py logs `rollout_metrics` and then the training `metrics`, both at + `total_steps + 1`. Only the second has a loss, which is why the branch keys + on one: without it the step reports twice -- each resending every series -- + and a throttled step ends up pending as the rollout half, losing the loss. + """ + logger = _make_logger() + logger.log_metrics(ROLLOUT_METRICS, step=1, prefix="train") + logger.log_metrics(TRAIN_METRICS, step=1, prefix="train") + + assert len(callback.train_steps) == 1 + assert callback.train_steps[0]["metrics"]["loss"] == 0.5 + + def test_every_train_scalar_is_forwarded(callback: _RecordingCallback) -> None: """Including the ones the old allow-list had no entry for.""" _make_logger().log_metrics(TRAIN_METRICS, step=0, prefix="train") @@ -432,6 +460,29 @@ def test_validation_with_nothing_usable_is_not_reported(callback: _RecordingCall assert callback.validations == [] +def test_a_validation_pass_without_a_loss_is_still_reported(callback: _RecordingCallback) -> None: + """No metric is privileged, so a loss is not what makes a pass worth recording. + + GRPO scores validation on rewards: its validate() returns `accuracy` and + `avg_length` and no loss at all. Gating on one dropped every validation pass + it ran -- not the loss curve, the whole pass. + """ + _make_logger().log_metrics(REWARD_VALIDATION_METRICS, step=10, prefix="validation") + + assert callback.validations[0]["metrics"] == REWARD_VALIDATION_METRICS + + +def test_a_loss_free_validation_leaves_the_best_metric_alone(callback: _RecordingCallback) -> None: + """Best-so-far tracks the validation loss, so a pass without one says nothing.""" + logger = _make_logger() + logger.log_metrics({"loss": 0.4}, step=10, prefix="validation") + logger.log_metrics(REWARD_VALIDATION_METRICS, step=20, prefix="validation") + + assert len(callback.validations) == 2 + assert logger._best_metric_value == 0.4 + assert logger._best_epoch == 1 + + def test_best_validation_loss_tracks_minimum(callback: _RecordingCallback) -> None: logger = _make_logger() logger.log_metrics({"loss": 0.5}, step=10, prefix="validation") From a27cd5af60054237a6308e9b80897ff487845cc0 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 14 Aug 2026 12:13:27 -0400 Subject: [PATCH 20/23] fix(customization): correct twelve defects across the progress reporting paths Everything a review of this branch turned up, in the shared callback and in all three backends that call it. Grouped below by the assumption each set breaks, because several are the same mistake wearing different clothes: a report says more than it observed, or arrives somewhere the code downstream did not expect. Seeding ------- fetch_current_metrics parsed the stored blob outside its try, so a status_details not shaped like {"metrics": {...}} raised AttributeError rather than seeding nothing. It runs from TrainingProgressCallback.__init__, which AutomodelRecipeWrapper and NemoRLLogger both build outside any try, so the training subprocess died at startup over data that should have cost it only its curves. Both the blob and its metrics entry are checked now. A read that failed and a read that found nothing both answered {}, leaving the caller no way to tell them apart. The server replaces a sent key wholesale, so seeding from {} after a failed read had the first train report replace the job's entire stored history with the one point this process held; a distributed launch racing a briefly unreachable Jobs service lands exactly there. Only a 404 now means "nothing stored"; every other failure warns and returns None, and the callback withholds the metrics key for the life of the process. This run's curves stall, which a later run recovers, rather than every previous run's being destroyed, which nothing recovers. Not addressed, and now written down where the seeding is: a series can hold points from more than one run. A replaced pod seeds itself from the old run's points and, nothing being able to restore a checkpoint, starts again at step one, so two values can sit at one step contradicting each other. Both fixes are bigger than they look -- a run identifier per point needs Studio to filter or overlay by run, and dropping the superseded points loses a failed attempt's history for good -- so the behaviour is left as it is, described in the module docstring and pinned by a test rather than left for someone to trip over. What a report states -------------------- report_training_start sent a literal step: 0. report_running derives percentage_done from a stated step and the merge is wholesale per key, so that 0 overwrote whatever progress the task had stored, in a field harder to spot than the metrics it already omits because the epoch beside it is not restated and goes on reading correctly. It fires before the first step and has no position to report, so it states none and lets the first train step do it. A report whose values were all unchartable resent every series anyway. They go in full or not at all, so a report that added no point has nothing to say about them -- up to 413 KB, by this module's own measurements, to say nothing. The four-line stamp-and-send tail every report repeated is now one _send. _default_backend is None for two of the three backends, so a change to it that missed a copy would have shown up only in unsloth's payload. Report cadence and cost ----------------------- _MAX_REPORTS_PER_RUN capped train reports and nothing else, so it did not bound a run. val_check_interval=1 is reachable, and it validates every step: a 20k-step run made 200 train reports and 20,000 validation ones, each resending every series in full and stored twice by the Jobs service -- the quadratic growth the cap exists to prevent, arriving through the other door. The same bound now applies to validation, counted in passes because passes come on their own cadence. It resolves to "every pass" for any ordinary configuration, so nothing in the existing regime moves. A withheld pass is held pending and flushed by close(), as the train path already did, since the final validation is the one worth having. Metric naming, per backend -------------------------- NeMo-RL's validate() logs once per dataloader, every call at the same step under `validation-`. Forwarded as-is, two datasets' loss landed as two points at one step in a single val_loss series -- the collision the _ rule exists to prevent, one level further down. The first prefix seen keeps the bare names, because NeMo-RL names the dataloader even when there is only one, and disambiguating unconditionally would rename the common case and take Studio's curve with it. Automodel's strip_val_prefix maps every name onto its unprefixed form, so a dict carrying both val_loss and loss collapsed them onto one key and the second silently replaced the first. The prefixed name wins now, and the collision is logged. Unsloth kept float() on loss and dropped it from learning_rate and grad_norm, so a value the Trainer logs as anything but a numbers.Real -- some paths log grad_norm as a 0-dim tensor rather than calling .item() -- failed the chartable filter and vanished from both the series and the report, with no log line saying why. Dead state ---------- _best_metric_value and _best_epoch go, with the two tests that pinned them. Nothing has ever read them; they predate this branch, which was extending write-only state. Test isolation -------------- Two module stubs outlived the files that installed them and leaked across the xdist worker. automodel's test_config set sys.modules["transformers"] to a MagicMock at import scope, and unsloth's bridge does `from transformers import TrainerCallback` at call time -- so HfTrainerProgressCallback became a mock subclass whose hooks did nothing, and two of its tests failed or passed vacuously depending on collection order. The RL logger test's nemo_rl stub had the same shape, which its own docstring conceded. Both are now scoped to the window that needs them. Signed-off-by: Albert Cui --- .../training/callbacks.py | 148 ++++++++------ .../customization_common/training/progress.py | 53 ++++- .../tests/training/test_callbacks.py | 178 ++++++++++------- .../tests/training/test_progress.py | 83 +++++++- .../tasks/training/backends/finetune.py | 24 ++- .../tasks/training/backends/test_callbacks.py | 3 +- .../tasks/training/backends/test_config.py | 23 ++- .../tasks/training/backends/test_finetune.py | 20 ++ .../backends/nemo_rl/nemo_rl_logger.py | 163 +++++++++++---- services/rl/tests/test_nemo_rl_logger.py | 186 ++++++++++++++---- .../training/backends/hf_trainer_callback.py | 31 ++- services/unsloth/tests/test_callbacks.py | 3 +- .../unsloth/tests/test_hf_trainer_callback.py | 52 +++++ 13 files changed, 741 insertions(+), 226 deletions(-) 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 ebdb65cc1b..b76350bcc2 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 @@ -52,10 +52,38 @@ job, so the write volume is double the numbers above. Deliberately accepted for batch training jobs, on the understanding that a -backend bounds its own report count -- NeMo-RL's ``resolve_log_interval`` caps a -run for exactly this reason. A backend that instead reports every step of a long -run pays quadratically, and the answer to that is delta appends in the transport -rather than trimming the series here. +backend bounds its own report count -- on *every* path that reports, since one +uncapped path is enough to make the total quadratic. NeMo-RL bounds both its +train steps and its validation passes for exactly this reason. A backend that +instead reports every step of a long run pays quadratically, and the answer to +that is delta appends in the transport rather than trimming the series here. + +Seeding, and what it inherits +---------------------------- +The accumulator lives in this process, so a new process taking over a task +would report from empty and -- the merge being wholesale per key -- replace +whatever the task had already stored. It therefore seeds itself from the +server, which is behaviour that predates the per-metric series: the reporter +has always read back the stored curves. + +Known issue: a series can end up holding points from more than one run, which +contradict each other. If a task's pod is replaced -- the Jobs service can +suspend and resume a Kubernetes Job, and the Volcano backend restarts on +``PodFailed`` when an execution profile raises ``maxRetry`` above its default +of zero -- the new process seeds itself from the old one's points and then +appends its own. Nothing restores a checkpoint, so it starts again at step one +and the series carries two values at each repeated step. Consumers see the +last one written at a given step, and a tail from the abandoned attempt beyond +wherever the new run has reached. + +Left alone deliberately. The fix is to identify which run a point came from, +so both can be kept and told apart, and that belongs with checkpoint restore: +without it there is no second curve worth the work to separate, and Studio's +loss chart keys a Map by step, so it would need to filter or overlay by run +before any of it were visible. Detecting the restart here and dropping the +superseded points is the other option, and it loses a failed attempt's history +permanently to keep one series single-valued -- a trade for the readers rather +than the data, and not one to make in passing. Backends subclass this and set :attr:`_default_backend`: unsloth stamps a ``backend`` field on each report (``"unsloth"``); automodel and NeMo-RL leave it @@ -146,12 +174,6 @@ def _point_step(point: object) -> float | int | None: return step if isinstance(step, (int, float)) else None -def _highest_step(series: Mapping[str, list[dict[str, float | int]]]) -> float | int: - """The furthest step any series in ``series`` has recorded, or 0 for none.""" - steps = [step for points in series.values() for point in points if (step := _point_step(point)) is not None] - return max(steps, default=0) - - class TrainingProgressCallback: """Report training progress to the Jobs service.""" @@ -162,20 +184,32 @@ class TrainingProgressCallback: def __init__(self, reporter: JobsServiceProgressReporter): self._reporter = reporter + seeded = reporter.fetch_current_metrics() + + #: True when the seed read failed outright, as opposed to finding nothing + #: stored. The accumulator is then known to be incomplete, and because + #: the server replaces a sent key wholesale, reporting it would overwrite + #: the very history that could not be read. So the ``metrics`` key is + #: withheld for the life of the process: this run's curves stall, which + #: is recoverable, rather than every previous run's being destroyed, + #: which is not. Current values and progress still report normally. + self._seed_unavailable = seeded is None + if self._seed_unavailable: + logger.warning( + "Could not read the stored metric series; this run will report progress and current " + "values but will not update the accumulated curves, to avoid overwriting them." + ) + #: series name -> [{step, epoch, value}], seeded from the server so a - #: resumed job continues its curves instead of restarting them. A seeded - #: point that records no step is dropped here: it can be placed neither - #: on a curve nor against a rewind, so carrying it only defers the - #: problem to whoever reads it. + #: process taking over a task continues its curves instead of replacing + #: them, which is what an empty accumulator would do. A seeded point that + #: records no step is dropped here: nothing can place it on a curve, so + #: carrying it only defers the problem to whoever reads it. self._series: dict[str, list[dict[str, float | int]]] = { name: [point for point in points if _point_step(point) is not None] - for name, points in reporter.fetch_current_metrics().items() + for name, points in (seeded or {}).items() } - #: Furthest step any series has reached, seeded points included. A report - #: below it means training resumed from a checkpoint and is replaying - #: steps that were already recorded; see :meth:`_discard_from`. - self._high_water_step: float | int = _highest_step(self._series) if any(self._series.values()): logger.info( "Seeded %d metric series from server (%d points): %s", @@ -187,6 +221,19 @@ def __init__(self, reporter: JobsServiceProgressReporter): def _resolve_backend(self, backend: str | None) -> str | None: return backend if backend is not None else self._default_backend + def _send(self, phase: str, details: dict[str, object], backend: str | None) -> None: + """Stamp the backend field, if there is one, and hand the report over. + + The tail every report shares, in one place because a change to it that + missed a copy would be close to invisible: ``_default_backend`` is + ``None`` for two of the three backends, so only unsloth's payload would + have shown the difference. + """ + resolved = self._resolve_backend(backend) + if resolved is not None: + details["backend"] = resolved + self._reporter.report_running(phase=phase, **details) + def _report_metrics( self, phase: str, @@ -203,9 +250,6 @@ def _report_metrics( what the Jobs service records as the task's phase. """ namespaced = _namespace(phase, metrics) - if step < self._high_water_step: - self._discard_from(step) - self._high_water_step = max(self._high_water_step, step) for name, value in namespaced.items(): self._series.setdefault(name, []).append({"step": step, "epoch": epoch, "value": value}) @@ -213,33 +257,13 @@ def _report_metrics( "step": step, "epoch": epoch, **namespaced, - "metrics": self._build_metrics_summary(), } - resolved = self._resolve_backend(backend) - if resolved is not None: - details["backend"] = resolved - self._reporter.report_running(phase=report_phase, **details) - - def _discard_from(self, step: int) -> None: - """Drop every recorded point at or after ``step``, across all series. - - A report below the high-water mark means training resumed from a - checkpoint and is replaying steps it already reported. The replayed - values are the real ones -- what is discarded belongs to work that was - rolled back -- so keeping both makes the curve double back on itself, - which is worse than the gap that removing them leaves. - - Every series, not just the one being appended to: validation runs on its - own cadence, so a val curve would otherwise carry its rolled-back points - until the next validation pass, potentially hundreds of steps later. - Points whose step cannot be read go too, there being no way to place - them relative to the rewind. - - Lists are truncated in place; :meth:`_build_metrics_summary` copies each - one before handing it over, so payloads already sent are unaffected. - """ - for points in self._series.values(): - points[:] = [point for point in points if (recorded := _point_step(point)) is not None and recorded < step] + # Sent in full or not at all, so a report that added no point has nothing + # to say about the curves: the stored copy already matches, and the merge + # leaves a key that is not mentioned standing. + if not self._seed_unavailable and namespaced: + details["metrics"] = self._build_metrics_summary() + self._send(report_phase, details, backend) def _build_metrics_summary(self) -> dict[str, list[dict[str, float | int]]]: """Build the accumulated metrics payload for inclusion in status_details. @@ -256,19 +280,23 @@ def report_training_start(self, max_steps: int, num_epochs: int, *, backend: str """Report that training has started with schedule information. Carries no ``metrics``: it fires before the first step, so it has nothing - to add to the curves, and sending the empty accumulator would replace a - resumed job's stored series with two empty lists. + to add to the curves, and sending the empty accumulator would replace the + stored series with two empty lists. + + Carries no ``step`` either, for the same reason one level along. + ``report_running`` turns a stated step into ``percentage_done``, and the + merge is wholesale per key, so a literal 0 here overwrites whatever + progress the task had stored -- in a field harder to spot than the series, + because the epoch beside it is not restated and goes on reading + correctly. Nothing has happened yet that this could truthfully report, so + it says nothing and lets the first train step state the position. """ self._reporter.configure_progress_tracking(max_steps, num_epochs) details: dict[str, object] = { - "step": 0, "max_steps": max_steps, "num_epochs": num_epochs, } - resolved = self._resolve_backend(backend) - if resolved is not None: - details["backend"] = resolved - self._reporter.report_running(phase="training", **details) + self._send("training", details, backend) def report_train_step( self, step: int, epoch: int, metrics: Mapping[str, object], *, backend: str | None = None @@ -324,10 +352,7 @@ def report_checkpoint_saved( } if checkpoint_path: details["checkpoint_path"] = checkpoint_path - resolved = self._resolve_backend(backend) - if resolved is not None: - details["backend"] = resolved - self._reporter.report_running(phase="checkpoint_saved", **details) + self._send("checkpoint_saved", details, backend) def report_epoch_end(self, step: int, epoch: int, *, backend: str | None = None) -> None: """Report that an epoch has completed.""" @@ -335,10 +360,7 @@ def report_epoch_end(self, step: int, epoch: int, *, backend: str | None = None) "step": step, "epoch": epoch, } - resolved = self._resolve_backend(backend) - if resolved is not None: - details["backend"] = resolved - self._reporter.report_running(phase="epoch_end", **details) + self._send("epoch_end", details, backend) def close(self) -> None: """Clean up resources.""" 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 917179ebf4..a5d14037ef 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 @@ -20,9 +20,11 @@ import logging import os +from collections.abc import Mapping from typing import Any, cast from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NotFoundError from nemo_platform_plugin.jobs.client import JobsClient from nemo_platform_plugin.jobs.schemas import PlatformJobStatus from nemo_platform_plugin.jobs.types import PlatformJobTaskUpdate @@ -87,18 +89,40 @@ def update_task( 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]]]: - """Read back every stored metric series, for resume seeding. + def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]] | None: + """Read back every stored metric series, to seed a new process. + + The only read this reporter makes, and it happens once per process. The + accumulator lives in memory, so a process taking over a task would report + from empty and replace what the task had stored; reading it back is what + keeps the curves continuous. Predates the per-metric series -- the two + loss curves were seeded the same way. - The only read this reporter makes, and it happens once per process. Because the server merges, no write path needs to know what is already stored -- a report that omits a field leaves it standing. 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 + they accumulate, and seeding only ``train_loss`` would silently restart + every other curve from empty. Non-list values are dropped so a malformed blob cannot poison the accumulator, and each list is copied so the caller's accumulator does not alias the response. + + Returns ``{}`` when there is nothing to seed from and ``None`` when the + read itself failed, which are not the same thing and must not look the + same to the caller. The merge is wholesale per key, so a caller that + treats a failed read as "nothing stored" sends its own partial + accumulator under ``metrics`` and destroys the history it could not read + -- see :class:`~nmp.customization_common.training.callbacks.TrainingProgressCallback`. + + Only a 404 means "nothing stored": every other failure is a transport or + server problem, and a distributed launch racing a briefly unreachable + Jobs service hits those, not the 404. + + A blob that reads back malformed is a successful read of unusable data, + so it seeds nothing and is *not* an error: overwriting it is the repair. + Neither branch may raise -- this runs from the callback's constructor, + which backends build outside any try, so raising here kills the training + process rather than costing it its seed. """ if not self._enabled: return {} @@ -111,13 +135,24 @@ def fetch_current_metrics(self) -> dict[str, list[dict[str, float | int]]]: job=self._job_ctx.job_id, step=self._job_ctx.step, ).data() - stored = cast(dict[str, Any], task.status_details or {}) - except Exception as e: + stored = task.status_details or {} + except NotFoundError: # Expected on a first run, where the task has no stored details yet. - logger.info(f"No stored status details to seed from: {e}") + logger.info("No stored status details to seed from: task not found") + return {} + except Exception as e: + logger.warning(f"Could not read stored metric series to seed from: {e}") + return None + + if not isinstance(stored, Mapping): + logger.warning(f"Stored status details are not an object ({type(stored).__name__}); seeding nothing") + return {} + + metrics = stored.get("metrics") or {} + if not isinstance(metrics, Mapping): + logger.warning(f"Stored metrics are not an object ({type(metrics).__name__}); seeding nothing") return {} - metrics = cast(dict[str, Any], stored.get("metrics", {}) or {}) return cast( dict[str, list[dict[str, float | int]]], {name: list(points) for name, points in metrics.items() if isinstance(points, list)}, diff --git a/packages/nmp_customization_common/tests/training/test_callbacks.py b/packages/nmp_customization_common/tests/training/test_callbacks.py index 6f48adb671..1268834161 100644 --- a/packages/nmp_customization_common/tests/training/test_callbacks.py +++ b/packages/nmp_customization_common/tests/training/test_callbacks.py @@ -29,7 +29,7 @@ def __init__(self, prior: dict[str, list[dict[str, Any]]] | None = None) -> None self.tracking: tuple[int, int] | None = None self.closed = False - def fetch_current_metrics(self) -> dict[str, list[dict[str, Any]]]: + def fetch_current_metrics(self) -> dict[str, list[dict[str, Any]]] | None: return self._prior def configure_progress_tracking(self, max_steps: int, num_epochs: int) -> None: @@ -42,6 +42,13 @@ def close(self) -> None: self.closed = True +class _UnreadableReporter(_RecordingReporter): + """A reporter whose one read failed, as opposed to finding nothing stored.""" + + def fetch_current_metrics(self) -> dict[str, list[dict[str, Any]]] | None: + return None + + @pytest.fixture def reporter() -> _RecordingReporter: return _RecordingReporter() @@ -168,13 +175,19 @@ def test_validation_without_a_loss_leaves_the_curve_empty(reporter: _RecordingRe def test_an_empty_metric_dict_still_reports_progress(reporter: _RecordingReporter) -> None: - """step/epoch are progress, not metrics; they land with or without a curve.""" + """step/epoch are progress, not metrics; they land with or without a curve. + + The series are left out rather than resent: nothing was added to them, the + stored copy already matches, and the merge leaves an unmentioned key + standing. Resending would pay the full payload for a report that changed + nothing -- up to 413 KB by this module's own measurements. + """ _make_callback(reporter).report_train_step(step=7, epoch=2, metrics={}) report = reporter.reports[-1] assert report["step"] == 7 assert report["epoch"] == 2 - assert report["metrics"] == {"train_loss": [], "val_loss": []} + assert "metrics" not in report # --------------------------------------------------------------------------- # @@ -348,67 +361,12 @@ def test_series_survive_a_resume_beyond_the_loss_curves() -> None: # --------------------------------------------------------------------------- # -# Resuming from a checkpoint rewinds the curves +# Seeded points, and the run they came from # --------------------------------------------------------------------------- # -def test_a_resumed_run_discards_the_steps_it_replays() -> None: - """Reporting runs ahead of checkpointing, so a resume replays reported steps. - - Checkpoint at 100, reports as far as 150, interruption. Training resumes - from the checkpoint and reports 110 again with a different value. Both - points kept, the curve doubles back on itself; the replayed one is the real - one, because 110-150 belong to work that was rolled back. - """ - prior = { - "train_loss": [{"step": step, "epoch": 1, "value": 1.0} for step in (100, 110, 120, 130, 140, 150)], - } - reporter = _RecordingReporter(prior) - _make_callback(reporter).report_train_step(step=110, epoch=1, metrics={"loss": 0.7}) - - assert reporter.reports[-1]["metrics"]["train_loss"] == [ - {"step": 100, "epoch": 1, "value": 1.0}, - {"step": 110, "epoch": 1, "value": 0.7}, - ] - - -def test_a_rewind_prunes_every_series_not_only_the_one_being_written() -> None: - """Validation has its own cadence, so its curve cannot wait to self-heal. - - Pruning only the series being appended to would leave val_loss holding its - rolled-back points until the next validation pass, which may be hundreds of - steps away -- and on a short run may never come. - """ - prior = { - "train_loss": [{"step": 100, "epoch": 1, "value": 1.0}, {"step": 150, "epoch": 2, "value": 0.9}], - "val_loss": [{"step": 100, "epoch": 1, "value": 1.1}, {"step": 150, "epoch": 2, "value": 1.2}], - "train_reward": [{"step": 150, "epoch": 2, "value": 0.3}], - } - reporter = _RecordingReporter(prior) - _make_callback(reporter).report_train_step(step=110, epoch=1, metrics={"loss": 0.7}) - - metrics = reporter.reports[-1]["metrics"] - assert [point["step"] for point in metrics["train_loss"]] == [100, 110] - assert [point["step"] for point in metrics["val_loss"]] == [100] - assert metrics["train_reward"] == [] - - -def test_forward_progress_never_prunes(reporter: _RecordingReporter) -> None: - """The ordinary path: every step is past the high-water mark, nothing moves.""" - callback = _make_callback(reporter) - for step in (1, 2, 3, 4): - callback.report_train_step(step=step, epoch=1, metrics={"loss": 1.0 / step}) - - assert [point["step"] for point in reporter.reports[-1]["metrics"]["train_loss"]] == [1, 2, 3, 4] - - def test_validation_and_training_at_the_same_step_both_land(reporter: _RecordingReporter) -> None: - """Not a rewind: NeMo-RL validates at step N before logging train N. - - The high-water comparison is strict for this reason -- treating "not ahead" - as a rewind would have the train report delete the validation point that - legitimately shares its step. - """ + """NeMo-RL validates at step N before logging train N; both belong on the curves.""" callback = _make_callback(reporter) callback.report_validation(step=10, epoch=1, metrics={"loss": 0.9}) callback.report_train_step(step=10, epoch=1, metrics={"loss": 0.8}) @@ -418,13 +376,27 @@ def test_validation_and_training_at_the_same_step_both_land(reporter: _Recording assert metrics["train_loss"] == [{"step": 10, "epoch": 1, "value": 0.8}] -def test_a_restart_from_scratch_clears_the_curves() -> None: - """A task that reruns without a checkpoint rewinds all the way to zero.""" +def test_a_restart_appends_to_the_previous_runs_points() -> None: + """Known issue, pinned so it cannot change without someone deciding to. + + A replaced pod -- a suspended and resumed Kubernetes Job, or a Volcano + restart where an execution profile raises maxRetry above zero -- seeds itself + from the old run's points and starts again at step one, nothing being able to + restore a checkpoint. The series then carries both runs, and two values can + sit at one step contradicting each other. + + Accepted rather than solved: telling the runs apart needs a run identifier on + each point, and dropping the superseded ones loses a failed attempt's history + for good. See the seeding note in callbacks.py. + """ prior = {"train_loss": [{"step": 50, "epoch": 1, "value": 1.0}]} reporter = _RecordingReporter(prior) _make_callback(reporter).report_train_step(step=1, epoch=1, metrics={"loss": 2.0}) - assert reporter.reports[-1]["metrics"]["train_loss"] == [{"step": 1, "epoch": 1, "value": 2.0}] + assert reporter.reports[-1]["metrics"]["train_loss"] == [ + {"step": 50, "epoch": 1, "value": 1.0}, + {"step": 1, "epoch": 1, "value": 2.0}, + ] def test_a_malformed_stored_point_cannot_raise_into_the_training_loop() -> None: @@ -449,14 +421,84 @@ def test_a_malformed_stored_point_cannot_raise_into_the_training_loop() -> None: ] -def test_a_malformed_stored_point_does_not_skew_the_high_water_mark() -> None: - """It is unplaceable, so it cannot be what a later report is compared against.""" - prior = {"train_loss": [{"step": "eight"}, {"step": 10, "epoch": 1, "value": 1.0}]} +def test_a_report_that_changes_no_curve_omits_the_series(reporter: _RecordingReporter) -> None: + """Nothing recorded means nothing to say: the stored copy already matches.""" + callback = _make_callback(reporter) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 0.5}) + callback.report_train_step(step=2, epoch=1, metrics={"grad_norm": float("nan")}) + + assert "metrics" in reporter.reports[0] + assert "metrics" not in reporter.reports[-1] + + +# --------------------------------------------------------------------------- # +# Training start states where the run actually starts +# --------------------------------------------------------------------------- # + + +def test_training_start_states_the_schedule_and_no_position(reporter: _RecordingReporter) -> None: + """It fires before the first step, so it has no position to report yet.""" + _make_callback(reporter).report_training_start(max_steps=940, num_epochs=2) + + report = reporter.reports[-1] + assert report["max_steps"] == 940 + assert report["num_epochs"] == 2 + assert "step" not in report + + +def test_training_start_does_not_reset_the_stored_progress() -> None: + """report_running derives percentage_done from a stated step, and it merges. + + A literal 0 here wrote 0% over whatever progress the task had stored, and + Studio rendered `2/2 (0%)` until the first train step landed -- many minutes + later at log_interval=10. Stating nothing leaves the stored value standing. + """ + prior = {"train_loss": [{"step": 470, "epoch": 2, "value": 0.5}]} reporter = _RecordingReporter(prior) + + _make_callback(reporter).report_training_start(max_steps=940, num_epochs=2) + + assert "step" not in reporter.reports[-1] + assert "percentage_done" not in reporter.reports[-1] + + +# --------------------------------------------------------------------------- # +# A seed that could not be read is not a seed that found nothing +# --------------------------------------------------------------------------- # + + +def test_a_failed_seed_read_withholds_the_series() -> None: + """The accumulator is known incomplete, and the merge replaces a sent key whole. + + Sending it would overwrite the stored history with the fraction of it this + process happens to have. Omitting the key leaves the history standing, which + costs this run's curves and is the recoverable half of the trade. + """ + reporter = _UnreadableReporter() callback = _make_callback(reporter) - callback.report_train_step(step=11, epoch=1, metrics={"loss": 0.7}) + callback.report_train_step(step=1, epoch=1, metrics={"loss": 0.5}) + callback.report_validation(step=1, epoch=1, metrics={"loss": 0.4}) + + assert all("metrics" not in report for report in reporter.reports) + + +def test_a_failed_seed_read_still_reports_progress_and_current_values() -> None: + """Only the accumulated curves are withheld; the scalars merge harmlessly.""" + reporter = _UnreadableReporter() + _make_callback(reporter).report_train_step(step=7, epoch=2, metrics={"loss": 0.5, "lr": 5e-06}) + + report = reporter.reports[-1] + assert report["step"] == 7 + assert report["epoch"] == 2 + assert report["train_loss"] == 0.5 + assert report["train_lr"] == 5e-06 + + +def test_an_empty_seed_read_still_reports_the_series(reporter: _RecordingReporter) -> None: + """The other half of the distinction: nothing stored is safe to write over.""" + _make_callback(reporter).report_train_step(step=1, epoch=1, metrics={"loss": 0.5}) - assert [point["step"] for point in reporter.reports[-1]["metrics"]["train_loss"]] == [10, 11] + assert reporter.reports[-1]["metrics"]["train_loss"] == [{"step": 1, "epoch": 1, "value": 0.5}] def test_series_are_snapshots_not_live_references(reporter: _RecordingReporter) -> None: diff --git a/packages/nmp_customization_common/tests/training/test_progress.py b/packages/nmp_customization_common/tests/training/test_progress.py index 7aedfe366f..a5a64d2a7a 100644 --- a/packages/nmp_customization_common/tests/training/test_progress.py +++ b/packages/nmp_customization_common/tests/training/test_progress.py @@ -14,7 +14,14 @@ from typing import Any +import httpx import pytest +from nemo_platform_plugin.client.errors import ( + InternalServerError, + NemoHTTPError, + NemoTransportError, + NotFoundError, +) from nmp.customization_common.training.progress import JobsServiceProgressReporter SERIES: dict[str, list[dict[str, Any]]] = { @@ -65,7 +72,7 @@ def __init__(self) -> None: class _Task: - def __init__(self, status_details: dict[str, Any]) -> None: + def __init__(self, status_details: Any) -> None: self.status_details = status_details def data(self) -> "_Task": @@ -77,8 +84,12 @@ class _Jobs: def __init__(self) -> None: self.sent: list[dict[str, Any]] = [] - self.stored: dict[str, Any] = {} + #: Typed loosely on purpose: several tests store a blob no annotation + #: describes honestly, because that is what a corrupted read returns. + self.stored: Any = {} self.fetches = 0 + #: Raised by the read instead of answering, for the failure cases. + self.fetch_error: Exception | None = None def client(self) -> Any: harness = self @@ -89,6 +100,8 @@ def update_job_step_task(self, **kwargs: Any) -> None: def get_job_step_task(self, **kwargs: Any) -> _Task: harness.fetches += 1 + if harness.fetch_error is not None: + raise harness.fetch_error return _Task(harness.stored) return _Client() @@ -186,12 +199,16 @@ 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"} + seeded = _reporter(jobs, stored).fetch_current_metrics() + + assert seeded is not None + assert set(seeded) == {"train_loss"} def test_fetch_current_metrics_copies_the_point_lists(jobs: _Jobs) -> None: """The caller appends to what it gets back; it must not alias the response.""" seeded = _reporter(jobs, STORED).fetch_current_metrics() + assert seeded is not None seeded["train_loss"].append({"step": 20, "epoch": 2, "value": 0.4}) assert len(SERIES["train_loss"]) == 1 @@ -202,6 +219,66 @@ def test_fetch_current_metrics_survives_an_unseeded_task(jobs: _Jobs) -> None: assert _reporter(jobs, {}).fetch_current_metrics() == {} +# --------------------------------------------------------------------------- # +# A read that failed is not a read that found nothing +# --------------------------------------------------------------------------- # + + +def _response(status: int) -> httpx.Response: + return httpx.Response(status, request=httpx.Request("GET", "http://jobs/task")) + + +def test_a_missing_task_seeds_from_nothing(jobs: _Jobs) -> None: + """A 404 is the one failure that genuinely means "nothing stored".""" + jobs.fetch_error = NotFoundError(_response(404)) + + assert _reporter(jobs, STORED).fetch_current_metrics() == {} + + +@pytest.mark.parametrize( + "error", + [ + NemoTransportError(httpx.ConnectError("connection refused")), + InternalServerError(_response(500)), + NemoHTTPError(_response(503)), + ], + ids=["transport", "500", "503"], +) +def test_a_failed_read_is_not_reported_as_an_empty_one(jobs: _Jobs, error: Exception) -> None: + """None, not {}: a caller told "nothing stored" overwrites what it could not read. + + A distributed launch racing a briefly unreachable Jobs service lands here, + not on the 404, and the merge replaces a sent key wholesale -- so the two + cases have to be distinguishable at the call site. + """ + jobs.fetch_error = error + + assert _reporter(jobs, STORED).fetch_current_metrics() is None + + +@pytest.mark.parametrize( + "stored", + [ + ["not", "an", "object"], + "junk", + {"metrics": ["not", "an", "object"]}, + {"metrics": "junk"}, + {"metrics": 3}, + ], + ids=["list-blob", "str-blob", "list-metrics", "str-metrics", "int-metrics"], +) +def test_a_malformed_blob_seeds_nothing_instead_of_raising(jobs: _Jobs, stored: Any) -> None: + """This runs from the callback's constructor, which backends build outside any try. + + Raising here kills the training process over a blob that only ever cost its + own points. ``{}`` rather than ``None``: the read succeeded and the data is + unusable, so overwriting it is the repair, not the damage. + """ + jobs.stored = stored + + assert _Reporter().fetch_current_metrics() == {} + + # --------------------------------------------------------------------------- # # Gating # --------------------------------------------------------------------------- # diff --git a/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py b/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py index db16cd29e6..31298844ca 100644 --- a/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py +++ b/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py @@ -86,8 +86,30 @@ def strip_val_prefix(metrics: Mapping[str, Any]) -> dict[str, Any]: ``train_bi_encoder`` adds ``val_acc1`` and ``val_mrr``. Only the callback should be deciding the phase, so the prefix comes off wherever the recipe happened to put one. + + Stripping can collide: a dict carrying both ``val_loss`` and ``loss`` maps + them onto one name, and whichever lands second silently replaces the other. + None of today's recipes do that, but nothing stops one from starting, and the + failure would read as a validation curve charting the wrong quantity. The + prefixed name wins, being the one the recipe marked as validation, and the + collision is logged rather than swallowed. """ - return {name.removeprefix("val_"): value for name, value in metrics.items()} + stripped: dict[str, Any] = {} + for name, value in metrics.items(): + bare = name.removeprefix("val_") + if bare not in stripped: + stripped[bare] = value + continue + # One of the two is the prefixed name -- they cannot both be, having come + # from one dict -- so whether this one wins is just whether it is that one. + prefixed, dropped = (name, bare) if bare != name else (f"val_{bare}", name) + logger.warning( + f"Validation metrics carry both {prefixed!r} and {bare!r}, which strip to one name; " + f"reporting the {prefixed!r} value and dropping {dropped!r}." + ) + if bare != name: + stripped[bare] = value + return stripped class AutomodelRecipeWrapper: diff --git a/services/automodel/tests/tasks/training/backends/test_callbacks.py b/services/automodel/tests/tasks/training/backends/test_callbacks.py index a106865cc8..6b8bb7075b 100644 --- a/services/automodel/tests/tasks/training/backends/test_callbacks.py +++ b/services/automodel/tests/tasks/training/backends/test_callbacks.py @@ -142,9 +142,10 @@ 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) + # No `step`: it fires before the first one, and a literal 0 would write + # 0% over whatever progress the task had stored. reporter.report_running.assert_called_once_with( phase="training", - step=0, max_steps=500, num_epochs=2, ) diff --git a/services/automodel/tests/tasks/training/backends/test_config.py b/services/automodel/tests/tasks/training/backends/test_config.py index 77391e4fe9..6f577897b4 100644 --- a/services/automodel/tests/tasks/training/backends/test_config.py +++ b/services/automodel/tests/tasks/training/backends/test_config.py @@ -16,7 +16,28 @@ sys.modules["nemo_automodel"] = MagicMock() sys.modules["nemo_automodel._transformers"] = MagicMock() sys.modules["nemo_automodel._transformers.registry"] = MagicMock() -sys.modules.setdefault("transformers", MagicMock()) + + +@pytest.fixture(autouse=True) +def _transformers_module(monkeypatch: pytest.MonkeyPatch) -> None: + """Supply `transformers` for the AutoConfig patches without leaking one. + + config.py imports transformers inside the functions that need it, so it is + only wanted while a test runs -- and only when the real package is absent, + which is the training image's arrangement, not necessarily the test env's. + + Installed at module scope with setdefault it won permanently whenever this + file happened to import before anything had loaded the real package, and the + MagicMock then leaked across the whole xdist worker. services/unsloth does + `from transformers import TrainerCallback` at call time, so its + HfTrainerProgressCallback became a mock subclass whose hooks did nothing -- + two of its tests passed vacuously or failed depending on collection order. + """ + try: + import transformers # noqa: F401 + except ImportError: + monkeypatch.setitem(sys.modules, "transformers", MagicMock()) + from nmp.automodel.tasks.training.backends.config import ( # noqa: E402 _configure_chat_dataset, diff --git a/services/automodel/tests/tasks/training/backends/test_finetune.py b/services/automodel/tests/tasks/training/backends/test_finetune.py index 27efd33726..0c6d78c4bb 100644 --- a/services/automodel/tests/tasks/training/backends/test_finetune.py +++ b/services/automodel/tests/tasks/training/backends/test_finetune.py @@ -12,6 +12,7 @@ from __future__ import annotations import importlib +import logging import sys from collections.abc import Iterator from types import ModuleType @@ -81,3 +82,22 @@ def test_only_a_leading_occurrence_is_removed(finetune: ModuleType) -> None: def test_an_empty_metric_dict_survives(finetune: ModuleType) -> None: assert finetune.strip_val_prefix({}) == {} + + +def test_a_collision_keeps_the_prefixed_value(finetune: ModuleType) -> None: + """Stripping can map two names onto one, and the wrong winner is undetectable. + + No recipe reports both today, but the loser vanishes silently and the + validation curve then charts whatever else was in the dict. The prefixed name + wins, being the one the recipe marked as validation. + """ + assert finetune.strip_val_prefix({"val_loss": 0.5, "loss": 0.7}) == {"loss": 0.5} + assert finetune.strip_val_prefix({"loss": 0.7, "val_loss": 0.5}) == {"loss": 0.5} + + +def test_a_collision_is_logged(finetune: ModuleType, caplog: pytest.LogCaptureFixture) -> None: + """Silently dropping a metric is how this would go unnoticed for a release.""" + with caplog.at_level(logging.WARNING): + finetune.strip_val_prefix({"val_loss": 0.5, "loss": 0.7}) + + assert "val_loss" in caplog.text 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 4b62ee1f64..f3e5a12ab7 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 @@ -22,10 +22,15 @@ # notion of a reporting cadence, so it is derived from val_period. _REPORTS_PER_VAL_PERIOD = 10 -# Upper bound on progress reports for one run, whatever the validation cadence -# asks for. Every train report resends every accumulated series in full and the -# Jobs service stores the blob twice, so the cost of a report is not constant -- -# see the payload note in nmp.customization_common.training.callbacks. +# Upper bound on progress reports for one run, per reporting path, whatever the +# validation cadence asks for. Every report resends every accumulated series in +# full and the Jobs service stores the blob twice, so the cost of a report is not +# constant -- see the payload note in nmp.customization_common.training.callbacks. +# +# Per path, so a run making full use of both is bounded by twice this. That is a +# bound either way; what matters is that no path is unbounded, and applying one +# budget across two cadences that fire independently would mean whichever ran +# first starved the other. _MAX_REPORTS_PER_RUN = 200 # NeMo-RL's metric dicts are forwarded whole. There is no allow-list: the @@ -59,6 +64,34 @@ def resolve_log_interval(val_period: int | None, max_steps: int) -> int: return max(per_val_period, per_run, 1) +def resolve_val_report_interval(val_period: int | None, max_steps: int) -> int: + """Validation passes between validation reports. + + The same bound as :func:`resolve_log_interval`, applied to the other report + path. A validation report costs exactly what a train report costs -- every + series resent in full, stored twice by the Jobs service -- so capping only + the train side bounds nothing: ``val_check_interval=1`` is reachable through + ``compute_val_check_interval``, and it produces one validation pass per step. + A 20k-step run then made 200 train reports and 20,000 validation ones. + + Returns 1 for any ordinary cadence, so nothing in the existing regime moves: + validating every 100 steps over 20k steps is 200 passes, which is already the + cap. Only a cadence that would exceed it thins the passes out. + """ + period = max(val_period or 0, 1) + passes = max(max_steps, 0) // period + return max((passes + _MAX_REPORTS_PER_RUN - 1) // _MAX_REPORTS_PER_RUN, 1) + + +def _val_dataset_name(prefix: str) -> str: + """The dataloader name NeMo-RL suffixed onto a validation prefix. + + ``validation`` carries none; ``validation-0`` and ``validation/nemo_gym`` + name their dataloader, the separator depending on the caller. + """ + return prefix[len("validation") :].lstrip("-/") + + 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: @@ -84,6 +117,7 @@ def __init__( log_interval: int = 10, max_steps: int | None = None, num_epochs: int | None = None, + val_report_interval: int = 1, ): """Initialize the NemoRL logger. @@ -93,35 +127,47 @@ def __init__( log_interval: Number of steps between progress updates. max_steps: Total number of training steps (optional, used for progress reporting). num_epochs: Total number of epochs (optional, used for progress reporting). + val_report_interval: Validation passes between validation reports. 1 reports + every pass, which is what any ordinary validation cadence resolves to. Raises: - ValueError: If ``steps_per_epoch`` or ``log_interval`` is < 1. Both are - used as divisors/moduli in ``log_metrics`` (epoch derivation and - log-interval throttling), so non-positive values are rejected up - front to fail fast instead of raising ZeroDivisionError mid-training. + ValueError: If ``steps_per_epoch``, ``log_interval`` or + ``val_report_interval`` is < 1. All three are used as + divisors/moduli in ``log_metrics`` (epoch derivation and the two + throttles), so non-positive values are rejected up front to fail + fast instead of raising ZeroDivisionError mid-training. """ if steps_per_epoch < 1: raise ValueError(f"steps_per_epoch must be >= 1, got {steps_per_epoch}") if log_interval < 1: raise ValueError(f"log_interval must be >= 1, got {log_interval}") + if val_report_interval < 1: + raise ValueError(f"val_report_interval must be >= 1, got {val_report_interval}") self._job_ctx = job_ctx or NMPJobContext.from_env() self._log_interval = log_interval + self._val_report_interval = val_report_interval self._max_steps = max_steps self._num_epochs = num_epochs self._steps_per_epoch = steps_per_epoch self._callback = TrainingProgressCallback(JobsServiceProgressReporter(self._job_ctx)) - # Track best metrics for monitoring - self._best_metric_value = float("inf") - 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 + # Validation passes seen so far, which is what the validation throttle + # counts against: passes arrive on their own cadence, so a step-based + # modulus would skip whole passes rather than thin them evenly. + self._val_passes = 0 + + # The prefix whose metrics keep the bare names; see _namespace_validation. + self._primary_val_prefix: str | None = None + + # Reports built but withheld by a throttle, at most one of each kind. + # Flushed on close() so a run's final train step and final validation pass + # are reported even when neither lands on its interval -- otherwise the + # last recorded values are stale. + self._pending: dict[str, dict[str, Any]] = {} _logger.info( f"Initialized NemoRLLogger with jobs_url={self._job_ctx.jobs_url}, " @@ -155,6 +201,7 @@ def for_schedule( 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), + val_report_interval=resolve_val_report_interval(val_period, max_steps), max_steps=max_steps, num_epochs=num_epochs, ) @@ -198,29 +245,63 @@ def log_metrics( # 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: - self._callback.report_train_step(**report) - self._pending_train_report = None + self._send("train", report) else: - self._pending_train_report = report + self._pending["train"] = report # Validation reports whatever the pass produced. There is one validation - # log per pass and no rollout twin to tell apart, so requiring a `loss` - # here bought nothing and cost whole passes: GRPO validates on - # `accuracy` and `avg_length` and reports no loss at all, so every - # validation it ran went unrecorded. The gate is only against a hollow - # report -- a pass whose metrics are all histograms says nothing. + # log per pass per dataloader and no rollout twin to tell apart, so + # requiring a `loss` here bought nothing and cost whole passes: GRPO + # validates on `accuracy` and `avg_length` and reports no loss at all, so + # every validation it ran went unrecorded. The gate is only against a + # hollow report -- a pass whose metrics are all histograms says nothing. elif prefix and prefix.startswith("validation"): if any(is_chartable(value) for value in metrics.values()): - self._callback.report_validation(step=step, epoch=epoch, metrics=dict(metrics)) - # Track best validation loss, when the algorithm reports one. - if is_chartable(metrics.get("loss")): - val_loss = float(metrics["loss"]) - if val_loss < self._best_metric_value: - self._best_metric_value = val_loss - self._best_epoch = epoch + self._val_passes += 1 + report = { + "step": step, + "epoch": epoch, + "metrics": self._namespace_validation(prefix, metrics), + } + if self._val_passes % self._val_report_interval == 0: + self._send("validation", report) + else: + self._pending["validation"] = report _logger.debug(f"log_metrics: step={step}, prefix={prefix}, metrics={metrics}") + def _namespace_validation(self, prefix: str, metrics: Mapping[str, Any]) -> dict[str, Any]: + """Fold the dataloader name into each metric name, past the first set. + + ``validate()`` loops over ``val_dataloader.items()`` and logs once per + dataset, every call at the same step under ``f"validation-{name}"``. + Forwarded as-is, two datasets' ``loss`` interleave as two points at one + step in a single ``val_loss`` series -- the collision the ``_`` + rule exists to prevent, one level further down. Automodel's + ``val_dataloaders`` loop has the same shape. + + The first prefix seen keeps the bare names, so ``val_loss`` stays + ``val_loss`` on the ordinary single-dataset run. NeMo-RL names that + dataloader too, so keying on "did a name arrive" would rename the common + case and take Studio's curve with it. Iteration order over the dataloader + dict is stable within a run and across a resume of the same config, so a + dataset keeps whichever namespace it started in. + """ + if self._primary_val_prefix is None: + self._primary_val_prefix = prefix + if prefix == self._primary_val_prefix: + return dict(metrics) + dataset = _val_dataset_name(prefix) or prefix + return {f"{dataset}_{name}": value for name, value in metrics.items()} + + def _send(self, kind: str, report: dict[str, Any]) -> None: + """Report one built payload, retiring anything withheld of that kind.""" + if kind == "train": + self._callback.report_train_step(**report) + else: + self._callback.report_validation(**report) + self._pending.pop(kind, None) + def log_hyperparams(self, params: Mapping[str, Any]) -> None: """Log hyperparameters and report training start. @@ -274,23 +355,25 @@ def close(self) -> None: if self._closed: return self._closed = True - self._flush_pending_train_report() + self._flush_pending() _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. + def _flush_pending(self) -> None: + """Report whatever the throttles withheld, oldest step first. + + Step order so a series' points are appended in the order they were + produced, rather than in whichever order the kinds happen to iterate. 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}") + pending, self._pending = self._pending, {} + for kind, report in sorted(pending.items(), key=lambda item: item[1]["step"]): + try: + self._send(kind, report) + except Exception as exc: # pragma: no cover - defensive, shutdown path + _logger.warning(f"Failed to flush final {kind} report: {exc}") def __del__(self): """Cleanup when the logger is destroyed.""" diff --git a/services/rl/tests/test_nemo_rl_logger.py b/services/rl/tests/test_nemo_rl_logger.py index f0c948cf0a..0f584a3701 100644 --- a/services/rl/tests/test_nemo_rl_logger.py +++ b/services/rl/tests/test_nemo_rl_logger.py @@ -13,7 +13,6 @@ from __future__ import annotations -import importlib.machinery import importlib.util import sys import types @@ -25,29 +24,28 @@ # 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. +# +# The stubs are removed again as soon as the import that needs them is done. Left in +# sys.modules they outlived this file and leaked a hollow `nemo_rl` package into +# every other test sharing the xdist worker -- a sibling test_dpo_config had already +# had to be written around exactly that. nemo_rl_logger binds LoggerInterface at +# import time, so nothing downstream needs the stub to still be there. +_stubbed: list[str] = [] 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") + _logger_mod = types.ModuleType("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) + for _name, _module in ( + ("nemo_rl", types.ModuleType("nemo_rl")), + ("nemo_rl.utils", types.ModuleType("nemo_rl.utils")), + ("nemo_rl.utils.logger", _logger_mod), + ): + if _name not in sys.modules: + sys.modules[_name] = _module + _stubbed.append(_name) 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 @@ -55,8 +53,12 @@ def _stub(name: str) -> types.ModuleType: NemoRLLogger, resolve_log_interval, resolve_steps_per_epoch, + resolve_val_report_interval, ) +for _name in _stubbed: + del sys.modules[_name] + class _RecordingCallback: """Stands in for TrainingProgressCallback, capturing what the logger forwards.""" @@ -65,6 +67,9 @@ def __init__(self) -> None: self.train_steps: list[dict[str, Any]] = [] self.validations: list[dict[str, Any]] = [] self.training_starts: list[dict[str, Any]] = [] + #: Report kinds in arrival order -- the real callback prunes against the + #: step of the last one it saw, so the sequence is part of the contract. + self.order: list[str] = [] self.closed = False def report_training_start(self, max_steps: int, num_epochs: int) -> None: @@ -72,9 +77,11 @@ def report_training_start(self, max_steps: int, num_epochs: int) -> None: def report_train_step(self, step, epoch, metrics, *, backend=None): self.train_steps.append({"step": step, "epoch": epoch, "metrics": metrics}) + self.order.append("train") def report_validation(self, step, epoch, metrics, *, backend=None): self.validations.append({"step": step, "epoch": epoch, "metrics": metrics}) + self.order.append("validation") def close(self) -> None: self.closed = True @@ -145,15 +152,16 @@ class _Histogram: # --------------------------------------------------------------------------- # -def test_module_stub_does_not_break_find_spec() -> None: - """The stub installed at import time outlives this module; it must be inert. +def test_the_import_stub_does_not_outlive_this_module() -> None: + """A stub left in sys.modules leaks a hollow nemo_rl across the xdist worker. - 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_dpo_config had to be rewritten around. + Worth asserting because the failure is silent and lands somewhere else: an + unrelated `import nemo_rl.algorithms.dpo` would get a module with no + attributes rather than a clean ImportError, and which tests broke would + depend on collection order. Vacuously true in the training image, where + nothing was stubbed because the real package is installed. """ - assert importlib.util.find_spec("nemo_rl") is not None + assert all(name not in sys.modules for name in _stubbed) # --------------------------------------------------------------------------- # @@ -324,6 +332,17 @@ def test_close_with_nothing_pending_reports_nothing(callback: _RecordingCallback assert callback.train_steps == [] +def test_close_still_flushes_a_step_ahead_of_the_last_validation(callback: _RecordingCallback) -> None: + """The ordinary interleaving: nothing overtook it, so the flush still runs.""" + logger = _make_logger(log_interval=10) + logger.log_metrics({"loss": 0.4}, step=13, prefix="validation") + logger.log_metrics(TRAIN_METRICS, step=14, prefix="train") + + logger.close() + + assert [r["step"] for r in callback.train_steps] == [14] + + # --------------------------------------------------------------------------- # # Schedule resolution — one formula for both drivers # --------------------------------------------------------------------------- # @@ -472,25 +491,122 @@ def test_a_validation_pass_without_a_loss_is_still_reported(callback: _Recording assert callback.validations[0]["metrics"] == REWARD_VALIDATION_METRICS -def test_a_loss_free_validation_leaves_the_best_metric_alone(callback: _RecordingCallback) -> None: - """Best-so-far tracks the validation loss, so a pass without one says nothing.""" +def test_a_loss_free_validation_is_still_a_reported_pass(callback: _RecordingCallback) -> None: + """A pass that scores on rewards alone reports, like any other.""" logger = _make_logger() logger.log_metrics({"loss": 0.4}, step=10, prefix="validation") logger.log_metrics(REWARD_VALIDATION_METRICS, step=20, prefix="validation") assert len(callback.validations) == 2 - assert logger._best_metric_value == 0.4 - assert logger._best_epoch == 1 -def test_best_validation_loss_tracks_minimum(callback: _RecordingCallback) -> None: +def test_one_dataset_keeps_the_bare_metric_names(callback: _RecordingCallback) -> None: + """NeMo-RL names the dataloader even when there is only one, so val_loss must survive.""" + logger = _make_logger() + logger.log_metrics({"loss": 0.4, "accuracy": 0.9}, step=10, prefix="validation-train_ds") + + assert callback.validations[0]["metrics"] == {"loss": 0.4, "accuracy": 0.9} + + +def test_a_second_dataset_does_not_share_the_first_ones_series(callback: _RecordingCallback) -> None: + """validate() logs once per dataloader, all at one step, all under `validation-*`. + + Passed through as-is, two datasets' `loss` land as two points at the same + step in one val_loss series -- the collision the _ prefix rule exists + to prevent, a level further down. + """ 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") + logger.log_metrics({"loss": 0.4}, step=10, prefix="validation-train_ds") + logger.log_metrics({"loss": 0.9}, step=10, prefix="validation-heldout") + + assert callback.validations[0]["metrics"] == {"loss": 0.4} + assert callback.validations[1]["metrics"] == {"heldout_loss": 0.9} + + +def test_a_dataset_keeps_the_namespace_it_started_in(callback: _RecordingCallback) -> None: + """Otherwise a curve would change names partway through the run.""" + logger = _make_logger() + for step in (10, 20): + logger.log_metrics({"loss": 0.4}, step=step, prefix="validation-train_ds") + logger.log_metrics({"loss": 0.9}, step=step, prefix="validation-heldout") + + assert [set(v["metrics"]) for v in callback.validations] == [ + {"loss"}, + {"heldout_loss"}, + {"loss"}, + {"heldout_loss"}, + ] + + +# --------------------------------------------------------------------------- # +# Validation reporting is bounded too +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + "val_period,max_steps,expected", + [ + (100, 100, 1), # 1 pass + (100, 20_000, 1), # 200 passes -- exactly the cap, still every pass + (1, 20_000, 100), # 20,000 passes -- thinned to 200 + (5, 20_000, 20), + (None, 20_000, 100), # treated as "every step" + (0, 0, 1), + ], +) +def test_resolve_val_report_interval(val_period: int | None, max_steps: int, expected: int) -> None: + assert resolve_val_report_interval(val_period, max_steps) == expected + + +def test_the_run_length_cap_bounds_validation_reports_too(callback: _RecordingCallback) -> None: + """Capping only the train side bounded nothing. + + `val_check_interval=1` is reachable, and it validates on every step: the + train reports were held to 200 while validation reported all 20,000, each + one resending every accumulated series in full. + """ + max_steps = 20_000 + logger = NemoRLLogger.for_schedule(max_steps=max_steps, num_epochs=1, val_period=1) + for step in _driver_steps(max_steps): + logger.log_metrics({"loss": 0.5}, step=step, prefix="validation") + + assert len(callback.validations) <= _MAX_REPORTS_PER_RUN + assert len(callback.validations) >= _MAX_REPORTS_PER_RUN // 2, "still a usable curve" + + +def test_an_ordinary_validation_cadence_reports_every_pass(callback: _RecordingCallback) -> None: + """Nothing in the existing regime moves: 200 passes is already the cap.""" + logger = NemoRLLogger.for_schedule(max_steps=20_000, num_epochs=1, val_period=100) + for step in range(100, 20_001, 100): + logger.log_metrics({"loss": 0.5}, step=step, prefix="validation") + + assert len(callback.validations) == 200 + + +def test_close_flushes_the_withheld_final_validation(callback: _RecordingCallback) -> None: + """The last pass is the one worth having; the throttle must not eat it.""" + logger = _make_logger(val_report_interval=10) + for step in (10, 20, 30): + logger.log_metrics({"loss": 0.5}, step=step, prefix="validation") + + assert callback.validations == [] + + logger.close() + + assert [v["step"] for v in callback.validations] == [30] + + +def test_pending_reports_flush_in_step_order(callback: _RecordingCallback) -> None: + """Both go to a callback that reads a report behind the last one as a rewind.""" + logger = _make_logger(log_interval=100, val_report_interval=100) + logger.log_metrics(TRAIN_METRICS, step=18, prefix="train") + logger.log_metrics({"loss": 0.5}, step=19, prefix="validation") + + logger.close() - assert logger._best_metric_value == 0.2 - assert logger._best_epoch == 2 + assert [r["step"] for r in callback.train_steps] == [18] + assert [v["step"] for v in callback.validations] == [19] + assert callback.order == ["train", "validation"] @pytest.mark.parametrize("prefix", ["validation", "validation-0", "validation/nemo_gym"]) diff --git a/services/unsloth/src/nmp/unsloth/tasks/training/backends/hf_trainer_callback.py b/services/unsloth/src/nmp/unsloth/tasks/training/backends/hf_trainer_callback.py index 016fd6b487..700cadeee7 100644 --- a/services/unsloth/src/nmp/unsloth/tasks/training/backends/hf_trainer_callback.py +++ b/services/unsloth/src/nmp/unsloth/tasks/training/backends/hf_trainer_callback.py @@ -14,6 +14,29 @@ def _epoch_from_value(raw_epoch: float | int, num_epochs: int) -> int: return max(1, min(num_epochs, math.ceil(float(raw_epoch)))) +def _as_float(value: Any) -> Any: + """Coerce a logged value to a float, leaving it alone if it will not go. + + The Trainer does not promise a builtin: some paths log ``grad_norm`` as a + 0-dim tensor rather than calling ``.item()`` on it, and a tensor is not a + ``numbers.Real``, so the callback's chartable filter drops it and the curve + silently never appears. ``float()`` converts anything with ``__float__``, + which covers those. + + A value it cannot convert is passed through rather than raised on or dropped + here: deciding what belongs in a series is the callback's job, and it already + drops what it cannot chart. ``None`` goes through untouched for the same + reason -- the trainer omits ``learning_rate`` and ``grad_norm`` on some + steps, and that is not an error. + """ + if value is None: + return None + try: + return float(value) + except (TypeError, ValueError): + return value + + def create_hf_trainer_progress_callback( progress_callback: TrainingProgressCallback, *, @@ -54,9 +77,9 @@ def on_log( # so the series is `train_lr` regardless of who reported it. The # callback drops whichever of these the trainer did not produce. metrics={ - "loss": float(logs["loss"]), - "lr": logs.get("learning_rate"), - "grad_norm": logs.get("grad_norm"), + "loss": _as_float(logs["loss"]), + "lr": _as_float(logs.get("learning_rate")), + "grad_norm": _as_float(logs.get("grad_norm")), }, backend=self._backend, ) @@ -76,7 +99,7 @@ def on_evaluate( self._progress.report_validation( step=int(state.global_step), epoch=_epoch_from_value(epoch_raw, self._num_epochs), - metrics={"loss": float(metrics["eval_loss"])}, + metrics={"loss": _as_float(metrics["eval_loss"])}, backend=self._backend, ) diff --git a/services/unsloth/tests/test_callbacks.py b/services/unsloth/tests/test_callbacks.py index 3adf700447..92737d6ea7 100644 --- a/services/unsloth/tests/test_callbacks.py +++ b/services/unsloth/tests/test_callbacks.py @@ -60,9 +60,10 @@ 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) + # No `step`: it fires before the first one, and a literal 0 would write + # 0% over whatever progress the task had stored. reporter.report_running.assert_called_once_with( phase="training", - step=0, max_steps=500, num_epochs=2, backend="unsloth", diff --git a/services/unsloth/tests/test_hf_trainer_callback.py b/services/unsloth/tests/test_hf_trainer_callback.py index 6f6c4d39f6..2b9cb593b3 100644 --- a/services/unsloth/tests/test_hf_trainer_callback.py +++ b/services/unsloth/tests/test_hf_trainer_callback.py @@ -13,6 +13,16 @@ ) +class _FloatLike: + """Converts under float() but is not a numbers.Real -- a 0-dim tensor's shape.""" + + def __init__(self, value: float) -> None: + self._value = value + + def __float__(self) -> float: + return self._value + + class TestEpochFromValue: def test_fractional_epoch_maps_to_one(self): assert _epoch_from_value(0.01314, 1) == 1 @@ -79,3 +89,45 @@ def test_on_evaluate_reports_validation(self, progress: tuple[TrainingProgressCa assert kwargs["phase"] == "validation" assert kwargs["val_loss"] == 1.75 assert kwargs["metrics"]["val_loss"][-1]["value"] == 1.75 + + def test_non_builtin_scalars_are_coerced(self, progress: tuple[TrainingProgressCallback, MagicMock]) -> None: + """Some Trainer paths log a 0-dim tensor rather than calling .item() on it. + + A tensor is not a numbers.Real, so the callback's chartable filter drops + it and the curve never appears -- with no log line saying why. Anything + with __float__ converts, which covers those. + """ + callback, reporter = progress + hf_callback = create_hf_trainer_progress_callback(callback) + + args = MagicMock(num_train_epochs=1) + state = MagicMock(max_steps=77, global_step=8, epoch=0.1) + hf_callback.on_train_begin(args, state, MagicMock()) + + hf_callback.on_log( + args, + state, + MagicMock(), + logs={"loss": 2.89, "learning_rate": _FloatLike(5e-5), "grad_norm": _FloatLike(10.6), "epoch": 0.1}, + ) + + kwargs = reporter.report_running.call_args.kwargs + assert kwargs["train_lr"] == 5e-5 + assert kwargs["train_grad_norm"] == 10.6 + + def test_a_value_that_will_not_convert_costs_only_its_own_curve( + self, progress: tuple[TrainingProgressCallback, MagicMock] + ) -> None: + """Deciding what belongs in a series is the callback's job, not the bridge's.""" + callback, reporter = progress + hf_callback = create_hf_trainer_progress_callback(callback) + + args = MagicMock(num_train_epochs=1) + state = MagicMock(max_steps=77, global_step=8, epoch=0.1) + hf_callback.on_train_begin(args, state, MagicMock()) + + hf_callback.on_log(args, state, MagicMock(), logs={"loss": 2.89, "grad_norm": "oops", "epoch": 0.1}) + + kwargs = reporter.report_running.call_args.kwargs + assert kwargs["train_loss"] == 2.89 + assert "train_grad_norm" not in kwargs From f21d67ec1459352b47a6ec5a935dc2a64eeda558 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 14 Aug 2026 15:02:25 -0400 Subject: [PATCH 21/23] fix(rl): derive the train report cadence from run length alone resolve_log_interval took the coarser of two floors: ~10 reports across one validation period, and _MAX_REPORTS_PER_RUN reports across the whole run. The first has no bearing on the question it was answering. How often someone wants the loss curve and the progress bar to move is unrelated to how often the run validates, and the val_period term decided the interval on nearly every real configuration, so the cap almost never engaged. Its effect was to hold the report count at ten per epoch at every scale. compute_val_check_interval returns steps_per_epoch when the user sets no val_check_interval, and returns it from an early branch that the later clamps never reach -- so on the default path val_period *was* the epoch. A one-epoch run drew its whole curve from ten points whether it was 32 steps or 20,000. report_running derives percentage_done from the step a train report states, so the progress bar advanced ten times too: once an hour on a 20k-step run at two seconds a step. The coupling also ran the wrong way. val_period is capped at steps_per_epoch, and a larger one meant a coarser interval, so choosing to validate less often -- which is what you do when validation is expensive -- made the training curve worse. _MAX_REPORTS_PER_RUN is now the only floor, and resolution is flat across three orders of magnitude: 32 steps -> 32 points, 3,125 -> 195, 20,000 -> 200, against ten for each of them before. Runs at or above 20 epochs were already governed by the cap and do not move, nor do runs shorter than the budget, which report every step either way. The cap keeps its value and gains the justification it should have had. Two hundred is roughly the number of points a chart a few hundred pixels wide can draw distinctly, and past which the extra points cost more than they show. That it also bounds a payload growing as the square of the report count is why it is a ceiling rather than a target -- and it is an argument that survives the transport learning to append deltas, where the cost argument alone would have retired with it. Signed-off-by: Albert Cui --- .../backends/nemo_rl/nemo_rl_logger.py | 64 +++++++++-------- services/rl/tests/test_nemo_rl_logger.py | 70 +++++++++++-------- 2 files changed, 76 insertions(+), 58 deletions(-) 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 f3e5a12ab7..24536be5cc 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 @@ -18,14 +18,18 @@ _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 - -# Upper bound on progress reports for one run, per reporting path, whatever the -# validation cadence asks for. Every report resends every accumulated series in -# full and the Jobs service stores the blob twice, so the cost of a report is not -# constant -- see the payload note in nmp.customization_common.training.callbacks. +# Progress reports for one run, per reporting path. Reporting is throttled to +# whatever cadence lands about this many across the run, so a curve has the same +# resolution whether the run is 300 steps or 30,000. Roughly the number of points a +# chart a few hundred pixels wide can draw distinctly, and past which the extra +# points cost more than they show. +# +# It is also what bounds the cost, which is why it is a ceiling and not a target. +# Every report resends every accumulated series in full and the Jobs service stores +# the blob twice, so upload grows as the square of the report count -- see the +# payload note in nmp.customization_common.training.callbacks. Capping the count is +# what keeps that finite; the resolution argument is what makes 200 the right +# number rather than merely a small one. # # Per path, so a run making full use of both is bounded by twice this. That is a # bound either way; what matters is that no path is unbounded, and applying one @@ -45,23 +49,27 @@ # one rides in the same dict as the scalars. -def resolve_log_interval(val_period: int | None, max_steps: int) -> int: - """Steps between progress reports: the coarser of two floors. +def resolve_log_interval(max_steps: int) -> int: + """Steps between progress reports: enough steps for ``_MAX_REPORTS_PER_RUN``. + + Run length is the only input. The cadence used to be derived from + ``val_period`` as well, targeting ~10 reports per validation period, and that + term decided the interval on nearly every real configuration. It had no + bearing on the question: how often someone wants the curve and the progress + bar to move is unrelated to how often the run validates. - The first targets ~10 reports per validation period, which is the cadence - someone watching the job expects. The second bounds the whole run at - ``_MAX_REPORTS_PER_RUN`` reports, because each report resends every series in - full -- so upload and stored-blob writes both grow as the square of the - report count, and it is the report count, not the step count, that drives it. + Its effect was to hold the report count at ten per epoch at every scale. + ``compute_val_check_interval`` returns ``steps_per_epoch`` when the user sets + no ``val_check_interval``, so on the default path ``val_period`` *was* the + epoch: a 20,000-step run drew its loss curve from ten points. The coupling + was inverted, too -- validating less often, which is what you do when + validation is expensive, made the training curve coarser. - The second floor is not a corner case guard: ``val_period`` is the user's - ``val_check_interval``, and any value below ``_REPORTS_PER_VAL_PERIOD`` -- - "validate every 5 steps" is an ordinary request -- floors the first one to - zero, which clamps to 1 and reports every step of an arbitrarily long run. + ``report_running`` derives ``percentage_done`` from the step a train report + states, so this sets the granularity of the progress bar as well as of the + chart. Ten reports across an eleven-hour run is ten movements of the bar. """ - per_val_period = (val_period or 0) // _REPORTS_PER_VAL_PERIOD - per_run = (max(max_steps, 0) + _MAX_REPORTS_PER_RUN - 1) // _MAX_REPORTS_PER_RUN - return max(per_val_period, per_run, 1) + return max((max(max_steps, 0) + _MAX_REPORTS_PER_RUN - 1) // _MAX_REPORTS_PER_RUN, 1) def resolve_val_report_interval(val_period: int | None, max_steps: int) -> int: @@ -187,20 +195,20 @@ def for_schedule( ) -> Self: """Build a logger from a NeMo-RL training schedule. - The arithmetic lives here rather than in each driver. DPO's copy read - ``(val_period // 10) + 1``, where the ``+1`` was a divide-by-zero guard - that also skewed every value it produced, and raised outright when - ``val_period`` was None. Owning it here fixes both and gives any further - algorithm one place to call. + The arithmetic lives here rather than in each driver, which is where DPO + derived its own ``log_interval`` from ``val_period`` and where any further + algorithm would have copied it. One place to call, and one place to fix. Args: steps_per_epoch: Authoritative value when the algorithm config carries one (DPO does); otherwise derived from max_steps and num_epochs. + val_period: Steps between validation passes. Sets the cadence of the + validation reports only; the train cadence is run length alone. """ 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), + log_interval=resolve_log_interval(max_steps), val_report_interval=resolve_val_report_interval(val_period, max_steps), max_steps=max_steps, num_epochs=num_epochs, diff --git a/services/rl/tests/test_nemo_rl_logger.py b/services/rl/tests/test_nemo_rl_logger.py index 0f584a3701..1844398100 100644 --- a/services/rl/tests/test_nemo_rl_logger.py +++ b/services/rl/tests/test_nemo_rl_logger.py @@ -349,45 +349,51 @@ def test_close_still_flushes_a_step_ahead_of_the_last_validation(callback: _Reco @pytest.mark.parametrize( - "val_period,max_steps,expected", + "max_steps,expected", [ - # Driven by val_period: ~10 reports across one validation period. - (100, 200, 10), - (10, 200, 1), - (5, 200, 1), # floors to 0 -> clamped - (1, 200, 1), - (0, 200, 1), - (None, 200, 1), # val_period is Optional - # Driven by the run-length cap, once val_period has stopped bounding - # anything. This is the regime a small val_check_interval lands in. - (5, 2_000, 10), - (5, 20_000, 100), - (None, 20_000, 100), - # Whichever floor is coarser wins; here it is val_period's. - (10_000, 20_000, 1_000), - # A run shorter than the cap is never throttled past its own length. - (0, 1, 1), + (0, 1), # degenerate, but the clamp keeps it a usable modulus + (1, 1), + (200, 1), # a run no longer than the budget reports every step + (313, 2), # 157 reports, not 313 + (2_000, 10), + (20_000, 100), + (20_001, 101), # ceiling division: never rounds down past the budget ], ) -def test_resolve_log_interval(val_period: int | None, max_steps: int, expected: int) -> None: - assert resolve_log_interval(val_period, max_steps) == expected +def test_resolve_log_interval(max_steps: int, expected: int) -> None: + assert resolve_log_interval(max_steps) == expected -def test_the_run_length_cap_bounds_the_report_count(callback: _RecordingCallback) -> None: - """val_period alone does not bound reporting, and the report is what costs. +@pytest.mark.parametrize("val_period", [None, 0, 1, 5, 100, 20_000]) +def test_the_train_cadence_does_not_move_with_the_validation_cadence( + callback: _RecordingCallback, val_period: int | None +) -> None: + """How often a run validates says nothing about how often it should report. - `val_check_interval=5` is an ordinary request, and it floors the - reports-per-validation-period term to zero -- so before the cap this run - reported all 20,000 steps. Each report resends every series in full, so that - is quadratic in upload and in stored-blob writes, not linear. + The two were coupled, and the coupling ran the wrong way: validating less + often -- what you do when validation is expensive -- made the training curve + coarser. Run length is the only thing that sets the train cadence now. + """ + logger = NemoRLLogger.for_schedule(max_steps=20_000, num_epochs=1, val_period=val_period) + + assert logger._log_interval == 100 + + +def test_a_long_run_on_the_default_cadence_draws_a_usable_curve(callback: _RecordingCallback) -> None: + """The regression the val_period term caused, at the scale where it mattered. + + `compute_val_check_interval` returns steps_per_epoch when the user sets no + `val_check_interval`, so a one-epoch run reached here with + `val_period == max_steps`. Targeting ~10 reports per validation period then + drew this whole curve from ten points, and moved the progress bar ten times, + however long the run was. """ max_steps = 20_000 - logger = NemoRLLogger.for_schedule(max_steps=max_steps, num_epochs=1, val_period=5) + logger = NemoRLLogger.for_schedule(max_steps=max_steps, num_epochs=1, val_period=max_steps) for step in _driver_steps(max_steps): logger.log_metrics(TRAIN_METRICS, step=step, prefix="train") - assert len(callback.train_steps) <= _MAX_REPORTS_PER_RUN - assert len(callback.train_steps) >= _MAX_REPORTS_PER_RUN // 2, "still a usable curve, not a throttle to nothing" + assert len(callback.train_steps) == _MAX_REPORTS_PER_RUN @pytest.mark.parametrize( @@ -406,10 +412,14 @@ def test_resolve_steps_per_epoch(max_steps: int, num_epochs: int | None, explici def test_for_schedule_builds_a_consistent_logger(callback: _RecordingCallback) -> None: - """DPO's own formula produced 11 here, via a `+1` that skewed every value.""" + """A 100-step run is shorter than the report budget, so every step reports. + + DPO's own formula produced 11 here -- one report for the whole run, from a + `+1` that was a divide-by-zero guard and skewed every value it produced. + """ logger = NemoRLLogger.for_schedule(max_steps=100, num_epochs=4, val_period=100) - assert logger._log_interval == 10 + assert logger._log_interval == 1 assert logger._steps_per_epoch == 25 assert logger._max_steps == 100 From 5d86886c98da387452232bc73e5107353c830494 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 14 Aug 2026 15:32:42 -0400 Subject: [PATCH 22/23] update docstrings to reflect seeding mechanism Signed-off-by: Albert Cui --- .../training/callbacks.py | 41 +++++++++++++------ .../tests/training/test_callbacks.py | 20 +++++---- 2 files changed, 41 insertions(+), 20 deletions(-) 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 b76350bcc2..2f98fbca14 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 @@ -71,19 +71,34 @@ suspend and resume a Kubernetes Job, and the Volcano backend restarts on ``PodFailed`` when an execution profile raises ``maxRetry`` above its default of zero -- the new process seeds itself from the old one's points and then -appends its own. Nothing restores a checkpoint, so it starts again at step one -and the series carries two values at each repeated step. Consumers see the -last one written at a given step, and a tail from the abandoned attempt beyond -wherever the new run has reached. - -Left alone deliberately. The fix is to identify which run a point came from, -so both can be kept and told apart, and that belongs with checkpoint restore: -without it there is no second curve worth the work to separate, and Studio's -loss chart keys a Map by step, so it would need to filter or overlay by run -before any of it were visible. Detecting the restart here and dropping the -superseded points is the other option, and it loses a failed attempt's history -permanently to keep one series single-valued -- a trade for the readers rather -than the data, and not one to make in passing. +appends its own. It takes one of two shapes, according to whether the backend +resumes from a checkpoint. + +Automodel and unsloth never do, and neither does a NeMo-RL run that has not +written a checkpoint yet to return to: ``save_period`` is ``val_period``, which +is ``steps_per_epoch``, so a single-epoch run saves only on its last step. +Training restarts at step one and the series carries both runs end to end. + +NeMo-RL otherwise does resume. ``dpo.setup()`` loads the latest checkpoint +unconditionally and the loop continues from the step recorded in it. Reporting +runs ahead of checkpointing -- the train cadence is set by run length, the save +cadence by ``val_period`` -- so the steps between that checkpoint and the +interruption had already been recorded, and the replayed points are appended +after the ones they supersede. The curve doubles back on itself across that +range rather than restarting. + +Either way, consumers see the last value written at a given step, and a tail +from the abandoned attempt beyond wherever the new run has reached. + +Left alone deliberately, and the two shapes are why. The clean fix is to +identify which run a point came from so both can be kept and told apart, but +Studio's loss chart keys a Map by step, so it would need to filter or overlay +by run before any of it were visible. Detecting the restart here and dropping +the superseded points is the other option, and it is right for one shape and +not the other: a replayed range is work that was rolled back and is no loss, +while a from-scratch restart's points are a failed attempt's entire history, +gone permanently to keep one series single-valued. Telling those two apart is +the work, and not a trade to make in passing. Backends subclass this and set :attr:`_default_backend`: unsloth stamps a ``backend`` field on each report (``"unsloth"``); automodel and NeMo-RL leave it diff --git a/packages/nmp_customization_common/tests/training/test_callbacks.py b/packages/nmp_customization_common/tests/training/test_callbacks.py index 1268834161..95fa553882 100644 --- a/packages/nmp_customization_common/tests/training/test_callbacks.py +++ b/packages/nmp_customization_common/tests/training/test_callbacks.py @@ -381,13 +381,19 @@ def test_a_restart_appends_to_the_previous_runs_points() -> None: A replaced pod -- a suspended and resumed Kubernetes Job, or a Volcano restart where an execution profile raises maxRetry above zero -- seeds itself - from the old run's points and starts again at step one, nothing being able to - restore a checkpoint. The series then carries both runs, and two values can - sit at one step contradicting each other. - - Accepted rather than solved: telling the runs apart needs a run identifier on - each point, and dropping the superseded ones loses a failed attempt's history - for good. See the seeding note in callbacks.py. + from the old run's points and appends its own. This is the shape it takes + when the backend does not resume: automodel and unsloth never do, and a + NeMo-RL run that has not written a checkpoint yet has none to return to. The + series then carries both runs end to end. + + The other shape is a NeMo-RL run that does resume -- dpo.setup() loads the + latest checkpoint unconditionally -- and replays the steps it had already + reported since that checkpoint, so its curve doubles back across that range + instead of restarting. Not pinned here. + + Accepted rather than solved either way: telling the runs apart needs a run + identifier on each point, and dropping the superseded ones is right for a + replay and destructive for a restart. See the seeding note in callbacks.py. """ prior = {"train_loss": [{"step": 50, "epoch": 1, "value": 1.0}]} reporter = _RecordingReporter(prior) From 25f30aeefac5ee113a9fee4ed5f520a65766b750 Mon Sep 17 00:00:00 2001 From: Albert Cui Date: Fri, 14 Aug 2026 18:21:21 -0400 Subject: [PATCH 23/23] refactor(customization): rename the metric-name qualifier off "namespace" _namespace read as one of four established meanings of the word before it read as the metric one. NamespacedModel and __schema_namespace__ sit one module away in the same package and also mean "prefix a name", but for pydantic schema class names rather than metric keys; the platform's own resource scoping (parse_resource_id("default/my-model")), pydantic's protected_namespaces, and argparse.Namespace are the other three. _qualify_metric_names says which names and what happens to them. The local at its call site goes from `namespaced` to `qualified`, which matters most at `if not self._seed_unavailable and qualified:` -- the condition being tested is "did anything survive the chartable filter", which the new name states and the old one did not. NemoRLLogger._namespace_validation becomes _qualify_by_dataset: the same verb one level down, where the phase qualifies a metric name and the dataset qualifies it further on a run with more than one validation dataloader. Renaming one and not the other would have been worse than leaving both alone. Four docstrings and two test docstrings follow the same word. The three remaining uses of "namespace" in these packages -- model_namespace on the model-entity path, NamespacedModel in the schema tests -- are the other meanings, and keeping them distinct is the point of the rename. No behaviour change: identifiers and prose only. Signed-off-by: Albert Cui --- .../training/callbacks.py | 22 +++++++++---------- .../tests/training/test_callbacks.py | 4 ++-- .../backends/nemo_rl/nemo_rl_logger.py | 8 +++---- 3 files changed, 17 insertions(+), 17 deletions(-) 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 2f98fbca14..99b17813c8 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 @@ -21,7 +21,7 @@ One rule, no exceptions: a metric is stored and reported as ``_``, where the backend supplies its framework's own ```` (``loss``, ``lr``, ``accuracy``) and the phase that produced it supplies the prefix. The same -namespaced name is used for the accumulated series and for the current-step +qualified name is used for the accumulated series and for the current-step value in ``status_details``. The prefix is load-bearing rather than cosmetic -- DPO reports ``accuracy`` in @@ -144,7 +144,7 @@ def is_chartable(value: Any) -> bool: return False -def _namespace(phase: str, metrics: Mapping[str, object]) -> dict[str, float | int]: +def _qualify_metric_names(phase: str, metrics: Mapping[str, object]) -> dict[str, float | int]: """The chartable subset of ``metrics``, keyed by ``_``. The single naming rule. A backend passes its framework's own metric names -- @@ -261,22 +261,22 @@ def _report_metrics( """Record every metric as a point and report them as current values. The one path both ``report_train_step`` and ``report_validation`` take. - ``phase`` namespaces the series (``train``/``val``); ``report_phase`` is - what the Jobs service records as the task's phase. + ``phase`` qualifies the series names (``train``/``val``); ``report_phase`` + is what the Jobs service records as the task's phase. """ - namespaced = _namespace(phase, metrics) - for name, value in namespaced.items(): + qualified = _qualify_metric_names(phase, metrics) + for name, value in qualified.items(): self._series.setdefault(name, []).append({"step": step, "epoch": epoch, "value": value}) details: dict[str, object] = { "step": step, "epoch": epoch, - **namespaced, + **qualified, } # Sent in full or not at all, so a report that added no point has nothing # to say about the curves: the stored copy already matches, and the merge # leaves a key that is not mentioned standing. - if not self._seed_unavailable and namespaced: + if not self._seed_unavailable and qualified: details["metrics"] = self._build_metrics_summary() self._send(report_phase, details, backend) @@ -328,9 +328,9 @@ def report_train_step( loss reports no ``train_loss``, and a name is stated only when it was observed, because a null charts as a real zero. - Taken as a dict rather than ``**kwargs`` so that the metric namespace and - this method's own parameters cannot collide. Backends forward whatever - their framework emits, and a framework is free to call something ``step``. + Taken as a dict rather than ``**kwargs`` so that the metric names and this + method's own parameters cannot collide. Backends forward whatever their + framework emits, and a framework is free to call something ``step``. """ self._report_metrics("train", "training", step, epoch, metrics, backend) diff --git a/packages/nmp_customization_common/tests/training/test_callbacks.py b/packages/nmp_customization_common/tests/training/test_callbacks.py index 95fa553882..20c967e1d8 100644 --- a/packages/nmp_customization_common/tests/training/test_callbacks.py +++ b/packages/nmp_customization_common/tests/training/test_callbacks.py @@ -65,7 +65,7 @@ def _make_callback(reporter: _RecordingReporter) -> TrainingProgressCallback: def test_the_phase_supplies_the_prefix(reporter: _RecordingReporter) -> None: - """Backends pass their framework's own names; the phase namespaces them.""" + """Backends pass their framework's own names; the phase qualifies them.""" callback = _make_callback(reporter) callback.report_train_step(step=1, epoch=1, metrics={"loss": 0.5, "lr": 5e-06, "grad_norm": 1.9}) callback.report_validation(step=1, epoch=1, metrics={"loss": 0.4, "accuracy": 0.9}) @@ -146,7 +146,7 @@ def test_reserved_names_cannot_be_reached_by_a_metric(reporter: _RecordingReport assert report["epoch"] == 1 assert isinstance(report["metrics"], dict), "the series payload survives" assert [r["phase"] for r in reporter.reports] == ["training", "validation"] - assert reporter.reports[0]["train_phase"] == 1.0, "the metric is kept, under its namespaced name" + assert reporter.reports[0]["train_phase"] == 1.0, "the metric is kept, under its qualified name" assert reporter.reports[0]["train_step"] == 2.0 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 24536be5cc..8a55efc28b 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 @@ -168,7 +168,7 @@ def __init__( # modulus would skip whole passes rather than thin them evenly. self._val_passes = 0 - # The prefix whose metrics keep the bare names; see _namespace_validation. + # The prefix whose metrics keep the bare names; see _qualify_by_dataset. self._primary_val_prefix: str | None = None # Reports built but withheld by a throttle, at most one of each kind. @@ -269,7 +269,7 @@ def log_metrics( report = { "step": step, "epoch": epoch, - "metrics": self._namespace_validation(prefix, metrics), + "metrics": self._qualify_by_dataset(prefix, metrics), } if self._val_passes % self._val_report_interval == 0: self._send("validation", report) @@ -278,7 +278,7 @@ def log_metrics( _logger.debug(f"log_metrics: step={step}, prefix={prefix}, metrics={metrics}") - def _namespace_validation(self, prefix: str, metrics: Mapping[str, Any]) -> dict[str, Any]: + def _qualify_by_dataset(self, prefix: str, metrics: Mapping[str, Any]) -> dict[str, Any]: """Fold the dataloader name into each metric name, past the first set. ``validate()`` loops over ``val_dataloader.items()`` and logs once per @@ -293,7 +293,7 @@ def _namespace_validation(self, prefix: str, metrics: Mapping[str, Any]) -> dict dataloader too, so keying on "did a name arrive" would rename the common case and take Studio's curve with it. Iteration order over the dataloader dict is stable within a run and across a resume of the same config, so a - dataset keeps whichever namespace it started in. + dataset keeps whichever naming it started with. """ if self._primary_val_prefix is None: self._primary_val_prefix = prefix