From d6a4a846a3b396313511a18a0dd072b9e74cd3e5 Mon Sep 17 00:00:00 2001 From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com> Date: Fri, 28 Aug 2026 18:28:20 -0400 Subject: [PATCH 1/3] fix(loss): min/max packed extrema instead of summing them SequencePackingLossWrapper folds per-sequence metric dicts into one. It special-cases extrema through a hardcoded allowlist naming only the four probs_ratio keys, so MseValueLossFn's values_min/values_max fall through to '+=' and are summed. Three packed sequences spanning -3..9 report values_min=4.0 -- a positive number for a critic whose predictions go negative -- and values_max=15.0. The error grows with packing density, so it is not a stable offset a reader could correct for. The loss function already says what these are: 'Min/max are per-MB; ppo.py takes min/max across MBs.' Five other sites apply the '_min'/'_max' suffix rule to this very dict, and one of them -- megatron_value_worker.py:611 -- is this wrapper's own direct consumer, skipping the divide because it is an extremum while the wrapper upstream has already summed it. Use the same rule here. That alone is not enough. MseValueLossFn returns 0.0 for a fully-masked sequence where ClippedPGLossFn returns +/-inf, and 0.0 is a plausible value that wins the min against an all-positive critic: the one-line version reports 0.0 where the truth is 3.0. sample_mask is loss_multiplier, which overlong_filtering zeroes per sample, and under packing one filtered sample in a pack is enough. So the sentinel moves to +/-inf and ppo.py skips it, matching what that file already does for probs_ratio at :1774 and :2757. Metrics only -- packed and unpacked losses are bit-identical. Reachable on ppo-qwen2.5-1.5b-gsm8k-1n8g-megatron-valuetp2sp-pp2cp2-pack, which nightly.txt runs. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com> --- nemo_rl/algorithms/loss/loss_functions.py | 12 +- nemo_rl/algorithms/loss/wrapper.py | 13 +- nemo_rl/algorithms/ppo.py | 8 +- .../test_packed_metric_aggregation.py | 148 ++++++++++++++++++ 4 files changed, 173 insertions(+), 8 deletions(-) create mode 100644 tests/unit/algorithms/test_packed_metric_aggregation.py diff --git a/nemo_rl/algorithms/loss/loss_functions.py b/nemo_rl/algorithms/loss/loss_functions.py index cafc6110b1c..4d438a75784 100755 --- a/nemo_rl/algorithms/loss/loss_functions.py +++ b/nemo_rl/algorithms/loss/loss_functions.py @@ -1898,12 +1898,20 @@ def __call__( ).item() # Min/max are per-MB; ppo.py takes min/max across MBs. + # +/-inf, not 0.0, for an empty mask: 0.0 is a plausible value and + # would win the min against an all-positive critic, silently + # flooring the reported range. ClippedPGLossFn uses the same + # sentinel for the same reason, and both consumers skip it. masked_values = values[mask.bool()] values_min = ( - masked_values.min().item() if masked_values.numel() > 0 else 0.0 + masked_values.min().item() + if masked_values.numel() > 0 + else float("inf") ) values_max = ( - masked_values.max().item() if masked_values.numel() > 0 else 0.0 + masked_values.max().item() + if masked_values.numel() > 0 + else float("-inf") ) # Explained variance sufficient statistics. diff --git a/nemo_rl/algorithms/loss/wrapper.py b/nemo_rl/algorithms/loss/wrapper.py index f1eb9dd7318..5b6dc9297ea 100644 --- a/nemo_rl/algorithms/loss/wrapper.py +++ b/nemo_rl/algorithms/loss/wrapper.py @@ -162,10 +162,15 @@ def __call__( # aggregate loss and metrics loss_accum += loss for k, v in metrics.items(): + # ``*_min``/``*_max`` are extrema, not additive quantities. + # Preserve explicitly registered non-suffix metrics while also + # covering new loss metrics that follow the suffix contract. + is_min = k in _SEQ_METRIC_MIN or k.endswith("_min") + is_max = k in _SEQ_METRIC_MAX or k.endswith("_max") if k not in metrics_accum: - if k in _SEQ_METRIC_MIN: + if is_min: metrics_accum[k] = float("inf") - elif k in _SEQ_METRIC_MAX: + elif is_max: metrics_accum[k] = float("-inf") else: metrics_accum[k] = 0 @@ -173,10 +178,10 @@ def __call__( val = v.item() if isinstance(v, torch.Tensor) and v.ndim == 0 else v # Skip inf/-inf sentinel values (from sequences with no valid tokens) - if k in _SEQ_METRIC_MIN: + if is_min: if not math.isinf(val): metrics_accum[k] = min(metrics_accum[k], val) - elif k in _SEQ_METRIC_MAX: + elif is_max: if not math.isinf(val): metrics_accum[k] = max(metrics_accum[k], val) else: diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 1f8479fbaff..36c8dff4c8f 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -1213,9 +1213,13 @@ def _compute_critic_metrics(value_results: dict[str, Any]) -> dict[str, Any]: if key in {"lr", "wd", "global_valid_seqs", "global_valid_toks", "grad_norm"}: critic_metrics[metric_name] = np.mean(value).item() elif key == "values_min": - critic_metrics[metric_name] = np.min(value).item() + # Skip the empty-mask sentinel, as the probs_ratio extrema are + # handled at :1774 and :2757. + finite = [x for x in value if not np.isinf(x)] + critic_metrics[metric_name] = np.min(finite).item() if finite else -1.0 elif key == "values_max": - critic_metrics[metric_name] = np.max(value).item() + finite = [x for x in value if not np.isinf(x)] + critic_metrics[metric_name] = np.max(finite).item() if finite else -1.0 elif isinstance(value, (np.ndarray, list)): critic_metrics[metric_name] = np.sum(value).item() else: diff --git a/tests/unit/algorithms/test_packed_metric_aggregation.py b/tests/unit/algorithms/test_packed_metric_aggregation.py new file mode 100644 index 00000000000..2c699479f78 --- /dev/null +++ b/tests/unit/algorithms/test_packed_metric_aggregation.py @@ -0,0 +1,148 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Packing must not change what a metric means. + +``SequencePackingLossWrapper`` folds per-sequence metric dicts into one. Sums +are right for globally normalized metrics and wrong for extrema, and the +workers that consume this dict downstream already tell the two apart by the +``_min``/``_max`` suffix (megatron_value_worker.py:611 and four sibling sites). +""" + +import pytest +import torch + +from nemo_rl.algorithms.loss.loss_functions import MseValueLossConfig, MseValueLossFn +from nemo_rl.algorithms.loss.wrapper import SequencePackingLossWrapper +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +def _value_prepare_fn(logits, data, loss_fn=None, **kwargs): + """Stand-in for megatron_value_worker._value_loss_prepare_fn. + + That function all-gathers across CP and shifts; neither applies here, so + this keeps only the part the wrapper's contract depends on -- the key the + loss is called with. + """ + del loss_fn, kwargs + return {"logits": logits}, data + + +def _batch(value_rows, sample_mask): + values = torch.tensor(value_rows, dtype=torch.float32) + data = BatchedDataDict( + { + "values": values.clone(), + "returns": torch.zeros_like(values), + "token_mask": torch.ones_like(values), + "sample_mask": torch.tensor(sample_mask, dtype=torch.float32), + } + ) + return data, values + + +def _packed_and_unpacked(value_rows, sample_mask): + """Run the same values both ways and return (unpacked_metrics, packed_metrics).""" + loss_fn = MseValueLossFn(MseValueLossConfig()) + data, values = _batch(value_rows, sample_mask) + logits = values.unsqueeze(-1) + global_valid_seqs = data["sample_mask"].sum() + global_valid_toks = (data["token_mask"] * data["sample_mask"].unsqueeze(-1)).sum() + + unpacked_loss, unpacked = loss_fn( + logits, data, global_valid_seqs, global_valid_toks + ) + + seq_len = values.shape[1] + cu_seqlens = torch.tensor( + [i * seq_len for i in range(len(value_rows) + 1)], dtype=torch.int32 + ) + wrapper = SequencePackingLossWrapper( + loss_fn=loss_fn, + prepare_fn=_value_prepare_fn, + cu_seqlens_q=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens, + ) + packed_loss, packed = wrapper( + logits.reshape(1, -1, 1), data, global_valid_seqs, global_valid_toks + ) + + # The loss must be untouched by anything here; if it moves, the test is + # measuring the wrong thing. + assert packed_loss.item() == pytest.approx(unpacked_loss.item(), abs=1e-6) + return unpacked, packed + + +def test_packing_reports_the_true_value_range(): + """Summed extrema are not extrema -- and the reported minimum flips sign. + + Three sequences spanning -3..9. Summing the per-sequence minima gives + -3 + 1 + 6 = 4, so ``critic/values_min`` reports a positive number for a + critic whose predictions go negative. It is the diagnostic operators read + to catch a value head drifting, and the error grows with packing density. + """ + unpacked, packed = _packed_and_unpacked( + [[-3.0, -1.0, 2.0], [1.0, 4.0, 2.0], [6.0, 9.0, 7.0]], [1, 1, 1] + ) + + assert unpacked["values_min"] == pytest.approx(-3.0) + assert unpacked["values_max"] == pytest.approx(9.0) + assert packed["values_min"] == pytest.approx(unpacked["values_min"]) + assert packed["values_max"] == pytest.approx(unpacked["values_max"]) + + +def test_a_fully_masked_sequence_does_not_floor_the_reported_minimum(): + """The empty-mask sentinel must not be a plausible value. + + ``sample_mask`` is ``loss_multiplier``, which ``overlong_filtering`` zeroes + per sample -- and under packing one filtered sample in a pack is enough. A + 0.0 sentinel would win the min against this all-positive critic and report + 0.0 instead of 3.0, so it has to be +/-inf and be skipped, as + ClippedPGLossFn's already is. + """ + unpacked, packed = _packed_and_unpacked( + [[3.0, 5.0, 4.0], [6.0, 9.0, 7.0], [8.0, 8.5, 8.2]], [1, 1, 0] + ) + + assert unpacked["values_min"] == pytest.approx(3.0) + assert packed["values_min"] == pytest.approx(3.0) + assert packed["values_max"] == pytest.approx(9.0) + + +def test_globally_normalized_metrics_are_still_summed(): + """Only extrema changed: everything else must still add up across sequences.""" + unpacked, packed = _packed_and_unpacked( + [[-3.0, -1.0, 2.0], [1.0, 4.0, 2.0], [6.0, 9.0, 7.0]], [1, 1, 1] + ) + + for key in ("values_mean", "returns_mean", "returns_sq_mean", "residual_sq_mean"): + assert packed[key] == pytest.approx(unpacked[key], abs=1e-5), key + + +def test_an_all_masked_microbatch_reports_the_sentinel_fallback(): + """Every sequence filtered: there is no value range, and inf must not leak.""" + import numpy as np + + from nemo_rl.algorithms.ppo import _compute_critic_metrics + + _, packed = _packed_and_unpacked([[3.0, 5.0, 4.0], [6.0, 9.0, 7.0]], [0, 0]) + assert np.isinf(packed["values_min"]) + + critic = _compute_critic_metrics( + { + "grad_norm": torch.tensor(0.0), + "loss": torch.tensor(0.0), + "all_mb_metrics": {"values_min": [packed["values_min"]]}, + } + ) + assert critic["critic/values_min"] == pytest.approx(-1.0) From aea9669699155175373dde2f326689d7a13beb71 Mon Sep 17 00:00:00 2001 From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com> Date: Sun, 30 Aug 2026 15:14:53 -0400 Subject: [PATCH 2/3] fix(loss): classify extrema by metric suffix Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com> --- nemo_rl/models/automodel/train.py | 2 +- .../policy/workers/dtensor_policy_worker.py | 2 +- .../policy/workers/megatron_policy_worker.py | 4 +- .../value/workers/megatron_value_worker.py | 2 +- .../test_packed_metric_aggregation.py | 38 +++++++++++++++++++ .../policy/test_megatron_split_parity.py | 4 +- 6 files changed, 45 insertions(+), 7 deletions(-) diff --git a/nemo_rl/models/automodel/train.py b/nemo_rl/models/automodel/train.py index f5932477e21..3bd638127f2 100644 --- a/nemo_rl/models/automodel/train.py +++ b/nemo_rl/models/automodel/train.py @@ -509,7 +509,7 @@ def automodel_forward_backward( ## scale by the number of global batches so we get the correct ## value when summing metrics across all microbatches for k in metrics.keys(): - if "_min" in k or "_max" in k: + if k.endswith(("_min", "_max")): continue metrics[k] /= num_global_batches diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker.py b/nemo_rl/models/policy/workers/dtensor_policy_worker.py index 29e12e3b8ad..7662fba62c4 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker.py @@ -947,7 +947,7 @@ def train( ## scale by the number of global batches so we get the correct ## value when summing metrics across all microbatches for k in loss_metrics.keys(): - if "_min" in k or "_max" in k: + if k.endswith(("_min", "_max")): continue loss_metrics[k] /= num_global_batches num_valid_samples = loss_metrics["num_valid_samples"] diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index f436b7424ae..5c8b3c8ac65 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -1202,7 +1202,7 @@ def train( for x in losses_reduced: loss_metrics = {} for k in x.keys(): - if "_min" in k or "_max" in k: + if k.endswith(("_min", "_max")): loss_metrics[k] = x[k] else: loss_metrics[k] = x[k] / num_global_batches @@ -2029,7 +2029,7 @@ def _scale_metric(name: str, value: Any) -> Any: for m in state["all_mb_metrics"]: out: dict[str, Any] = {} for k, v in m.items(): - if "_min" in k or "_max" in k: + if k.endswith(("_min", "_max")): out[k] = v else: out[k] = _scale_metric(k, v) diff --git a/nemo_rl/models/value/workers/megatron_value_worker.py b/nemo_rl/models/value/workers/megatron_value_worker.py index ea546575592..e05c0edeeca 100644 --- a/nemo_rl/models/value/workers/megatron_value_worker.py +++ b/nemo_rl/models/value/workers/megatron_value_worker.py @@ -614,7 +614,7 @@ def train( for x in losses_reduced: loss_metrics = {} for k in x.keys(): - if "_min" in k or "_max" in k: + if k.endswith(("_min", "_max")): loss_metrics[k] = x[k] else: loss_metrics[k] = x[k] / num_global_batches diff --git a/tests/unit/algorithms/test_packed_metric_aggregation.py b/tests/unit/algorithms/test_packed_metric_aggregation.py index 2c699479f78..14cd571031b 100644 --- a/tests/unit/algorithms/test_packed_metric_aggregation.py +++ b/tests/unit/algorithms/test_packed_metric_aggregation.py @@ -83,6 +83,44 @@ def _packed_and_unpacked(value_rows, sample_mask): return unpacked, packed +class _AdditiveControlLossFn: + def __call__( + self, + logits: torch.Tensor, + data: BatchedDataDict, + global_valid_seqs: torch.Tensor, + global_valid_toks: torch.Tensor, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + del logits, global_valid_seqs, global_valid_toks + value = data["values"].sum() + return value.new_zeros(()), { + "latency_minibatch": value, + "latency_maximum": value, + } + + +def test_embedded_extrema_substrings_do_not_make_metrics_extrema() -> None: + """Only suffixes classify extrema; ordinary metric names remain additive.""" + data, values = _batch([[1.0, 2.0], [3.0, 4.0]], [1, 1]) + cu_seqlens = torch.tensor([0, 2, 4], dtype=torch.int32) + wrapper = SequencePackingLossWrapper( + loss_fn=_AdditiveControlLossFn(), + prepare_fn=_value_prepare_fn, + cu_seqlens_q=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens, + ) + + _, metrics = wrapper( + values.reshape(1, -1, 1), + data, + data["sample_mask"].sum(), + data["token_mask"].sum(), + ) + + assert metrics["latency_minibatch"] == pytest.approx(10.0) + assert metrics["latency_maximum"] == pytest.approx(10.0) + + def test_packing_reports_the_true_value_range(): """Summed extrema are not extrema -- and the reported minimum flips sign. diff --git a/tests/unit/models/policy/test_megatron_split_parity.py b/tests/unit/models/policy/test_megatron_split_parity.py index e7fee4a9802..f2d5d528df5 100644 --- a/tests/unit/models/policy/test_megatron_split_parity.py +++ b/tests/unit/models/policy/test_megatron_split_parity.py @@ -152,9 +152,9 @@ def _run_split( def _reduce_metric(key: str, values: list) -> float: """Collapse a per-microbatch metric list the way grpo.py's reducer does.""" - if "_min" in key: + if key.endswith("_min"): return float(np.min(values)) - if "_max" in key: + if key.endswith("_max"): return float(np.max(values)) if key in ("lr", "wd", "global_valid_seqs", "global_valid_toks"): return float(np.mean(values)) From 4a5af89017bd2bf51a9fa4cd32ba81ec59de7f20 Mon Sep 17 00:00:00 2001 From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com> Date: Sat, 12 Sep 2026 14:06:51 -0400 Subject: [PATCH 3/3] test(loss): preserve registered packed extrema Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com> --- tests/unit/algorithms/test_packed_metric_aggregation.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/unit/algorithms/test_packed_metric_aggregation.py b/tests/unit/algorithms/test_packed_metric_aggregation.py index 14cd571031b..ded771a6740 100644 --- a/tests/unit/algorithms/test_packed_metric_aggregation.py +++ b/tests/unit/algorithms/test_packed_metric_aggregation.py @@ -14,9 +14,9 @@ """Packing must not change what a metric means. ``SequencePackingLossWrapper`` folds per-sequence metric dicts into one. Sums -are right for globally normalized metrics and wrong for extrema, and the -workers that consume this dict downstream already tell the two apart by the -``_min``/``_max`` suffix (megatron_value_worker.py:611 and four sibling sites). +are right for globally normalized metrics and wrong for extrema. Most extrema +use the ``_min``/``_max`` suffix; explicitly registered exceptions must keep +their reduction as new metrics are added. """ import pytest @@ -96,6 +96,7 @@ def __call__( return value.new_zeros(()), { "latency_minibatch": value, "latency_maximum": value, + "opd_full_decomposition_error": value, } @@ -119,6 +120,7 @@ def test_embedded_extrema_substrings_do_not_make_metrics_extrema() -> None: assert metrics["latency_minibatch"] == pytest.approx(10.0) assert metrics["latency_maximum"] == pytest.approx(10.0) + assert metrics["opd_full_decomposition_error"] == pytest.approx(7.0) def test_packing_reports_the_true_value_range():