Make packed sequence alignment FP8-safe

[jinghanyao1-hub]

SkyRL FP8-Safe Sequence Alignment PR

Summary

This PR fixes Transformer Engine FP8 failures caused by sequence lengths that are valid for SkyRL/Megatron TP or CP layout, but invalid for TE FP8 GEMMs.

Main currently pads sequence lengths only to satisfy Megatron layout constraints:

• non-packed remove_left_padding: pad to tp_size when tp_size > 1;
• packed preprocess_packed_seqs: pad each subsequence to tp_size, or tp_size * cp_size * 2 when CP is enabled.

That is not enough for FP8. TE/cublasLt FP8 GEMMs require the local token dimension used by each rank to be 16-aligned. Without this, FP8 forward/backward can fail with:

Assertion failed: ret.lda % 16 == 0. Leading dimension requirement on A for FP8 GEMM. Caller must pad.

For CP, the subtle part is that preprocess_packed_seqs pads the global subsequence length and then shards it by cp_size. So aligning the global length to 16 is not sufficient: after global_len // cp_size, the local CP-rank slab can still be 8, 4, 2, etc.

Root Cause

There are two failure modes.

1. TP-only or CP-disabled path:

   Example: tp=8, cp=1

   Main can emit token lengths such as 8 or 24. These are divisible by TP, but not valid for the cublasLt FP8 leading-dimension requirement.

2. Packed CP path:

   Example: tp=2, cp=2

   Main pads a raw length of 8 to global length 8, then CP shards it:

   local_len = 8 // 2 = 4
   

   That local length is invalid for TE FP8.

   Example: tp=1, cp=4

   global_len = 8
   local_len = 8 // 4 = 2
   

   Also invalid.

Fix

The PR makes packed sequence alignment satisfy both Megatron layout constraints and TE FP8 local-rank constraints.

For packed sequences:

layout_align = tp_size * cp_size * 2 if cp_size > 1 else tp_size
align_size = math.lcm(layout_align, 16 * cp_size)

Why 16 * cp_size:

• TE needs each rank's local token slab to be 16-aligned.
• CP local length is global_padded_length // cp_size.
• Therefore the global padded length must be divisible by 16 * cp_size.

For non-packed left-padding removal:

align_size = math.lcm(16, tp_size)

remove_left_padding already asserts cp_size == 1, so no CP multiplier is needed there.

The same packed alignment formula is applied in both places that must stay in lockstep:

• skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py
• skyrl/train/dataset/collators.py

This keeps controller-side packed row offsets and worker-side preprocess_packed_seqs offsets identical.

H100 Verification

NVIDIA H100 80GB HBM3
torch 2.11.0+cu128
Transformer Engine 2.11.0

The verification used remove_left_padding and preprocess_packed_seqs outputs, then fed those resulting local token lengths into a TE FP8 Linear forward/backward.

Upstream Main Fails

Results:

remove_left_padding raw=8  tp=8 cp=1: len=8  -> FAIL cublasLt lda % 16
remove_left_padding raw=17 tp=8 cp=1: len=24 -> FAIL cublasLt lda % 16

preprocess_packed_seqs raw=8  tp=8 cp=1: local_total=8  maxq=8  -> FAIL cublasLt lda % 16
preprocess_packed_seqs raw=17 tp=8 cp=1: local_total=24 maxq=24 -> FAIL cublasLt lda % 16

preprocess_packed_seqs raw=8 tp=2 cp=2: local_total=4 maxq=8 -> FAIL TE FP8 dimension assertion
preprocess_packed_seqs raw=8 tp=1 cp=4: local_total=2 maxq=8 -> FAIL TE FP8 dimension assertion

Patched Branch Passes

Results:

remove_left_padding raw=8  tp=8 cp=1: len=16 -> PASS
remove_left_padding raw=17 tp=8 cp=1: len=32 -> PASS

preprocess_packed_seqs raw=8  tp=8 cp=1: local_total=16 maxq=16 -> PASS
preprocess_packed_seqs raw=17 tp=8 cp=1: local_total=32 maxq=32 -> PASS

preprocess_packed_seqs raw=8 tp=2 cp=2: local_total=16 maxq=32 -> PASS
preprocess_packed_seqs raw=8 tp=1 cp=4: local_total=16 maxq=64 -> PASS

Unit Tests

13 passed
tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_cp.py
tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_multiseq.py

16 passed, 1 warning
tests/train/test_packing_round_trip.py
tests/train/test_sft_packing_collate.py

Why This Fix Covers All Observed Failures

The observed FP8 failures all reduce to one invariant:

the token dimension consumed locally by TE FP8 GEMMs must be 16-aligned

For cp=1, local length equals the padded sequence length, so lcm(tp_size, 16) is enough.

For packed cp>1, local length equals:

global_padded_subsequence_length // cp_size

So the global packed subsequence length must be divisible by:

16 * cp_size

The patched packed formula:

math.lcm(tp_or_tp_cp_layout_align, 16 * cp_size)

simultaneously preserves:

• TP sequence-parallel divisibility;
• CP zigzag chunk divisibility;
• TE FP8 16-aligned local GEMM dimensions;
• controller/worker packed-row offset consistency.

[gemini-code-assist]

Code Review

This pull request updates the sequence alignment and padding logic in Megatron utilities and dataset collators to support Transformer Engine FP8 GEMM requirements, ensuring that local token slabs are 16-aligned. The feedback suggests refactoring the duplicated alignment helper function _get_packed_seq_align_size into a single public utility to prevent code duplication and potential logic drift across the codebase and tests.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[erictang000]

PR looks good, just 1 comment

[erictang000]

is it possible to only do 16 alignment if fp8 is enabled? here in the worker we can likely access the megatron model config to check if fp8 is enabled, and in collators.py, we should just pass through megatron_config.transformer_config_kwargs["fp8"] to see if it's enabled and we need to pad subsequences to 16*cp_size.

otherwise we will do unnecessary padding for the bf16 case, which could create additional memory pressure, especially when we have cp_size > 1

[jinghanyao1-hub]

Fixed