[dyurk-lila]

What does this PR do?

When we set a batch size we always want to make it an exact multiple of our number of DP replicas so that each replica gets an equal amount of training data. However, if we enable sequence packing, a fixes number of inputs turns into a variable number of packed sequences with no guarantee that this will align to our DP count. This can result in DP replicas sitting idle during the end of the batch.

The idea of this PR is to allow the batch size to vary slightly in order to attempt to make the number of packed sequences divisible by the DP count as closely as we can. The algorithm guarantees that any individual batch cannot differ from the baseline target by more than some fixed tolerance (default 5%) to avoid large distortions in training dynamics.

How it works

Each training step logs train/actual_batch_size to W&B: the global count of original examples, counting sequences within packed bins and excluding dummy padding rows.

The sampler reserves a real sequence per DP rank for future steps when the epoch has enough examples. Short tails use valid one-token, zero-loss rows for missing ranks; every real example is still consumed exactly once.

The batch sampler and collator share make_sft_sequence_packer() for MFFD selection, capacity, and TP/CP/FP8 alignment. The sampler omits DP padding when counting natural bins.

DPAlignedPackingBatchSampler looks for nearby boundaries whose unpadded packing result is DP-divisible. It preserves sample order and every example, bounds each batch’s size variation, and repays cumulative drift at the run’s final scheduled step, including when it falls in a later epoch.

Checkpoint state includes the sampler cursor and carry buffer so resume reproduces the same pending examples. Evaluation keeps its configured batch size.

Usage

Set align_packing_bins_to_dp=true with use_sequence_packing=true. packing_batch_size_allowed_variation defaults to 0.05 and must be between zero and one. Custom samplers are unsupported; validation also rejects a minimum allowed batch smaller than the DP world size.

Packing results on public data

A recorded CPU packing simulation compared MFFD with ordinary DP bin-count rounding against the same MFFD packer with bounded DP alignment. Each geometry used 677 simulated optimizer steps and the default 5% batch-size tolerance, with the same ordered example stream in both runs.

Global batch / bin capacity / DP	Baseline bins	DP-aligned bins	Fewer bins	Token fill, before → after
1024 / 64K / 4	141,096	140,084	0.717%	98.914% → 99.629%
128 / 4K / 8	81,176	78,944	2.750%	94.958% → 97.643%

The length pool contained 3,072 tokenized sequences: 256 from each of the 12 released Nemotron 3.5 Lightning SFT blend entries. The simulation resampled sources uniformly and lengths with replacement using seed 42, then clipped lengths to each geometry’s capacity. Tokenization used nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-Base-BF16 and the Lightning 3.5 chat template. Per-sequence and aggregate alignment multiples were both 1 for the 4K case and 16 for the 64K case. All released subsets were included instead of applying the recipe’s global nano_v3 filter; this is a sampled, uniformly weighted study rather than the full recipe mixture.

Token fill is valid tokens divided by total bin capacity. These results measure packing efficiency; they do not measure training throughput or sampler CPU overhead.

Sample order was preserved, and no example moved by more than one optimizer step. Actual batch sizes ranged from 1005–1040 around a nominal 1024, and 122–134 around a nominal 128. At the 4K geometry, per-step token-count coefficient of variation increased from 2.39% to 3.46%, so the packing gain comes with somewhat more variation in each step’s token load.

Test coverage

Training-loop tests verify the logged example count for DP-aligned packing, fixed-size packing, and unpacked tail batches with both synchronous and asynchronous collation.

Short-tail tests verify that one real example is logged as one example and three tokens despite four DP rows, and that a real one-token example is distinguishable from dummy rows.

Epoch-tail regression checks cover DP=4 and DP=8 with the default 5% tolerance, including fixed-step resets before and at the tail; every emitted batch is passed through the packed collator.

Configuration and dataloader tests cover DP divisibility, sample order/count, bounded cumulative drift, carry-buffer restoration, fixed-step boundaries, dataset paths, and StatefulDataLoader resume.

Validation

Before the short-tail metrics follow-up, focused CPU suites passed in both mirrors: 415 tests across collation, DP sampling, worker batching, layout, active spans, and Megatron correctness (13 Megatron-only skips); 40 bin-packer tests; and 3 unpacked active-span guard tests with megatron-core installed. Ruff and Black passed. Seven targeted Triton value/gradient tests passed on one H100; the full GPU suite and benchmarks were not rerun.

After the metrics follow-up, 39 focused SFT CPU tests passed in each mirror, including synchronous and asynchronous short-tail logging; Ruff, Black, and secret detection passed on the changed files.

Review order

Part 4 of 9 in a linked series. The stack is based on main; each later PR targets the preceding layer. Review and merge from bottom to top.

Depends on #2371. Review only this layer.

