feat(rl): thread seq_lens through transport, packer, and trainer forward

Mirror of the SFT-side contract: packed RL micro-batches carry explicit
per-sample boundaries end to end — prepare_sample stamps them,
_materialize_bin accumulates them across the bin (asserting they sum to
the packed length), pad_micro_batch records padding as its own segment,
and the RL dataloader tensorizes them into the micro-batch for
forward(). Under CP the trainer passes seq_lens=None for now, matching
the SFT side, until global boundary semantics land.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: hubert-marek

src/prime_rl/trainer/batch.py

@@ -214,6 +214,7 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch
         rl_weights=rl_weights,
         ce_weights=ce_weights,
         ref_kl_weights=ref_kl_weights,
+        seq_lens=[len(input_ids)],
     )
 
 
@@ -291,6 +292,7 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch:
     ref_logprobs: list[float] | None = [] if has_ref_logprobs else None
     mm_token_type_ids: list[int] | None = [] if has_mm_token_type_ids else None
     streams: dict[str, list[float] | None] = {name: ([] if has_stream[name] else None) for name in STREAM_FILL}
+    seq_lens: list[int] = []
     routed_experts: RoutedExperts | None = None
     lora_num_tokens = [0] * num_loras
 
@@ -322,10 +324,12 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch:
                 assert routed_experts.shape[1:] == sample.routed_experts.shape[1:]
                 routed_experts.data += sample.routed_experts.data
                 routed_experts.shape[0] += sample.routed_experts.shape[0]
+        seq_lens.extend(sample.seq_lens if sample.seq_lens is not None else [sample_len])
         lora_num_tokens[lora_idx] += sample_len
 
     sequence_lengths = [len(sample.input_ids) for _, sample in bin_content.samples]
     assert sum(sequence_lengths) == len(input_ids), (sequence_lengths, len(input_ids))
+    assert sum(seq_lens) == len(input_ids), (seq_lens, len(input_ids))
     first_sample = bin_content.first_sample
 
     return MicroBatch(
@@ -345,6 +349,7 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch:
         rl_weights=streams["rl_weights"],
         ce_weights=streams["ce_weights"],
         ref_kl_weights=streams["ref_kl_weights"],
+        seq_lens=seq_lens,
     )
 
 
@@ -464,6 +469,8 @@ def pad_micro_batch(micro_batch: MicroBatch, pad_to_multiple_of: int) -> MicroBa
         )
     if micro_batch.mm_token_type_ids is not None:
         micro_batch.mm_token_type_ids.extend([0] * padding_size)
+    if micro_batch.seq_lens is not None:
+        micro_batch.seq_lens.append(padding_size)
     if micro_batch.routed_experts is not None:
         _pad_routed_experts(micro_batch, padding_size)
     micro_batch.env_names.extend([""] * padding_size)

src/prime_rl/trainer/rl/data.py

@@ -34,6 +34,7 @@ class TensorMicroBatch(TypedDict):
 
     # Batch level
     lora_num_tokens: Int[Tensor, "n_loras"]
+    seq_lens: Int[Tensor, "segments"] | None
 
     # MoE router replay
     routed_experts: Int[Tensor, "batch seq layers topk"] | None
@@ -129,6 +130,7 @@ def _get_sample_micro_batch(self, generator: torch.Generator) -> TensorMicroBatc
             "sequence_lengths": sequence_lengths,
             "loss_mask": loss_mask.unsqueeze(0),
             "lora_num_tokens": lora_num_tokens,
+            "seq_lens": torch.tensor(sequence_lengths, dtype=torch.long),
             "routed_experts": None,
             "mm_kwargs": None,
             "mm_token_type_ids": None,
@@ -161,6 +163,7 @@ def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch:
             "sequence_lengths": [self.seq_len],
             "loss_mask": torch.ones(self.seq_len, dtype=torch.bool).unsqueeze(0),
             "lora_num_tokens": lora_num_tokens,
+            "seq_lens": torch.tensor([self.seq_len], dtype=torch.long),
             "routed_experts": None,
             "mm_kwargs": None,
             "mm_token_type_ids": None,
@@ -256,6 +259,7 @@ def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch:
             env_names=micro_batch.env_names,
             sequence_lengths=micro_batch.sequence_lengths,
             lora_num_tokens=torch.tensor(micro_batch.lora_num_tokens, dtype=torch.int32),
+            seq_lens=torch.tensor(micro_batch.seq_lens, dtype=torch.long) if micro_batch.seq_lens is not None else None,
             mm_kwargs=mm_kwargs,
             mm_token_type_ids=torch.tensor(micro_batch.mm_token_type_ids, dtype=torch.long).unsqueeze(0)
             if micro_batch.mm_token_type_ids is not None

src/prime_rl/trainer/rl/train.py

@@ -421,13 +421,18 @@ def load_run_checkpoint(_optimizer, idx: int) -> None:
                 else None
             )
 
+            seq_lens = micro_batch["seq_lens"].to("cuda") if micro_batch.get("seq_lens") is not None else None
+
             labels = shift_tensor_left(input_ids)
 
             # VLM + CP is not supported: MRoPE requires global positions but CP shards the sequence
             if cp_enabled and mm_kwargs is not None:
                 raise NotImplementedError("Context parallelism is not supported with VLM/multimodal training")
 
             if cp_enabled:
+                # seq_lens spans the pre-shard sequence; consuming it under CP needs
+                # global-boundary handling that isn't wired up here yet.
+                seq_lens = None
                 input_ids, forward_position_ids = setup_cp_params(
                     input_ids, position_ids, cp_rank, cp_size, cp_group, cp_style=config.model.cp_style
                 )
@@ -465,6 +470,7 @@ def load_run_checkpoint(_optimizer, idx: int) -> None:
                     temperature=temperatures,
                     mm_kwargs=mm_kwargs,
                     mm_token_type_ids=mm_token_type_ids,
+                    seq_lens=seq_lens,
                     routed_experts=routed_experts,
                 )
 

src/prime_rl/transport/types.py

@@ -108,3 +108,6 @@ class MicroBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True):
     # Packer-derived metadata used for run-local token exports.
     run_id: str | None = None
     run_step: int | None = None
+
+    # Packed sample boundaries. Sum equals len(input_ids) when present.
+    seq_lens: list[int] | None = None

tests/unit/orchestrator/test_batch.py

@@ -161,6 +161,7 @@ def test_randomized_packing_invariants():
         for batch in flat_batches:
             assert len(batch.input_ids) <= seq_len
             assert sum(batch.sequence_lengths) == len(batch.input_ids)
+            assert batch.seq_lens == batch.sequence_lengths
             assert sum(batch.lora_num_tokens) == len(batch.input_ids)
             assert len(batch.env_names) == len(batch.input_ids)
 
@@ -176,6 +177,7 @@ def test_pad_micro_batch_preserves_explicit_sequence_lengths():
 
     assert len(padded.input_ids) == 6
     assert padded.sequence_lengths == [4, 2]
+    assert padded.seq_lens == [4, 2]
     assert sum(padded.sequence_lengths) == len(padded.input_ids)
     assert padded.loss_mask[-2:] == [False, False]
 
@@ -271,6 +273,7 @@ def test_prepare_batch_packs_different_temperatures(make_training_example):
     # Second sample (4 tokens): all get temp 1.1
     assert flat_batches[0].temperatures[4:8] == [1.1, 1.1, 1.1, 1.1]
     assert flat_batches[0].env_names == ["env-a"] * 4 + ["env-b"] * 4
+    assert flat_batches[0].seq_lens == [4, 4]
 
 
 def test_prepare_sample_propagates_weight_streams(make_training_example):