refactor(data): limit GPT SFT packing to microbatch loaders

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

docs/training/packed-sequences.md

@@ -46,8 +46,8 @@ schemas: prompt/completion (`GPTSFTDataset`) and chat
 read, parse, and tokenize only the rows selected for the current logical
 microbatch. Collation then concatenates those rows into one physical THD batch
 row; it does not materialize an offline dataset or load the full source into
-RAM. With GPT-SFT's default global-batch sampler, Bridge first divides the
-global batch into logical microbatches and packs each one independently.
+RAM. Use a microbatch-yielding `single` or `cyclic` dataloader; GPT-SFT
+in-batch packing does not support the global-batch `batch` dataloader.
 
 ## When to Use It
 
@@ -120,14 +120,16 @@ config and use a logical micro-batch larger than one:
 
 ```text
 dataset.enable_in_batch_packing=true
+dataset.dataloader_type=single
 train.micro_batch_size=4
 ```
 
 The collator preserves each sample's prompt/completion or chat loss mask and
 emits current MCore packed metadata (`cu_seqlens_q`, `cu_seqlens_kv`, and the
 corresponding padded boundaries when CP/SP alignment is required). The model
 sees one physical THD row. `enable_in_batch_packing` and
-`enable_offline_packing` are mutually exclusive.
+`enable_offline_packing` are mutually exclusive. The `batch` dataloader is not
+supported for GPT-SFT in-batch packing; use `single` or `cyclic`.
 
 ## Stable Constraints
 
@@ -136,6 +138,8 @@ The durable constraints for packed sequences in Bridge are:
 - offline packed SFT requires configured `micro_batch_size == 1`
 - GPT-SFT/Direct-HF/VLM in-batch packing requires configured `micro_batch_size > 1`;
   collation flattens those input rows into one physical THD batch row
+- GPT-SFT in-batch packing requires `dataloader_type="single"` or `"cyclic"`;
+  the global-batch `"batch"` dataloader is not supported
 - Energon online packing currently supports the eager Qwen-VL collator path,
   requires physical `micro_batch_size == 1`, the generic `vlm_step`, per-token loss, and
   `ddp.average_in_collective=False`

scripts/training/README.md

@@ -208,12 +208,15 @@ host memory. For example:
 --dataset local-jsonl \
 dataset.dataset_root=/data/sft \
 dataset.enable_in_batch_packing=true \
+dataset.dataloader_type=single \
 train.micro_batch_size=4
 ```
 
 The logical micro-batch must be larger than one. Collation emits one physical
 THD row and derives per-sequence CP/SP alignment from the resolved topology.
-Do not combine `dataset.enable_in_batch_packing=true` with offline packing.
+Use the `single` or `cyclic` dataloader; the global-batch `batch` dataloader is
+not supported. Do not combine `dataset.enable_in_batch_packing=true` with
+offline packing.
 
 ### VLM data
 

skills/nemo-mbridge-perf-sequence-packing/SKILL.md

@@ -108,15 +108,16 @@ 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. With the default `dataloader_type="batch"`, the DataLoader first
-divides each per-DP-rank global batch into logical microbatches and packs each
-logical microbatch into its own physical THD row.
+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:
@@ -239,23 +240,14 @@ def prepare_padded_or_packed_sequence_batch(
         return
 ```
 
-GPT-SFT direct-row packing and global-batch microbatch preservation:
+GPT-SFT direct-row packing:
 
 ```627:671:src/megatron/bridge/data/datasets/gpt_sft.py
 def _collate_in_batch(self, batch):
     ...
     return build_mcore_thd_sequence_batch_from_rows(...)
 ```
 
