fix(rl): GRPO training metrics - #1264
Draft
albcui wants to merge 7 commits into
Draft
Conversation
The NeMo-RL logger was written against DPO and forwards almost nothing from a
GRPO run. Three problems, all in the translation layer rather than upstream.
Validation was reported for DPO only. The branch gated on a `loss` key, but
GRPO's validation dict is {accuracy, avg_length} and carries no loss at all
(nemo_rl/algorithms/grpo.py builds it in validate()). Every GRPO validation
pass was therefore dropped silently -- no report, no best-metric update, no
log line saying anything had been skipped. `val_loss` is now optional: report
whenever anything usable arrived, and omit the key rather than sending null,
which would chart as a real zero.
The metric key list was DPO-shaped. `preference_loss` and
`rewards_rejected_mean` do not exist under GRPO, so a GRPO run forwarded only
loss/lr/grad_norm and three token counts -- no reward, no advantages, no KL,
no truncation rate. Reward is the metric that says whether RL is converging,
and it never left the pod. The list is now a union across algorithms, selected
by presence, covering the reward/advantage block, policy-optimization health,
and NeMo-Gym's rollout metrics.
Widening the list meant `has_metric_value` had to stop assuming numbers.
NeMo-RL's metric dicts interleave non-scalars with the scalars:
calculate_single_metric emits a `<key>/histogram` 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 three, so the old
None-check would have crashed mid-training on the first widened key rather
than dropping the value.
Separately, RL never accumulated a metric time series. `report_running`
REPLACES the task's status_details blob, so reporting only the current step
leaves no history -- the loss curve was unrecoverable and fetch_current_metrics
always came back empty, so resume could not seed either. The shared
customization callback has done this since it was written; the RL copy simply
had not. train_loss/val_loss now accumulate in the same {step, epoch, value}
shape Studio already renders as CustomizationMetricValue[], and every report
path carries the payload, including report_training_start -- which would
otherwise blank a resumed job's seeded series.
Only those two series accumulate. The wider metric set rides along as
current-step scalars: every series is resent in full on every update, so the
payload grows with series count times step count. Putting the whole GRPO
surface on that wire needs the Jobs metrics transport reworked, which is
deliberately out of scope here.
Both files were byte-identical to main before this change.
Signed-off-by: Albert Cui <albcui@nvidia.com>
test_tokenizer_omits_chat_template_when_none asserted that a model with no
chat template omits the key, and established the premise indirectly: the
fixture model dir has no tokenizer, so resolve_chat_template was expected to
fall through to None.
It does not, once another suite has run in the same session.
services/automodel/tests/tasks/training/backends/test_config.py installs
`sys.modules.setdefault("transformers", MagicMock())` at module scope and never
removes it, so resolve_chat_template's AutoTokenizer.from_pretrained(...) call
returns a Mock whose `.chat_template` is truthy. The key is then present and the
assertion fails -- but only when the automodel tests are collected first, which
is why it passed per-service and failed in a combined run.
Patch resolve_chat_template directly instead, mirroring the neighbouring
test_tokenizer_keeps_chat_template_when_present. That states the actual premise
rather than arranging for it, and is immune to what else is in sys.modules.
Does not address the leak itself; two unsloth hf_trainer_callback tests fail
the same way (`from transformers import TrainerCallback` yields a Mock base
class) and are untouched here.
Signed-off-by: Albert Cui <albcui@nvidia.com>
…ation RL carried a standalone TrainingProgressCallback rather than subclassing the shared one, and the previous commit made that worse by copying ~25 lines of accumulation logic into it. Every other customization service (unsloth, automodel) subclasses the shared class; RL was the outlier, and the one that had drifted ahead in features. The three RL-only behaviours are additive, so they move up into the shared class rather than justifying a fork: `**additional_metrics` for backend-specific current-step scalars, and an optional `val_loss` for algorithms that do not produce one. RL's copy becomes a two-line subclass, leaving `_default_backend` as None so no `backend` key is added and its status-detail shape is unchanged. Promoting the accumulation exposed that the shared class had the same data-loss hole this branch just closed for RL. report_training_start, report_epoch_end and report_checkpoint_saved all omitted the `metrics` payload, and report_running REPLACES status_details rather than merging, so any of them erased the accumulated series from stored status until the next train step resent it -- and lost it outright if the job died in that window. automodel calls both report_epoch_end and report_checkpoint_saved mid-training, so this was reachable in practice, not theoretical. Two automodel tests and one unsloth test pinned the buggy payload with exact-kwargs assertions; they now assert the series survives instead. 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. Two defects fixed while in here: 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(), the only end-of-run hook available (`step_finished` is per-step). The flush is reachable from __del__, so it swallows and logs rather than raising. The two drivers derived the same two parameters with different formulas: at val_period=100 DPO computed a log_interval of 11 and GRPO 10, and steps_per_epoch was read from config in one and derived in the other. Both now call NemoRLLogger.for_schedule, which owns the arithmetic and still prefers an explicit steps_per_epoch when the algorithm config carries one (DPO does). DPO's reporting cadence changes slightly as a result -- its `+1` was a divide-by-zero guard that also skewed every value. Behaviour-preserving elsewhere: no wire-shape change for any backend beyond the `metrics` payload now being present on reports that previously dropped it. Signed-off-by: Albert Cui <albcui@nvidia.com>
The metric series survived mid-training reports as of the previous commit, but
not the end of the run. `status_details` is REPLACED by the Jobs service, and
the last three writes of any training job come from the runner process, not the
training driver:
runner report_running("processing_checkpoint") <- no metrics
runner report_completed("Training completed") <- no metrics
runner report_error(...) <- blanks the whole blob
TrainingProgressCallback resends the series on every report it makes, which is
why the fix for report_epoch_end/report_checkpoint_saved worked. It cannot help
here: the runner is a *different process* from the driver that accumulated the
series (backend.execute_training spawns the driver as a subprocess), so it holds
nothing to resend. Studio reads the curve straight out of status_details
(CustomizationDetailsPanel), so in practice the chart was populated while a job
ran and empty the moment it stopped -- and emptiest on the failure path, where
the partial curve is worth the most.
Preservation therefore has to live below the callback, at update_task, the one
choke point every write path shares. When an update does not carry `metrics`,
read the stored series back and re-attach it. The server is already the source
of truth for the series -- that is how resume seeding works -- so this reuses
fetch_current_metrics rather than introducing a second notion of "current".
Only `metrics` is carried over. The rest of the blob is deliberately a
current-state snapshot (phase, step, lr, ...); merging that would leave a
completed task advertising a mid-training step.
Costs one GET per update that omits `metrics`, which is the handful the runner
makes per job. Per-step training reports always carry their own series and skip
the fetch, so the hot path is unchanged.
Applies to all three customization services, since they share this reporter.
Adds the first test coverage for progress.py.
Signed-off-by: Albert Cui <albcui@nvidia.com>
The previous commit added a pending-report flush on NemoRLLogger.close() to stop
the last training step being dropped by the log_interval throttle. Nothing calls
close().
Both drivers append the logger to `logger_inst.loggers` and never tear 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 grpo_train/dpo_train never call
finish() either; the only caller upstream is the single-controller path, which
these drivers do not use.
So the flush ran only from __del__, at GC or interpreter shutdown, where both
_flush_pending_train_report and the reporter's update_task swallow exceptions.
On a clean return refcounting probably got there. On SIGTERM, an unhandled
exception or a cancelled job it did not -- which is the case the flush exists
for.
Two hooks, because neither alone is sufficient:
finish() aliases close() under the name the composite dispatches, so the
flush happens even if a driver forgets.
drivers call close() explicitly from a finally, because dpo_train and
grpo_train never trigger the composite's finish() at all.
DPO's driver had no try/finally around dpo_train; it has one now. GRPO's runs
before the existing environment teardown, which is slow and can itself raise.
Tested at the seam that broke: one test performs the composite's exact
`getattr(logger, "finish", None)` lookup, so renaming the method fails loudly.
The driver calls are 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 <albcui@nvidia.com>
`log_metrics` opened with `step = step + 1 # ...we start counting from 1`, but both callers already count from 1: grpo.py and dpo.py log `total_steps + 1`, where total_steps is 0-based and incremented *after* the log. The logger added a second increment on top. A 23-step run therefore recorded steps 2..24 against max_steps=23. Until the previous commit that only skewed a `step` field; now that the series is the x-axis of a rendered loss curve, the whole curve sat one step right of the truth and percentage_done saturated a step early. It also interacted with the log_interval throttle. `step % log_interval == 0` was evaluated on the inflated step, so reports landed on true steps 9, 19, 29 -- and the final step was withheld even when max_steps *was* a multiple of the interval (at max_steps=20, interval=10 the run's last reported loss was step 19's). The pending-report flush was covering for this; it now handles only the case it was written for. Epoch derivation is fixed by the same change, since it was reading the same inflated step and flipping an epoch early at the boundary. It now clamps at zero: step 0 does arrive, from the validate-at-start path both algorithms run before training, and it belongs to epoch 1 rather than epoch 0. Rides with this branch because it is the same wire values, and this branch already changes DPO's reporting cadence -- correcting the labels separately would mean reviewing that blast radius twice. The tests drove on 0-indexed steps, which is why this survived them. They now generate the sequence a real N-step run produces, via a helper that says so, and pin both ends of the range plus the epoch boundaries. Signed-off-by: Albert Cui <albcui@nvidia.com>
Independent one-liners, grouped so they are easy to drop; none change behaviour
on a path exercised today.
callbacks: splat `**additional_metrics` first in report_train_step, so a backend
metric named `metrics` or `train_loss` cannot silently replace the accumulated
series or the step's own loss. report_validation already had this order; the two
now agree. Every other colliding name is a real parameter and so already errors
at the call site.
nemo_rl_logger: add `rewards_chosen_mean` alongside `rewards_rejected_mean`.
Forwarding one half of DPO's reward pair makes it hard to read.
test_nemo_rl_logger: give the nemo_rl module stubs a real `__spec__`. find_spec
consults sys.modules before the finders and raises on a `__spec__` of None, so
the bare ModuleType turned any later `find_spec("nemo_rl")` in the same session
into a ValueError. The stub is installed at import time and never torn down --
exactly the shape of leak the chat_template test on this branch had to be
rewritten around, so it should not be left as a trap for the next one.
test_grpo_config: `lambda *_, **__` for the resolve_chat_template patch. The
call site passes keywords today; the keyword-only stub breaks silently if that
ever changes.
Signed-off-by: Albert Cui <albcui@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
We previously built the Nemo RL Logger against DPO only, which lacks coverage for GRPO specific metrics, e.g.
reward/advantages/*etc (see grpo.py). Grounded on the RL source code, this PR expands the current list of metric names to support both DPO and GRPO runs.Separately, RL's job progress reporter callback implementation didn't accumulate metrics as time series. Each
report_runningcall replaced the task's entire status_details blob, so reporting only the current step. This issue is not in Automodel nor Unsloth because they use a differentTrainingProgressCallbackinnmp_customization_common-- this PR refactors RL to use the shared one, which does a copy-then-write.Related Issue
Changes
Type of Change
Quality Gates
Verification
Signed-off-by:traileruv run pre-commit run -apasses, or any blocked checks are identified belowTargeted validation: