nemo-mbridge-perf-sequence-packing

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.