[dyurk-lila]

What does this PR do?

Account for the physical aligned footprint of packed sequences, applying CP layout alignment per sequence and TP/FP8 tail padding once per packed slab.

TLDR: packing capacity checks, collation, Megatron preprocessing, replay metadata, and token-based batching use the same alignment model.

How it works

• sequence_length_multiple accounts for constraints on each sequence, including CP/SP layout gaps.
• packed_length_multiple accounts for TP/FP8 alignment after sequence footprints are combined.
• FFD keeps a hard capacity; MFFD and balanced place an oversized aligned sequence alone as a soft-cap singleton. All three check physical footprint before combining sequences.
• Collation, packed preprocessing, token metadata, targets, and masks agree on segment boundaries and aggregate padding.
• Worker token batching counts each controller-packed segment’s alignment gap before applying aggregate padding.

Charging TP/FP8 alignment once per packed slab avoids padding every short sequence to a constraint that applies only after packing. Per-sequence CP constraints still apply before placement, preventing an apparently valid bin from overflowing after layout padding.

Relationship to the general FP8 training PR

#1898 adds broad end-to-end FP8 training and rollout support. Its packing changes add FP8-recipe-specific tile requirements and make SFT bin packing account for aligned lengths, but combine CP/SP, TP, and FP8 into one alignment paid independently by every packed sequence.

This PR addresses a different scope question: CP/SP layout gaps are genuinely per-sequence, while TP and FP8 constrain the aggregate local token slab and should be paid once per packed bin or microbatch. It also moves this two-level accounting into the generic packers, covering non-FP8 CP/TP layouts and worker-side balanced token batching in addition to controller-side FP8 SFT packing.

The changes are complementary rather than replacements. When the branches are reconciled, #1898’s recipe-specific blockwise/MXFP8 tile requirement should feed packing_align_size_total, while this PR’s packing_align_size_sequence versus packing_align_size_total split determines where that padding is charged.

Validation

At the current series tip, 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.

Review order

Part 2 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 #2369. 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

High Risk
Changes packed THD layout, bin budgets, and loss normalization across Megatron CP/TP/FP8; misalignment would corrupt gradients silently, though tests cover collation, preprocess, and batching parity.

Overview
Makes Megatron sequence packing alignment-aware by splitting padding into per-sequence layout gaps (CP/SP) versus a single aggregate tail pad (TP/FP8), and threading that model through collation, bin packing, preprocessing, and worker microbatching.

Replaces get_packed_seq_align_size with get_packing_align_size_sequence and get_packing_align_size_total. Generic packers (SeqPacker, FFD/MFFD/balanced) gain sequence_length_multiple and packed_length_multiple so capacity checks use physical footprint; MFFD can place an aligned singleton alone when it exceeds the token budget. PackedDataCollator bins on max_tokens_per_microbatch with those multiples; preprocess_packed_seqs, token metadata, valid masks, and per-token loss counts mirror the same segment offsets and one-shot aggregate padding. TokenBasedBatchIterator counts sub_seq_lengths alignment gaps and passes the same multiples from MegatronWorker.

Docs now describe controller MFFD bin capacity as max_tokens_per_microbatch (must be >= max_length) instead of implying per-sequence max_length bins or strict multiples of max_length.

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

[gemini-code-assist]

Code Review

This pull request refactors the sequence packing and alignment logic to separate per-sequence layout alignment constraints (for CP/SP) from aggregate token slab alignment constraints (for TP/FP8). Instead of padding every sub-sequence to the global alignment size, sub-sequences are padded to their individual layout alignment constraints, and the aggregate TP/FP8 alignment padding is applied once at the end of the packed token slab. Feedback on these changes suggests pre-aligning token counts to packed_length_multiple in TokenBasedBatchIterator when sub_seq_lengths is present to prevent underestimating aggregate padding, as well as adding defensive checks to ensure tensors are not empty before indexing their last elements in preprocess_packed_seqs and build_token_metadata_layout.

[gemini-code-assist]

When sub_seq_lengths is not None, each row represents an already-packed bin with its own independent aggregate tail padding. Since these bins are processed as separate rows in the microbatch rather than concatenated, their padding overhead is paid independently. Pre-aligning self._token_counts to packed_length_multiple and setting packed_length_multiple = 1 prevents the packer from under-estimating the aggregate padding when grouping multiple bins, which could otherwise lead to over-packing and potential OOMs.

