PoR: [VLM] Refactor: clean up VLM data pipeline (processing / batching / packing / masking) #4041

Summary

The current VLM data path in Megatron-Bridge has accumulated several design
issues that make it hard to add new VLM models, share code between the Energon
and Hugging Face data sources, and reason about packing / CP behavior.

This issue tracks a staged refactor to:

1. Unify VLM data processing between the Energon TaskEncoder and the HF
   processor path.
2. Move batching / padding / packing out of the VLM forward step and into
   the processing / collate layer.
3. Replace the current text-search masking with a more reliable strategy
   (Tony-style token-search or HF generation tag).
4. Investigate and fix the M-RoPE issue that surfaces when Qwen2.5-VL is
   wired into the VLM step.

It came out of a design review between @yaoyu-33 and @cuichenx.

Current state (problems)

1. Three layers are tangled together

Conceptually the VLM data path has three layers:

The code today does not respect this layering:

• Some processing (apply_chat_template, masking, tokenization) lives in
  tasks/vlm/task_encoder.py for the Energon path, and a near-duplicate
  copy lives in the HF encoder task encoder. These should be a single piece
  of code with one TaskEncoder-style API.
• The VLM forward step (forward_step / get_batch) currently does
  additional padding-then-unpadding work on samples that were already
  batched and padded by the Energon batch collator. The flow is roughly:
  pad in collator -> remove pad in step -> repack in step, which is both
  wasteful and confusing.

2. "Packing" in the VLM step is not real packing

What is currently called "packing" in the VLM step is in-microbatch
packing: at micro-batch time we concatenate everything in the microbatch
into one THD sequence with no bin size and no packing algorithm. That is
not equivalent to LLM-side offline sequence packing.

Implications:

• We cannot pack across microbatches without changing the data path.
• BSHD already pads to the longest sample in the microbatch; the current
  in-batch packing only saves the padding inside that microbatch.
• We need a real packing path eventually; see "Long-term options" below.

3. Masking is fragile

The current find_pattern_indices-style masking searches for tokenized
answer text inside input_ids. This is unreliable across tokenizers
(leading-space tokens, BOS handling, etc.) and forces the user to manually
craft answer strings that survive tokenization.

We have at least two better options:

