feat(grpo): deduplicate multimodal payloads

[aroshanghias-nvd]

Deduplicate multimodal GRPO payloads

Stacked on #3414 at f27dc6d6. Candidate commit:
d95bb1ae4a4ff2d0207bd0edc1da94a6ba7f44a0.

What this changes

• Represents repeated multimodal rows as logical row-to-segment mappings over a
  compact physical tensor store, while retaining normal logical batch semantics.
• Preserves media identity through GRPO G-expansion, slicing, concatenation,
  dynamic batching, sequence packing, DP sharding, async replay, DAPO assembly,
  and replay-buffer checkpoint restore.
• Re-interns equivalent physical media after shard formation so DP-local Ray
  transfers retain savings even when packing or sharding splits prompt groups.
• Covers image, video, and audio through shared primitives. Model-level
  qualification is currently Nemotron image and Qwen2.5-VL image; audio/video
  behavior is covered by focused data-structure tests rather than model runs.
• Supports native NeMo-RL rollout and NeMo Gym rollout, sync and async GRPO.
  Gym request payloads are not expected to shrink because Gym receives each
  request independently; Gym return, replay, and policy-training boundaries do.
• Adds opt-in exact Ray-boundary payload metrics and forwards them to W&B.
• Pre-casts policy-bound pixel_values from fp32 to bf16 before Ray transfer,
  matching Omni MR !12 and halving physical pixel-tensor bytes. Preprocessing,
  Gym return, and replay storage retain fp32; both dedup modes use the same cast.
• Keeps shared VLM defaults off. The maintained Nemotron VLM/Omni recipes enable
  deduplication; payload metrics remain debug-only and disabled in recipes.

Representation and safety contract

PackedTensor stores physical segments plus CSR-like row offsets and segment
indices. Each physical segment also carries opaque provenance. Provenance is a
candidate equality partition, never proof of media equality: re-interning still
checks tensor value, dtype, device, and shape. Missing media rows are explicit,
and correlated media keys preserve per-row logical counts and ordering.

Flag-off retains the legacy expanded logical representation. Both modes use
bf16 pixel_values for policy-bound transfers, so dedup A/B runs do not differ
in dtype. Legacy pickles/checkpoints
remain readable, and compact replay checkpoints can be restored with the flag
off because sparse media normalization follows the data representation rather
than the current flag.

Quantified G=16 evidence

The W&B campaigns below quantify deduplication before the bf16 pre-cast was
ported. They remain the structural off/on evidence. The pre-cast applies equally
to both modes and independently halves physical policy-bound pixel_values
bytes; the focused protocol-5 test above validates the current combined path.

Nemotron Omni 30B, Megatron, async NeMo Gym

4 nodes, 8 GPUs/node, TP=2, EP=16, CP=2, 10 steps. Dedup off/on jobs
14856877/14856876 completed with exact logical trajectory, reward, and token
parity. TMPE mean/max was 1.022114/1.028493 off and
1.021894/1.027261 on.

Exact Ray boundary	Serialized off/on reduction
Gym request	1.00x (expected: requests are independent)
Gym return	284.60x
Replay push	15.49x
Replay sample	15.55x
Rollout to policy	14.90x
Policy logprobs	3.94x
Reference logprobs	3.93x
Policy train	3.91x

W&B: off,
on.

An independently repeated one-step final-telemetry pair (14869430/14869431)
matched all 16 logical rows, rewards, token counts, masks, advantages, content,
and agent references. TMPE was 1.014925 off and 1.014564 on. Its reductions
were 620.02x at Gym return, 15.78x at replay push/sample, 13.41x rollout-to-policy,
and 3.77-3.85x across policy shards.

Qwen2.5-VL 3B, Megatron, native rollout

1 node, TP=1, DP=8, G=16, 10 steps. Jobs 14858320/14858321 matched logical
trajectories, rewards, and generated tokens exactly. Mean loss was
0.0748708278 off and 0.0748722717 on; maximum paired loss delta was
1.21668e-5. TMPE mean/max was 1.022007/1.025749 off and
1.021798/1.024828 on. Serialized reductions were 14.97x at prompt repeat,
17.05x at vLLM generation, 15.06x rollout-to-policy, and approximately 1.99x
across policy shards.

W&B: off,
on.

Qwen2.5-VL 3B, AutoModel, native rollout

1 node, DP=8, G=16, 10 sampled steps. Jobs 14979492/14979493 completed all
steps. The first complete generation-to-training step matched exactly; later
independently updated sampled trajectories diverged, so this qualifies the
AutoModel data path and multi-step health, not strict multi-step A/B identity.
Serialized reductions were 14.96x prompt repeat, 17.05x vLLM generation,
15.29x rollout-to-policy, and approximately 1.99x at policy boundaries.

W&B: off,
on.

Gemma 3 4B, AutoModel, transport-only diagnostic

A 10-step native-rollout G=16 pair on pre-rebase candidate 5496708e also
exercised the generic AutoModel transport path. The first-step trajectory hash
matched exactly. Averaged across ten steps, serialized payloads fell 43.12x at
the exact sharded vLLM-generation Ray boundary, 15.68x in the rollout-to-policy
batch, and approximately 2.00x at the sharded policy logprob, reference-logprob,
and training Ray boundaries. The 2x shard-local result is expected because 16
logical generations were distributed over eight AutoModel DP workers.

