cp: feat: VLM pretokenized data pipeline with neat packing

[HuiyingLi]

Summary

VLM pretokenized data pipeline with greedy knapsack packing, cherry-picked and adapted from @ZhiqiLi-Nvidia's zhiqi-dev branch. This PR lands the data pipeline components needed for efficient packed-sequence VLM fine-tuning.

What's included

Packing engine (CP from zhiqi-dev zhiqil@nvidia.com)

• Greedy knapsack bin-packing with configurable packing_ratio and pack_size
• LengthGroupedSampler for token-aware distributed sampling
• configure_packing monkey-patch for flash_attention_2 and sdpa packed masks
• _get_unpad_data patch for flash_attn_varlen_func with per-document cu_seqlens

Data loading (CP from zhiqi-dev zhiqil@nvidia.com)

• Meta-dataset loader: ShareGPT format, LMDB images, video support
• RobustDatasetWrapper with retry logic and fake image injection for text-only samples
• Template-based label building (works across processor variants)
• VideoMetadata refactor for timestamp handling
• precompute_tokens.py offline tokenization script

Collate functions & integration (CP zhiqi-dev zhiqil@nvidia.com + new)

• Attention-aware collaters for Qwen2.5-VL, Qwen3-VL, Qwen3.5, Kimi-VL, Kimi-K2.5-VL
• packed_sequence top-level YAML config section (replaces nested dataset.packing)
• Recipe integration in finetune.py: pretokenize → pack → train with FA2

Recipes

• qwen3_5/qwen3_5_4b_neat_packing.yaml — Qwen3.5-4B with train + validation
• qwen3/qwen3_vl_4b_neat_packing.yaml — Qwen3-VL-4B-Thinking
• qwen3/qwen3_vl_moe_30b_neat_packing.yaml — Qwen3-VL-30B-A3B MoE with EP=8

Cherry-pick origin

Most feature commits are cherry-picked from zhiqi-dev (authored by @ZhiqiLi-Nvidia). Key mappings:

This branch	zhiqi-dev	Description
63eaf11	a5aa050	Greedy knapsack packing
2372959	24126fe	LengthGroupedSampler
58a9f69	36222b6	Meta-dataset loading
f5f3df3	d3e440b	RobustDatasetWrapper + fake image
e6ca6d7	e45a7cb	Template-based label building
e7769cb	8b78f18	VideoMetadata refactor
dc5e060	f02d78e	precompute_tokens script
7b59b29	multiple	Wire up packing + collaters

Post-CP commits (lint, tests, recipes, fixes) are by @HuiyingLi.

Validation

Recipe	Config	Result
Qwen3.5-4B	qwen3_5_4b_neat_packing.yaml	https://wandb.ai/Nemo-automodel/vlm-pack/runs/bgt2dw6n?nw=nwuserhuiyingl
Qwen3-VL-4B	qwen3_vl_4b_neat_packing.yaml	https://wandb.ai/Nemo-automodel/vlm-pack/runs/j0pm5m2o?nw=nwuserhuiyingl
Qwen3-VL-30B MoE EP=8	qwen3_vl_moe_30b_neat_packing.yaml	https://wandb.ai/Nemo-automodel/vlm-pack/runs/2ja5q4ut?nw=nwuserhuiyingl

Test plan

• Full convergence run with wandb logging

🤖 Generated with Claude Code

[HuiyingLi]

/ok to test 8597ffa

[HuiyingLi]

/claude review

[claude]

Bug: _resolve_lmdb_image is defined in nemo_automodel.components.datasets.vlm.utils, not in datasets. This import will raise ImportError at runtime.

Suggested change

	from nemo_automodel.components.datasets.vlm.datasets import (
	_convert_sharegpt_to_conversation,
	_load_json_or_jsonl,
	_resolve_lmdb_image,
	)
	from nemo_automodel.components.datasets.vlm.datasets import (
	_convert_sharegpt_to_conversation,
	_load_json_or_jsonl,
	)
	from nemo_automodel.components.datasets.vlm.utils import _resolve_lmdb_image

[claude]

Bug: fake_indices is computed from the original examples list, but conversations may have been shortened by _drop_overlong_samples above. After dropping, the batch tensors have fewer rows than len(examples), so these indices can be out-of-bounds or point to the wrong samples.

Same issue exists in default_collate_fn (around line 1115) and qwen2_5_collate_fn (around line 467).

You need to either:

1. Track which examples survived filtering and remap indices, or
2. Move fake-image injection to happen before _drop_overlong_samples so it's already baked into the conversation, and track the _injected_fake flag on the conversation (not on examples).

[claude]

These two sum(...) expressions compute values but discard them (not assigned to any variable). If they're needed for the summary, assign them; otherwise remove them.

Suggested change

	if has_stats:
	total_images = sum(t.get("n_images", 0) for t in timings.values())
	total_text_tokens = sum(t.get("n_text_tokens", 0) for t in timings.values())
	total_media_tokens = sum(t.get("n_media_tokens", 0) for t in timings.values())

[HuiyingLi]

/claude review

[HuiyingLi]

/claude review

[claude]

Bug: processor_type is reassigned to "default" on line 375 before the warning on line 376, so the warning always logs "You are using default with default collate function" instead of the actual processor type name.

Suggested change

	if processor_type not in COLLATE_FNS:
	processor_type = "default"
	logging.warning(f"You are using {processor_type} with default collate function.")
	if processor_type not in COLLATE_FNS:
	logging.warning(f"You are using {processor_type} with default collate function.")
	processor_type = "default"

[claude]

Light review complete. One bug found (inline comment posted): the warning message in finetune.py always logs "default" instead of the actual processor type because the variable is reassigned before the log call.

The rest of the PR — packing engine, collate functions, dataset loading, retry logic, and sharding — looks correct. The autoregressive label shift is safe across all collate functions (shapes diverge before the input-trimming loop). Tests cover the key paths well.

Nice work on the knapsack packing and robust dataset wrapper.

[HuiyingLi]

/ok to test 358b5f7

[HuiyingLi]

/ok to test d1e450a

[HuiyingLi]

/ok to test eb01df7

[HuiyingLi]

/ok to test 65d42d6

[HuiyingLi]

/ok to test 849ba79