• Token-search (Tony's approach): user supplies start/end token IDs
  (e.g. assistant turn boundaries) and we mask between them. Works for any
  number of (start, end) pairs, also covers tool-calling regions.
• HF chat-template generation tag: apply_chat_template(..., return_assistant_tokens_mask=True) already returns the assistant mask
  if the chat template marks it. Cleanest when the template supports it,
  but not all model templates do.

LLaMA-Factory uses the token-search variant with a fallback. We should
match a similar two-layer approach.

4. M-RoPE / Qwen2.5-VL on the VLM step

Qwen2.5-VL "runs" on the VLM step today but produces small numerical
errors. M-RoPE is not used by the VLM step path itself (that path goes
through GPTModel), so the error is in how positional encodings get
threaded through after the vision encoder. Needs a follow-up
investigation; not blocking the refactor above.

5. Vision-encoder side artifacts

image_grid_thw and similar artifacts must be passed through to the
vision encoder. They currently go via forward_args.update in
normalize_for_model, which is a Bridge-internal seam. We should make
this contract explicit so new VLMs (and Omni later) can plug in without
rediscovering it.

Proposed refactor

Step 1 — Unify processing

• Make HF-path "processing" and Energon-path "task encoding" share a single
  per-model class with the same API (encode-sample / collate).
• Default behavior: call HF AutoProcessor for image+text, apply chat
  template, tokenize, build labels mask. Per-model overrides only when the
  community processor does not cover something (e.g. extra metadata).

Step 2 — Strip the VLM forward step

• VLM forward_step / get_batch should only do:
  • read already-collated batch from the data iterator,
  • rename / route fields into forward_args,
  • skip fields on ranks that don't need them (PP/EP non-input ranks).
• Remove the in-step pad-then-unpad-then-repack logic.
• Move all batching / padding / in-batch packing into the collator
  (batch / collate_fn).

Step 3 — Fix masking

• Adopt token-search masking (multiple (start_token_id, end_token_id)
  pairs) as the default.
• If the chat template supplies the generation tag, prefer
  return_assistant_tokens_mask=True and use that directly.
• Document the trade-off and the fallback chain in the VLM data skill.

Step 4 — Long-term packing

In-batch packing stays as the default, but we open a path to either:

• Offline packing: run only the processor (no vision encoder) to get
  per-sample token counts (text tokens + image-placeholder tokens via
  image_grid_thw-based mapping), then build an offline pack index, the
  same way LLM-side packing works. This is feasible for both Energon and
  HF datasets.
• DP-group online packing: gather samples across the DP group inside
  a step (analogous to the MCore-dev rerun_data_iterator /
  sequence-packing scheduler) and pack at that granularity. The MCore
  scheduler implementation lives in the dev branch and is not yet on
  main — we would need to either upstream it or vendor it.

We pick one (likely offline packing) once Step 1–3 are landed.

Step 5 — Investigate M-RoPE on VLM step

Track separately. Goal: get Qwen2.5-VL bit-exact on the VLM step path,
or document why it cannot be and gate it appropriately.

Order of work

1. Generate a current-state map of where each piece of logic lives
   (loading / processing / batching / packing / masking) in
   src/megatron/bridge/data/ and src/megatron/bridge/training/ —
   primarily the VLM task encoder, the HF encoder task encoder, and the
   VLM forward step. (blocking — code is hard to navigate today)
2. Step 1: unify processing.
3. Step 2: strip the VLM step.
4. Step 3: fix masking.
5. Step 4: design + implement offline packing for VLMs.
6. Step 5: M-RoPE follow-up.

Steps 1–3 unblock other VLM work and should land first. Step 4 is the
real long-term win.

Out of scope

• Omni / cross-attention models. Token-level LLaVA-style integration is
  the only assumption we are making in this refactor.
• Dynamic CP for VLMs.

Notes / open questions

• Naming of the new shared processing class — VLMTaskEncoder vs
  something HF-flavored. Today the HF-side and Energon-side classes have
  diverged names; pick one in Step 1.
• Whether to vendor MCore's sequence-packing scheduler from the dev
  branch or wait for it to land upstream before doing online DP-group
  packing.
• The "first ~200 steps slow, then suddenly fast" behavior reproduces in
  NeMo as well — likely orthogonal to this refactor but worth a separate
  issue.

cc @yaoyu-33 @cuichenx

Progress log — status as of 2026-06-25

Added as a tracking update. Step↔PR mapping inferred from PR titles and merge order; nothing above this line was changed.

Step 1 — Unify processing — mostly done (HF/VLM path)

• ✅ refactor(data): move VLM collators to model modules #4166 refactor(data): move VLM collators to model modules (merged 2026-06-06)
• ✅ feat(data): add text chat collate for unified HF datasets #4307 feat(data): add text chat collate for unified HF datasets (merged 2026-06-24)
• Shared HF-path encoder is data/energon/hf_task_encoder.py::HFTaskEncoder; per-model VLM collates plug into it (qwen_vl, gemma_vl, glm_vl, kimi_vl, ministral3, nemotron_vl).
• ⚠️ Audio models are NOT on the unified encoder. NemotronOmniTaskEncoder (data/energon/nemotron_omni_task_encoder.py) and the qwen_audio path keep bespoke encoders (mel-feature extraction / audio tokens not covered by HF AutoProcessor). Consistent with the "Out of scope: Omni" note above, but flagged here for visibility.
• TODO: resolve the shared-class naming (VLMTaskEncoder vs HF-flavored) called out in Notes.

Step 2 — Strip the VLM forward step (batching/packing → collate) — done

• ✅ fix(data): route HF VLM sequence packing by collate support #4255 fix(data): route HF VLM sequence packing by collate support (merged 2026-06-10)
• ✅ refactor(data): move VLM sequence batching to collate #4315 refactor(data): move VLM sequence batching to collate (merged 2026-06-25)
  • New collate-time helper data/sequence_batching.py::pad_or_pack_sequence; ~300 lines of pad/truncate/pack removed from training/vlm_step.py.
  • All 9 registered collates updated; defer_in_batch_packing_to_step escape hatch added for Qwen3-VL.
  • Nemotron Omni: although Omni is out-of-scope for Step 1, its collate WAS swept into Step 2 — refactor(data): move VLM sequence batching to collate #4315 deleted Omni's ~40-line inline packer and routed it through pad_or_pack_sequence (renamed pack_sequences → enable_in_batch_packing; now packs at B==1).
• ⚠️ TODO (parity, from review of refactor(data): move VLM sequence batching to collate #4315): on the HFTaskEncoder + Qwen path, min_pixels/max_pixels now fall back to the HF processor defaults instead of the prior explicit 200704/1003520 — verify/restore before relying on parity.
• Note: pad_or_pack_sequence is still a pad-then-pack round-trip (collate pads to [B, S], packer strips padding and re-flattens). Correct but wasteful; the native-flat-packing fix is Step 4 below.

Step 3 — Fix masking — in progress

• ✅ fix(data): anchor VLM assistant loss masks to chat templates #4169 fix(data): anchor VLM assistant loss masks to chat templates (merged 2026-06-24)
• ✅ fix(data): add Ministral3 assistant mask boundaries #4477 fix(data): add Ministral3 assistant mask boundaries (merged 2026-06-24)
• 🔲 fix(data): preserve Kimi thinking in VLM collate #4478 fix(data): preserve Kimi thinking in VLM collate (OPEN, stacked on refactor(data): move VLM sequence batching to collate #4315) — Kimi reasoning_content / <think> handling + assistant-boundary inference.
• TODO: per-model masking fixes will keep landing; document the token-search vs generation-tag fallback chain in the VLM data skill.

Step 4 — Long-term packing — not started

• 🔲 No PR yet. Build flattened packed tensors directly in the packing collates (offline pack index, or DP-group online packing) to remove the Step-2 pad-then-pack round-trip. Partly gated on the MCore-dev sequence-packing scheduler not yet on main (see Notes).

Related / possibly superseded

• 🔲 Refactor dataloader creation in VLM #1127 Refactor dataloader creation in VLM (DRAFT, stale since 2025-10) — confirm whether superseded by this stack or still intended.