perf-sequence-packing

Par nvidia · skills

Validez et utilisez les séquences empaquetées ainsi que l'entraînement sur de longs contextes dans Megatron-Bridge, en distinguant le packed SFT hors ligne pour les LLM du packing en batch pour les VLM, et en appliquant les bonnes contraintes CP.

npx skills add https://github.com/nvidia/skills --skill perf-sequence-packing

Compétence Packing de Séquences

Pour les informations de base et le niveau de recommandation, voir :

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

Activation

SFT offline avec packing pour le finetuning LLM :

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

cfg.train.micro_batch_size = 1
cfg.dataset.seq_length = 4096
cfg.model.seq_length = 4096
cfg.dataset.dataset_kwargs = {"pad_to_max_length": True}
cfg.dataset.packed_sequence_specs = PackedSequenceSpecs(
    packed_sequence_size=4096,
    pad_seq_to_mult=1,
)

Si CP est activé :

cfg.model.context_parallel_size = 2
cfg.model.calculate_per_token_loss = True
cfg.ddp.average_in_collective = False
cfg.dataset.packed_sequence_specs.pad_seq_to_mult = cfg.model.context_parallel_size * 2

# Si sequence_parallel est aussi activé, utiliser lcm(2*CP, CP*TP):
# import math
# cfg.dataset.packed_sequence_specs.pad_seq_to_mult = math.lcm(2 * CP, CP * TP)
# Voir src/megatron/bridge/training/vlm_step.py pour la logique de référence.

Si CUDA graphs est activé pour ce chemin avec packing :

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

Note : pad_cu_seqlens = True requiert aussi un fichier de métadonnées JSON aux côtés du dataset packed (vérifié dans src/megatron/bridge/data/datasets/sft.py). Les datasets packed personnalisés qui omettent le fichier de métadonnées déclencheront une assertion lors de l'initialisation du dataset.

Packing in-batch pour le finetuning VLM :

cfg.dataset.pack_sequences_in_batch = True
cfg.train.micro_batch_size = 2

Baseline long-context :

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

Anchors de Code

Surface de configuration SFT packed LLM :

if packed_sequence:
    dataset_kwargs = {"pad_to_max_length": True}
    packed_sequence_specs = PackedSequenceSpecs(packed_sequence_size=seq_length, pad_seq_to_mult=pad_seq_to_mult)
else:
    dataset_kwargs = {}
    packed_sequence_specs = None

Validation Bridge :

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 ... packed_sequence_size > 0 and self.train.micro_batch_size > 1:
    raise ValueError(...)
...
if getattr(self.dataset, "pack_sequences_in_batch", False) and self.train.micro_batch_size == 1:
    raise ValueError(...)

Runtime in-batch VLM :

if enable_packing:
    ...
    ) = pack_batch_sequences(
        ...
        pad_token_id=0,
        pad_to_multiple_of=cp_size * 2 if cp_size > 1 else 1,
    )

Contrainte runtime packed THD :

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)")

Écueils

  1. Le SFT offline avec packing et le packing in-batch VLM sont des fonctionnalités différentes avec des règles de micro-batch opposées.
  2. Quand CP est activé, les longueurs de séquence packées doivent respecter la divisibilité par 2 * context_parallel_size.
  3. Pour le finetuning avec CP, calculate_per_token_loss=True et ddp.average_in_collective=False sont obligatoires.
  4. pad_cu_seqlens=True requiert aussi pad_to_max_length=True.
  5. Le support du packing est spécifique à la famille de modèles. Qwen3-Next, GLM-4.5 et Qwen3.5-VL contiennent des opt-outs explicites dans différents chemins.
  6. Le finetuning MTP est documenté comme incompatible avec les séquences packées.

Vérification

Utiliser la couverture d'unit tests vérifiée :

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 pack_sequences_in_batch or context_parallel_seq_length_divisibility or context_parallel_finetuning_validations" -v && \
uv run python -m pytest tests/unit_tests/training/test_vlm_step.py -k "enable_packing" -v

Critères de succès :

  • la première commande rapporte 8 passed
  • la deuxième commande rapporte 14 passed
  • la troisième commande rapporte 2 passed

Skills similaires