feat(sft): pad to a multiple of cp and remove fp8 indexer dynamism

src/prime_rl/orchestrator/batch.py

@@ -696,11 +696,10 @@ def _distribute_group(
 
 def pad_micro_batch(micro_batch: MicroBatch, pad_to_multiple_of: int) -> MicroBatch:
     """
-    Pad a micro batch with the given padding size sample
-    Return the padded micro batch.
+    Pad a micro batch with loss-masked tokens up to the next multiple of `pad_to_multiple_of`.
     Args:
         micro_batch: The micro batch to pad.
-        padding_size: The number of padding tokens to add.
+        pad_to_multiple_of: The multiple to pad the sequence length up to.
     Returns:
         The padded micro batch.
     """

src/prime_rl/trainer/models/kernels/fp8_indexer.py

@@ -12,6 +12,7 @@
 FP8_MAX = 448.0
 FP8_MIN = -448.0
 FP8_EPS = 1e-10
+KEY_ALIGNMENT = 16
 
 
 @triton.jit
@@ -88,17 +89,6 @@ def per_token_group_quant_fp8(
     return x_q, x_s
 
 
-@triton.autotune(
-    configs=[
-        triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=4, num_stages=2),
-        triton.Config({"BLOCK_M": 64, "BLOCK_N": 128}, num_warps=4, num_stages=2),
-        triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=8, num_stages=2),
-        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=8, num_stages=2),
-        triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=8, num_stages=3),
-        triton.Config({"BLOCK_M": 64, "BLOCK_N": 128}, num_warps=8, num_stages=3),
-    ],
-    key=["S_Q", "S_K_BUCKET", "H", "D"],
-)
 @triton.jit
 def _triton_fp8_indexer_kernel(
     Q_fp8,
@@ -116,7 +106,6 @@ def _triton_fp8_indexer_kernel(
     stride_ws,
     H: tl.constexpr,
     D: tl.constexpr,
-    S_K_BUCKET: tl.constexpr,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
 ):
@@ -201,6 +190,8 @@ def fp8_indexer(
 
     q_flat = q.reshape(S_q * H, D).contiguous()
     q_fp8, q_scales = per_token_group_quant_fp8(q_flat, group_size=D, use_ue8m0=True)
+    S_k_aligned = triton.cdiv(S_k, KEY_ALIGNMENT) * KEY_ALIGNMENT
+    k = torch.nn.functional.pad(k, (0, 0, 0, S_k_aligned - S_k))
     k_fp8, k_scales = per_token_group_quant_fp8(k.contiguous(), group_size=D, use_ue8m0=True)
 
     q_fp8 = q_fp8.view(S_q, H, D).permute(1, 0, 2).contiguous()
@@ -209,12 +200,10 @@ def fp8_indexer(
     q_scales = q_scales.view(S_q, H)
     w = w * q_scales
 
-    logits = torch.empty(S_q, S_k, dtype=torch.float32, device=device)
+    logits = torch.empty(S_q, S_k_aligned, dtype=torch.float32, device=device)
 
-    grid = lambda meta: (
-        triton.cdiv(S_q, meta["BLOCK_M"]),
-        triton.cdiv(S_k, meta["BLOCK_N"]),
-    )
+    BLOCK_M, BLOCK_N = 64, 128
+    grid = (triton.cdiv(S_q, BLOCK_M), triton.cdiv(S_k_aligned, BLOCK_N))
     _triton_fp8_indexer_kernel[grid](
         q_fp8,
         k_fp8,
@@ -224,16 +213,17 @@ def fp8_indexer(
         ks,
         ke,
         S_q,
-        S_k,
+        S_k_aligned,
         q_fp8.stride(0),
         q_fp8.stride(1),
         k_fp8.stride(0),
         w.stride(0),
         H=H,
         D=D,
-        # Packed document tails change S_k between batches. Reuse tuning within
-        # a size bucket; exact S_k still controls strides and all bounds masks.
-        S_K_BUCKET=triton.next_power_of_2(S_k),
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        num_warps=4,
+        num_stages=2,
     )
 
     actual_topk = min(topk, S_k)

src/prime_rl/trainer/parallel_dims.py

@@ -112,7 +112,7 @@ def _build_mesh_with_ep(self) -> DeviceMesh:
 
         # Create all the submesh here to ensure all required process groups are
         # initialized:
-        # Mesh for data loading (no communication on this mesh)
+        # Mesh for data loading and per-step token counts
         dp_mesh_dim_names = []
         # Mesh for param sharding
         dp_shard_cp_mesh_dim_names = []
@@ -172,7 +172,7 @@ def _build_mesh_without_ep(self) -> DeviceMesh:
 
         # Create all the submesh here to ensure all required process groups are
         # initialized:
-        # Mesh for data loading (no communication on this mesh)
+        # Mesh for data loading and per-step token counts
         dp_mesh_dim_names = []
         # Mesh for param sharding
         dp_shard_cp_mesh_dim_names = []
@@ -263,13 +263,6 @@ def fsdp_gradient_divide_factor(self) -> int:
     def non_data_parallel_size(self):
         return self.cp * self.pp
 
-    @cached_property
-    def seq_len_divisor(self):
-        # Context Parallel requires that seq_len be divisible by 2 * CP degree,
-        # when load balancing is enabled (by default).
-        # https://github.com/pytorch/pytorch/blob/4f62dcc/torch/distributed/tensor/experimental/_attention.py#L1246
-        return self.cp * 2
-
     @cached_property
     def logger(self):
         return get_logger()
@@ -309,7 +302,7 @@ def resolve_ep(config: ModelConfig) -> None:
     get_logger().info(f"EP auto: world_size={world_size}, dp_replicate={dp_replicate} -> resolved ep={resolved_ep}")
 
 
-def get_parallel_dims(config: ModelConfig, seq_len: int | None = None) -> ParallelDims:
+def get_parallel_dims(config: ModelConfig) -> ParallelDims:
     assert isinstance(config.ep, int), (
         f"config.ep must be resolved to an int before get_parallel_dims; got {config.ep!r}. "
         "Call resolve_ep(config) first."
@@ -325,12 +318,4 @@ def get_parallel_dims(config: ModelConfig, seq_len: int | None = None) -> Parall
         world_size=dist.get_world_size(),
     )
 
-    # Validate sequence length against parallel dimensions requirements
-    if seq_len is not None and seq_len % parallel_dims.seq_len_divisor != 0:
-        raise ValueError(
-            f"Sequence length ({seq_len}) must be divisible by "
-            f"seq_len_divisor ({parallel_dims.seq_len_divisor}) for the given parallel dimensions. "
-            f"This requirement comes from context parallel (CP={config.cp})."
-        )
-
     return parallel_dims

src/prime_rl/trainer/rl/train.py

@@ -339,7 +339,9 @@ def train(config: TrainerConfig):
         cp_size = parallel_dims.cp
 
         step_tokens_per_expert = 0
+        step_local_num_tokens = 0
         for micro_step, micro_batch in enumerate(micro_batch_workers(prepare, micro_batches)):
+            step_local_num_tokens += micro_batch["input_ids"].shape[1]
             input_ids = micro_batch["input_ids"].to("cuda")
             position_ids = micro_batch["position_ids"].to("cuda")
             advantages = micro_batch["advantages"].to("cuda")
@@ -627,9 +629,10 @@ def train(config: TrainerConfig):
         if is_moe_model:
             tensor_stats.update(get_expert_load_stats(step_tokens_per_expert, dp_cp_group))
 
-        # Compute step metrics
-        num_local_tokens = seq_len * batch_size
-        num_tokens = parallel_dims.get_mesh("dp").size() * num_local_tokens
+        # The dp mesh excludes cp, whose ranks hold the same rows.
+        global_num_tokens = torch.tensor(step_local_num_tokens, dtype=torch.int64, device="cuda")
+        dist.all_reduce(global_num_tokens, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group())
+        num_tokens = global_num_tokens.item()
         progress.total_tokens += num_tokens
         progress.total_samples += batch_size
         perf_counter = get_perf_counter(model, seq_len)

src/prime_rl/trainer/sft/data.py

@@ -474,12 +474,13 @@ def __iter__(self):
 
 
 class CatDataset(StatefulIterableDataset):
-    """Concatenate text and multimodal samples into one fixed-length row."""
+    """Pack samples into rows of at most `seq_len` tokens, padded to a multiple of `pad_to_multiple_of`."""
 
-    def __init__(self, dataset: StatefulIterableDataset, seq_len: int):
+    def __init__(self, dataset: StatefulIterableDataset, seq_len: int, pad_to_multiple_of: int):
         self.logger = get_logger()
         self.dataset = dataset
         self.seq_len = seq_len
+        self.pad_to_multiple_of = pad_to_multiple_of
         self.pending_sample: Sample | None = None
 
     def state_dict(self) -> dict:
@@ -583,7 +584,7 @@ def _finalize_pack(self, packed: dict[str, Any], seq_len: int) -> dict:
             if kept > 0:
                 result["seq_lens"].append(kept)
             remaining -= kept
-        pad_len = seq_len - len(result["input_ids"])
+        pad_len = -len(result["input_ids"]) % self.pad_to_multiple_of
         if pad_len > 0:
             result["input_ids"].extend([0] * pad_len)
             result["position_ids"].extend(range(pad_len))
@@ -753,8 +754,8 @@ def setup_dataset(
         raise ValueError(f"Invalid dataset type: {config.type}")
 
 
-def setup_dataloader(dataset: StatefulIterableDataset, config: DataConfig) -> StatefulDataLoader:
-    packing_dataset = CatDataset(dataset, config.seq_len * config.micro_batch_size)
+def setup_dataloader(dataset: StatefulIterableDataset, config: DataConfig, cp: int) -> StatefulDataLoader:
+    packing_dataset = CatDataset(dataset, config.seq_len * config.micro_batch_size, pad_to_multiple_of=cp)
     return StatefulDataLoader(
         packing_dataset,
         batch_size=1,

src/prime_rl/trainer/sft/train.py

@@ -113,7 +113,7 @@ def train(config: SFTConfig):
     resolve_ep(config.model)
 
     # Initialize parallel dimensions
-    parallel_dims = get_parallel_dims(config.model, config.data.seq_len)
+    parallel_dims = get_parallel_dims(config.model)
 
     total_micro_batches = config.data.batch_size * config.model.cp
     micro_batches_per_step = world.world_size * config.data.micro_batch_size
@@ -192,7 +192,7 @@ def train(config: SFTConfig):
         processor=processor,
         multimodal=multimodal,
     )
-    dataloader = setup_dataloader(dataset, config.data)
+    dataloader = setup_dataloader(dataset, config.data, config.model.cp)
 
     val_raw_dataset = None
     if config.val is not None:
@@ -257,7 +257,7 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]:
 
         if cp_enabled:
             # CP requires the sequence length to be divisible by cp_size. CatDataset
-            # pads every pack to seq_len; shard_for_cp raises on violations.
+            # pads every pack to a multiple of it; shard_for_cp raises on violations.
             defer_vlm_cp_to_model = (
                 mm_kwargs is not None and "image_grid_thw" in mm_kwargs and config.model.cp_style == "ulysses"
             )
@@ -345,7 +345,7 @@ def run_validation(step: int) -> None:
             processor=processor,
             multimodal=multimodal,
         )
-        val_dataloader = setup_dataloader(val_dataset, config.val.data)
+        val_dataloader = setup_dataloader(val_dataset, config.val.data, config.model.cp)
 
         # No train/eval switch: no dropout in these models, and toggling would trigger torch.compile recompilation
         mean_loss, nan_count = run_eval_loop(val_dataloader)
@@ -451,7 +451,9 @@ def is_online_eval_step(step: int) -> bool:
             )
 
         step_tokens_per_expert = 0
+        step_local_num_tokens = 0
         for micro_step, micro_batch in enumerate(micro_batches):
+            step_local_num_tokens += micro_batch["input_ids"].shape[1]
             if config.log.log_data:
                 print_sample(
                     micro_batch["input_ids"].flatten().tolist(), micro_batch["loss_mask"].flatten().tolist(), tokenizer
@@ -549,14 +551,10 @@ def is_online_eval_step(step: int) -> bool:
         if memory_profiler is not None:
             memory_profiler.step()
 
-        # Compute step metrics. CP shards the same sequences across cp ranks
-        # (sequence-sharded data parallelism on the seq dim), so the unique
-        # training tokens per step is dp_size * (batch_per_dp_rank * seq).
-        # The `dp` mesh excludes cp by construction (parallel_dims.py), mirroring
-        # the RL trainer's accounting (rl/train.py).
-        dp_size = parallel_dims.get_mesh("dp").size()
-        num_local_tokens = config.data.seq_len * (config.data.batch_size // dp_size)
-        num_tokens = dp_size * num_local_tokens
+        # The dp mesh excludes cp, whose ranks hold the same rows.
+        global_num_tokens = torch.tensor(step_local_num_tokens, dtype=torch.int64, device="cuda")
+        dist.all_reduce(global_num_tokens, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group())
+        num_tokens = global_num_tokens.item()
         progress.total_tokens += num_tokens
         dataset_progress = get_dataset_progress(dataloader)
         progress.total_samples = dataset_progress["step"]

tests/unit/train/sft/test_dataloader.py

@@ -19,7 +19,7 @@ def setup_fake_dataloader(config: FakeDataConfig, non_dp_size: int = 1):
         seed=config.seed,
         non_dp_size=non_dp_size,
     )
-    return dataset, setup_dataloader(dataset, config)
+    return dataset, setup_dataloader(dataset, config, cp=1)
 
 
 def test_fake_dataset_single_rank_state():
@@ -88,7 +88,7 @@ def test_fake_dataset_single_rank_resume():
 
     # Reload dataloader
     state_dict = dataloader.state_dict()
-    dataloader = setup_dataloader(dataset, config)
+    dataloader = setup_dataloader(dataset, config, cp=1)
     dataloader.load_state_dict(state_dict)
     dataiter = iter(dataloader)
 
@@ -110,7 +110,8 @@ def test_fake_dataset_single_rank_state_with_packing():
         micro_batch = next(dataiter)
         num_packed_examples = len(micro_batch["input_ids"][micro_batch["loss_mask"]].unique())
         step += num_packed_examples
-        assert micro_batch["input_ids"].shape == (1, 128)
+        assert micro_batch["input_ids"].shape[0] == 1
+        assert micro_batch["input_ids"].shape[1] <= 128
         assert micro_batch["seq_lens"].sum() == micro_batch["input_ids"].shape[1]
         worker_state = dataloader.state_dict()["_snapshot"]["_worker_snapshots"]["worker_0"]["dataset_state"]
         pending_sample = worker_state.get("pending_sample")
@@ -186,7 +187,7 @@ def setup_epoch_dataloader():
                 seq_len=config.seq_len,
                 non_dp_size=non_dp_size,
             )
-            return setup_dataloader(dataset, config)
+            return setup_dataloader(dataset, config, cp=1)
 
         dataloader = setup_epoch_dataloader()
         dataiter = iter(dataloader)

tests/unit/train/sft/test_sft_dataset.py

@@ -443,6 +443,7 @@ def test_cat_dataset_packs_multimodal_samples():
             ),
         ],
         seq_len=5,
