fix(loss): min/max packed extrema instead of summing them - #3901
Open
tianyi-zhang-02 wants to merge 3 commits into
Open
fix(loss): min/max packed extrema instead of summing them#3901tianyi-zhang-02 wants to merge 3 commits into
tianyi-zhang-02 wants to merge 3 commits into
Conversation
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>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
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
2 participants
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.
What does this PR do?
Fixes packed min/max diagnostics that were being summed as additive metrics.
SequencePackingLossWrapperand the worker aggregators now recognize extrema only by the explicit_min/_maxsuffix contract. Names that merely contain those substrings in the middle remain additive. Fully masked value samples keep the existing sentinel behavior, and an all-sentinel critic step maps to-1.0like policy metrics.Losses and gradients are unchanged.
Validation
Final SHA:
43f85aec43dd692eedb855f107643dae5bafba5f, based on upstreammainatccbcd4cc5.The gated Megatron parity test also passed on two real H100s. It runs two optimizer steps through freshly initialized sync and split policies, then compares the loss curve, grad norm, and every reduced per-microbatch metric.
Environment: 2× NVIDIA H100 80 GB (from a 4-GPU Runpod host),
nvcr.io/nvidia/nemo-rl:v0.7.0, Python 3.13, PyTorch 2.11.0+cu130, CUDA 13, andNCCL_NVLS_ENABLE=0. The source checkout and Megatron-Bridge submodule were pinned to the PR's exact SHA; this was necessary because the pinned v0.7.0 image's bundled Bridge predates currentmain.