feat(trainer): remove CP token padding and duplicated RL microbatches

docs/scaling.md

@@ -122,6 +122,16 @@ cp = 2                       # CP degree
 cp_style = "ulysses"         # "ring"
 ```
 
+RL distributes the original packed tokens into balanced CP shards. For example, nine tokens across four ranks produce lengths `[3, 2, 2, 2]`; rows shorter than the CP degree include empty shards. The same partition applies to positions, labels, temperatures, sampling masks, and router replay. Packed sample boundaries remain independent of rank boundaries. Ulysses retains its head divisibility requirements.
+
+The packer splits at whole-sample boundaries to balance distributed iterations. When a lane has no sample for an iteration, it executes an explicit zero-token slot. FSDP parameter hooks, expert communication, and vision participation still run, so other lanes can use parameters and experts owned by that rank. Empty slots contribute no loss members or sample annotations. A globally empty training step is rejected before updating the optimizer.
+
+These behaviors are enabled by `model.cp_unpadded` and `model.inactive_micro_batches`, both defaulting to `true`. Managed RL configures the orchestrator accordingly; a standalone orchestrator must use matching trainer settings. `perf/input_tokens`, `perf/real_micro_batches`, and `perf/inactive_micro_batches` report the executed data separately from empty slots. Sampling-set `-1` entries, kernel workspace alignment, and SFT dataset packing retain their own storage contracts.
+
+Changing token counts can create additional compiled variants, particularly for zero- and one-token shards. CUDA Graph capture depends on the chosen compile mode and kernel support. Measure useful-token throughput and memory on the target workload; removing input tokens does not guarantee lower step latency.
+
+Ring CP gathers a combined buffer of real keys and values before applying causal packed attention. Its temporary memory scales with the full row's KV width. Recurrent layers redistribute features or gather their real sequence as required by their kernel. Compilation still follows each kernel's existing restrictions; FA4 and multi-LoRA contain graph breaks in the current dependencies and require the default `fullgraph=false` setting.
+
 ### Activation Checkpointing and Offloading
 
 | Knob | Memory ↓ | Throughput ↓ |

packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py

@@ -661,6 +661,9 @@ class OrchestratorConfig(BaseConfig):
     pad_to_multiple_of: int = Field(1, ge=1)
     """Pad each packed micro batch to a multiple of this value (the trainer's cp degree). Auto-filled by the ``rl`` entrypoint; set explicitly for standalone orchestrator runs with cp > 1."""
 
+    inactive_micro_batches: bool = True
+    """Send explicit zero-token slots for distributed iteration alignment. Managed RL copies this setting from the trainer."""
+
     max_steps: int | None = None
     """Maximum training steps. If None, runs indefinitely."""
 

packages/prime-rl-configs/src/prime_rl/configs/rl.py

@@ -686,7 +686,8 @@ def validate_mooncake_offload_requires_slurm(self):
 
     @model_validator(mode="after")
     def auto_setup_deployment(self):
-        self.orchestrator.pad_to_multiple_of = self.trainer.model.cp
+        self.orchestrator.pad_to_multiple_of = 1 if self.trainer.model.cp_unpadded else self.trainer.model.cp
+        self.orchestrator.inactive_micro_batches = self.trainer.model.inactive_micro_batches
         if self.deployment.type == "single_node":  # single-node
             # set num_train_workers to the number of data replicas
             non_data_parallel_size = self.trainer.model.cp

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -315,6 +315,12 @@ class ModelConfig(BaseModelConfig):
     cp: int = 1
     """Context parallelism degree. 1 disables CP."""
 
+    cp_unpadded: bool = True
+    """Distribute original tokens into balanced CP shards, including empty shards for short rows."""
+
+    inactive_micro_batches: bool = True
+    """Align distributed iterations with zero-token slots that retain parameter and expert participation. Requires cp_unpadded when CP is enabled."""
+
     cp_style: Literal["ring", "ulysses"] = "ring"
     """CP communication style. ``ring`` uses ring-attention all-gather/reduce-scatter (requires custom kernels per attention type). ``ulysses`` uses all-to-all to redistribute Q/K/V from sequence-sharded to head-sharded, runs vanilla attention locally on the full sequence, then all-to-all back — works out-of-the-box with any attention kernel (softmax FA, linear attention, mamba, etc.)."""
 
@@ -369,6 +375,8 @@ def vlm_cp_requires_ulysses(self):
 
     @model_validator(mode="after")
     def validate_cp(self):
+        if self.inactive_micro_batches and self.cp > 1 and not self.cp_unpadded:
+            raise ValueError("inactive_micro_batches with CP requires cp_unpadded")
         if self.cp > 1 and self.attn not in ["flash_attention_2", "flash_attention_3", "flash_attention_4", "auto"]:
             raise ValueError("CP is only supported with flash attention 2, 3, or 4")
         if self.cp > 1 and self.impl not in ("custom", "auto"):

src/prime_rl/orchestrator/packing.py

@@ -13,6 +13,7 @@ def __init__(self, config: OrchestratorConfig):
         self.seq_len = config.seq_len
         self.num_train_workers = config.num_train_workers
         self.pad_to_multiple_of = config.pad_to_multiple_of
+        self.inactive_micro_batches = config.inactive_micro_batches
         try:
             model_config = AutoConfig.from_pretrained(
                 config.model.name, trust_remote_code=config.tokenizer.trust_remote_code
@@ -32,4 +33,5 @@ def pack(self, samples: list[TrainingSample]) -> list[list[MicroBatch]]:
             num_train_workers=self.num_train_workers,
             bin_cost=self.bin_cost,
             pad_to_multiple_of=self.pad_to_multiple_of,
+            inactive_micro_batches=self.inactive_micro_batches,
         )

src/prime_rl/trainer/batch.py

@@ -667,11 +667,19 @@ def packed_samples_into_micro_bs(
             bins.append(_MicroBatchBin.from_sample(sample))
 
     if num_train_workers > 1:
-        target_count = max(
-            ((len(bins) + num_train_workers - 1) // num_train_workers) * num_train_workers,
-            num_train_workers,
-        )
-        _expand_bins_by_splitting(bins, target_count, bin_cost)
+        while True:
+            count = len(bins)
+            groups = [[], []]
+            for bin_content in bins:
+                has_images = any(_is_multimodal_sample(sample) for sample in bin_content.samples)
+                groups[int(has_images)].append(bin_content)
+            for group in groups:
+                target_count = ((len(group) + num_train_workers - 1) // num_train_workers) * num_train_workers
+                _expand_bins_by_splitting(group, target_count, bin_cost)
+            bins = groups[1] + groups[0]
+            # Splitting a mixed bin can create a text-only bin in another execution group.
+            if len(bins) == count:
+                break
 
     return [_materialize_bin(bin_content) for bin_content in bins]
 
@@ -809,11 +817,28 @@ def _make_dummy_batch(source: MicroBatch) -> MicroBatch:
     return dummy
 
 
-def _pad_group_for_distribution(group: list[MicroBatch], num_train_workers: int) -> list[MicroBatch]:
+def _pad_group_for_distribution(
+    group: list[MicroBatch], num_train_workers: int, inactive_micro_batches: bool = False
+) -> list[MicroBatch]:
     """Pad a group of micro batches so its length is divisible by num_train_workers."""
     num_padding = -len(group) % num_train_workers
     if num_padding > 0 and len(group) > 0:
-        dummy = _make_dummy_batch(group[0])
+        dummy = (
+            MicroBatch(
+                input_ids=[],
+                loss_mask=[],
+                advantages=[],
+                inference_logprobs=[],
+                position_ids=[],
+                sequence_lengths=[],
+                temperatures=[],
+                env_names=[],
+                seq_lens=[],
+                inactive=True,
+            )
+            if inactive_micro_batches
+            else _make_dummy_batch(group[0])
+        )
         group.extend([dummy] * num_padding)
     return group
 
@@ -824,6 +849,7 @@ def prepare_batch(
     num_train_workers: int,
     bin_cost: Callable[[Sequence[int]], int],
     pad_to_multiple_of: int = 1,
+    inactive_micro_batches: bool = True,
 ) -> list[list[MicroBatch]]:
     """
     Prepare a batch of problems for each GPU. Each batch is a list of micro batches.
@@ -844,8 +870,8 @@ def prepare_batch(
     text_batches = [b for b in micro_batches if not _is_multimodal_sample(b)]
 
     # Pad each group independently so its count is divisible by num_train_workers
-    mm_batches = _pad_group_for_distribution(mm_batches, num_train_workers)
-    text_batches = _pad_group_for_distribution(text_batches, num_train_workers)
+    mm_batches = _pad_group_for_distribution(mm_batches, num_train_workers, inactive_micro_batches)
+    text_batches = _pad_group_for_distribution(text_batches, num_train_workers, inactive_micro_batches)
 
     # Alignment check after distribution padding so the dummy batches are covered too
     for micro_batch in (*mm_batches, *text_batches):

src/prime_rl/trainer/distributed/collectives.py

@@ -3,6 +3,8 @@
 import prime_kernels
 import torch
 import torch.distributed as dist
+import torch.distributed._functional_collectives as funcol
+from torch._dynamo.comptime import comptime
 from torch.distributed import ProcessGroup
 
 
@@ -254,8 +256,34 @@ def all_gather(x: torch.Tensor, dim: int, group: ProcessGroup) -> torch.Tensor:
     return _all_gather(x, dim, group.size(), group.group_name)
 
 
+def all_to_all_cp(x: torch.Tensor, total_tokens: int, group: ProcessGroup, sequence_to_head: bool) -> torch.Tensor:
+    """Exchange flattened token rows using the balanced CP partition."""
+    degree = group.size()
+    # Unit and empty shards need concrete broadcast strides in compiled backward.
+    if total_tokens < 2 * degree:
+        comptime.force_static(total_tokens)
+    lengths = [(total_tokens + degree - 1 - rank) // degree for rank in range(degree)]
+    repeated = [lengths[group.rank()]] * degree
+    outputs, inputs = (lengths, repeated) if sequence_to_head else (repeated, lengths)
+    exchanged = torch.ops._c10d_functional.all_to_all_single(x.contiguous(), outputs, inputs, group.group_name)
+    return funcol.wait_tensor(exchanged)
+
+
+def all_gather_cp(x: torch.Tensor, dim: int, total_tokens: int, group: ProcessGroup) -> torch.Tensor:
+    """Gather balanced token shards; backward sums every peer's contribution."""
+    if total_tokens % group.size() == 0:
+        rows = x.movedim(dim, 0).contiguous()
+        gathered = funcol.wait_tensor(funcol.all_gather_single(rows, 0, group))
+        return gathered.movedim(0, dim).contiguous()
+    rows = x.movedim(dim, 0)
+    repeated = rows.repeat(group.size(), *([1] * (rows.ndim - 1)))
+    return all_to_all_cp(repeated, total_tokens, group, True).movedim(0, dim).contiguous()
+
+
 __all__ = [
     "all_gather",
+    "all_gather_cp",
+    "all_to_all_cp",
     "all_to_all_single",
     "all_to_all_single_equal",
     "mxfp8_all_to_all_combine",

src/prime_rl/trainer/model.py

@@ -271,10 +271,13 @@ def get_load_balance_stats(
         num_routed_tokens = tokens_per_expert.sum() / block_mlp.router.top_k
         tokens_per_expert = tokens_per_expert.sort(dim=0, descending=True).values[block_mlp.router.top_k :]
         balanced_load = tokens_per_expert.mean()
-        max_vio = (tokens_per_expert.max() - balanced_load) / balanced_load
+        max_vio = (tokens_per_expert.max() - balanced_load) / balanced_load.clamp_min(1e-12)
+        max_vio = torch.where(num_routed_tokens > 0, max_vio, torch.nan)
         per_layer_max_vio.append(max_vio.detach())
 
-        routing_confidence = routing_confidence_sum / num_routed_tokens
+        routing_confidence = torch.where(
+            num_routed_tokens > 0, routing_confidence_sum / num_routed_tokens.clamp_min(1), torch.nan
+        )
         per_layer_routing_confidence.append(routing_confidence.detach())
 
         block_mlp.tokens_per_expert.zero_()
@@ -295,12 +298,15 @@ def get_global_moe_stats(
     for name, values in get_load_balance_stats(model, group=ep_group).items():
         if values is None:
             continue
-        value = values.max() if name == "max_vio" else values.mean()
+        # Empty routing groups have no observation to contribute to the global statistic.
+        valid = ~values.isnan()
         if name == "max_vio":
+            value = values.masked_fill(~valid, 0.0).max()
             dist.all_reduce(value, op=dist.ReduceOp.MAX, group=dp_cp_group)
         else:
-            dist.all_reduce(value, op=dist.ReduceOp.SUM, group=dp_cp_group)
-            value /= dist.get_world_size(dp_cp_group)
+            totals = torch.stack((values.masked_fill(~valid, 0.0).sum(), valid.sum()))
+            dist.all_reduce(totals, op=dist.ReduceOp.SUM, group=dp_cp_group)
+            value = totals[0] / totals[1].clamp_min(1)
         stats[name] = value.to("cpu")
     return stats
 
@@ -1095,6 +1101,7 @@ def forward(
     # True when seq_lens holds the full pre-CP-shard document boundaries
     # (kept global because documents can straddle the shard cut).
     seq_lens_are_pre_shard: bool = False,
+    cp_total_tokens: int | None = None,
 ) -> PrimeLmOutput:
     kwargs = {
         "input_ids": input_ids,
@@ -1124,6 +1131,8 @@ def forward(
     if isinstance(model, PreTrainedModelPrimeRL):
         kwargs["seq_lens"] = seq_lens
         kwargs["seq_lens_are_pre_shard"] = seq_lens_are_pre_shard
+        if cp_total_tokens is not None:
+            kwargs["cp_total_tokens"] = cp_total_tokens
 
     if routed_experts is not None:
         kwargs["routed_experts"] = routed_experts