1. perf(sft): use modified first-fit decreasing packing #2369 — use modified first-fit decreasing packing
2. perf(train): make sequence packing alignment-aware #2370 — make sequence packing alignment-aware
3. refactor(megatron): share packed segment layout derivation #2371 — share packed segment layout derivation
4. perf(sft): align packed bin counts to data parallelism #2372 — align packed bin counts to data parallelism
5. perf(megatron): autotune fused LM-head Triton kernels #2373 — autotune fused LM-head Triton kernels
6. feat(megatron): fuse LM-head entropy for RL #2374 — fuse LM-head entropy for RL
7. feat(megatron): add block-sparse fused LM-head forward #2375 — add block-sparse fused LM-head forward
8. perf(megatron): add block-sparse fused LM-head backward #2376 — add block-sparse fused LM-head backward
9. perf(megatron): fuse sample-support candidate projection with Triton #2521 — fuse sample-support candidate projection with Triton

---

Note

Medium Risk
Changes training batch composition, dataloader checkpoint state, and packed collator behavior near epoch tails; bounded variation limits drift but affects Megatron packing throughput paths.

Overview
Adds optional DP-aligned sequence packing for Megatron SFT so training batches can vary slightly (default ±5%) instead of padding packed-bin counts up to dp_size.

When align_packing_bins_to_dp=true, a stateful DPAlignedPackingBatchSampler nudges contiguous batch boundaries so unpadded MFFD bin counts are DP-divisible, preserves sample order, checkpoints a carry buffer, and repays sample-count drift on the run’s final scheduled step. The collator and sampler share make_sft_sequence_packer(); packing can keep empty DP padding bins (allow_empty_bins) and emit zero-loss placeholder rows on short epoch tails. Training logs train/actual_batch_size and real token counts from packed-batch metadata; config validation rejects custom samplers and unsafe variation vs. DP size. VLM training forces the flag off.

^{Reviewed by Cursor Bugbot for commit 2cb8e35. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request introduces the align_packing_bins_to_dp feature, which dynamically adjusts training batch boundaries within a configurable variation limit to ensure packed bins are divisible by the data parallel (DP) size, thereby avoiding unnecessary padding. It implements the DPAlignedPackingBatchSampler to manage these boundaries while preserving sample order, updates the collator and logging to track actual batch sizes and tokens, and adds validation and tests. The review feedback suggests two important improvements: bounding the reserve calculation in the sampler to prevent violating the allowed variation on short datasets, and adding validation to ensure the configured variation allows at least one sample of variation to avoid silent no-ops.

[gemini-code-assist]

The current logic for calculating reserve can severely violate the configured packing_batch_size_allowed_variation limit on short datasets or epochs. When remaining_samples is slightly larger than (future_steps + 1) * dp_size, the sampler reserves future_steps * dp_size samples, which can force the current batch size to drop well below the allowed minimum batch size (low).

To prevent this, the reservation should be bounded so that it does not force the current batch size below low while still reserving at least future_steps samples to ensure all future steps have at least one sample.

            remaining_samples = len(self.batch_sampler.sampler) - self.samples_emitted
            # Reserve a real row per DP rank in each future step when the
            # epoch has enough examples; otherwise reserve one per step and
            # let the collator supply zero-loss rows on short tails.
            future_steps = total_batches - self.batch_index - 1
            reserve = max(
                future_steps,
                min(future_steps * self.batch_sampler.dp_size, remaining_samples - low),
            )

[gemini-code-assist]

If math.floor(cfg.batch_size * variation) is 0, the allowed variation in terms of number of samples is 0. This means align_packing_bins_to_dp will silently act as a no-op because the batch size cannot vary at all. We should explicitly validate that the variation allows at least 1 sample of variation.

        variation = cfg.packing_batch_size_allowed_variation
        if not 0 < variation < 1:
            raise ValueError("packing_batch_size_allowed_variation must be in (0, 1), " f"got {variation}.")
        if math.floor(cfg.batch_size * variation) < 1:
            raise ValueError(
                f"packing_batch_size_allowed_variation={variation} with batch_size={cfg.batch_size} "
                f"allows 0 samples of variation. Increase packing_batch_size_allowed_variation "
                f"to at least {1 / cfg.batch_size:.4f} or increase batch_size."
            )
        min_batch_size = math.ceil(cfg.batch_size * (1 - variation))

[greptile-apps]

Confidence Score: 4/5

[High risk] Modifies core training data sampling and packing logic.

The PR appears safe to merge, though batch-boundary search may reduce throughput for larger batches or tolerances.

Findings

1.  Repeated Packing Slows Batch Selection ▶

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Ordered sample indices] --> B[Candidate batch prefixes]
  B --> C[Unpadded MFFD bin counts]
  C --> D[Select batch boundary]
  D --> E[Packed collator]
  E --> F[DP-divisible bins and real-count metadata]

Loading

_{Reviews (1) · Last reviewed commit: "fix(sft): exclude dummy packing rows fro..."}

[greptile-apps]

Repeated packing slows batch selection

Each candidate batch size triggers a fresh packing run on the dataloader’s main-process path. With a batch size of 1024, the default 5% tolerance can mean checking up to 103 candidates per step, and larger tolerances allow more. This extra CPU work could offset the throughput gain from reducing DP padding. Consider reusing packing results or limiting the search work.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit da01aa6. Configure here.}