Skip to content

feat(data): add PackedTensor preprocessing modes - #4082

Open
rohitrango wants to merge 1 commit into
NVIDIA-NeMo:rohit/sft_v2_stage2from
rohitrango:rohit/packedtensor_padding
Open

feat(data): add PackedTensor preprocessing modes#4082
rohitrango wants to merge 1 commit into
NVIDIA-NeMo:rohit/sft_v2_stage2from
rohitrango:rohit/packedtensor_padding

Conversation

@rohitrango

@rohitrango rohitrango commented Sep 10, 2026

Copy link
Copy Markdown
Contributor

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

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

image
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)
image

@rohitrango
rohitrango requested review from a team as code owners September 10, 2026 18:22
@copy-pr-bot

copy-pr-bot Bot commented Sep 10, 2026

Copy link
Copy Markdown

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.

@github-actions github-actions Bot added the Documentation Improvements or additions to documentation label Sep 10, 2026
@rohitrango
rohitrango changed the base branch from main to rohit/sft_v2_stage2 September 10, 2026 18:25
@rohitrango
rohitrango force-pushed the rohit/packedtensor_padding branch 3 times, most recently from 07aae6b to 9aa7c9d Compare September 10, 2026 18:41
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant