perf(train): make sequence packing alignment-aware

[dyurk-lila]

Summary

• distinguish per-sequence layout alignment from aggregate packed-slab alignment
• make FFD capacity checks account for the physical aligned footprint used by packed SFT batches
• make balanced token batching use the same alignment-aware accounting on Megatron workers
• apply aggregate TP/FP8 tail padding once, rather than padding every sequence to that multiple
• keep the collator, packed preprocessing, token metadata, and target layouts in lockstep

Alignment model

Two constraints have different scopes:

• packing_align_size_sequence: alignment required independently by each sequence. This is 2 * tp_size * cp_size when context parallelism is active, because each sequence must support the CP/SP layout; otherwise it is 1.
• packing_align_size_total: alignment required by the final packed token slab. This covers TP divisibility and the Transformer Engine FP8 requirement (16 * cp_size) and is paid once after the per-sequence footprints are summed.

The packers expose matching sequence_length_multiple and packed_length_multiple inputs. They first align each sequence to the former, then round each candidate bin total to the latter for capacity checks. This avoids both underestimating CP layout gaps and overestimating TP/FP8 padding by charging only CP/SP alignment per sequence and TP/FP8 alignment once per packed bin.

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.

Affected paths

• controller-level SFT packing (PackedDataCollator, FFD)
• worker-side token-based microbatching (TokenBasedBatchIterator, balanced packing)
• Megatron packed-sequence preprocessing and CP sharding
• packed token metadata, targets, masks, and GPU parity references

Testing

• pytest -q tests/backends/skyrl_train/distributed/test_bin_packing.py (24 passed)
• added focused tests for CP per-sequence gaps, aggregate TP/FP8 tails, FFD capacity, balanced batching, packed metadata, and round trips
• Ruff lint
• Black formatting check

---

Note

High Risk
Changes packed THD layout, collator offsets, and microbatch token budgets across SFT/RL Megatron paths; misalignment can corrupt loss/grads without obvious failures, though coverage is broad.

Overview
Refactors Megatron packed-sequence alignment so CP/SP layout padding is applied per sub-sequence, while TP and FP8 padding is applied once on the aggregate token slab instead of rounding every segment to the same multiple.

Replaces get_packed_seq_align_size with get_packing_align_size_sequence (CP-only per-sequence gaps; 1 when CP is off) and get_packing_align_size_total (TP + optional FP8 on the full packed length). preprocess_packed_seqs, PackedDataCollator, token_metadata, and packed valid masks are updated to share that model so collator row layout and worker preprocessing stay matched.

Bin packing and token microbatching now take sequence_length_multiple and packed_length_multiple so FFD/balanced capacity checks budget the aligned physical footprint (including one aggregate tail pad per bin/microbatch). Megatron workers pass the same multiples from _packed_sequence_length_multiples into get_microbatch_iterator when remove_microbatch_padding is enabled.

Tests and parity references are refreshed for contiguous sub-sequences without CP, smaller rows when TP/FP8 only affect the tail, and CP rows that still insert per-sequence gaps.

^{Reviewed by Cursor Bugbot for commit 0d1f5c8. 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 alignment logic for Megatron-based training by separating the per-sequence layout alignment from the aggregate total alignment. This optimization applies TP and FP8 tail padding once to the final sequence instead of to every individual sequence, reducing padding overhead. The review feedback identifies potential IndexError crashes in megatron_utils.py and token_metadata.py when indexing the last element of empty tensors (e.g., during empty input batches) and suggests adding safety guards.

[gemini-code-assist]

There's a potential IndexError here if seqlens_in_batch_padded is an empty tensor. This can happen if seqlens_in_batch is empty, for example with an empty input batch. Accessing seqlens_in_batch_padded[-1] would cause a crash. It's safer to add a guard to handle this edge case.

Suggested change

	seqlens_in_batch_padded = seqlens_in_batch + 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.
	aggregate_pad_size = (-seqlens_in_batch_padded.sum()) % packing_align_size_total
	seqlens_in_batch_padded[-1] += aggregate_pad_size
	seqlens_in_batch_padded = seqlens_in_batch + 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

[dyurk-lila]

SFT rejects an empty dataloader explicitly at skyrl/train/sft_trainer.py:1890, and RL padding microbatches deliberately contain one valid token at skyrl/backends/skyrl_train/workers/worker_utils.py:316. We don't need an extra defensive check here.

[gemini-code-assist]

Similar to another comment, there's a potential IndexError on line 71 if padded_sequence_lengths_tensor is empty. This could occur if sequence_lengths_tensor is empty (e.g., from an empty attention_mask). Adding a guard would make this more robust.

Suggested change

	padded_sequence_lengths_tensor = sequence_lengths_tensor + (-sequence_lengths_tensor % packing_align_size_sequence)
	aggregate_pad_size = (-padded_sequence_lengths_tensor.sum()) % packing_align_size_total
	padded_sequence_lengths_tensor[-1] += aggregate_pad_size
	padded_sequence_lengths_tensor = sequence_lengths_tensor + (-sequence_lengths_tensor % packing_align_size_sequence)
	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

[dyurk-lila]

Same as above