perf(bagel): batch trainside UniGRPO rollout forwards

examples/unified_model/bagel_trainside_unigrpo.yaml

@@ -26,12 +26,9 @@
 # Full fine-tuning: the MoT decoder experts are trained directly (no LoRA). FSDP
 # shards params + grads + optimizer states across all ranks.
 #
-# NOTE (rebased onto main): the rollout is per-sample navit bs=1 (the branch's
-# pack-B / forward_batch_size rollout packing is dropped here — it conflicts with
-# main's KV-in-conditions diffusion design; re-add as a follow-up if rollout
-# throughput matters). With bs=1 the rollout and replay geometry match, so the
-# image side uses old_logp_source=rollout (the on-policy ratio is 1 without a
-# pre-update replay).
+# Rollout packing is opt-in. forward_batch_size=1 preserves the per-sample path;
+# larger values block-diagonally pack text-only thinking chains and same-shape
+# diffusion images while conditions continue to hold per-sample KV contexts.
 #
 # Launch (multi-node; the launcher sets num_devices = NUM_NODES * GPUS_PER_NODE):
 #   PYTHONPATH=$PWD python -m unirl.train_unified_model \
@@ -83,6 +80,9 @@ pipeline:
   trajectory_precision: fp32
   logprob_precision: fp32
   shift: 3.0
+  # Pack at most B thinking chains / images into each navit forward. Keep 1 as
+  # the conservative default; tune per GPU after parity and memory validation.
+  forward_batch_size: ${oc.env:BAGEL_FWD_BS,1}
   strategy:
     _target_: unirl.sde.kernels.FlowSDEStrategy
 
@@ -127,7 +127,7 @@ backend:
   # transformer (the decoder blocks the bundle unfroze).
 
 # Single TRAINSIDE rollout (the M=1 / UniGRPO mode). stage_attrs eval-scopes BOTH
-# trainable stages (the same shared transformer). Per-sample navit bs=1.
+# trainable stages (the same shared transformer).
 rollout:
   _target_: unirl.rollout.engine.trainside.engine.TrainsideRolloutEngine
   stage_attrs: [diffusion, ar]
@@ -167,15 +167,14 @@ algorithm:
     clip_range: 1.0e-6                 # flow trust region (clip range)
     clip_schedule: constant
     # Rollout anchor: μ_old (sde_means) and π_old (sde_logp) are the rollout's own
-    # per-SDE-step recordings. Valid because the rollout is per-sample navit bs=1 — the
-    # same geometry as the bs=1 train replay — so on-policy μ_old == μ_θ → ratio = 1.
-    # (Replay anchor would only be needed if rollout ran a different geometry, e.g. pack-B.)
+    # per-SDE-step recordings. Block-diagonal packing isolates every sequence, so
+    # each row retains the same replay geometry and on-policy ratio semantics.
     old_logp_source: rollout
     mse_weight: 1.5e-5                 # velocity-MSE weight (replaces the latent KL)
     # GRPO-Guard RatioNorm: per-SDE-step normalize the flow ratio (it is otherwise
     # left-shifted, mean < 1, so clipping never engages). Only bites on the off-policy
     # update(s), i.e. needs num_updates_per_batch >= 2. μ_old comes from segment.sde_means
-    # as recorded by the rollout (bs=1, matching the bs=1 replay geometry).
+    # as recorded by the rollout.
     ratio_norm: true
     grad_reweight: false              # GRPO-Guard's optional 2nd part (x 1/dt); off by default
     conditions_cls:

unirl/models/bagel/ar.py

@@ -30,6 +30,7 @@
 from contextlib import nullcontext
 from dataclasses import dataclass
 from dataclasses import field as dc_field
+from functools import partial
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
 import torch