Suggested change

	sub_seq_lengths = data.get("sub_seq_lengths")
	if sub_seq_lengths is None:
	self._token_counts = data["attention_mask"].sum(dim=1).cpu().tolist()
	else:
	self._token_counts = [
	sum(length + (-length % sequence_length_multiple) for length in row.tolist()) for row in sub_seq_lengths
	]
	sequence_length_multiple = 1 # The segment footprints are already aligned.
	sub_seq_lengths = data.get("sub_seq_lengths")
	if sub_seq_lengths is None:
	self._token_counts = data["attention_mask"].sum(dim=1).cpu().tolist()
	else:
	self._token_counts = [
	sum(length + (-length % sequence_length_multiple) for length in row.tolist()) for row in sub_seq_lengths
	]
	self._token_counts = [
	length + (-length % packed_length_multiple) for length in self._token_counts
	]
	sequence_length_multiple = 1 # The segment footprints are already aligned.
	packed_length_multiple = 1 # The row footprints are already fully aligned.

[gemini-code-assist]

Add a defensive check to ensure seqlens_in_batch_padded is not empty before attempting to index its last element, preventing potential runtime index errors on empty or degenerate batches.

Suggested change

	# TP and FP8 operate on the aggregate local token slab. Attach their tail
	# padding once to the final sequence rather than to every sequence.
	aggregate_pad_size = (-seqlens_in_batch_padded.sum()) % packing_align_size_total
	seqlens_in_batch_padded[-1] += aggregate_pad_size
	# TP and FP8 operate on the aggregate local token slab. Attach their tail
	# padding once to the final sequence rather than to every sequence.
	if seqlens_in_batch_padded.numel() > 0:
	aggregate_pad_size = (-seqlens_in_batch_padded.sum()) % packing_align_size_total
	seqlens_in_batch_padded[-1] += aggregate_pad_size

[gemini-code-assist]

Add a defensive check to ensure padded_sequence_lengths_tensor is not empty before indexing its last element, preventing potential runtime index errors on empty or degenerate batches.

Suggested change

	aggregate_pad_size = (-padded_sequence_lengths_tensor.sum()) % packing_align_size_total
	padded_sequence_lengths_tensor[-1] += aggregate_pad_size
	if padded_sequence_lengths_tensor.numel() > 0:
	aggregate_pad_size = (-padded_sequence_lengths_tensor.sum()) % packing_align_size_total
	padded_sequence_lengths_tensor[-1] += aggregate_pad_size

[greptile-apps]

Confidence Score: 4/5

[High risk] Refactors sequence packing alignment logic across training pipeline.

The PR appears safe to merge, with non-blocking documentation and parity-test corrections recommended.

Findings

1.  Capacity omits alignment minimum ▶
2.  Parity bins use different capacity ▶

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Sequence lengths] --> B[Per-sequence CP/SP alignment]
  B --> C[Controller bin packing]
  C --> D[Aggregate TP/FP8 tail padding]
  D --> E[Collated rows]
  E --> F[Worker microbatch packing]
  F --> G[Megatron packed layout]

Loading

_{Reviews (1) · Last reviewed commit: "test(packing): cover MFFD soft singleton..."}

[greptile-apps]

Capacity omits alignment minimum

This says bins have capacity max_tokens_per_microbatch, but the collator raises that capacity to the aggregate alignment when it is larger. For example, TP4 with blockwise FP8 can use a 512-token capacity even when the configured budget is lower. This makes memory and batch-shape planning misleading. The SFT README gives the same capacity description; both guides should explain the alignment-driven minimum.

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!

[greptile-apps]

Parity bins use different capacity

_recompute_flat_bins uses the configured budget without the collator's alignment-driven minimum, and its aggregate alignment calculation omits fp8_recipe. When FP8 alignment exceeds the budget, the collator can combine short sequences in a bin that this helper splits. The resulting example-to-bin mapping differs, so the parity comparison no longer checks the intended examples. Compute capacity and recipe the same way as the collator.

[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 8a1b45d. Configure here.}