nemo-mbridge-perf-sequence-packing

シーケンスパッキングスキル

安定した背景情報と推奨レベルについては、以下を参照してください。

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

有効化

LLM ファインチューニング向けのオフラインパック SFT：

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

オフラインパック長を選択する

テキストのみの LLM SFT と PEFT の検証では、モデルのコンテキスト制限、メモリ、モデルファミリーのサポートが許す場合、まず 8192 トークンのオフラインパックから始めてください。オプティマイザーステップあたりのトークンスロット数を同じにして、パック長をベンチマークします。

token_slots_per_step = packed_sequence_size * global_batch_size

たとえば、2K/GBS32、4K/GBS16、8K/GBS8 は、それぞれステップあたり 65,536 個のトークンスロットを確保します。長いパックでは、各物理 MBS1 行により多くのソース例をまとめられるため、勾配蓄積とステップあたりのオーバーヘッドを削減できます。一方で、アクティベーションメモリが増加し、カーネル幅の制約が表面化する可能性もあります。そのため、長いほど常に高速だと決めつけず、測定した構成のうち適合する最大のものを選択してください。

オフラインパッキングには MBS1 が必要です。global_batch_size % data_parallel_size == 0 および global_batch_size >= data_parallel_size を満たす必要があります。そのため、8K/GBS8 のワークロードでは DP を 8 以下にする必要があります。model.seq_length、dataset.seq_length、packed_sequence_size を同じ値にし、いずれかを変更した後は新しいパックデータ出力ルートを使用して、セットアップ後に解決された設定を確認してください。

トークンスロット数が同じでも、パック長が異なれば数値的に同一になるわけではありません。ターゲットが長くなると、切り捨てとパックの構成が変わります。検証済みのエビデンスを置き換える前に、有限損失、スキップ／NaN がないこと、および収束のセンチネルを再実行してください。

CP を有効にしてファインチューニングする場合：

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

SFT と PEFT には同じアラインメント式を使用してください。SP が無効な TP1/CP1 では 1、SP が有効な TP4/CP1 では 4 になります。オフラインパッキングではこの値が自動的に導出されないため、明示的に固定し、トポロジーの変更後にパックデータを再構築してください。

ディスパッチャーまたはカーネルで最終トークン幅を固定する必要がある場合：

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

packed_sequence_size は、カーネルの倍数条件を満たす値にしてください。たとえば、HybridEP で 128 トークンの結合チャンクを使用する場合、幅は 128 で割り切れる必要があります。これは、各構成シーケンスを CP/SP に合わせる pad_seq_to_mult とは別の設定です。

このパック処理で CUDA グラフを有効にする場合、トークン幅を固定する必要があり、パックのメタデータも静的な形状でなければなりません。

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

注： pad_cu_seqlens = True の場合、パックデータセットと同じ場所にメタデータ JSON ファイルも必要です（src/megatron/bridge/data/datasets/sft.py でアサートされます）。メタデータファイルがないカスタムのパックデータセットでは、データセットの初期化時にアサーションが発生します。

GPT SFT およびサポート対象の VLM ファインチューニングでのインバッチパッキング：

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

ローカルまたはマテリアライズ済みの GPT-SFT JSONL では、既存の mmap ベースのデータセットを維持し、トークナイズを遅延実行します。プロンプト／補完形式（GPTSFTDataset）とチャット形式（GPTSFTChatDataset）の両方で、損失マスクのセマンティクスが維持されます。各 DataLoader の出力が 1 つの論理マイクロバッチになるよう、dataloader_type="single" または "cyclic" を使用してください。GPT-SFT のインバッチパッキングは、グローバルバッチの "batch" データローダーをサポートしていません。

