nemo-mbridge-perf-sequence-packing

Files (6)

• BENCHMARK.md
• SKILL.md
• card.yaml
• evals.json · evals/
• skill-card.md
• skill.oms.sig

SKILL.md (first lines)

# 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:

```python
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:

```text
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:

```python
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:

```python
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:

```python
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:

```python
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:

```python

From NVIDIA/skills, Apache-2.0.

More from NVIDIA/skills