-```18:57:src/megatron/bridge/data/batch_utils.py
-def collate_finetuning_microbatches(samples, *, micro_batch_size, collate_fn):
-    ...
-```
-
-Do not flatten a GPT-SFT `batch`-sampler global batch into one THD row. The
-training schedule still consumes the configured number of logical
-microbatches; each one must retain independent packed metadata.
-
 Packed THD runtime constraint:
 
 ```94:108:src/megatron/bridge/training/gpt_step.py
@@ -275,17 +267,18 @@ if cu_seqlens.dim() > 1 and cu_seqlens.size(0) != 1:
 ## Pitfalls
 
 1. 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.
-2. When CP is enabled, packed sequence lengths must respect `2 * context_parallel_size` divisibility.
-3. For finetuning with CP, `calculate_per_token_loss=True` and `ddp.average_in_collective=False` are required.
-4. `pad_cu_seqlens=True` also requires `pad_to_max_length=True`.
-5. Packing support is model-family-specific. `Qwen3-Next`, `GLM-4.5`, and `Qwen3.5-VL` contain explicit opt-outs in different paths.
-6. MTP finetuning is documented as incompatible with packed sequences.
-7. Synthetic padding rows, including negative indices remapped through `samples_mapping`, must retain an all-zero loss mask.
-8. `global_batch_size` must be divisible by and no smaller than data parallel size when offline packing uses MBS1.
-9. Derive `pad_seq_to_mult` from CP/TP/SP for both SFT and PEFT; do not hardcode different values by workload type.
-10. `pad_to_max_length` controls final pack width and is conditional on fixed-shape execution requirements.
-11. Energon `packing_buffer_size` is per worker and also affects validation; global/eval batch counts refer to physical packs rather than source conversations.
-12. Exact Energon loader resume requires unchanged shards/splits, DP world size, worker counts, shuffle settings/seed, processor, sequence length, topology, and packing-buffer size.
+2. GPT-SFT in-batch packing requires `dataloader_type="single"` or `"cyclic"`; it does not support `"batch"`.
+3. When CP is enabled, packed sequence lengths must respect `2 * context_parallel_size` divisibility.
+4. For finetuning with CP, `calculate_per_token_loss=True` and `ddp.average_in_collective=False` are required.
+5. `pad_cu_seqlens=True` also requires `pad_to_max_length=True`.
+6. Packing support is model-family-specific. `Qwen3-Next`, `GLM-4.5`, and `Qwen3.5-VL` contain explicit opt-outs in different paths.
+7. MTP finetuning is documented as incompatible with packed sequences.
+8. Synthetic padding rows, including negative indices remapped through `samples_mapping`, must retain an all-zero loss mask.
+9. `global_batch_size` must be divisible by and no smaller than data parallel size when offline packing uses MBS1.
+10. Derive `pad_seq_to_mult` from CP/TP/SP for both SFT and PEFT; do not hardcode different values by workload type.
+11. `pad_to_max_length` controls final pack width and is conditional on fixed-shape execution requirements.
+12. Energon `packing_buffer_size` is per worker and also affects validation; global/eval batch counts refer to physical packs rather than source conversations.
+13. 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
 
@@ -295,9 +288,8 @@ 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 or batch_dataloader" -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/data/test_batch_utils.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 && \
@@ -310,6 +302,6 @@ 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
-- a GPT-SFT global batch is pre-collated as one THD row per logical microbatch
+- 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

skills/nemo-mbridge-perf-sequence-packing/card.yaml

@@ -79,6 +79,7 @@ known_constraints:
   - 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.

src/megatron/bridge/data/batch_utils.py

@@ -14,61 +14,11 @@
 
 """Batch handling utilities for finetuning data."""
 
-from collections.abc import Callable
 from typing import Any, Iterator
 
 import torch
 
 
-_PRECOLLATED_MICROBATCHES_KEY = "__precollated_microbatches__"
-
-
-def collate_finetuning_microbatches(
-    samples: list[Any],
-    *,
-    micro_batch_size: int,
-    collate_fn: Callable[[list[Any]], dict[str, Any]],
-) -> dict[str, list[dict[str, Any]]]:
-    """Collate a global batch as independently packed logical microbatches.
-
-    The ``batch`` finetuning sampler deliberately sends a full per-DP-rank
-    global batch to one DataLoader collate call. In-batch packing must happen
-    after that input is divided into logical microbatches; otherwise the whole
-    global batch becomes one physical THD row and cannot be consumed by the
-    configured number of pipeline microbatches.
-
-    Args:
-        samples: Samples assigned to one DP rank for the current global batch.
-        micro_batch_size: Number of logical samples to pack together.
-        collate_fn: Dataset collator that converts one logical microbatch to a
-            physical packed batch.
-
-    Returns:
-        A pin-memory-compatible mapping containing the pre-collated
-        microbatches.
-
-    Raises:
-        ValueError: If the sample count cannot form complete microbatches or a
-            collate result is not a mapping.
-    """
-    if micro_batch_size <= 0:
-        raise ValueError("micro_batch_size must be greater than 0.")
-    if not samples or len(samples) % micro_batch_size != 0:
-        raise ValueError(
-            f"Global batch sample count {len(samples)} must be a positive multiple of "
-            f"micro_batch_size {micro_batch_size}."
-        )
-
-    microbatches = []
-    for start in range(0, len(samples), micro_batch_size):
-        microbatch = collate_fn(samples[start : start + micro_batch_size])
-        if not isinstance(microbatch, dict):
-            raise ValueError("Finetuning collate_fn must return a batch dictionary.")
-        microbatches.append(microbatch)
-
-    return {_PRECOLLATED_MICROBATCHES_KEY: microbatches}
-
-
 def split_batch_into_microbatches(
     batch: dict[str, Any], num_microbatches: int, enforce_divisible: bool = True
 ) -> list[dict[str, Any]]:
@@ -149,9 +99,8 @@ def prepare_finetuning_batch(
 
     This function handles the finetuning-specific data flow:
     1. Gets the full global batch from the iterator
-    2. Reuses pre-collated logical microbatches when the DataLoader performed
-       in-batch packing, otherwise splits the dense global batch
-    3. Extracts the dynamic sequence length from the resulting microbatches
+    2. Extracts the dynamic sequence length from the batch
+    3. Splits the batch into microbatches with consistent sequence length
     4. Returns an iterator over microbatches and the extracted sequence length
 
     Args:
@@ -182,27 +131,6 @@ def prepare_finetuning_batch(
     # Get full global batch from dataloader
     global_batch = next(data_iterator)
 
-    precollated_microbatches = (
-        global_batch.get(_PRECOLLATED_MICROBATCHES_KEY) if isinstance(global_batch, dict) else None
-    )
-    if precollated_microbatches is not None:
-        if not isinstance(precollated_microbatches, list) or not all(
-            isinstance(microbatch, dict) for microbatch in precollated_microbatches
-        ):
-            raise ValueError("Pre-collated finetuning microbatches must be a list of batch dictionaries.")
-        if len(precollated_microbatches) != num_microbatches:
-            raise ValueError(
-                f"Expected {num_microbatches} pre-collated microbatches, got {len(precollated_microbatches)}."
-            )
-
-        sequence_lengths = [
-            microbatch[seq_key].size(1)
-            for microbatch in precollated_microbatches
-            if isinstance(microbatch.get(seq_key), torch.Tensor) and microbatch[seq_key].dim() >= 2
-        ]
-        seq_length = max(sequence_lengths, default=default_seq_length)
-        return iter(precollated_microbatches), seq_length
-
     # Extract dynamic seq_length from the full batch
     seq_length = default_seq_length
     if seq_key in global_batch and isinstance(global_batch[seq_key], torch.Tensor):

src/megatron/bridge/data/builders/gpt_sft.py

@@ -197,6 +197,10 @@ def validate(self) -> None:
             raise ValueError("enable_offline_packing and enable_in_batch_packing are mutually exclusive.")
         if self.in_batch_packing_pad_to_multiple_of <= 0:
             raise ValueError("in_batch_packing_pad_to_multiple_of must be greater than 0.")
+        if self.enable_in_batch_packing and self.dataloader_type == "batch":
+            raise ValueError(
+                "GPT-SFT in-batch packing does not support dataloader_type='batch'; use 'single' or 'cyclic'."
+            )
         if self.enable_offline_packing and self.offline_packing_specs is None:
             raise ValueError("offline_packing_specs must be set when enable_offline_packing=True.")
         if self.offline_packing_specs is not None and not self.enable_offline_packing:

src/megatron/bridge/data/datasets/direct_sft.py

@@ -73,10 +73,6 @@ def __init__(
         self._base_examples = base_examples
         self._length = int(max(0, target_length))
         self._processor = processor
-        # The global-batch sampler uses this runtime capability to collate each
-        # logical microbatch independently before THD packing. Step-owned packing
-        # must continue through the ordinary dense global-batch path.
-        self.enable_in_batch_packing = enable_in_batch_packing and not defer_in_batch_packing_to_step
         # Choose collate implementation by processor type name when not provided
         collate_key = type(processor).__name__ if processor is not None else "default"
         explicit_collate_impl = collate_impl is not None
@@ -92,7 +88,7 @@ def __init__(
             "pad_to_multiple_of": pad_to_multiple_of,
             # Active deferral user: Qwen3-VL. Other VLM/HF collates should pack
             # here when enable_in_batch_packing is set.
-            "enable_in_batch_packing": self.enable_in_batch_packing,
+            "enable_in_batch_packing": enable_in_batch_packing and not defer_in_batch_packing_to_step,
             "in_batch_packing_pad_to_multiple_of": in_batch_packing_pad_to_multiple_of,
         }
         if explicit_collate_impl:

src/megatron/bridge/data/samplers.py

@@ -3,15 +3,12 @@
 """Dataloaders."""
 
 import random
-from functools import partial
 from typing import Any, Callable, Iterator, Optional
 
 import numpy as np
 import torch
 from torch.utils.data import DataLoader, Dataset
 
-from megatron.bridge.data.batch_utils import collate_finetuning_microbatches
-
 
 def build_pretraining_data_loader(
     dataset: Dataset,
@@ -113,23 +110,13 @@ def build_pretraining_data_loader(
     else:
         raise Exception("{} dataloader type is not supported.".format(dataloader_type))
 
-    effective_collate_fn = collate_fn
-    if dataloader_type == "batch" and getattr(dataset, "enable_in_batch_packing", False) is True:
-        if collate_fn is None:
-            raise ValueError("In-batch packing with dataloader_type='batch' requires a dataset collate_fn.")
-        effective_collate_fn = partial(
-            collate_finetuning_microbatches,
-            micro_batch_size=micro_batch_size,
-            collate_fn=collate_fn,
-        )
-
     # Torch dataloader.
     return DataLoader(
         dataset,
         batch_sampler=batch_sampler,
         num_workers=num_workers,
         pin_memory=pin_memory,
-        collate_fn=effective_collate_fn,
+        collate_fn=collate_fn,
         persistent_workers=persistent_workers,
         worker_init_fn=worker_init_fn,
     )

tests/unit_tests/data/builders/test_gpt_sft_config.py

@@ -83,6 +83,7 @@ def test_in_batch_config_round_trip_is_declarative_and_serializable(tmp_path):
         preprocessing=PromptCompletionSFTPreprocessingConfig(),
         enable_in_batch_packing=True,
         in_batch_packing_pad_to_multiple_of=8,
+        dataloader_type="single",
     )
 
     serialized = ConfigContainer._convert_value_to_dict(config)
@@ -99,12 +100,25 @@ def test_config_rejects_nonpositive_in_batch_packing_multiple(tmp_path):
         dataset_root=tmp_path,
         enable_in_batch_packing=True,
         in_batch_packing_pad_to_multiple_of=0,
+        dataloader_type="single",
     )
 
     with pytest.raises(ValueError, match="in_batch_packing_pad_to_multiple_of must be greater than 0"):
         config.validate()
 
 
+def test_config_rejects_batch_dataloader_with_in_batch_packing(tmp_path):
+    config = GPTSFTDatasetConfig(
+        seq_length=128,
+        dataset_root=tmp_path,
+        enable_in_batch_packing=True,
+        dataloader_type="batch",
+    )
+
+    with pytest.raises(ValueError, match="does not support dataloader_type='batch'"):
+        config.validate()
+
+
 @pytest.mark.parametrize(
     "dataset_kwargs",
     [
@@ -128,6 +142,7 @@ def test_config_rejects_dense_attention_mask_with_in_batch_packing(tmp_path):
         seq_length=128,
         dataset_root=tmp_path,
         enable_in_batch_packing=True,
+        dataloader_type="single",
         dataset_kwargs={"get_attention_mask_from_fusion": False},
     )
 
@@ -140,6 +155,7 @@ def test_config_rejects_offline_and_in_batch_packing(tmp_path):
         seq_length=128,
         dataset_root=tmp_path,
         enable_in_batch_packing=True,
+        dataloader_type="single",
         enable_offline_packing=True,
         offline_packing_specs=PackedSequenceSpecs(packed_sequence_size=128),
     )
@@ -663,6 +679,7 @@ def test_builder_owns_runtime_materialization_and_shared_construction(monkeypatc
     config = _hf_config(tmp_path)
     config.enable_in_batch_packing = True
     config.in_batch_packing_pad_to_multiple_of = 8
+    config.dataloader_type = "single"
     materialize_mock = []
     dataset_calls = []