Nemo Mbridge Perf Sequence Packing

NVIDIAOfficialPublished by NVIDIA

Skill
Tool integration
Software development
Apache-2.0
Medium risk

From the author

Validate and use packed sequences and long-context training in Megatron-Bridge, including offline LLM packing, collate-time VLM packing, Energon online packing, and CP constraints.

— the description by NVIDIA

AutoGineer summary

Written by AutoGineer from the skill’s own files, not by the author

This skill provides the configuration patterns and validation requirements for implementing sequence packing in Megatron-Bridge. It covers three distinct packing methodologies: offline LLM packing, collate-time in-batch packing, and Energon online packing for VLMs. The guide details how to align packing specifications with context-parallel (CP), tensor-parallel (TP), and sequence-parallel (SP) topologies, ensuring that memory usage and kernel constraints are managed effectively during long-context training.

What it does

Configures offline packed SFT for LLM finetuning with specific sequence length and padding requirements.
Enables in-batch packing for GPT-SFT and VLM workloads to optimize data throughput.
Implements Energon online packing for Qwen-VL models using per-worker candidate buffers.
Provides alignment formulas for CP/TP/SP topologies to ensure valid sequence lengths.
Enforces configuration constraints to prevent invalid combinations of packing methods and dataloaders.
Includes a comprehensive suite of unit tests to verify packing metadata and loss mask integrity.

When to use it

When optimizing training throughput for long-context LLM or VLM finetuning.
When memory constraints require efficient packing of multiple sequences into single micro-batches.
When deploying models with context-parallelism that require specific sequence length divisibility.
When using Energon-based data loading for high-resolution multimodal models like Qwen-VL.

Advantages

Reduces gradient accumulation overhead by aggregating source examples into physical micro-batches.
Optimizes memory utilization for long-context training through structured sequence packing.
Provides clear validation logic to prevent common configuration errors in complex parallel topologies.
Supports diverse packing strategies tailored to specific model families and data loading requirements.

Usage examples

“How do I configure offline packing for an 8192-token LLM finetuning job?”: The agent provides a Python configuration snippet setting enable_offline_packing to True, defining PackedSequenceSpecs with appropriate padding multiples, and ensuring micro-batch size is set to 1.
“What are the requirements for enabling in-batch packing in GPT-SFT?”: The agent explains that enable_in_batch_packing must be True, dataloader_type must be set to single or cyclic, and micro-batch size must be greater than 1.

Install

Download the skill as a ZIP:
https://mp-staging.autogineer.com/agent-skills/nvidia/nemo-mbridge-perf-sequence-packing/download

Unzip it into the cross-client skills folder of your project:

mkdir -p .agents/skills && unzip nemo-mbridge-perf-sequence-packing.zip -d .agents/skills

Documentation by NVIDIA

Sequence Packing Skill

For stable background and recommendation level, see:

@docs/training/packed-sequences.md
@skills/nemo-mbridge-perf-sequence-packing/card.yaml

Enablement

Offline packed SFT for LLM finetuning:

import math

from megatron.bridge.data.datasets.packed_sequence import PackedSequenceSpecs

cfg.train.micro_batch_size = 1
cfg.train.global_batch_size = 8
cfg.dataset.seq_length = 8192
cfg.model.seq_length = 8192
cfg.dataset.enable_offline_packing = True

cp_size = cfg.model.context_parallel_size
tp_size = cfg.model.tensor_model_parallel_size
cp_multiple = 2 * cp_size if cp_size > 1 else 1
sp_multiple = cp_size * tp_size if cfg.model.sequence_parallel and tp_size > 1 else 1
cfg.dataset.offline_packing_specs = PackedSequenceSpecs(
    packed_sequence_size=8192,
    pad_seq_to_mult=math.lcm(cp_multiple, sp_multiple),
)

Choose the offline pack length

For text-only LLM SFT and PEFT verification, start with an 8192-token offline pack when the model context limit, memory, and model-family support allow it. Benchmark pack lengths at equal token slots per optimizer step:

token_slots_per_step = packed_sequence_size * global_batch_size

