feat(data): add PackedTensor preprocessing modes

[rohitrango]

What does this PR do ?

Adds configurable PackedTensor preprocessing and patchifies Nemotron Omni pixels before materialization.

• Replaces the top-level pad_to_max_shape boolean with preprocess_mode and preprocess_kwargs.
• Adds native-resolution patchification into per-image (C_i, P²) blocks, packs them on dimension 0, and returns (1, total_C, P²).
• Carries preprocessing settings through slicing, concatenation, wire transport, and materialization.
• Enables patchification with patch size 16 for supported Nemotron Omni processors.
• Adds NemotronH_Omni_Reasoning_V3Processor to the placeholder-style processor set.

Issues

None.

Usage

PackedTensor(
    pixel_values,
    dim_to_pack=0,
    preprocess_mode="patchify",
    preprocess_kwargs={"patch_dim": 16},
)

Before your PR is "Ready for review"

• Read and followed the contributor guidelines.
• Added or updated unit coverage for preprocessing and patchification.
• Full pre-commit lint suite passed on Slurm job 18281365 (Ruff lint, import sorting, Ruff format, Pyrefly, and config checks).
• Affected unit suite passed on Slurm job 18281365: 90 passed, 1 skipped (megatron.bridge unavailable).
• GRPO and SFT GPU recipe validation passed on Slurm job 18281365.
• Updated relevant inline documentation.

Additional Information

• Stacked on feat: add Energon-owned SFT sequence packing #4105 via the rohit/sft_v2_stage2 base branch. Retarget this PR to main after feat: add Energon-owned SFT sequence packing #4105 merges.

Scripts validated

uv run examples/run_vlm_grpo.py --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-2n8g-megatron-tp8ep8.v1-tq_mooncake.yaml logger.wandb.name=grpo-nemotron-with-patches logger.wandb_enabled=true logger.wandb.project=sft-dev cluster.num_nodes=1

uv run examples/run_sft_v2.py --config examples/configs/recipes/vlm/vlm_sft-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8-energon.v1.packing.yaml logger.wandb.name=sft-nemotron-with-patches logger.wandb_enabled=true logger.wandb.project=sft-dev sft.max_steps=50

(brown: baseline, green: with patches)

Previous PR and self-review

Previous PR: #4082. A self-review was performed on the previous PR, and the resulting feedback was addressed in this PR.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[terrykong]

This PR moves Nemotron Omni pixel_values from pad-to-max-shape to patchify, so each image is cut into vision patches at its native resolution and packed into one sequence instead of being padded up to the largest image in the batch. It adds preprocess_mode / preprocess_kwargs to PackedTensor, a _patchify_segments helper, and a get_preprocess policy function that picks the mode per processor and per field.

The patchify layout is not invented here — it's what Bridge's own Nemotron Omni SFT collate produces (collate_fn.py L342-L346), and the per-frame kernel does the same reshape/permute (nemotron_omni_utils.py L63-L67). Three lanes checked independently that _patchify_segments is bit-for-bit identical to Bridge's _patchify_dynamic_images on mixed-resolution input. _preprocess_spec is also a genuinely good call: it replaces what would otherwise be ten hand-written copies of the same keyword pair across every copy/slice/concat path, and an audit of all of them found none missed.

The one blocking problem is that get_preprocess decides patchify from the processor alone, with no idea which backend will read the result. Megatron takes pre-patchified input by design; the AutoModel/DTensor forward takes 4-D only, so two nightly VLM lanes now hit an unpack error raised inside HuggingFace code — details in the inline comment on get_preprocess.

One note on the description: it says a self-review on the previous PR (#4082) "was addressed in this PR", but the change this PR contributes is identical to #4082's — every added and removed line matches across all 10 files, with only hunk offsets moved by the rebase — so no code changed as a result. If that review happened somewhere off GitHub, could you link it? A pending review would be invisible to us, so I can't tell from here whether one exists.

FYI, for a later PR and not this one: the next PR up the stack calls PackedTensor(..., pad_to_max_shape=True) in nemo_rl/data/energon/multimodal/task_encoders/nemotron_multimodal.py while carrying this PR's __init__, which takes no such keyword and no **kwargs — that'll raise TypeError once this lands.

Reviewed by a team of agents.

Generated by Claude Code

[rohitrango]

Validated end to end:

uv run examples/run_vlm_grpo.py --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml
uv run examples/run_vlm_grpo.py --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v2.yaml

[rohitrango]

/ok to test 2fc6596

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

[terrykong]

/ok to test 972e9f4

[rohitrango]

/ok to test 1cba512

[rohitrango]

/ok to test 30821cc