feat(data): enable native Energon packing for Qwen3-VL by yaoyu-33 · Pull Request #5285 · NVIDIA-NeMo/Megatron-Bridge

yaoyu-33 · 2026-08-03T23:53:40Z

Summary

enable Energon-native online sequence packing for Qwen-VL through Energon's select_samples_to_pack and pack_selected_samples APIs
use exact post-processor token lengths and first-fit-decreasing selection over packing_buffer_size candidates per worker
emit current MCore THD metadata (cu_seqlens_q/kv, optional padded boundaries, and max_seqlen_q/kv) with visual tensors, labels, loss masks, and position IDs
preserve exact per-image/per-video-frame ViT FLOPS accounting when several conversations share one physical pack
add strict configuration guardrails, recipe metadata, a runnable example, tests, and a complete prepare/connect/train tutorial

Background and behavior

The existing Qwen-VL path could pack only within a collator microbatch. That lifecycle does not use Energon's buffered selection, checkpoint, or restore APIs. This change makes packing_buffer_size the sole selector for native Energon packing on supported Qwen-VL recipes.

packing_buffer_size is a candidate count per worker, not bytes or packed tokens. Each source sample is processed once before selection, then Energon chooses compatible samples from the buffer and emits one physical MBS1 pack containing multiple independent conversations. THD boundaries prevent cross-conversation attention; terminal labels and aligned gaps are masked so loss does not cross conversations.

This is online packing. Users prepare the normal Energon WebDataset once; no offline packed-data rewrite is required.

The legacy enable_in_batch_packing and defer_in_batch_packing_to_step fields remain because generic HF, Nemotron Omni, dataset conversion, and the legacy Qwen step still use them. Native packing leaves those flags at their defaults and rejects conflicting owners.

Supported scope and guardrails

eager Qwen-VL with the generic vlm_step
TP/SP and CP, including automatic CP/SP alignment derivation
physical MBS1, per-token loss, and DDP sum reduction (average_in_collective=False)
resumable Energon dataloader state, pending pack groups, multi-worker loading, and finite partial-buffer flush
exact additive vision FLOPS for multiple images and temporal video frames, including Qwen3-VL deepstack mergers

Configuration fails early for unsupported task encoders/models, legacy packing conflicts, MBS greater than 1, MTP, language or vision CUDA graphs, DistTrain, and pipeline parallelism. With EP greater than one, native packing emits fixed-width physical rows so other dispatchers are not rejected by this validation layer; only standard eager alltoall has the EP8 runtime evidence below. Requested MoE EP communication overlap or delayed weight-gradient compute is disabled with a warning before communication-overlap setup. Current MCore excludes fixed-width padding from z-loss, auxiliary-loss, and expert-bias statistics but may still dispatch it, so expert-capacity/token-dropping configurations remain unvalidated. Consecutive overlength native samples use Energon bounded failure tolerance instead of spinning indefinitely.

Tutorial

The expanded Energon tutorial covers:

background and when to use native packing
preparation of example and production WebDataset metadata
connection of an existing dataset and Qwen processor
starting LoRA training through the public Qwen3-VL recipe alias
buffer sizing, worker memory, packing-quality metrics, validation behavior, checkpoint/resume requirements, and current limitations

The Qwen3-VL example and packed-sequence documentation link to the same workflow.

Validation

Final review and targeted tests

uv run pre-commit run --all-files: all hooks passed
latest independent CW functional delta: 55 selected tests passed in job 15185140 (packing/native/text 30, packed model legacy-SP-CP-mask 6, transformer propagation/checkpoint 3, native config 13, overlap/delay fallback 3)
final CW native delta: 22 passed (encoder, builder, config, tutorial contract)
final CW VLM/FLOPS suite: 154 passed
post-deepstack CW FLOPS rerun: 128 passed
broader affected unit groups earlier in the branch: 514 passed; one unrelated existing metadata collision on origin/main remains
three independent reviewers completed multiple API/restore, training/math, config/backward-compatibility, documentation, and post-fix adversarial rounds with no remaining P0/P1/P2 findings
Energon API compatibility checked against 7.0.0 through 7.4.1