For example, 2K/GBS32, 4K/GBS16, and 8K/GBS8 each expose 65,536 token slots per step. Longer packs aggregate more source examples into each physical MBS1 row and can reduce gradient accumulation and per-step overhead. They also increase activation memory and may expose kernel-width constraints, so select the largest measured configuration that fits rather than assuming longer is always faster.

Offline packing requires MBS1. Require global_batch_size % data_parallel_size == 0 and global_batch_size >= data_parallel_size; an 8K/GBS8 workload therefore needs DP no larger than 8. Keep model.seq_length, dataset.seq_length, and packed_sequence_size equal, use a fresh packed-data output root after changing any of them, and inspect the resolved post-setup configuration.

Equal token slots do not make different pack lengths numerically identical: the longer target changes truncation and pack membership. Rerun finite-loss, no-skip/NaN, and convergence sentinels before replacing verified evidence.

For finetuning with CP enabled:

cfg.model.context_parallel_size = 2
cfg.model.calculate_per_token_loss = True
cfg.ddp.average_in_collective = False

Use the same alignment formula for SFT and PEFT. It produces 1 for TP1/CP1 with SP disabled and 4 for TP4/CP1 with SP enabled. Offline packing does not derive the value automatically, so pin it explicitly and rebuild packed data after a topology change.

If a dispatcher or kernel requires a fixed final token width:

cfg.dataset.dataset_kwargs = {
    **(cfg.dataset.dataset_kwargs or {}),
    "pad_to_max_length": True,
}

Choose packed_sequence_size to satisfy the kernel multiple. For example, HybridEP with a 128-token combine chunk requires a width divisible by 128. This is separate from pad_seq_to_mult, which aligns each constituent sequence for CP/SP.

If CUDA graphs are enabled for this packed path, fixed token width is required and packed metadata must also have a static shape:

cfg.dataset.offline_packing_specs.pad_cu_seqlens = True
cfg.dataset.dataset_kwargs["pad_to_max_length"] = True

Note: pad_cu_seqlens = True also requires a metadata JSON file alongside the packed dataset (asserted in src/megatron/bridge/data/datasets/sft.py). Custom packed datasets that omit the metadata file will hit an assertion at dataset initialization.

In-batch packing for GPT SFT and supported VLM finetuning:

cfg.dataset.enable_in_batch_packing = True
cfg.dataset.dataloader_type = "single"
cfg.train.micro_batch_size = 4

For local or materialized GPT-SFT JSONL, this keeps the existing mmap-backed dataset and performs tokenization lazily. Both prompt/completion (GPTSFTDataset) and chat (GPTSFTChatDataset) preserve their loss-mask semantics. Use dataloader_type="single" or "cyclic" so every DataLoader yield is one logical microbatch; GPT-SFT in-batch packing does not support the global-batch "batch" dataloader.

Energon online packing for Qwen-VL uses Energon's per-worker candidate buffer instead of limiting selection to one collator micro batch:

cfg.dataset.packing_buffer_size = 16
cfg.dataset.micro_batch_size = 1
cfg.train.micro_batch_size = 1
cfg.model.calculate_per_token_loss = True
cfg.ddp.average_in_collective = False

packing_buffer_size is the sole native-packing selector; leave the legacy collator- and step-owned packing flags at their defaults. Use vlm_step. The buffer size counts prepared candidate samples per worker, not bytes or packed tokens. Since prepared image/video patch tensors remain in host memory until selection, start at 8-16 for high-resolution or video data and measure worker RSS, first-batch latency, and bin fill before increasing it. This path does not write offline packs; the source WebDataset shards remain unchanged. It supports eager Qwen-VL with MBS1 and rejects MTP, CUDA graphs, Qwen3-VL DistTrain, and PP. Requested MoE expert-parallel communication overlap is disabled with a warning. Standard eager alltoall EP has functional coverage for Qwen3.6-35B-A3B at TP1/PP1/EP8 with overlap disabled; this is not performance evidence. Other EP dispatchers are accepted with fixed-width native packs but do not yet have equivalent runtime evidence. The Qwen-VL model derives a MoE padding mask from logical and physical THD boundaries so fixed-width gaps do not enter auxiliary-loss, z-loss, or expert-bias statistics. Current MCore may still dispatch padded positions; expert-capacity/token-dropping configurations lack native-packing runtime coverage.