Qwen-VL の Energon オンラインパッキングでは、選択対象を 1 つの collator マイクロバッチに限定せず、Energon のワーカーごとの候補バッファーを使用します。

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 が唯一のネイティブパッキング選択項目です。従来のコレーターおよびステップが所有するパッキングフラグはデフォルトのままにしてください。vlm_step を使用してください。バッファーサイズは、ワーカーごとの準備済み候補サンプル数を表し、バイト数やパック済みトークン数ではありません。準備済みの画像／動画パッチテンソルは選択されるまでホストメモリーに保持されるため、高解像度または動画データでは 8〜16 から始め、増やす前にワーカーの RSS、最初のバッチのレイテンシ、ビン充填率を測定してください。 このパスではオフラインパックは書き込まれません。ソースの WebDataset シャードは変更されません。MBS1 での eager Qwen-VL をサポートし、MTP、CUDA グラフ、Qwen3-VL DistTrain、PP は拒否します。要求された MoE エキスパート並列通信のオーバーラップは、警告とともに無効になります。標準の eager alltoall EP は、オーバーラップを無効にした TP1/PP1/EP8 の Qwen3.6-35B-A3B で機能面のカバレッジがあります。これは性能の根拠ではありません。ほかの EP ディスパッチャーは固定幅のネイティブパックで受け付けられますが、同等のランタイム検証はまだありません。Qwen-VL モデルは論理および物理 THD 境界から MoE パディングマスクを導出するため、固定幅のギャップが補助損失、z-loss、エキスパートバイアスの統計に入りません。現在の MCore はパディング位置を引き続きディスパッチする可能性があります。エキスパート容量／トークン破棄の構成には、ネイティブパッキングのランタイムカバレッジがありません。

長文脈のベースライン：

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

コードアンカー

LLM パック済み SFT の設定項目：

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

共有のテキストデータセットヘルパーは、現在、固定幅パックを有効にしています。これはヘルパーのデフォルトとして扱い、オフラインパッキングの普遍的なランタイム要件とは見なさないでください。選択したディスパッチャー、カーネル、または CUDA グラフのパスで静的幅が必要な場合は、その設定を維持してください。

ブリッジの検証：

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(...)

VLM プロバイダーが使用するコレーション時のインバッチランタイム：

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 の直接行パッキング：

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

パック済み THD のランタイム制約：

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

注意点

1. オフラインのパック済み SFT、ランタイムのインバッチパッキング、Energon のオンラインパッキングは、それぞれ異なる機能です。オフラインパッキングと Energon パッキングでは物理 MBS1 を使用し、ランタイムのインバッチパッキングでは MBS が 1 より大きくなります。
2. GPT-SFT のインバッチパッキングには dataloader_type="single" または "cyclic" が必要です。"batch" はサポートされません。
3. CP が有効な場合、パック済みシーケンス長は 2 * context_parallel_size で割り切れる必要があります。
4. CP を使ったファインチューニングでは、calculate_per_token_loss=True と ddp.average_in_collective=False が必要です。
5. pad_cu_seqlens=True には pad_to_max_length=True も必要です。
6. パッキングのサポートはモデルファミリー固有です。Qwen3-Next、GLM-4.5、Qwen3.5-VL では、異なるパスに明示的なオプトアウトが含まれています。
7. MTP ファインチューニングは、パック済みシーケンスとの互換性がないと記載されています。
8. samples_mapping を通じて再マッピングされる負のインデックスを含む合成パディング行では、損失マスクをすべてゼロにする必要があります。
9. オフラインパッキングで MBS1 を使用する場合、global_batch_size はデータ並列サイズで割り切れ、かつそれ以上である必要があります。
10. SFT と PEFT の両方で CP/TP/SP から pad_seq_to_mult を導出してください。ワークロード種別ごとに異なる値をハードコードしないでください。
11. pad_to_max_length は最終パック幅を制御し、固定形状での実行要件に応じて設定します。
12. Energon の packing_buffer_size はワーカーごとの値であり、検証にも影響します。グローバル／評価バッチ数が表すのはソースの会話数ではなく、物理パック数です。
13. Energon ローダーを正確に再開するには、シャード／分割、DP ワールドサイズ、ワーカー数、シャッフル設定／シード、プロセッサー、シーケンス長、トポロジー、パッキングバッファーサイズを変更しないでください。

検証

チェックイン済みのユニットテストを使用してください：

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

成功条件：

• 選択したすべてのテストが成功する
• オフラインとインバッチの設定検証が相互排他的に維持される
• パック済みメタデータが MCore THD 形式でトレーニングステップに渡される
• GPT-SFT のインバッチパッキングが、グローバルバッチの "batch" データローダーを拒否する
• ネイティブ Energon パッキングが保留中のグループを正確に復元し、有限の部分バッファーをサンプルを破棄せずにフラッシュする
• マッピングされたパディング行が損失に寄与しない