+        pad_to_multiple_of=1,
     )
 
     packed = next(iter(dataset))
@@ -467,23 +468,24 @@ def test_cat_dataset_packs_text_and_multimodal_samples_together():
                 mm_token_type_ids=[0, 1],
             ),
             _sft_sample([4]),
-            _sft_sample([5, 6]),
+            _sft_sample([5, 6, 7, 8, 9, 10]),
         ],
-        seq_len=5,
+        seq_len=9,
+        pad_to_multiple_of=3,
     )
 
     dataiter = iter(dataset)
     packed = next(dataiter)
     text_pack = next(dataiter)
 
-    assert packed["input_ids"] == [1, 2, 3, 4, 0]
-    assert packed["loss_mask"] == [True, True, True, True, False]
-    assert packed["seq_lens"] == [1, 2, 2]
+    assert packed["input_ids"] == [1, 2, 3, 4, 0, 0]
+    assert packed["loss_mask"] == [True, True, True, True, False, False]
+    assert packed["seq_lens"] == [1, 2, 3]
     assert packed["mm_kwargs"] is not None
-    assert packed["mm_token_type_ids"] == [0, 0, 1, 0, 0]
-    assert text_pack["input_ids"] == [5, 6, 0, 0, 0]
-    assert text_pack["loss_mask"] == [True, True, False, False, False]
-    assert text_pack["seq_lens"] == [5]
+    assert packed["mm_token_type_ids"] == [0, 0, 1, 0, 0, 0]
+    assert text_pack["input_ids"] == [5, 6, 7, 8, 9, 10]
+    assert text_pack["loss_mask"] == [True] * 6
+    assert text_pack["seq_lens"] == [6]
     assert text_pack["mm_kwargs"] is None
     assert text_pack["mm_token_type_ids"] is None