EOS real-model EP8 functional validation

final-head rerun at Bridge bfa6d894d1ca11c3437938feb266bd51bd72b15a and MCore 6513e3e23d6b5eda6a1c934990b15e804237732b with Qwen3.6-35B-A3B revision 995ad96eacd98c81ed38be0c5b274b04031597b0, 256 experts, TP1/PP1/CP1/EP8, standard eager alltoall, EP overlap and delayed wgrad disabled
real Flickr8k data at fixed revision c7a02cb487da0eca729e133194b0f842cb578c08 through an Energon cache
packing OFF and ON each completed 12/12 forward/backward steps with finite loss and zero skipped/NaN iterations
final-head packing OFF W&B and packing ON W&B; this is functional evidence, not a performance comparison

CW interactive training validation

real Qwen3-VL processor plus real Energon 7.4 PackingDataset: 876 physical tokens from four source sequences; checkpoint/restore produced identical next-pack keys and tensors
1-GPU LoRA, sequence length 1024, buffer size 4, MBS1/GBS2: two iterations completed with losses 4.552594 and 4.611749, no skipped iterations or NaNs
fresh save/resume: restored model and Energon dataloader at iteration 1 and continued at iteration 2 with matching loss 4.611749
CP=2 through the public recipe alias with no explicit step function: iteration completed with loss 4.589449, no skipped iteration or NaN

The fixed-width follow-up changes affect the EP execution path: physical tail padding is masked from MoE loss/statistics, while current MCore may still dispatch padded positions. The existing W&B runs validate standard eager alltoall, overlap disabled; they do not validate allgather, flex, capacity dropping, or performance. No dependency or CI workflow changes are included.

Known reporting caveat

progress.txt Tokens and time/tokens still represent configured token capacity (consumed_train_samples * model.seq_length), not exact partial-pack occupancy or useful unpadded tokens. Making that cumulative counter resume-correct requires a separate backward-compatible TrainState checkpoint migration.

copy-pr-bot · 2026-08-03T23:53:44Z

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.

svcnvidia-nemo-ci · 2026-08-06T00:46:35Z

/nvskills-ci

claude · 2026-08-06T00:56:23Z

Review — Energon online sequence packing for Qwen3-VL

Solid, cohesive change. The native-packing path is well-gated by config validation, the FLOP refactor cleanly preserves per-media attention boundaries, and test coverage is thorough (config gating, task-encoder pack/batch, savable-loader restore + partial-buffer flush, exact-vs-legacy ViT FLOP stats). Observations:

Correctness / behavior:

resolve_global_flops_stats flips seqlen_sum/seqlen_squared_sum to None when seqlen_sum <= 0 while still returning vision stats; the wrapper resolve_global_flops_seqlen_stats preserves the old 3-value contract (confirmed by test_legacy_vision_patch_count_remains_backward_compatible).
train.py gates on has_exact_vision_stats = vision_patch_sum > 0, so a pure-text step falls back to the legacy num_vision_patches path (also 0) — consistent, no double-count.
batch() rejects mixed packed/unpacked samples and enforces physical MBS1; overlength guards in both encode_sample and batch are defensive and tested.

Docs / comments:

encode_sample docstring 'Output format' still describes only the deferred (unpacked) behavior; the native-packing branch now returns example=None with an eager prepared_sequence. Flagged inline.

Nits (non-blocking):

QwenVLTaskEncoder keeps both self.seq_length and self.seq_len aliased to max_padding_length; new packing code uses self.seq_length while batch() still reads self.seq_len. Pre-existing dual-field footgun.

Suggested test cases:

