[data] tracking: VLM SFT benchmark — high-res & interleaved multi-image dataset support #3133

Summary

Track end-to-end support for high-resolution image and interleaved multi-image datasets in the VLM SFT benchmark pipeline (Energon-based).

Driven by user requests to cover corner cases where encoder bottlenecks appear (e.g., 10+ images per sample, very large image sizes).

Phase 1 — Functional Test (throughput + loss convergence)

Qwen-VL model
Verify data loading via Energon, confirm images are properly processed, throughput holds up, and loss converges. (needs data sample balance for vision encoder)

• TIGER-Lab/Mantis-Instruct — interleaved multi-image (10+ images/sample), throughput verification + Energon image processing
• Ahren09/InfoVQA — high-res infographic images, throughput + loss convergence
• lmms-lab/LLaVA-Video-178K — video pipeline functional check

Functional test checklist (per dataset)

• datasets.load_dataset() loads successfully
• Data format parseable by current VLM data pipeline (Energon)
• Images/video decode correctly (especially large images, multi-image, video)
• Single forward pass without OOM
• 10-50 step training run completes without errors, loss converges
• For video: confirm on-the-fly loading support

Changes note:

• Adding proper support for video
• Adding control and processing logic for max_num_images, max_num_frames, max_visual_tokens
• Fixing issue when sequence truncation truncating visual tokens (e.g. when running with PP>1, sequence length needs to constant)
• Adding wrapper QwenVLEnergonProvider to bring VLM-specific knobs to CLI
  => General setup:
  • max_pixels: upper-bound for media height/width, resize if exceeding
  • max_num_images, max_num_frames: upper-bound for number of medias, dropping if exceeding
  • max_visual_tokens: upper-bound for number of media tokens, dropping if exceeding

Records:

• Wandb convergence runs for 8b:
  • infovqa: https://wandb.ai/nvidia/vlm_energon_naver/runs/r3ocpr6e?nw=nwuserhuvu
  • mantis: https://wandb.ai/nvidia/vlm_energon_naver/runs/5tpirdqv?nw=nwuserhuvu
  • llava_video: https://wandb.ai/nvidia/vlm_energon_naver/runs/g7mbdx6t?nw=nwuserhuvu
• Wandb benchmarking runs for 30b and 235b:
  • https://wandb.ai/nvidia/vlm_energon_naver/reports/Qwen3-VL-Energon-benchmarking--VmlldzoxNjkzOTcwNw

Nemotron-Omni

• TIGER-Lab/Mantis-Instruct — interleaved multi-image (10+ images/sample), throughput verification + Energon image processing
• Ahren09/InfoVQA — high-res infographic images, throughput + loss convergence
• lmms-lab/LLaVA-Video-178K — video pipeline functional check

Functional test checklist (per dataset)

• datasets.load_dataset() loads successfully
• Data format parseable by current VLM data pipeline (Energon)
• Images/video decode correctly (especially large images, multi-image, video)
• Single forward pass without OOM
• 10-50 step training run completes without errors, loss converges
• For video: confirm on-the-fly loading support

Records:

• Wandb convergence runs for 8b:
  • infovqa: https://wandb.ai/nvidia/vlm_energon_naver/runs/pxdv6i81?nw=nwuserhuvu
  • mantis: https://wandb.ai/nvidia/vlm_energon_naver/runs/cg39bjqw?nw=nwuserhuvu
  • llava_video: https://wandb.ai/nvidia/vlm_energon_naver/runs/bz1c38zw?nw=nwuserhuvu

Lastest updates:

• Full verified with video dataset (lmms-lab/LLaVA-Video-178K). Including: fine-tuning, loss curve inspection, running inference. Note: video VLM always uses none-dynamic resolution path.
• Full verified with images dataset (TIGER-Lab/Mantis-Instruct, Ahren09/InfoVQA) on dynamic resolution path. Including: fine-tuning, loss curve inspection, running inference.
• Updated Nemotron-Omni's README.md adding more details/instructions on input processing modes (dynamic/non-dynamic resolution).
• Exposing existing issues:
  • Resume training broke with checkpoint optimizer loading OOM: Reported to Chen, fixed with fix(ckpt): reduce distributed optimizer resume memory #4936.
  • Context parallel divergence: Reported to Chen, fixed with [model] fix: preserve Nemotron runtime config #4974.

Phase 2 — Benchmark + Perf Optimization

After functional tests pass, move to performance reproduction and benchmark accuracy.

• TIGER-Lab/Mantis-Instruct — interleaved multi-image (10+ images/sample), throughput verification + Energon image processing
• Ahren09/InfoVQA — (continued) benchmark accuracy for high-res
• lmms-lab/LLaVA-Video-178K — (continued) video perf benchmark
• ~~lmms-lab/LLaVA-NeXT-Interleave-Bench~~~~ — perf reproduction for interleaved multi-image (17 GB, 39K rows)~~
• ~~HuggingFaceFV/finevideo~~~~ — stretch: long video + audio~~

Tokens stats and data processing time:

Link to measured performance: https://docs.google.com/spreadsheets/d/1MlPPv2UA-rDC-VhrniRN384-XwcYn5S5YkReJ1c4YVU/edit?usp=sharing

Since these datasets are vision heavy (either containing multiple images, videos or high-resolution), the main gain comes from improving the vision module efficiency. One method is by enabling data parallelism for vision module across CP ranks. The input images/frames tensors will be sharded along the number of images dimension (data parallel). Each shard will run on a CP rank. The results will then gathered before sending to LLM. This mechanism (vision data parallel across CP ranks) are similar to LLava implementation in Megatron-LM.

Related changes regarding enabling DP for vision module:

• Code implementation PR: feat(model): DP of Nemotron Omni vision encoder across CP ranks #5645
• Docs + skills PR: docs(model): document Nemotron Omni vision_dp_over_cp #5646

Dataset Summary

Context

• All benchmarking and validation should be conducted within an Energon-based setup.
• Users specifically care about samples with 10+ image paths per sample, which expose encoder bottlenecks.
• MP-DocVQA and MINT-1T-PDF were considered but excluded since Mantis-Instruct and LLaVA-NeXT-Interleave-Bench already cover 10+ image interleaved samples.