Long-context baseline:

cfg.model.seq_length = 16384
cfg.dataset.seq_length = 16384
cfg.model.context_parallel_size = 2

Code Anchors

LLM packed SFT config surface:

dataset_kwargs = {}
offline_packing_specs = None
if enable_offline_packing:
    dataset_kwargs["pad_to_max_length"] = True
    offline_packing_specs = PackedSequenceSpecs(packed_sequence_size=seq_length, pad_seq_to_mult=pad_seq_to_mult)

return _text_hf_dataset_config(
    source=HFDatasetSourceConfig(dataset_name="squad"),
    preprocessing=PromptCompletionSFTPreprocessingConfig(separator=" "),
    seq_length=seq_length,
    enable_offline_packing=enable_offline_packing,
    offline_packing_specs=offline_packing_specs,
    dataset_kwargs=dataset_kwargs,
    val_proportion=0.1,
    num_workers=1,
)

The shared text-dataset helper currently opts into fixed-width packs. Treat that as a helper default, not a universal offline-packing runtime requirement; preserve it when the selected dispatcher, kernel, or CUDA-graph path requires static width.

Bridge validation:

enable_in_batch_packing = getattr(self.dataset, "enable_in_batch_packing", False)
enable_offline_packing = getattr(self.dataset, "enable_offline_packing", False)
offline_packing_specs = getattr(self.dataset, "offline_packing_specs", None)

if enable_offline_packing and enable_in_batch_packing:
    raise ValueError("enable_offline_packing and enable_in_batch_packing are mutually exclusive.")
if enable_offline_packing and offline_packing_specs is None:
    raise ValueError("offline_packing_specs must be set when enable_offline_packing=True.")
...
if enable_in_batch_packing:
    ...
    cp_multiple = 2 * cp_size if cp_size > 1 else 1
    sp_multiple = cp_size * tp_size if has_sp and tp_size > 1 else 1
    self.dataset.in_batch_packing_pad_to_multiple_of = math.lcm(cp_multiple, sp_multiple)

if self.model.context_parallel_size > 1:
    assert self.model.seq_length % (self.model.context_parallel_size * 2) == 0, ...
    if isinstance(self.dataset, FinetuningDatasetConfig):
        assert self.model.calculate_per_token_loss, ...
        assert not self.ddp.average_in_collective, ...
...
if enable_offline_packing and self.train.micro_batch_size > 1:
    raise ValueError(...)
...
if enable_in_batch_packing and self.train.micro_batch_size == 1:
    raise ValueError(...)

Collate-time in-batch runtime used by VLM providers:

def prepare_padded_or_packed_sequence_batch(
    batch,
    *,
    sequence_length,
    ...
    enable_in_batch_packing=False,
    in_batch_packing_pad_to_multiple_of=1,
    ...
):
    ...
    if enable_in_batch_packing:
        pack_right_padded_sequence_batch_to_mcore_thd(
            batch,
            sequence_length=sequence_length,
            pad_to_multiple_of=in_batch_packing_pad_to_multiple_of,
            ...
        )
        return

GPT-SFT direct-row packing:

def _collate_in_batch(self, batch):
    ...
    return build_mcore_thd_sequence_batch_from_rows(...)

Packed THD runtime constraint:

if batch.get("cu_seqlens_q") is not None:
    cu_seqlens = batch.get("cu_seqlens_q_padded")
    if cu_seqlens is None:
        cu_seqlens = batch["cu_seqlens_q"]
    if cu_seqlens.dim() > 1 and cu_seqlens.size(0) != 1:
        raise ValueError("Packed THD batches expect micro-batch size 1 for context-parallel slicing (THD layout)")
    return cu_seqlens.squeeze()

cu_seqlens = batch["cu_seqlens"]
if cu_seqlens.dim() > 1 and cu_seqlens.size(0) != 1:
    raise ValueError("Packed THD batches expect micro-batch size 1 for context-parallel slicing (THD layout)")

Pitfalls

