feat(rl): orchestrator.pad_to_seq_len, and pad with real tokens to avoid MoE hot spots

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

@@ -539,6 +539,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."""
 
+    pad_to_seq_len: bool = False
+    """Pad every packed micro batch to ``seq_len`` instead of to a multiple of ``pad_to_multiple_of``. Rows then have one static shape, which avoids recompilation on dynamic sequence lengths (torch.compile, DeepGEMM JIT kernels, ...) at the cost of extra padding."""
+
     max_steps: int | None = None
     """Maximum training steps. If None, runs indefinitely."""
 

src/prime_rl/orchestrator/batch.py

@@ -19,6 +19,7 @@
     TopLogprobs,
     TrainingSample,
 )
+from prime_rl.utils.sequence import repeat_to_length
 
 # Backfill value per component weight stream when a packed sample doesn't
 # carry it: absent rl means weight 1.0 on the loss mask, absent ce/ref_kl
@@ -156,7 +157,7 @@ def _pad_routed_experts(micro_batch: MicroBatch, padding_size: int) -> None:
     routed_experts = micro_batch.routed_experts
     assert routed_experts is not None
     row_size = _routed_experts_row_size(routed_experts)
-    routed_experts.data += b"\0" * (padding_size * row_size)
+    routed_experts.data += repeat_to_length(routed_experts.data, padding_size * row_size)
     routed_experts.shape[0] += padding_size
 
 
@@ -628,7 +629,12 @@ def pad_micro_batch(micro_batch: MicroBatch, pad_to_multiple_of: int) -> MicroBa
     if not (pad_to_multiple_of > 1 and padding_size > 0):
         return micro_batch
 
-    micro_batch.input_ids.extend([1] * padding_size)
+    # Identical pad tokens route to the same MoE experts and can overload one EP rank, so text rows pad with their
+    # own tokens. Multimodal rows keep a constant id: a repeated image placeholder would have no feature to receive.
+    if micro_batch.mm_refs is None:
+        micro_batch.input_ids.extend(repeat_to_length(micro_batch.input_ids, padding_size))
+    else:
+        micro_batch.input_ids.extend([1] * padding_size)
     micro_batch.advantages.extend([0.0] * padding_size)
     micro_batch.loss_mask.extend([False] * padding_size)
     micro_batch.position_ids.extend(list(range(padding_size)))
@@ -787,7 +793,8 @@ class BatchPacker:
     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
+        # Rows hold at most `seq_len` tokens, so padding to a multiple of it pads every row to exactly `seq_len`
+        self.pad_to_multiple_of = config.seq_len if config.pad_to_seq_len else config.pad_to_multiple_of
         # Built as the trainer builds it. The flash-attention version selects kernels, not modules.
         model_config = ModelConfig(
             name=config.model.name, attn="flash_attention_3", trust_remote_code=config.tokenizer.trust_remote_code

src/prime_rl/trainer/sft/data.py

@@ -22,6 +22,7 @@
 from prime_rl.trainer.world import get_world
 from prime_rl.utils.chat_template import deserialize_tool_calls, normalize_messages
 from prime_rl.utils.logger import get_logger
+from prime_rl.utils.sequence import repeat_to_length
 from prime_rl.utils.utils import format_time
 
 
@@ -592,7 +593,13 @@ def _finalize_pack(self, packed: dict[str, Any], seq_len: int) -> dict:
             remaining -= kept
         pad_len = -len(result["input_ids"]) % self.pad_to_multiple_of
         if pad_len > 0:
-            result["input_ids"].extend([0] * pad_len)
+            # Identical pad tokens route to the same MoE experts and can overload one EP rank, so text packs pad with
+            # their own tokens. Multimodal packs keep a constant id: a repeated image placeholder would have no feature
+            # to receive.
+            if packed["mm_kwargs"] is None:
+                result["input_ids"].extend(repeat_to_length(result["input_ids"], pad_len))
+            else:
+                result["input_ids"].extend([0] * pad_len)
             result["position_ids"].extend(range(pad_len))
             result["loss_mask"].extend([False] * pad_len)
             result["target_ids"].extend([0] * pad_len)

src/prime_rl/utils/sequence.py

@@ -50,3 +50,9 @@ def get_cu_seqlens_from_seq_lens(seq_lens: torch.Tensor, total_tokens: int | Non
     cu_seqlens[0] = 0
     cu_seqlens[1:] = seq_lens.cumsum(dim=0, dtype=torch.int32)
     return cu_seqlens, int(seq_lens.max().item())
+
+
+def repeat_to_length(values: list | bytes, length: int) -> list | bytes:
+    """Repeat `values` end to end and truncate to `length` elements."""
+    num_repeats = length // len(values) + 1
+    return (values * num_repeats)[:length]