@@ -75,7 +76,12 @@ def __init__(self, *, temperature: float = 1.0, top_p: float = 1.0, top_k: int =
         self.top_p = float(top_p)
         self.top_k = int(top_k)
 
-    def step(self, logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    def step(
+        self,
+        logits: torch.Tensor,
+        *,
+        generators: Optional[List[torch.Generator]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         if logits.dim() != 2:
             raise ValueError(f"BagelARStep.step: expected logits shape [B, vocab], got {tuple(logits.shape)}")
 
@@ -107,7 +113,18 @@ def step(self, logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
             scaled = torch.full_like(scaled, float("-inf")).scatter(-1, sorted_idx, sorted_vals)
 
         probs = F.softmax(scaled, dim=-1)
-        token_id = torch.multinomial(probs, num_samples=1).squeeze(-1)
+        if generators is None:
+            token_id = torch.multinomial(probs, num_samples=1).squeeze(-1)
+        else:
+            if len(generators) != int(probs.shape[0]):
+                raise ValueError(f"BagelARStep.step: got {len(generators)} generators for batch {int(probs.shape[0])}.")
+            token_id = torch.cat(
+                [
+                    torch.multinomial(probs[row], num_samples=1, generator=generator)
+                    for row, generator in enumerate(generators)
+                ],
+                dim=0,
+            )
         log_prob = log_probs_full.gather(-1, token_id.unsqueeze(-1)).squeeze(-1)
         return token_id, log_prob
 
@@ -128,10 +145,17 @@ def __init__(
         autocast_precision: str = "bf16",
         logprob_precision: str = "fp32",
         replay_mode: str = "train",
+        forward_batch_size: Optional[int] = 1,
     ) -> None:
         self.model = model
         self.autocast_dtype = parse_torch_dtype(autocast_precision, field_name="BagelARStage.autocast_precision")
         self.logprob_dtype = parse_torch_dtype(logprob_precision, field_name="BagelARStage.logprob_precision")
+        forward_batch_size = 1 if forward_batch_size is None else int(forward_batch_size)
+        require(
+            forward_batch_size >= 1,
+            f"BagelARStage.forward_batch_size must be >= 1; got {forward_batch_size!r}.",
+        )
+        self.forward_batch_size = forward_batch_size
         # Replay scorer for the GRPO ratio's new_logp:
         #   "train"     — one grad forward_train per sample (nested mask: image full +
         #                 text causal); the und path INCLUDING the image is trained.
@@ -199,6 +223,21 @@ def _resolve_stop_ids(self, params: Optional[BagelARParams], sampling_params: AR
         ids.append(int(self.model.new_token_ids["eos_token_id"]))  # <|im_end|>, as in the vendored gen_text
         return list(dict.fromkeys(ids))
 
+    @staticmethod
+    def _batched_text_ids(prompt_splits: List[List[Dict[str, Any]]]) -> Optional[List[torch.Tensor]]:
+        """Concatenate each sample's text splits, or decline packing for ViT/empty inputs."""
+        out: List[torch.Tensor] = []
+        for splits in prompt_splits:
+            ids: List[torch.Tensor] = []
+            for split in splits:
+                if split.get("kind") != "text":
+                    return None
+                ids.append(split["ids"].reshape(-1).to(dtype=torch.long))
+            if not ids:
+                return None
+            out.append(torch.cat(ids, dim=0))
+        return out
+
     # ------------------------------------------------------------------
     # Rollout
     # ------------------------------------------------------------------
@@ -225,22 +264,57 @@ def autoregress(
         stop_ids = self._resolve_stop_ids(params, sampling_params)
         start_id = int(self.model.new_token_ids["bos_token_id"])
 
+        text_id_lists = self._batched_text_ids(conditions.prompt_splits)
+        use_batched = self.forward_batch_size > 1 and text_id_lists is not None and len(text_id_lists) > 1
+
         generated: List[List[int]] = []
         logps: List[List[float]] = []
         with torch.no_grad(), self._autocast_ctx(device):
-            for splits in conditions.prompt_splits:
-                ctx = self._prefill(splits, device=device)
-                tokens_i, logps_i = rl_ops.decode_text(
-                    bagel,
-                    ctx,
-                    start_token_id=start_id,
-                    sample_fn=step.step,
-                    max_new_tokens=int(sampling_params.max_new_tokens),
-                    stop_ids=stop_ids,
-                    device=device,
-                )
-                generated.append(tokens_i)
-                logps.append(logps_i)
+            if use_batched:
+                for start in range(0, len(text_id_lists), self.forward_batch_size):
+                    id_chunk = text_id_lists[start : start + self.forward_batch_size]
+                    generator_chunk: Optional[List[torch.Generator]] = None
+                    if step.temperature > 0.0 and len(id_chunk) > 1:
+                        seeds = torch.randint(
+                            0,
+                            (1 << 63) - 1,
+                            (len(id_chunk),),
+                            dtype=torch.int64,
+                            device=device,
+                        ).cpu()
+                        generator_chunk = []
+                        for seed in seeds.tolist():
+                            generator = torch.Generator(device=device)
+                            generator.manual_seed(int(seed))
+                            generator_chunk.append(generator)
+                    ctx = rl_ops.prefill_text_batched(bagel, id_chunk, device=device)
+                    tokens, token_logps = rl_ops.decode_text_batched(
+                        bagel,
+                        ctx,
+                        start_token_id=start_id,
+                        sample_fn=(
+                            partial(step.step, generators=generator_chunk) if generator_chunk is not None else step.step
+                        ),
+                        max_new_tokens=int(sampling_params.max_new_tokens),
+                        stop_ids=stop_ids,
+                        device=device,
+                    )
+                    generated.extend(tokens)
+                    logps.extend(token_logps)
+            else:
+                for splits in conditions.prompt_splits:
+                    ctx = self._prefill(splits, device=device)
+                    tokens_i, logps_i = rl_ops.decode_text(
+                        bagel,
+                        ctx,
+                        start_token_id=start_id,
+                        sample_fn=step.step,
+                        max_new_tokens=int(sampling_params.max_new_tokens),
+                        stop_ids=stop_ids,
+                        device=device,
+                    )
+                    generated.append(tokens_i)
+                    logps.append(logps_i)
 
         return TextSegment.pack(
             tokens=[torch.tensor(t, dtype=torch.long, device=device) for t in generated],