fix(loss): min/max packed extrema instead of summing them

[tianyi-zhang-02]

What does this PR do?

SequencePackingLossWrapper folds the per-sequence metric dicts into one. It already knows extrema are not additive — it just decides which ones by a hardcoded allowlist naming only the four probs_ratio_* keys (wrapper.py:147-164). MseValueLossFn's values_min/values_max (loss_functions.py:1334-1335) are not on it, so they fall through to += and get summed.

Measured on three packed sequences spanning -3..9:

metric	truth (unpacked)	packed, on main
values_min	-3.0	4.0
values_max	9.0	15.0

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 saturating or diverging, and because the error is a sum it grows with packing density rather than being a stable offset someone could mentally correct for.

This is a stale allowlist, not a design call. The loss function already states the contract in a comment at loss_functions.py:1304 — "Min/max are per-MB; ppo.py takes min/max across MBs." — and ppo.py really does np.min/np.max on those two keys in _compute_critic_metrics. Five sites already apply the "_min" in k or "_max" in k rule to a loss's metric dict: megatron_value_worker.py:611, megatron_policy_worker.py:1034 and :1779, dtensor_policy_worker.py:940, automodel/train.py:531. Only the first sees MseValueLossFn's dict — the others aggregate the policy loss — but it is this wrapper's own direct consumer — it deliberately skips dividing values_min by num_global_batches because it is an extremum, while the wrapper upstream has already summed it. The wrapper is the odd one out.

Why it is three hunks and not one

Adding the two keys to the existing allowlist fixes the sum and introduces a quieter bug. MseValueLossFn returns 0.0 for a fully-masked sequence (loss_functions.py:1306-1311) where ClippedPGLossFn returns ±inf (:759-762), and 0.0 is a plausible value that wins the min against an all-positive critic. I measured it: with values [3,5,4] [6,9,7] [8,8.5,8.2] and the third sequence masked, the one-line version reports values_min = 0.0 where the truth is 3.0.

That case is not hypothetical — sample_mask is loss_multiplier, which overlong_filtering zeroes per sample (ppo.py:1467-1472), and under packing one filtered sample in a pack is enough, where today an entire microbatch would have to be masked. So the sentinel moves to ±inf and ppo.py skips it, copying the shape that file already uses for probs_ratio at :1774 and :2757.

The suffix rule is safe here: across every LossFunction in the repo the only metric keys containing _min/_max are the four probs_ratio_* and these two, all extrema by intent.

Validation

• Four CPU tests in a new tests/unit/algorithms/test_packed_metric_aggregation.py. Mutation matrix, each hunk reverted independently:

  state	result
  upstream/main	3 failed, 1 passed
  suffix rule only	2 failed, 2 passed
  + sentinel, no consumer filter	1 failed, 3 passed
  all three hunks	4 passed

  test_globally_normalized_metrics_are_still_summed passes in every state — it is the control that pins that values_mean, returns_mean, returns_sq_mean and residual_sq_mean still add up.

• 129 passed / 51 skipped across test_loss_functions.py, test_ppo.py, test_sequence_packing_fusion.py, test_draft_loss_wrapper.py and the new file.

• ruff, ruff format clean; pyrefly unchanged from main on both touched files (48 → 48), neither of which is in the checked set.

Scope

Metrics only. Packed and unpacked losses are bit-identical in both test cases — I assert that inside the helper so the tests cannot pass while measuring the wrong thing. No gradient, no training behavior, nothing branches on these values.

Reachable on examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-1n8g-megatron-valuetp2sp-pp2cp2-pack.yaml, which sets value.sequence_packing.enabled: true and is run by tests/test_suites/nightly.txt:307. megatron/train.py rejects fuse_loss for a custom prepare_fn, and the value worker supplies one, so value + packing lands on this wrapper.

To be precise about which worker sees what: the aggregation fix is Megatron-value-only, since the DTensor value worker has no packing path. The sentinel change is inside MseValueLossFn, so it is visible to both. Both workers feed the same _compute_critic_metrics, which now filters it, so neither can surface a ±inf.

One behavior change worth naming: on a step where every microbatch is fully masked, critic/values_min and critic/values_max now log -1.0 instead of 0.0. That matches what probs_ratio_min already does in the same situation, and 0.0 was never a real measurement there.

---

Corrections (pushed to the description only; no code change). Three citations in an earlier revision were off:

• The Min/max are per-MB comment is at loss_functions.py:1304, not :1302.
• The np.min/np.max consumer moved to ppo.py:1177 once this PR's own hunk added lines above it; the original citation was the pre-change position.
• "Five other sites apply the rule to this same dict" was too strong — only megatron_value_worker.py:611 sees MseValueLossFn's dict; the other four aggregate the policy loss's. The argument is unchanged, since that one site is this wrapper's direct consumer.

I also sharpened the Scope section: the sentinel change is not packing-gated, so it is visible to the DTensor value worker too, though it cannot surface there.

---

Second correction (description only). I claimed the DTensor value worker's num_valid_samples > 0 gate (dtensor_value_worker_v2.py:395) already prevents a ±inf from reaching aggregation there. That is not true on current main: MseValueLossFn reports num_valid_samples as the microbatch row count, so a fully-masked microbatch reports 2, not 0, and the gate does not fire. I checked by running it.

It becomes true only after my #3850, which changes that metric to sample_mask.sum().

This PR does not depend on that, and is not stacked on it. I checked by running the fully-masked case against main: the loss reports values_min = inf, the DTensor gate does not fire (num_valid_samples = 2), the sentinel reaches _compute_critic_metrics, and hunk 3 filters it to -1.0. No inf leaks on either worker with or without #3850. #3850 only adds a second, earlier place the same case gets dropped. Merge them in whichever order suits you.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.