Offline packed SFT, runtime in-batch packing, and Energon online packing are different features. Offline and Energon packing use physical MBS1; runtime in-batch packing uses MBS greater than one.
GPT-SFT in-batch packing requires dataloader_type="single" or "cyclic"; it does not support "batch".
When CP is enabled, packed sequence lengths must respect 2 * context_parallel_size divisibility.
For finetuning with CP, calculate_per_token_loss=True and ddp.average_in_collective=False are required.
pad_cu_seqlens=True also requires pad_to_max_length=True.
Packing support is model-family-specific. Qwen3-Next, GLM-4.5, and Qwen3.5-VL contain explicit opt-outs in different paths.
MTP finetuning is documented as incompatible with packed sequences.
Synthetic padding rows, including negative indices remapped through samples_mapping, must retain an all-zero loss mask.
global_batch_size must be divisible by and no smaller than data parallel size when offline packing uses MBS1.
Derive pad_seq_to_mult from CP/TP/SP for both SFT and PEFT; do not hardcode different values by workload type.
pad_to_max_length controls final pack width and is conditional on fixed-shape execution requirements.
Energon packing_buffer_size is per worker and also affects validation; global/eval batch counts refer to physical packs rather than source conversations.
Exact Energon loader resume requires unchanged shards/splits, DP world size, worker counts, shuffle settings/seed, processor, sequence length, topology, and packing-buffer size.

Verification

Use the checked-in unit coverage:

uv run python -m pytest tests/unit_tests/training/utils/test_packed_seq_utils.py -v && \
uv run python -m pytest tests/unit_tests/training/test_config.py -k "packed_sequence or enable_in_batch_packing or offline_and_in_batch_packing_are_mutually_exclusive or context_parallel_seq_length_divisibility or context_parallel_finetuning_validations" -v && \
uv run python -m pytest tests/unit_tests/data/packing/test_in_batch.py -v && \
uv run python -m pytest tests/unit_tests/data/datasets/test_gpt_sft.py -k "in_batch_packing" -v && \
uv run python -m pytest tests/unit_tests/data/builders/test_gpt_sft_config.py -v && \
uv run python -m pytest tests/unit_tests/training/test_vlm_step.py -k "deferred_in_batch_packing or packed_metadata" -v && \
uv run python -m pytest tests/unit_tests/models/qwen_vl/data/test_energon.py tests/unit_tests/data/builders/test_energon_builder.py -v && \
uv run python -m pytest tests/unit_tests/tutorials/test_multimodal_data_tutorials.py -k "native_packing_loader" -v && \
uv run python -m pytest tests/unit_tests/data/datasets/test_packed_parquet.py -k "negative_index_zeroes_loss_mask" -v && \
uv run python -m pytest tests/unit_tests/data/datasets/test_sft.py -k "mapped_padding_rows_do_not_contribute_to_loss" -v

Success criteria:

all selected tests pass
offline and in-batch configuration validation remains mutually exclusive
packed metadata reaches the training step in MCore THD form
GPT-SFT in-batch packing rejects the global-batch "batch" dataloader
native Energon packing restores pending groups exactly and flushes finite partial buffers without dropping samples
mapped padding rows do not contribute to loss

Files

ATTRIBUTION.mdAdded by AutoGineer867 B
BENCHMARK.md4.8 KiB
LICENSELicence file added by AutoGineer11.2 KiB
SKILL.md13.0 KiB
card.yaml7.2 KiB
evals/evals.json2.5 KiB
skill-card.md4.1 KiB
skill.oms.sig4.7 KiB

card.yaml

title: packed_sequences_long_context
validated_on: "2026-08-05"
summary: >
  Megatron-Bridge supports offline text SFT packing, collate-time VLM packing,
  and Energon online packing for Qwen-VL. Long-context training is primarily
  expressed through context parallelism and memory tradeoff knobs.
validation_status:
  offline_packed_sft_runtime:
    - code_verified
  vlm_in_batch_packing_runtime:
    - code_verified
  qwen3_vl_energon_online_packing_runtime:
    - code_verified
  qwen36_moe_ep8_energon_online_packing_runtime:
    - functional_verified
  cp_and_packing_validation_rules:
    - code_verified
  packed_seq_helper_behavior:
    - code_verified
  vlm_packing_helper_behavior:
    - code_verified
  mapped_padding_loss_mask_behavior:
    - code_verified
  packed_cp_functional_smoke_in_tree:
    - recipe_verified
  long_context_llama_recipe_coverage:
    - recipe_verified
  public_cp_backend_guidance:
    - doc_only
  long_context_perf_claims:
    - unclear