tests/unit_tests/data/builders/test_energon_builder.py::test_qwen_factory_enables_native_energon_packing_from_buffer_size
tests/unit_tests/data/builders/test_energon_builder.py::test_native_energon_packing_rejects_incompatible_modes
tests/unit_tests/data/builders/test_energon_builder.py::test_native_energon_packing_rejects_unsupported_task_encoder
tests/unit_tests/data/builders/test_energon_builder.py::test_qwen_factory_preserves_limits_and_deferred_packing
tests/unit_tests/models/qwen_vl/data/test_energon.py::TestQwenVLTaskEncoderNativePacking::test_select_samples_to_pack_uses_aligned_lengths_without_drops
tests/unit_tests/models/qwen_vl/data/test_energon.py::TestQwenVLTaskEncoderNativePacking::test_pack_and_batch_emit_canonical_thd_metadata_and_visual_order
tests/unit_tests/models/qwen_vl/data/test_energon.py::TestQwenVLTaskEncoderNativePacking::test_batch_rejects_malformed_overlength_packed_group
tests/unit_tests/models/qwen_vl/data/test_energon.py::TestQwenVLTaskEncoderNativePacking::test_encode_sample_preserves_restore_key_without_buffering_raw_media
tests/unit_tests/models/qwen_vl/data/test_energon.py::TestQwenVLTaskEncoderNativePacking::test_overlength_native_sample_counts_toward_energon_failure_tolerance
tests/unit_tests/training/test_config.py::test_native_energon_packing_sets_variable_sequences_and_cp_sp_alignment
tests/unit_tests/training/test_config.py::test_native_energon_packing_requires_per_token_loss
tests/unit_tests/training/test_config.py::test_native_energon_packing_requires_non_averaged_collective
tests/unit_tests/training/test_config.py::test_native_energon_packing_rejects_unsupported_execution_modes
tests/unit_tests/training/test_config.py::test_native_energon_packing_rejects_qwen_dist_train
tests/unit_tests/training/test_config.py::test_native_energon_packing_does_not_require_model_capability_flag
tests/unit_tests/training/test_vlm_step.py::test_forward_step_preserves_independent_image_flops_boundaries
tests/unit_tests/training/utils/test_flop_utils.py::test_grid_thw_preserves_equal_image_attention_boundaries
tests/unit_tests/training/utils/test_flop_utils.py::test_grid_thw_treats_video_frames_as_independent_attention_sequences
tests/unit_tests/training/utils/test_flop_utils.py::test_grid_thw_supports_thinker_nested_vision_config
tests/unit_tests/training/utils/test_flop_utils.py::test_deepstack_visual_indexes_add_one_merger_each
tests/unit_tests/training/utils/test_flop_utils.py::test_exact_vision_stats_accumulate_and_request_global_reduce
tests/unit_tests/training/utils/test_flop_utils.py::test_zero_vision_stats_still_request_matching_dp_collective
tests/unit_tests/training/utils/test_flop_utils.py::test_legacy_vision_patch_count_remains_backward_compatible
tests/unit_tests/training/utils/test_flop_utils.py::test_exact_vision_stats_share_integer_all_reduce_across_dp
tests/unit_tests/tutorials/test_multimodal_data_tutorials.py::test_qwen_native_packing_loader_restores_pending_groups_and_flushes_partial_buffer (num_workers 0 and 2)
tests/unit_tests/scripts/training/test_recipe_metadata.py::test_benchmark_recipe_metadata_selects_task_and_step (qwen3_vl_8b_peft_energon_config and qwen3_vl_8b_peft_1gpu_h100_bf16_energon_config map to vlm_step)
tests/unit_tests/scripts/training/test_run_recipe.py (qwen3_vl_8b_peft_energon_config and qwen3_vl_8b_peft_1gpu_h100_bf16_energon_config map to lora / vlm_step)
tests/unit_tests/recipes/qwen_vl/test_qwen3_vl_recipes.py::test_qwen3_vl_8b_peft_energon_dataset_params

No perf/recipe launcher configs under scripts/performance/configs/ were touched — no perf tests impacted.

svcnvidia-nemo-ci · 2026-08-06T02:21:47Z

/nvskills-ci

svcnvidia-nemo-ci · 2026-08-06T03:17:47Z

/nvskills-ci

svcnvidia-nemo-ci · 2026-08-06T04:38:59Z

/nvskills-ci