feat(vlm): route VLM GRPO through TQ trainer when data_plane.enabled - #2957
Open
ZhiyuLi-Nvidia wants to merge 17 commits into
Open
feat(vlm): route VLM GRPO through TQ trainer when data_plane.enabled#2957ZhiyuLi-Nvidia wants to merge 17 commits into
ZhiyuLi-Nvidia wants to merge 17 commits into
Conversation
Contributor
Author
|
/ok to test 14e1105 |
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_vlm
branch
from
July 27, 2026 18:37
14e1105 to
38ef87c
Compare
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_vlm
branch
from
July 29, 2026 00:25
88aef2c to
9977c05
Compare
Contributor
Author
|
/ok to test 9977c05 |
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_vlm
branch
from
July 29, 2026 21:41
9977c05 to
6afd638
Compare
Contributor
Author
|
/ok to test 6afd638 |
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_vlm
branch
from
July 29, 2026 22:16
6afd638 to
456b985
Compare
Contributor
Author
|
/ok to test 456b985 |
zyzhou5
reviewed
Jul 30, 2026
| from nemo_rl.algorithms.grpo_sync import grpo_train_sync | ||
|
|
||
| print("🚀 Running synchronous VLM GRPO training (TransferQueue)") | ||
| return grpo_train_sync |
Contributor
There was a problem hiding this comment.
do we have a plan to update the file name? the current naming is quite confusing
grpo_train_sync -> TQ path while grpo_train -> legacy path
Contributor
Author
There was a problem hiding this comment.
Yeap, I'd expect this issue would be fixed once legacy path is retired.
|
|
||
| # Packed per-sample: jagged; ``PackedTensor`` in-memory, wire form is | ||
| # ``torch.nested`` parent + ``<key>__lengths`` companion. | ||
| PACKED_MULTIMODAL_FIELDS = frozenset( |
Contributor
There was a problem hiding this comment.
Suggested by agents - Nemotron-Omni needs imgs_sizes here
Contributor
Author
There was a problem hiding this comment.
Added. Also added num_frames since processors.py creates it right next to imgs_sizes and they're coupled — would have failed on the next field otherwise.
Contributor
There was a problem hiding this comment.
can you add a note/TODO to refactor this later into a single source of truth (i.e. ProcessorInterface)
Contributor
Author
There was a problem hiding this comment.
nemotron-omni now supported.
Also added TODO comment.
| # optional routed_experts under R3 replay. | ||
| present_multimodal = _LP_MULTIMODAL_FIELDS & set(meta.fields or ()) | ||
| lp_fields = fields_with_optional_routed_experts( | ||
| [*LP_SEED_FIELDS, *present_multimodal], |
Contributor
There was a problem hiding this comment.
These image fields are included for logprobs, but train_from_meta() later fetches only DP_TRAIN_FIELDS, which has no image fields. So logprobs see the image, while the GRPO training update does not.
Contributor
Author
There was a problem hiding this comment.
Confirmed, fixed. Train now fetches the same multimodal fields as logprob. train/probs_ratio_max and train/probs_ratio_min are normal now.
ZhiyuLi-Nvidia
added a commit
that referenced
this pull request
Aug 18, 2026
Addresses PR #2957 review comments. 1. train_from_meta / train_microbatches_from_meta fetched only the static text-only DP_TRAIN_FIELDS, so a VLM training forward ran image-blind while prev/ref logprobs were computed *with* images. Both now ship DP_TRAIN_FIELDS + _present_multimodal_fields(meta), the same per-batch add-on the logprob dispatch uses (renamed _LP_MULTIMODAL_FIELDS -> _WIRE_MULTIMODAL_FIELDS since it is no longer logprob-only). The Qwen3.5-35B-A3B geo3k run shows the fingerprint: train/ probs_ratio_max 216-2661 and probs_ratio_min 0.0 on every step, with probs_ratio_clamped pinned to the 1.2/0.8 clip bounds, while the non-DP baseline holds probs_ratio_max == probs_ratio_min == 1.0 (the strictly-on-policy invariant for a single inner step). Mean probs_ratio stays ~0.998 because only the image tokens are wrong, which is why the metrics compared in the PR body -- all computed upstream of the training forward -- did not move. 2. BatchedDataDict.truncate_tensors narrowed every >=2-D tensor on dim 1 to the microbatch seqlen. The in-memory PackedTensor form is skipped by torch.is_tensor, but the data-plane wire form is a plain tensor whose dim 1 is patch/image count, so dynamic batching would silently corrupt images (or raise narrow: length > size). Skip PACKED_MULTIMODAL_FIELDS and their __lengths companions; per-token fields stay sequence-aligned and still truncate. 3. Register imgs_sizes (Nemotron-Omni, packed along dim 0 per data/processors.py) in PACKED_MULTIMODAL_FIELDS. Tests: train/logprob multimodal field parity + text-only stays empty (tests/unit/data_plane/test_multimodal_wire_roundtrip.py), and truncate_tensors leaving the wire form intact while truncating mm_token_type_ids (tests/unit/data/test_multimodal_dict.py). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_vlm
branch
from
August 18, 2026 09:28
456b985 to
231b2ad
Compare
Contributor
Author
|
/ok to test 231b2ad |
Contributor
Author
|
/ok to test c0ae0bd |
ZhiyuLi-Nvidia
added a commit
that referenced
this pull request
Aug 19, 2026
Addresses PR #2957 review comments. 1. train_from_meta / train_microbatches_from_meta fetched only the static text-only DP_TRAIN_FIELDS, so a VLM training forward ran image-blind while prev/ref logprobs were computed *with* images. Both now ship DP_TRAIN_FIELDS + _present_multimodal_fields(meta), the same per-batch add-on the logprob dispatch uses (renamed _LP_MULTIMODAL_FIELDS -> _WIRE_MULTIMODAL_FIELDS since it is no longer logprob-only). The Qwen3.5-35B-A3B geo3k run shows the fingerprint: train/ probs_ratio_max 216-2661 and probs_ratio_min 0.0 on every step, with probs_ratio_clamped pinned to the 1.2/0.8 clip bounds, while the non-DP baseline holds probs_ratio_max == probs_ratio_min == 1.0 (the strictly-on-policy invariant for a single inner step). Mean probs_ratio stays ~0.998 because only the image tokens are wrong, which is why the metrics compared in the PR body -- all computed upstream of the training forward -- did not move. 2. BatchedDataDict.truncate_tensors narrowed every >=2-D tensor on dim 1 to the microbatch seqlen. The in-memory PackedTensor form is skipped by torch.is_tensor, but the data-plane wire form is a plain tensor whose dim 1 is patch/image count, so dynamic batching would silently corrupt images (or raise narrow: length > size). Skip PACKED_MULTIMODAL_FIELDS and their __lengths companions; per-token fields stay sequence-aligned and still truncate. 3. Register imgs_sizes (Nemotron-Omni, packed along dim 0 per data/processors.py) in PACKED_MULTIMODAL_FIELDS. Tests: train/logprob multimodal field parity + text-only stays empty (tests/unit/data_plane/test_multimodal_wire_roundtrip.py), and truncate_tensors leaving the wire form intact while truncating mm_token_type_ids (tests/unit/data/test_multimodal_dict.py). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_vlm
branch
from
August 19, 2026 00:01
c0ae0bd to
d265750
Compare
Contributor
Author
|
/ok to test d265750 |
Contributor
|
can this PR be validated on The implementation from #3290 explicitly forbids models like Nano-Omni from working with the DataPlane, see |
rohitrango
requested changes
Aug 24, 2026
|
|
||
| # Packed per-sample: jagged; ``PackedTensor`` in-memory, wire form is | ||
| # ``torch.nested`` parent + ``<key>__lengths`` companion. | ||
| PACKED_MULTIMODAL_FIELDS = frozenset( |
Contributor
There was a problem hiding this comment.
can you add a note/TODO to refactor this later into a single source of truth (i.e. ProcessorInterface)
| # the padding could fix them. Padding to the *global* batch max | ||
| # (not a per-shard max) is deliberate: every DP rank then sees | ||
| # identical trailing dims for the forward. | ||
| if self.pad_to_max_shape: |
Contributor
There was a problem hiding this comment.
why not flatten the tensor and unflatten after? this can allow jagged tensors at this stage instead of having multiple points to pad the tensor, and is easier to maintain for future ops for mm tensors. maybe im missing something?
Contributor
There was a problem hiding this comment.
can this PR be validated on nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 ?
The implementation from #3290 explicitly forbids models like Nano-Omni from working with the DataPlane, see nemo_rl/data_plane/worker_mixin.py:137 . The check was added because PackedTensor objects could not be passed through the TQ, which is now implemented here.
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_vlm
branch
from
August 30, 2026 05:23
d265750 to
f67dd38
Compare
Mirror the `run_grpo.py` launcher pattern in `run_vlm_grpo.py` so the VLM entrypoint dispatches to `grpo_train_sync` (TransferQueue) when `data_plane.enabled=true` and otherwise stays on the legacy `grpo_train`. The policy factory is also selected at the launcher level (`TQPolicy` when enabled) so the legacy trainer remains data-plane-agnostic per the architecture invariant. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Mirror `grpo.py`'s `_initial_policy_generation_stale` check that the legacy trainer uses to seed `POLICY_GENERATION_STALE`. The sync trainer hardcoded `POLICY_GENERATION_STALE = True`, forcing a refit at iter 1 even when setup() had already synced weights (`synchronizer.is_stale` is `False`). The redundant refit resets vLLM's CUDA-graph capture and prefix-cache / KV-cache state, so vLLM's step-1 `generation_logprobs` disagree with Megatron's re-scoring on the same tokens by ~ln(40) nats/token (`train/token_mult_prob_error` ~ 40 at step 1, converging by step 3). Loss is unaffected (`probs_ratio = 1` throughout), but the diagnostic metric spike is cosmetic noise and gets masking-thresholded in some configs. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
The TQ sync trainer's ``sync_rollout_actor`` silently dropped
PackedTensor multimodal fields (pixel_values, image_grid_thw,
mm_token_type_ids for Qwen2.5-VL / Qwen3-VL, token_type_ids for
Gemma3) because the write path filtered by
``isinstance(v, torch.Tensor)`` — ``PackedTensor`` is not a
``torch.Tensor`` subclass. Consequence: trainer's ``prev_logprobs``
were computed without image embeddings while vLLM's
``generation_logprobs`` used them, producing on Qwen3.5-A3B-Base +
geometry3k:
* train/token_mult_prob_error avg ~176 (should be ~1.02)
* train/sampling_importance_ratio ~0.982 uniform (should be ~1.000)
* train/policy_kl_error avg ~175 (should be ~0.02)
Loss was on-policy (probs_ratio = 1) so training didn't crash, but
the diagnostic metrics flagged the mismatch clearly.
Fix threads multimodal fields end-to-end via a binary wire dispatch
(torch.Tensor | np.ndarray[object]) with these pieces:
* nemo_rl/data_plane/field_registry.py — new FieldSpec (alignment /
container / stages) collapses five sprinkled schema constants
(DP_TRAIN_FIELDS, LP_SEED_FIELDS, TOKEN_ALIGNED_FIELDS,
ADDITIONAL_OPTIONAL_KEY_TENSORS, LP_STAGE_EXCLUDED) into one
registry. Structural container check catches mis-registrations
loudly at write time.
* nemo_rl/experience/sync_rollout_actor.py — write-side converts each
PackedTensor to torch.nested.NestedTensor (jagged layout) +
companion ``<key>__lengths`` int32 tensor before crossing the wire
boundary. None entries are padded with zero-length tensors to keep
the nested batch dim aligned with lengths (mixed text/image
batches). SmolVLM (dim_to_pack=1) fails loudly with
NotImplementedError — ragged_idx threading is a follow-up.
* nemo_rl/distributed/batched_data_dict.py::get_multimodal_dict —
read-side reconstructs the per-sample PackedTensor by slicing the
materialize-padded rectangular tensor at ``__lengths[i]``.
Defensive assertion on batch-dim alignment between the parent and
companion.
* nemo_rl/data_plane/codec.py — pack_jagged_fields dispatch is binary
(Tensor + ndarray[object]); errors loudly on any other type. No
packed_tensor container category.
* nemo_rl/data_plane/column_io.py::kv_first_write — filter matches
the codec's binary universe (Tensor incl. torch.nested; ndarray[object]).
* nemo_rl/models/policy/tq_policy.py::_logprob_dispatch — fetches
via field_registry.fields_for_stage("logprob" | "ref_lp") instead
of a hardcoded include list. New multimodal keys drop into the
registry as one line each; no per-stage constant updates needed.
Diagnostic sub-timers (rollout / flatten / bulk_assemble /
kv_first_write / finish_gen) emit under timing/rollout/sub_* so the
driver-side ``generation`` timer can be decomposed for perf analysis.
Empirical verification on Qwen3.5-A3B-Base + geometry3k (2N x 8G,
container nemo-rl:nightly-07242026):
* SIR 1.00005 / 1.00002 / 0.99998 (bit-for-bit legacy parity)
* tmpe step 1..3 = 1.019 / 1.020 / 1.057
* mean(3-9) total_step_time: TQ 122.0s vs Legacy 125.3s — **TQ is
3.3s faster at steady state**. sub_kv_first_write dropped from
~2.0s (np.ndarray[object] of PackedTensor wrappers, per-object
pickle) to ~0.64s (native torch.nested storage buffer). Ray RPC
(driver ↔ actor) ~2.8s, TQ put_samples ~0.64s, net TQ tax on
generation ~3.4s, offset by ~5s savings in policy_training from
per-worker TQ fetch. Net TQ throughput ~2.6% faster than legacy.
Known limitations documented in code:
* SmolVLM ``dim_to_pack=1`` support requires ragged_idx handling
(follow-up).
* mm_token_type_ids inclusion in the logprob fetch is what
actually fixes Qwen2.5-VL / Qwen3-VL 3D RoPE positional encoding
for image tokens on the trainer side.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Adds three tests exercising the VLM multimodal wire path added by the preceding commit: - `test_get_multimodal_dict_full_vlm_wire_roundtrip` — build a realistic batch (pixel_values / image_grid_thw as PackedTensor with a None entry, mm_token_type_ids as per-token rectangular tensor, plus non-multimodal input_ids/token_mask). Encode via `to_nested_wire`, simulate materialize (`to_padded_tensor`), decode via `from_nested_wire`. Assert every packed field's `.as_tensor()` matches pre-wire, per-token fields pass through unchanged, and non-multimodal keys are silently skipped. Covers both `as_tensors=True` and `as_tensors=False` paths. - `test_get_multimodal_dict_missing_companion_asserts` — wire-form parent without its `__lengths` companion must raise with a wire-contract message, not a bare KeyError deep in trainer forward. - `test_get_multimodal_dict_empty_batch_skips_wire_field` — 0-row DP shard: `from_nested_wire` returns None and the read side skips the field (no crash on `PackedTensor.__init__`'s len>0 assert). Also adds a static registry check: the field names used by the tests must be in `PACKED_MULTIMODAL_FIELDS` / `PER_TOKEN_MULTIMODAL_FIELDS` — guards against a future rename that would silently drop the field. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
End-to-end verification through the data-plane ABC contract: build a rollout batch with PackedTensor pixel_values + image_grid_thw + mm_token_type_ids → simulate sync_rollout_actor's write loop → kv_first_write to NoOpDataPlaneClient → read_columns (materialize) → assert get_multimodal_dict on the fetched batch matches the pre-wire .as_tensor() output. Guards the silent-drop regression class this PR was written to fix, plus the pad_to_seqlen exclusion for multimodal fields (asserts pixel_values shape stays [B, max_patches, ...] and doesn't inflate to [B, seqlen, ...]). Complements the two BatchedDataDict-layer defensive tests in tests/unit/data/test_multimodal_dict.py (companion-missing and empty-batch guards). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Addresses PR #2957 review comments. 1. train_from_meta / train_microbatches_from_meta fetched only the static text-only DP_TRAIN_FIELDS, so a VLM training forward ran image-blind while prev/ref logprobs were computed *with* images. Both now ship DP_TRAIN_FIELDS + _present_multimodal_fields(meta), the same per-batch add-on the logprob dispatch uses (renamed _LP_MULTIMODAL_FIELDS -> _WIRE_MULTIMODAL_FIELDS since it is no longer logprob-only). The Qwen3.5-35B-A3B geo3k run shows the fingerprint: train/ probs_ratio_max 216-2661 and probs_ratio_min 0.0 on every step, with probs_ratio_clamped pinned to the 1.2/0.8 clip bounds, while the non-DP baseline holds probs_ratio_max == probs_ratio_min == 1.0 (the strictly-on-policy invariant for a single inner step). Mean probs_ratio stays ~0.998 because only the image tokens are wrong, which is why the metrics compared in the PR body -- all computed upstream of the training forward -- did not move. 2. BatchedDataDict.truncate_tensors narrowed every >=2-D tensor on dim 1 to the microbatch seqlen. The in-memory PackedTensor form is skipped by torch.is_tensor, but the data-plane wire form is a plain tensor whose dim 1 is patch/image count, so dynamic batching would silently corrupt images (or raise narrow: length > size). Skip PACKED_MULTIMODAL_FIELDS and their __lengths companions; per-token fields stay sequence-aligned and still truncate. 3. Register imgs_sizes (Nemotron-Omni, packed along dim 0 per data/processors.py) in PACKED_MULTIMODAL_FIELDS. Tests: train/logprob multimodal field parity + text-only stays empty (tests/unit/data_plane/test_multimodal_wire_roundtrip.py), and truncate_tensors leaving the wire form intact while truncating mm_token_type_ids (tests/unit/data/test_multimodal_dict.py). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
run_vlm_grpo passes processor= to whichever trainer _select_trainer
returns, but grpo_train_sync did not accept it, so every VLM run with
data_plane.enabled=true died after full model load with:
TypeError: grpo_train_sync() got an unexpected keyword argument
'processor'
grpo_train uses processor for exactly one thing:
attach_initial_nemo_gym_image_payloads, gated on
grpo.deduplicate_multimodal_data. That flag is already rejected for
data_plane.enabled=true by _validate_multimodal_dedup_capability, so the
sync trainer can never need the object -- accept it for signature parity
and assert the invariant instead of ignoring the argument, so relaxing
that upstream guard fails loudly rather than silently dropping image
payloads.
Adds test_sync_trainer_is_call_compatible_with_legacy_trainer, which
diffs the two signatures. The e2e that surfaced this costs two nodes and
~12 minutes of setup before it fails; the signature check is instant.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
The DTensor/AutoModel backend has its own sequence-dim pre-flight,
check_sequence_dim, which is the analog of megatron's
get_and_validate_seqlen. Only the megatron one was taught to skip
multimodal fields, so a TQ VLM run on the automodel backend died at
step 1:
AssertionError: Dim 1 must be the sequence dim, expected dim 1=2432
but got shape torch.Size([32, 1, 3])
[32, 1, 3] is image_grid_thw (batch, num_images, t/h/w). Packed
multimodal fields are never sequence-aligned -- dim 1 is
num_images/num_patches -- and their <key>__lengths wire companions are
1-D. In-memory these ride as PackedTensor and are skipped by
torch.is_tensor, but the data-plane wire form is a plain tensor.
The skip goes inside check_sequence_dim rather than onto its existing
skip_keys parameter: all seven call sites need it, and none of them
should have to know the wire format.
Also collapses three byte-identical inline copies of the same check in
the v1 DTensor worker into calls to check_sequence_dim. v1 is the
default worker (dtensor_cfg._v2 defaults to false), so it had the same
latent bug; sharing the helper keeps the skip in one place.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Multi-agent review of the VLM/TQ path. Each item below is either
reproduced on hardware or verified against source; severities are
post-adversarial-review.
1. mooncake_cpu was dead on arrival for any VLM run. The
`<key>__lengths` companions minted by `PackedTensor.to_nested_wire`
are dense int32[B], and `_promote_1d_leaves` rejects any dense 1-D
field not declared in PROMOTE_1D_FIELDS. Reproduced on 1 node:
ray::SyncRolloutActor.rollout_to_tq()
ValueError: Mooncake field 'pixel_values__lengths' is a dense 1D
tensor but is not declared in data_plane.schema.PROMOTE_1D_FIELDS
Derived the companion names from PACKED_MULTIMODAL_FIELDS rather
than hand-listing them, so a new packed modality is covered
automatically. Fixes the read side too -- `_from_wire` squeezes only
declared fields. Verified: same job now runs both steps clean.
Missed because every e2e used backend=simple, where the transform
is a no-op.
2. The rollout write dropped `pixel_dtype`, which every legacy analog
passes (grpo.py:2106/:3195, grpo_sync.py:644) and no worker
re-applies. pixel_values are deliberately fp32 out of the processor
(see data/processors.py), so they crossed the wire fp32 where legacy
shipped bf16 -- 2x the largest column. Verified by inspection of all
call sites, not measured.
3. `to_nested_wire` concatenated a logical row's segments before
applying pad_to_max_shape, inverting `as_tensor`'s order; a
multi-segment dedup row with differing trailing dims raised in
torch.cat before the padding could run. Now pads segments first.
Latent today (dedup is rejected under data_plane) -- trap removal.
4. `from_nested_wire` turned zero-length rows into empty tensors rather
than None, so an image-free DP shard produced pixel_values of shape
(0, ...) where legacy gives None, and logical_segment_counts_by_row
reported 1 instead of 0. Also added a companion/row length check --
a short companion would silently misalign images against samples.
5. fp8 QKV calibration filtered on DP_CALIB_INPUT_FIELDS, which names a
`multi_modal_inputs` column that is never written, so a VLM run
calibrated image-blind. Newly reachable: run_vlm_grpo only started
routing to grpo_train_sync in this PR.
6. Docs/typing: the `_logprob_dispatch` docstring stated the opposite of
what the code does (a reader following it would restore the
image-blind bug); the registry comment claimed to be the "single
source of truth" when data/processors.py performs the actual
classification; a comment cited a test path that does not exist;
TOKEN_ALIGNED_FIELDS hand-duplicated PER_TOKEN_MULTIMODAL_FIELDS;
`encode_multimodal_for_wire` was unannotated in a pyrefly-checked
file.
Not addressed here, both needing a product decision: SmolVLM cannot run
with data_plane.enabled (get_dim_to_pack_along returns 1, which
to_nested_wire rejects) and wants a setup()-time gate; and VLM+TQ has no
nightly coverage because tests/test_suites/llm/common-tq.env gates on
^(grpo|dapo|prorlv2)-, which no vlm_grpo-* recipe can match.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Closes the coverage gaps the review found. 218 passed on GPU (tests/unit/data_plane/ + unit/data/test_multimodal_dict.py). Encoder / wire boundary (test_multimodal_dict.py): - First direct coverage of `encode_multimodal_for_wire`: parent + __lengths emission, per-token passthrough, all-empty skip, unregistered-field KeyError, wrong-type guards. - `to_nested_wire` guard rails: dim_to_pack != 0 -> NotImplementedError, all-None -> (None, None), pad_to_max_shape rank mismatch -> ValueError. - Regression for the pad-before-concat fix: a dedup row spanning 2x4 and 4x2 segments round-trips instead of raising in torch.cat. - Regressions for the companion length check and for empty rows matching legacy None semantics. - truncate_tensors leaves wire-form multimodal intact while still truncating per-token maps. Dispatch (test_multimodal_wire_roundtrip.py): - The roundtrip helper now calls the production `encode_multimodal_for_wire` instead of reimplementing its branches -- the reimplementation was blind by construction to exactly the drift its docstring claimed to catch. - ref_lp and the SC `train_microbatches_from_meta` path each carry their own copy of the multimodal add-on and had none; the SC path could have regressed image-blind while the sync path stayed green. - The train/logprob parity assertion now compares against the full registry; the previous third assertion was implied by the two above it. - materialize's pad_to_seqlen exclusion was never actually exercised -- pad_to_seqlen comes from meta.extra_info, which the test never stamped, so the guard against the ~40x blow-up was untested despite the module docstring claiming otherwise. - `_stub_tq_policy` used `__del__ = None`, which does not suppress the destructor: CPython installs tp_finalize whenever __del__ is in the class dict, then calls None() -> TypeError, surfacing as PytestUnraisableExceptionWarning charged to whichever test is at GC. Architecture invariants (test_architecture_invariants.py): - run_vlm_grpo's `_select_trainer` copy is now pinned; only run_grpo's was, and the VLM launcher is the one that shipped the processor= TypeError. - The signature-compat check now binds the launcher's actual call shape rather than demanding full parameter parity, which would have forced every future grpo_train parameter into grpo_train_sync as dead weight. Seq-dim skip (test_automodel_train.py): - check_sequence_dim's multimodal skip, plus a negative case proving per-token fields are still validated. The matching get_and_validate_seqlen tests are not included: the container ships neither megatron.bridge nor transformer_engine, so they could not be verified here (the pre-existing TestGetAndValidateSeqlen cannot run in it either). They want the mcore CI shard. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Two type errors from the Lint check, both introduced by the preceding review-fix commit: - `row_segments` typed as list[list[Tensor]] but built by a comprehension whose `is not None` filter does not narrow Optional[Tensor], so it inferred list[list[Tensor | None]]. Rebuilt with an explicit loop. - `encode_multimodal_for_wire` yielded the `__lengths` companion while only guarding `nested is None`; `to_nested_wire` returns Optional for both, so the yield did not match the newly declared Iterator[tuple[str, Tensor]]. Guard both. The second is a direct consequence of annotating the function in that commit -- the annotation is what exposed the looseness. ruff was clean throughout, which is why this only surfaced in CI. Verified with pyrefly in the nightly container (errors shown: 0); the repo venv here cannot install dev deps. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
…anions TQ already tracks per-row shape and RL was overriding it in two places. Both overrides existed to undo something RL itself did, so remove the cause rather than the compensation. 1. Multimodal `<key>__lengths` companions `codec.materialize` ran `to_padded_tensor` on every nested leaf, rectangularizing packed multimodal fields to `[B, max_rows, ...]` and destroying the row boundaries. The `<key>__lengths` companion field existed solely to recover them. TQ stores one entry per row (`storage/managers/base.py::_generate_values` unbinds nested fields), so the value it hands back already carries the true per-row shapes. Multimodal fields now skip `to_padded_tensor` and stay nested; `PackedTensor.from_wire` unbinds instead of slicing by companion length. That removes a wire field per packed modality, the derived `PROMOTE_1D_FIELDS` union, the `_WIRE_MULTIMODAL_FIELDS` companion tier, and the `.endswith(LENGTHS_SUFFIX)` skips in both seq-dim validators and `truncate_tensors`. `pad_to_max_shape` is deliberately kept. It is not a transport artifact: `as_tensor` pads to the same batch max, and a deduplicated logical row can span segments with differing trailing dims that `torch.cat` cannot join otherwise. 2. `PROMOTE_1D_FIELDS` `transfer_queue.metadata.extract_field_schema` rebinds a *local* for 1-D inputs (`value = value.unsqueeze(-1)`) and derives the sample shape from it, while `_generate_values` iterates the *original* `(N,)` tensor into N 0-d rows. The schema claims `(1,)`; storage holds `()`. Only the KV path notices, because it is the only one that reconstructs from the schema -- `SimpleStorage` fetches stored objects and never consults it. RL compensated by reshaping the payload to match the wrong schema. `_patch_scalar_field_schema` fixes the schema instead, so the reported shape matches the stored rows and the column stacks back to a dense `(N,)`. It rebinds in all three importing modules (both storage managers bind the name at import time) and is self-verifying: it probes `KVStorageManager._generate_values` and refuses to install if a TQ revision ever starts storing 1-D fields as `(1,)` rows. This covers every dense 1-D field rather than a declared allowlist, so `PROMOTE_1D_FIELDS`, `_promote_1d_leaves` and `_from_wire`'s squeeze branch are all deleted. Tested: 255 passed / 11 skipped across tests/unit/data_plane/ plus test_multimodal_dict.py; 398 passed on test_batched_data_dict.py, tests/unit/data/ and test_automodel_train.py. The 11 skips are the mooncake_cpu fixtures (no RDMA on the test host), so the KV path the scalar-schema patch targets is NOT yet covered by a run -- that and Nemotron-Omni's pad_to_max_shape path remain unverified end to end. Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Three defects kept VLM GRPO from running on the TransferQueue data plane, each reproduced before fixing: * ``BatchedDataDict.slice`` indexed a raw ``torch.nested`` value on its ragged dim. A nested tensor satisfies ``isinstance(v, Tensor)`` and has a valid ``shape[0]``, so it cleared both guards and failed only at the indexing. ``codec.materialize`` now reassembles packed multimodal fields into ``PackedTensor`` at the decode boundary, restoring the invariant that no raw nested value reaches a consumer. * ``_from_wire`` densified any nested field whose rows happened to share a shape, which for a packed field is a data-dependent accident (every sample carrying one image) that discards row boundaries. * ``_validate_multimodal_dedup_capability`` rejected every ``data_plane.enabled=true`` config, blocking all six Nemotron-Omni recipes. The gap it guards is NeMo-Gym specific -- ``grpo_train_sync`` does not call ``attach_initial_nemo_gym_image_payloads``, itself gated on ``should_use_nemo_gym`` -- so the check now names that combination. ``to_wire`` additionally flattens each segment to 1-D. Rows then vary only in dim 0, so ``torch.jagged`` accepts ragged trailing dims and mixed rank alike; the batch-max padding it used to materialize is gone from the wire and from TQ storage, and TQ never falls back to the deprecated strided layout. Padding moves to ``as_tensor``, where a rectangle is actually required -- it is a model input constraint, not a transport one. The shapes flattening removes travel on ``KVBatchMeta.tags``, the transport's per-sample channel, projected with the rows by ``subset``/``slice``/``concat``. ``shard_meta_for_dp`` did not propagate tags and would have dropped them on every per-rank fetch. Nothing in ``nemo_rl/data_plane`` interprets them. ``_global_pad_shape`` pins the batch-wide pad target so every DP rank sees identical trailing dims; a rank-local max let the logprob and training passes encode the same media at different widths. It propagates through all nine ``PackedTensor`` constructors -- ``concat`` dropping it silently reverted ``shard_by_batch_size`` to a per-shard max. Tested: * Qwen3.5-35B-A3B geo3k 2n8g, 20/20 steps. This recipe trains one inner step per rollout, so ``probs_ratio`` is an identity: measured exactly 1.000000 on every step, across three successive wire formats. * Nemotron-Omni-30B-A3B clevr 1n8g, 10/10 steps. ``token_mult_prob_error`` 1.0138-1.0153 against 1.0138-1.0155 for the same recipe with ``data_plane.enabled=false``. * 245 passed / 11 skipped on tests/unit/data_plane and test_multimodal_dict. Nightly gates match what each recipe can actually assert: Qwen3.5 gates ``probs_ratio`` exactly; Nemotron-Omni trains 16 inner steps, where that metric measures policy drift and its max ranges 5.85-29.21 run to run on identical code (the non-data-plane path included), so it gates ``token_mult_prob_error`` instead. Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_vlm
branch
from
August 30, 2026 23:01
f67dd38 to
8d5e0fd
Compare
Contributor
Author
|
/ok to test 8d5e0fd |
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_vlm
branch
from
August 30, 2026 23:17
bde7809 to
530eca4
Compare
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
The pad width is scratch the model discards -- mcore crops it via imgs_sizes before patchification, and the AutoModel path rejects mixed-resolution batches outright -- so no consumer needs shards to agree on it. Drop the transported max from the row tags and the _global_pad_shape plumbing that carried it; as_tensor now computes the max over the rows it holds. Row tags keep shapes (payload geometry, unrecoverable after to_wire flattens) and pad (the field's policy flag), so the data plane no longer carries a pad target it cannot interpret. Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
cspades
added a commit
to cspades/RL
that referenced
this pull request
Sep 1, 2026
Signed-off-by: Cory Ye <cye@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
3 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.
Summary
Route VLM GRPO through the TransferQueue data plane when
data_plane.enabled=true, and thread VLM multimodal payloads (pixel_values,image_grid_thw,mm_token_type_ids,imgs_sizes,num_frames) through the TQ wire.examples/run_vlm_grpo.pymirrors therun_grpo.pylauncher pattern: dispatch togrpo_train_syncwhen the data plane is enabled, otherwise stay on the legacygrpo_train.torch.nestedparent plus a<key>__lengthscompanion.PackedTensor.to_nested_wire/from_nested_wireare the only boundary between the domain wrapper and the wire.materialize's pad-to-seqlen skip,truncate_tensors, both per-backend seq-dim validators, the TQ fetch list) dispatch through two registries innemo_rl/data/multimodal_utils.py.data_plane.enabled=false(default).Test plan
End-to-end.
vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16withdata_plane.enabled=truereachesMAX_STEPS=10. Job 15968308, 2N x 8G, containernemo-rl:nightly-08102026.squashfs, 49m17s.Wandb: https://wandb.ai/nvidia/nemorl-dataplane-zhiyul/runs/wjhk1cc3
Recipe switched from the megatron variant for two reasons:
vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-megatron-ep16is intests/test_suites/disabled.txt(pre-existing Qwen3.5 + Megatron + EP hang insample_tokens), and the automodel variant inheritsdynamic_batching.enabled=true, which is the path review comment #3 asked to exercise.train/probs_ratio_maxandtrain/probs_ratio_min= exactly 1.0 on all 10 steps. On a single-inner-step on-policy update the training forward must reproduceprev_logprobs, so this ratio is an invariant. Before the fix below it was 216–2661 / 0.0 with the clamped values pinned to the 1.2/0.8 clip bounds, against a non-DP baseline of exactly 1.0/1.0.train/rewardmean 0.444, range [0.370, 0.532].validation/accuracyat step 10 = 0.535 vs 0.544 non-DP baseline — within noise.train/token_mult_prob_errorspikes at steps 4 (48.3) and 9 (6.99), 1.022–1.062 elsewhere; the non-DP baseline had a comparable spike (9.5 at step 2). It compares rollout vsprev_logprobs, both computed with images, so it is upstream of this change, andprobs_ratiostaying at 1.0 shows no train/logprob mismatch remains.Mooncake backend.
data_plane.backend=mooncake_cpu: https://wandb.ai/nvidia/nemorl-dataplane-zhiyul/runs/qtn9sl4uReview fixes
Addresses the four review comments plus findings from a follow-up review:
train_from_metafetched only the static text-onlyDP_TRAIN_FIELDSwhile the logprob dispatch shipped the multimodal columns. Both now use_present_multimodal_fields(meta).truncate_tensorsnarrowed wire-form multimodal tensors on dim 1 under dynamic batching; the AutoModel backend'scheck_sequence_dimneeded the same skip as megatron'sget_and_validate_seqlen.imgs_sizesregistered (plusnum_frames, its coupled partner).grpo_train_syncacceptsprocessorfor launcher signature parity.pixel_dtype, matching the legacy analogs —pixel_valueswere crossing the wire in fp32 where legacy shipped bf16.to_nested_wirepads segments before concatenating (mirroringas_tensor) and walks logical rows underPackedTensordeduplication;from_nested_wirerestoresNonefor empty rows and validates companion length.