feature_meaning:
  offline_packed_sft: >
    Pre-tokenized packed finetuning datasets built through PackedSequenceSpecs
    and consumed through THD packed-sequence metadata.
  vlm_in_batch_packing: >
    Collate-time batch concatenation path for supported VLM training flows,
    controlled by enable_in_batch_packing=True and emitted in MCore THD form.
  energon_online_packing: >
    Per-worker candidate-buffer packing for Qwen-VL WebDataset streams,
    controlled by packing_buffer_size and emitted as one physical MBS1 THD row.
  long_context_training: >
    Training at longer sequence lengths, primarily enabled through context
    parallelism plus recipe-specific long-context presets and memory tuning
    knobs.
recommended_path:
  llm_packed_sft:
    target_seq_length: 8192_when_model_and_memory_support_it
    train.micro_batch_size: 1
    train.global_batch_size: preserve_equal_token_slots
    dataset.enable_offline_packing: true
    dataset.offline_packing_specs.packed_sequence_size: match_seq_length
    dataset.offline_packing_specs.pad_seq_to_mult: lcm(cp_multiple, sp_multiple)
    sizing_rule: packed_sequence_size * global_batch_size = token_slots_per_step
  fixed_width_execution:
    when: dispatcher_kernel_or_cuda_graph_requires_static_token_width
    dataset.dataset_kwargs.pad_to_max_length: true
    packed_sequence_size: divisible_by_kernel_chunk
  cuda_graph_packing:
    dataset.offline_packing_specs.pad_cu_seqlens: true
    requires_pack_metadata: true
  cp_finetuning:
    model.calculate_per_token_loss: true
    ddp.average_in_collective: false
    dataset.offline_packing_specs.pad_seq_to_mult: 2 * context_parallel_size
  vlm_in_batch_packing:
    dataset.enable_in_batch_packing: true
    train.micro_batch_size: ">1"
  qwen3_vl_energon_online_packing:
    dataset.packing_buffer_size: 16
    train.micro_batch_size: 1
    model.calculate_per_token_loss: true
    ddp.average_in_collective: false
    step_func: vlm_step
known_constraints:
  - seq_length must be divisible by 2 * context_parallel_size when CP > 1.
  - Offline packed SFT requires micro_batch_size == 1.
  - With offline MBS1 packing, global_batch_size must be divisible by and no smaller than data_parallel_size.
  - pad_seq_to_mult is derived from CP/TP/SP and follows the same rule for SFT and PEFT.
  - pad_to_max_length is conditional on a fixed-width dispatcher, kernel, or CUDA-graph path.
  - VLM in-batch packing requires micro_batch_size > 1.
  - GPT-SFT in-batch packing requires dataloader_type single or cyclic; batch is unsupported.
  - Energon online packing requires physical micro_batch_size == 1 and counts candidate samples per worker.
  - Energon online packing is mutually exclusive with collate-time in-batch packing and the deferred Qwen step.
  - Energon online packing does not support MTP, CUDA graphs, or pipeline parallelism.
  - Requested MoE expert-parallel communication overlap is disabled with a warning for Energon online packing.
  - EP dispatchers other than standard alltoall use fixed-width native packs but lack equivalent runtime evidence.
  - Current MCore may dispatch fixed-width padding; expert-capacity/token-dropping configurations lack runtime coverage.
  - For finetuning with CP > 1, calculate_per_token_loss must be true.
  - For finetuning with CP > 1, ddp.average_in_collective must be false.
  - pad_cu_seqlens=true also requires pad_to_max_length=true.
  - Synthetic padding rows must retain an all-zero loss mask after samples_mapping remaps their indices.
  - Fine-tuning sequence packing is documented as unsupported with MTP.
