perf(train): make sequence packing alignment-aware

skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py

@@ -38,7 +38,8 @@
 from megatron.core.utils import get_attr_wrapped_model, unwrap_model
 
 from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import (
-    get_packed_seq_align_size,
+    get_packing_align_size_sequence,
+    get_packing_align_size_total,
     get_unpacked_seq_align_size,
 )
 
@@ -438,12 +439,9 @@ def preprocess_packed_seqs(
       per row. This is the historical SkyRL behavior used by the RL path
       and the existing SFT path without mini-batch packing.
     - ``sub_seq_lengths is not None``: each row may contain multiple
-      sub-sequences concatenated end-to-end. ``sub_seq_lengths[r]`` lists
-      the per-sub-sequence valid token counts for row ``r``. Tokens
-      ``input_ids[r, :sum(sub_seq_lengths[r])]`` are assumed to be the
-      concatenated sub-sequences in order; any trailing tokens in the row
-      are pad. ``cu_seqlens`` enumerates every sub-sequence across every
-      row.
+      sub-sequences. They are laid out in order with any required CP gaps
+      between them; aggregate TP/FP8 padding follows the final sequence.
+      ``cu_seqlens`` enumerates every sub-sequence across every row.
 
     CP splits sequence into CP*2 chunks, and each GPU gets 2 chunks (GPU0
     gets first and last chunks, GPU1 gets second and second last chunks,
@@ -453,7 +451,8 @@ def preprocess_packed_seqs(
     tp_size = mpu.get_tensor_model_parallel_world_size()
     cp_size = mpu.get_context_parallel_world_size()
     cp_rank = mpu.get_context_parallel_rank()
-    align_size = get_packed_seq_align_size(tp_size, cp_size, fp8_enabled=fp8_enabled)
+    packing_align_size_sequence = get_packing_align_size_sequence(tp_size, cp_size)
+    packing_align_size_total = get_packing_align_size_total(tp_size, cp_size, fp8_enabled=fp8_enabled)
 
     batch_size = input_ids.shape[0]
 
@@ -467,7 +466,7 @@ def preprocess_packed_seqs(
         # Per-row, per-sub-seq starting column within the original padded row.
         # We need this to gather sub-seq tokens from the padded input_ids.
         # NOTE: the controller-side collator (``PackedDataCollator``)
-        # advances ``row_offset += round_up(length, align_size)`` between
+        # advances ``row_offset += round_up(length, packing_align_size_sequence)`` between
         # consecutive sub-sequences in the same row so that flash-attn varlen
         # sees TP/CP-aligned segment boundaries. We MUST mirror that here —
         # otherwise sub-seq i (for i > 0) would be read starting inside the
@@ -481,9 +480,11 @@ def preprocess_packed_seqs(
                 flat_seqlens.append(length_int)
                 row_index_of_subseq.append(r)
                 intra_row_offset_of_subseq.append(running)
-                # Pad each sub-seq independently to align_size, matching the
+                # Pad each sub-seq independently to packing_align_size_sequence, matching the
                 # collator's row layout.
-                pad = (align_size - length_int % align_size) % align_size
+                pad = (
+                    packing_align_size_sequence - length_int % packing_align_size_sequence
+                ) % packing_align_size_sequence
                 running += length_int + pad
 
         seqlens_in_batch = torch.tensor(flat_seqlens, dtype=torch.int32, device=input_ids.device)
@@ -492,8 +493,14 @@ def preprocess_packed_seqs(
         seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
         num_subseqs = batch_size
 
-    pad_size = (align_size - seqlens_in_batch % align_size) % align_size
+    pad_size = (
+        packing_align_size_sequence - seqlens_in_batch % packing_align_size_sequence
+    ) % packing_align_size_sequence
     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
 
     cu_seqlens = torch.zeros(num_subseqs + 1, dtype=torch.int32, device=input_ids.device)
     cu_seqlens[1:] = torch.cumsum(seqlens_in_batch, dim=0)

skyrl/backends/skyrl_train/distributed/megatron/packing_utils.py

@@ -9,8 +9,15 @@ def is_fp8_enabled(fp8: Any) -> bool:
     return bool(fp8)
 
 
-def get_packed_seq_align_size(tp_size: int, cp_size: int, fp8_enabled: bool = False) -> int:
-    """Return global per-subsequence padding needed for TP/CP layout."""
+def get_packing_align_size_sequence(tp_size: int, cp_size: int) -> int:
+    """Return the alignment required independently for each packed sequence."""
+    if cp_size > 1:
+        return tp_size * cp_size * 2
+    return 1
+
+
+def get_packing_align_size_total(tp_size: int, cp_size: int, fp8_enabled: bool = False) -> int:
+    """Return the alignment required for the aggregate packed token slab."""
     if cp_size > 1:
         layout_align = tp_size * cp_size * 2
     else:

skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py

@@ -5,7 +5,8 @@
 import torch
 
 from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import (
-    get_packed_seq_align_size,
+    get_packing_align_size_sequence,
+    get_packing_align_size_total,
     get_unpacked_seq_align_size,
 )
 
@@ -63,8 +64,11 @@ def build_token_metadata_layout(
         )
 
     cp_size = mpu.get_context_parallel_world_size()
-    align_size = get_packed_seq_align_size(tp_size, cp_size, fp8_enabled=fp8_enabled)
-    padded_sequence_lengths_tensor = sequence_lengths_tensor + (-sequence_lengths_tensor % align_size)
+    packing_align_size_sequence = get_packing_align_size_sequence(tp_size, cp_size)
+    packing_align_size_total = get_packing_align_size_total(tp_size, cp_size, fp8_enabled=fp8_enabled)
+    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 = padded_sequence_lengths_tensor.tolist()
     cu_seqlens_padded = torch.cat(
         (

skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py

@@ -104,8 +104,8 @@ def _build_packed_valid_mask(
 ) -> torch.Tensor:
     """Build a ``[1, T]`` real-token mask aligned to the packed (THD) logits layout.
 
-    1.0 for real tokens, 0.0 for the per-segment alignment padding that ``preprocess_packed_seqs``
-    inserts between sub-sequences. This is the packed counterpart of the ``[batch, seq]``
+    1.0 for real tokens, 0.0 for the layout gaps and aggregate tail padding that
+    ``preprocess_packed_seqs`` inserts. This is the packed counterpart of the ``[batch, seq]``
     ``attention_mask`` the decoupled MTP draft loss uses to mask invalid positions; mirrors the
     index math of :func:`_build_packed_targets` but scatters ones instead of token ids.
     """

skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py

@@ -37,6 +37,11 @@
     get_megatron_optimizer_param_scheduler,
     init_megatron_optim_config,
 )
+from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import (
+    get_packing_align_size_sequence,
+    get_packing_align_size_total,
+    is_fp8_enabled,
+)
 from skyrl.backends.skyrl_train.inference_servers.remote_inference_client import (
     SKYRL_LORA_ADAPTER_NAME,
 )
@@ -330,6 +335,19 @@ def extract_weights(self, dtype: torch.dtype):
 
 
 class MegatronWorker:
+    def _packed_sequence_length_multiples(self) -> tuple[int, int]:
+        """Return per-sequence and aggregate packed THD alignments."""
+        if not self.cfg.remove_microbatch_padding:
+            return 1, 1
+        model_config = get_model_config(self.actor_module[0])
+        tp_size = mpu.get_tensor_model_parallel_world_size()
+        cp_size = mpu.get_context_parallel_world_size()
+        return get_packing_align_size_sequence(tp_size, cp_size), get_packing_align_size_total(
+            tp_size,
+            cp_size,
+            fp8_enabled=is_fp8_enabled(model_config.fp8),
+        )
+
     def _maybe_setup_fake_int4_qat(self):
         """Wire up INT4-served training and return the BF16 bridge-weights path.
 
@@ -644,10 +662,13 @@ def _forward_logprobs(self, data: TrainingInputBatch) -> torch.Tensor:
         use_token_batching = self.cfg.max_tokens_per_microbatch > 0
 
         if use_token_batching:
+            sequence_length_multiple, packed_length_multiple = self._packed_sequence_length_multiples()
             microbatch_iterator = get_microbatch_iterator(
                 data,
                 micro_batch_size=self.cfg.micro_forward_batch_size_per_gpu,
                 max_tokens_per_microbatch=self.cfg.max_tokens_per_microbatch,
+                sequence_length_multiple=sequence_length_multiple,
+                packed_length_multiple=packed_length_multiple,
             )
         else:
             microbatch_iterator = None
@@ -1150,10 +1171,13 @@ def forward_backward(
         use_token_batching = self.cfg.max_tokens_per_microbatch > 0
 
         if use_token_batching:
+            sequence_length_multiple, packed_length_multiple = self._packed_sequence_length_multiples()
             microbatch_iterator = get_microbatch_iterator(
                 data,
                 micro_batch_size=self.cfg.micro_train_batch_size_per_gpu,
                 max_tokens_per_microbatch=self.cfg.max_tokens_per_microbatch,
+                sequence_length_multiple=sequence_length_multiple,
+                packed_length_multiple=packed_length_multiple,
             )
         else:
             microbatch_iterator = None

skyrl/backends/skyrl_train/workers/worker_utils.py

@@ -257,11 +257,17 @@ def __init__(
         self,
         data: TrainingInputBatch,
         max_tokens_per_microbatch: int,
+        sequence_length_multiple: int = 1,
+        packed_length_multiple: int = 1,
     ):
         """
         Args:
             data: The training input batch to chunk.
             max_tokens_per_microbatch: Maximum number of tokens per microbatch.
+            sequence_length_multiple: Round each sequence's packing footprint
+                up to this multiple.
+            packed_length_multiple: Round the aggregate microbatch footprint
+                up to this multiple.
         """
         super().__init__(data)
         self._max_tokens_per_microbatch = max_tokens_per_microbatch
@@ -273,7 +279,12 @@ def __init__(
         # Create microbatches based on token count. The "balanced" packer treats
         # the token budget as a soft cap: a sequence longer than the budget gets
         # its own (over-budget) microbatch rather than raising.
-        packer = make_seq_packer("balanced", bin_capacity=self._max_tokens_per_microbatch)
+        packer = make_seq_packer(
+            "balanced",
+            bin_capacity=self._max_tokens_per_microbatch,
+            sequence_length_multiple=sequence_length_multiple,
+            packed_length_multiple=packed_length_multiple,
+        )
         self._microbatches = packer.pack(self._token_counts)
 
         # Synchronize the number of microbatches across all DP workers
@@ -426,19 +437,32 @@ def reorder_and_combine_batches(self, batches: List[TensorBatch]) -> TensorBatch
 
 
 def get_microbatch_iterator(
-    data: TrainingInputBatch, micro_batch_size: int, max_tokens_per_microbatch: int
+    data: TrainingInputBatch,
+    micro_batch_size: int,
+    max_tokens_per_microbatch: int,
+    sequence_length_multiple: int = 1,
+    packed_length_multiple: int = 1,
 ) -> BaseBatchIterator:
     """Factory function to get the appropriate microbatch iterator.
 
     Args:
         data: The training input batch.
         micro_batch_size: Number of samples per microbatch (used if max_tokens_per_microbatch <= 0).
         max_tokens_per_microbatch: Maximum tokens per microbatch. If > 0, uses token-based batching.
+        sequence_length_multiple: Round each sequence's packing footprint up
+            to this multiple for token-based batching.
+        packed_length_multiple: Round each aggregate microbatch footprint up
+            to this multiple for token-based batching.
 
     Returns:
         A BaseBatchIterator instance.
     """
     if max_tokens_per_microbatch > 0:
-        return TokenBasedBatchIterator(data, max_tokens_per_microbatch=max_tokens_per_microbatch)
+        return TokenBasedBatchIterator(
+            data,
+            max_tokens_per_microbatch=max_tokens_per_microbatch,
+            sequence_length_multiple=sequence_length_multiple,
+            packed_length_multiple=packed_length_multiple,
+        )
     else:
         return SampleBasedBatchIterator(data, sample_batch_size=micro_batch_size, drop_last=False)

skyrl/train/dataset/bin_packing.py

@@ -47,25 +47,38 @@ def __init__(
         bin_capacity: int,
         min_bin_count: Optional[int] = None,
         bin_count_multiple: Optional[int] = None,
+        sequence_length_multiple: int = 1,
+        packed_length_multiple: int = 1,
     ):
         if min_bin_count is not None and min_bin_count < 0:
             raise ValueError("min_bin_count must be nonnegative")
         if bin_count_multiple is not None and bin_count_multiple < 1:
             raise ValueError("bin_count_multiple must be positive")
+        if sequence_length_multiple < 1:
+            raise ValueError("sequence_length_multiple must be positive")
+        if packed_length_multiple < 1:
+            raise ValueError("packed_length_multiple must be positive")
 
         self.bin_capacity = bin_capacity
         self.min_bin_count = min_bin_count
         self.bin_count_multiple = bin_count_multiple
+        self.sequence_length_multiple = sequence_length_multiple
+        self.packed_length_multiple = packed_length_multiple
 
     @abstractmethod
     def _pack_implementation(self, sequence_lengths: List[int]) -> List[List[int]]:
         """Pack sequences into bins. Override in sub-class."""
 
     def _validate_sequence_lengths(self, sequence_lengths: List[int]) -> None:
         for length in sequence_lengths:
-            if length > self.bin_capacity:
+            if self._packed_length(length) > self.bin_capacity:
                 raise ValueError(f"Sequence length {length} exceeds bin capacity {self.bin_capacity}")
 
+    def _packed_length(self, sequence_length_sum: int) -> int:
+        """Return the physical footprint after aggregate tail padding."""
+        multiple = self.packed_length_multiple
+        return sequence_length_sum + (-sequence_length_sum % multiple)
+
     def _adjust_bin_count(self, bins: List[List[int]]) -> List[List[int]]:
         """Pad the bin list to a multiple of ``bin_count_multiple``.
 
@@ -123,8 +136,10 @@ def _adjust_bin_count(self, bins: List[List[int]]) -> List[List[int]]:
         return adjusted_bins
 
     def pack(self, sequence_lengths: List[int]) -> List[List[int]]:
-        """Pack ``sequence_lengths`` into bins and apply DP-symmetry adjustment."""
-        bins = self._pack_implementation(sequence_lengths)
+        """Pack layout-aligned sequences and apply aggregate tail padding."""
+        multiple = self.sequence_length_multiple
+        packing_lengths = [length + (-length % multiple) for length in sequence_lengths]
+        bins = self._pack_implementation(packing_lengths)
         bins = self._adjust_bin_count(bins)
         return bins
 
@@ -144,19 +159,20 @@ def _pack_implementation(self, sequence_lengths: List[int]) -> List[List[int]]:
         indexed.sort(reverse=True)
 
         bins: List[List[int]] = []
-        bin_remaining: List[int] = []
+        bin_lengths: List[int] = []
 
         for length, idx in indexed:
             placed = False
-            for i, remaining in enumerate(bin_remaining):
-                if remaining >= length:
+            for i, bin_length in enumerate(bin_lengths):
+                new_length = bin_length + length
+                if self._packed_length(new_length) <= self.bin_capacity:
                     bins[i].append(idx)
-                    bin_remaining[i] -= length
+                    bin_lengths[i] = new_length
                     placed = True
                     break
             if not placed:
                 bins.append([idx])
-                bin_remaining.append(self.bin_capacity - length)
+                bin_lengths.append(length)
 
         return bins
 
@@ -199,7 +215,7 @@ def _pack_implementation(self, sequence_lengths: List[int]) -> List[List[int]]:
             if bin_tokens_heap:
                 bin_tokens, bin_idx = bin_tokens_heap[0]
                 new_bin_tokens = bin_tokens + length
-                if new_bin_tokens <= self.bin_capacity:
+                if self._packed_length(new_bin_tokens) <= self.bin_capacity:
                     bins[bin_idx].append(idx)
                     heapq.heapreplace(bin_tokens_heap, (new_bin_tokens, bin_idx))
                     placed = True
@@ -222,6 +238,8 @@ def make_seq_packer(
     bin_capacity: int,
     min_bin_count: Optional[int] = None,
     bin_count_multiple: Optional[int] = None,
+    sequence_length_multiple: int = 1,
+    packed_length_multiple: int = 1,
 ) -> SeqPacker:
     """Factory returning a configured :class:`SeqPacker` instance.
 
@@ -234,6 +252,10 @@ def make_seq_packer(
             ``dp_size``).
         bin_count_multiple: Force the total bin count to be a multiple of
             this value (typically ``dp_size``).
+        sequence_length_multiple: Round each sequence's packing footprint up
+            to this multiple before placement.
+        packed_length_multiple: Round each bin's aggregate footprint up to
+            this multiple when checking capacity.
     """
     if isinstance(algorithm, str):
         try:
@@ -250,4 +272,6 @@ def make_seq_packer(
         bin_capacity=bin_capacity,
         min_bin_count=min_bin_count,
         bin_count_multiple=bin_count_multiple,
+        sequence_length_multiple=sequence_length_multiple,
+        packed_length_multiple=packed_length_multiple,
     )