This is transport-only evidence. Gemma is not included in model correctness or
performance qualification because its unrelated vLLM/AutoModel token-logprob
mismatch was catastrophic with deduplication both disabled and enabled.

W&B: off,
on.

Current-#3414 validation

• Implementation-equivalent full unit campaign: 956 passed, 6 skipped
  (SLURM 15214209). This included full GRPO, PPO, distillation, Gym/replay,
  batching, AutoModel data, vLLM utils, config validation, serialization, and
  payload-metric suites.
• Megatron CP=2 model tests: 2 passed (SLURM 15214753), including model-ingress,
  logprob/loss/gradient parity and distributed-checkpoint round-trip.
• The dedup patch was replayed without conflicts from c3b8dacb onto the latest
  feat: Multimodal nemo gym compatible grpo pipeline #3414 head f27dc6d6; git range-diff confirms the product patch is identical.
  The intervening feat: Multimodal nemo gym compatible grpo pipeline #3414 delta only changes PPO documentation/configuration,
  ppo.py, and dataset utilities, with no file overlap with this patch.
• On implementation commit 91959244, Ruff format/check passed and Pyrefly
  reported zero shown errors across 219 modules (255 pre-existing ignored
  diagnostics) in SLURM 15222427. The focused
  GRPO/PPO/data/Gym/replay/backend/config/payload suite passed on GPU with
  1004 passed, 6 skipped in SLURM 15222850.
• Bf16 policy-transport commit d95bb1ae ports Omni MR !12 to GRPO without
  materializing deduplicated rows. Focused validation: 196 generic GRPO/data/
  payload tests passed (SLURM 15234623), 39 AutoModel data tests passed
  (15234855), and 41 Megatron data tests passed with one expected 2-GPU skip
  (15234856). Ruff format/check passed; Pyrefly reported zero errors on the
  changed whitelisted module. The payload test verifies exactly 2x fewer
  physical pixel_values bytes plus lower protocol-5 serialized size.
• Scope-cleanup commit 29de7b52 removes the standalone design document and
  redundant per-recipe debug defaults. The resulting complete configuration set
  passed with 547 passed, 6 skipped in SLURM 15224236.
• An independent review of the implementation-equivalent c3b8dacb patch found
  no code-level merge blockers.
• Implementation-equivalent Qwen2.5-VL 3B AutoModel native G=16 jobs
  15216656/15217828 completed two steps each. The first trajectory hash
  matched exactly; later independently sampled trajectories diverged after
  optimizer/refit, as expected for this non-frozen smoke. Aggregate serialized
  reductions were 14.99x prompt repeat, 17.06x vLLM generation, 15.30x
  rollout-to-policy, 1.994x policy/reference logprobs, and 1.991x policy train.
  W&B: off,
  on.
• Implementation-equivalent Nemotron Omni 30B Megatron async Gym G=16, TP=2,
  EP=16, CP=2 job 15219896 completed two steps per side on four 8-GPU nodes. All
  losses were finite and TMPE stayed healthy: 1.01897-1.01950 off and
  1.02197-1.02418 on. Aggregate serialized reductions were 406.68x at Gym
  return, 15.67x replay push, 15.69x replay sample, 15.28x rollout-to-policy,
  3.97x policy logprobs, and 3.96x policy train; Gym request was exactly 1.00x.
  These were independent sampled runs, so the frozen 10-step campaign above,
  rather than this smoke, supplies the strict parity claim. W&B:
  off,
  on.

Evaluation limits

• Gemma 3 is excluded from correctness and performance qualification: its
  vLLM/AutoModel token-logprob mismatch is catastrophic with dedup both off and
  on. Its transport-only diagnostic is reported separately above.
• Audio/video model runs are not claimed; shared media primitives have focused
  coverage and are intentionally media-type agnostic.
• AutoModel VLM CP>1 is rejected by the upstream worker independently of this
  change. Megatron CP=2 is qualified; CP>2 is not claimed.
• Payload instrumentation is opt-in because protocol-5 serialization is itself
  nontrivial work. It is intended for qualification/debugging, not production.

Test plan

• PackedTensor logical/physical invariants, copy, pickle, legacy restore,
  empty rows, missing keys, cross-key alignment, and differing equivalence groups.
• Repeat/slice/concat/from-batches/dynamic-batch/sequence-pack/DP-shard
  normalization and shard-local re-interning.
• Native and Gym sync/async integrations, multi-turn stale-media handling,
  replay/DAPO/checkpoint restore, and flag-off behavior.
• Non-GRPO primitive blast radius: full PPO and distillation unit suites.
• AutoModel CP=1 worker materialization and Megatron CP=2 model tests.
• Exact sharded Ray-argument telemetry and W&B aggregation.
• Protected CI after PR submission.

[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.

[aroshanghias-nvd]

Superseded by #3534, which has the identical tested tree on a clean one-commit branch based directly on main. This avoids carrying the squash-merged #3414 commit history without force-pushing.