known_limitations:
  - Packing support is model-family-specific rather than universal.
  - Qwen3-Next SFT disables packed sequences.
  - GLM-4.5 SFT and PEFT disable packed sequences.
  - Qwen3.5-VL disables enable_in_batch_packing.
  - The repo does not contain checked-in benchmark results validating long-context throughput claims.
  - Prepared visual tensors remain resident in every Energon worker's packing buffer, so buffer size must be memory-tuned.
  - Qwen3.6-35B-A3B has functional TP1/PP1/EP8 coverage with the standard eager alltoall dispatcher and EP
    communication overlap disabled; this does not establish a performance result.
evidence:
  - docs/training/packed-sequences.md
  - docs/performance-guide.md
  - docs/training/multi-token-prediction.md
  - docs/models/llama/llama3.md
  - docs/training/hierarchical-context-parallel.md
  - src/megatron/bridge/data/datasets/packed_sequence.py
  - src/megatron/bridge/data/datasets/packed_parquet.py
  - src/megatron/bridge/data/datasets/sft.py
  - src/megatron/bridge/training/utils/packed_seq_utils.py
  - src/megatron/bridge/training/gpt_step.py
  - src/megatron/bridge/data/sequence_batching.py
  - src/megatron/bridge/training/vlm_step.py
  - src/megatron/bridge/models/qwen_vl/data/energon.py
  - src/megatron/bridge/data/builders/energon.py
  - src/megatron/bridge/training/config.py
  - src/megatron/bridge/recipes/utils/dataset_utils.py
  - src/megatron/bridge/recipes/common.py
  - src/megatron/bridge/recipes/llama/h100/llama3.py
  - src/megatron/bridge/recipes/qwen/h100/qwen3_next.py
  - src/megatron/bridge/recipes/glm/h100/glm45.py
  - src/megatron/bridge/recipes/qwen_vl/h100/qwen35_vl.py
  - src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
  - tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py
  - tests/unit_tests/data/datasets/test_packed_parquet.py
  - tests/unit_tests/data/datasets/test_sft.py
  - tests/unit_tests/data/packing/test_in_batch.py
  - tests/unit_tests/training/utils/test_packed_seq_utils.py
  - tests/unit_tests/training/test_config.py
  - tests/unit_tests/training/test_vlm_step.py
  - tests/unit_tests/models/qwen_vl/data/test_energon.py
  - tests/unit_tests/data/builders/test_energon_builder.py
  - tests/unit_tests/tutorials/test_multimodal_data_tutorials.py
  - scripts/performance/utils/overrides.py
follow_up_validation:
  - Run the checked-in packed-plus-CP functional test and record whether it still passes on current infrastructure.
  - Add a tiny no-download end-to-end smoke test for offline packed SFT.
  - Add a checked-in long-context training smoke for at least one 16K or 64K recipe.
  - Cross-link public packing docs to model-family opt-outs and the MTP incompatibility note.
  - Benchmark Qwen-VL Energon buffer sizes against worker RSS, first-batch latency, and bin fill.

Requirements and compatibility

Works with: Claude Code, Claude apps, OpenAI Codex, GitHub Copilot, Cursor, Gemini CLI, Kiro, OpenCode, Goose, Amp, Roo Code, JetBrains Junie, Factory Droid, Devin (ex-Windsurf), Mistral Vibe, OpenHands, Letta

“Works with” means the tool reads the open Agent Skills format; a skill written for one tool may use features others lack.

Needs: Needs Python

Security

What our scanner noticed· Medium risk

Large base64-encoded blobs

A static scan is not a guarantee. Skills are third-party content: review a skill before you let an agent use it.

Attribution and licence

Published by

NVIDIA

Publisher

NVIDIA (Official)

Source

NVIDIA/skills · skills/nemo-mbridge-perf-sequence-packing · commit fce3cca

Licence

Apache License 2.0 (redistribution allowed)

LICENSE (added by AutoGineer from LICENSE-APACHE)
ATTRIBUTION.md (written by AutoGineer)

Synced

2026-11-08

AutoGineer is not affiliated with or endorsed by NVIDIA.

Download ZIP

The ZIP holds the skill’s files, its licence files and an ATTRIBUTION.md. Download nemo-mbridge-perf-sequence-packing.zip