fix(rollout): train safely on incomplete groups

areal/api/cli_args.py

@@ -77,6 +77,11 @@ class NormConfig:
         default=1, metadata={"help": "Group size for group-level normalization"}
     )
 
+    @property
+    def uses_group_statistics(self) -> bool:
+        """Whether normalization derives statistics from prompt groups."""
+        return self.mean_level == "group" or self.std_level == "group"
+
     def __post_init__(self):
         """Validate normalization configuration."""
         valid_levels = {"batch", "group", None}
@@ -1675,6 +1680,17 @@ class PPOActorConfig(TrainEngineConfig):
         default=None, metadata={"help": "Normalization configuration for advantages."}
     )
 
+    # Partial rollout groups
+    min_usable_group_size: int | None = field(
+        default=None,
+        metadata={
+            "help": "Minimum usable rollout slots a prompt group must keep to stay "
+            "trainable when some slots fail or are filtered. None derives the "
+            "minimum from reward_norm/adv_norm: 2 when either uses group "
+            "statistics (1 for a singleton target group), else 1."
+        },
+    )
+
     # KL Control
     kl_ctl: float = field(default=0.1, metadata={"help": "KL divergence coefficient"})
     kl_estimator: str = field(
@@ -1766,6 +1782,30 @@ class PPOActorConfig(TrainEngineConfig):
         metadata={"help": "Maximum number of new tokens to generate"},
     )
 
+    def resolve_min_usable_group_size(self, target_group_size: int) -> int:
+        """Minimum usable rollout slots a group must keep to stay trainable.
+
+        An explicit ``min_usable_group_size`` wins. Otherwise group-relative
+        normalization needs at least two group members before partial groups
+        become a hazard; a singleton target group is complete by definition,
+        so it keeps the minimum of one.
+        """
+        if self.min_usable_group_size is not None:
+            return self.min_usable_group_size
+        for normalization in (self.reward_norm, self.adv_norm):
+            if normalization is None:
+                continue
+            if isinstance(normalization, (dict, DictConfig)):
+                uses_group_statistics = (
+                    normalization.get("mean_level") == "group"
+                    or normalization.get("std_level") == "group"
+                )
+            else:
+                uses_group_statistics = normalization.uses_group_statistics
+            if uses_group_statistics:
+                return min(2, target_group_size)
+        return 1
+
     def should_compute_prox_logp(self) -> bool:
         """Determine if forward pass is needed for proximal log-probabilities.
 

areal/api/engine_api.py

@@ -202,6 +202,7 @@ def rollout_batch(
         group_size: int = 1,
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
+        min_usable_group_size: int = 1,
     ) -> list[dict[str, Any]]:
         """Submit a batch of requests and wait for results.
 
@@ -221,12 +222,16 @@ def rollout_batch(
         group_size : int, optional
             Number of times to run the workflow per input and concatenate results.
             Default is 1 (no grouping).
+        min_usable_group_size : int, optional
+            Estimator-owned minimum number of usable logical rollout slots. Must be
+            between 1 and ``group_size``. Default is 1.
 
         Returns
         -------
         list[dict[str, Any]]
             A list of trajectory dictionaries, one per accepted rollout result.
-            Each trajectory contains tensors with shape [group_size, seqlen, ...].
+            Each trajectory contains tensors whose leading dimension is the number
+            of usable slots, between ``min_usable_group_size`` and ``group_size``.
         """
         raise NotImplementedError()
 
@@ -241,6 +246,7 @@ def prepare_batch(
         dynamic_bs: bool = False,
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
+        min_usable_group_size: int = 1,
     ) -> list[dict[str, Any]]:
         """Prepare a batch of data for training from a dataloader.
 
@@ -261,6 +267,9 @@ def prepare_batch(
             If True, enables dynamic batch sizing. The method will stop collecting
             when (accepted + rejected) >= batch_size, returning only accepted results.
             This results in variable-sized batches of valid data. Default is False.
+        min_usable_group_size : int, optional
+            Estimator-owned minimum number of usable logical rollout slots. Must be
+            between 1 and ``group_size``. Default is 1.
 
         Returns
         -------
@@ -741,6 +750,7 @@ def submit(
         is_eval: bool = False,
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
+        min_usable_group_size: int = 1,
     ) -> int:
         """Submit a request to the inference engine and return immediately.
 
@@ -773,6 +783,9 @@ def submit(
         is_eval : bool, optional
             Whether this is an evaluation workflow. Affects variables like trajectory dump path
             and statistics keys. By default False.
+        min_usable_group_size : int, optional
+            Estimator-owned minimum number of usable logical rollout slots. Must be
+            between 1 and ``group_size``. Default is 1.
 
         Returns
         -------
@@ -846,6 +859,7 @@ def rollout_batch(
         group_size: int = 1,
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
+        min_usable_group_size: int = 1,
     ) -> list[dict[str, Any]]:
         """Submit a batch of requests to the inference engine and wait for the results.
 
@@ -873,6 +887,9 @@ def rollout_batch(
         group_size : int, optional
             Number of times to run the workflow per input and concatenate results.
             Default is 1 (no grouping).
+        min_usable_group_size : int, optional
+            Estimator-owned minimum number of usable logical rollout slots. Must be
+            between 1 and ``group_size``. Default is 1.
 
         Returns
         -------
@@ -893,6 +910,7 @@ def prepare_batch(
         dynamic_bs: bool = False,
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
+        min_usable_group_size: int = 1,
     ) -> list[dict[str, Any]]:
         """Asynchronously submit and wait until a full batch is ready with controlled staleness.
 
@@ -902,9 +920,10 @@ def prepare_batch(
 
             This method caches an internal data generator on the first call.
             The ``dataloader``, ``workflow``, ``workflow_kwargs``, ``group_size``,
-            and ``should_accept_fn`` parameters are captured at the first invocation
-            and reused in all subsequent calls. Passing different arguments in
-            later calls will **not** take effect.
+            ``reward_normalization``, ``drop_incomplete_group``,
+            ``min_usable_group_size``, and ``should_accept_fn`` parameters are captured
+            at the first invocation and reused in all subsequent calls. Passing
+            different arguments in later calls will **not** take effect.
 
             If you need to switch configurations mid-training, consider:
 
@@ -936,6 +955,9 @@ def prepare_batch(
             If True, enables dynamic batch sizing. The method will stop collecting
             when (accepted + rejected) >= batch_size, returning only accepted results.
             This results in variable-sized batches of valid data. Default is False.
+        min_usable_group_size : int, optional
+            Estimator-owned minimum number of usable logical rollout slots. Must be
+            between 1 and ``group_size``. Default is 1.
 
         Returns
         -------

areal/api/workflow_api.py

@@ -22,6 +22,21 @@ async def arun_episode(
         ----
         Returning `None` implies that this trajectory is rejected and will not be used for training.
 
+        When group-relative reward or advantage normalization is enabled
+        (``mean_level='group'`` or ``std_level='group'`` in
+        ``actor.reward_norm`` / ``actor.adv_norm``) and rollouts are grouped
+        (``n_samples >= 2``), each ``arun_episode`` call must contribute
+        exactly one training sample: a tensor dict with batch size 1, or a
+        dict containing a single interaction. Returning multiple samples per
+        episode (e.g. one row per turn, or tree-search branches) raises a
+        non-retryable ``WorkflowContractError`` that terminates training,
+        because group statistics would otherwise treat same-episode rows as
+        independent group members. Ungrouped rollouts (``n_samples=1``)
+        install no group wrapper and are not checked; a multi-sample episode
+        is then normalized as its own group of same-episode rows. To train
+        multi-sample workflows, use batch-level normalization or merge each
+        episode into a single sequence.
+
         See concrete example implementations under the `areal/workflow` directory.
 
         Parameters

areal/engine/core/__init__.py

@@ -4,12 +4,14 @@
 
 from areal.engine.core.train_engine import (
     aggregate_eval_losses,
+    compute_microbatch_loss_weight,
     compute_total_loss_weight,
     reorder_and_pad_outputs,
 )
 
 __all__ = [
     "aggregate_eval_losses",
+    "compute_microbatch_loss_weight",
     "compute_total_loss_weight",
     "reorder_and_pad_outputs",
 ]

areal/engine/core/train_engine.py

@@ -14,19 +14,37 @@
 
 from areal.infra.platforms import current_platform
 from areal.utils.data import (
+    TRANSPORT_DUMMY_KEY,
     MicroBatchList,
     pad_and_stack_tensors_along_first_dim,
     reorder_list,
     unpack_sequence,
 )
 
 __all__ = [
+    "compute_microbatch_loss_weight",
     "compute_total_loss_weight",
     "aggregate_eval_losses",
     "reorder_and_pad_outputs",
 ]
 
 
+def compute_microbatch_loss_weight(
+    microbatch: dict[str, Any],
+    loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor],
+) -> torch.Tensor:
+    """Return zero without invoking an objective on transport-only data."""
+    if microbatch.get(TRANSPORT_DUMMY_KEY) is not True:
+        return loss_weight_fn(microbatch)
+    reference = next(
+        (value for value in microbatch.values() if isinstance(value, torch.Tensor)),
+        None,
+    )
+    if reference is None:
+        raise ValueError("Transport micro-batch does not contain a tensor")
+    return torch.zeros((), dtype=torch.float32, device=reference.device)
+
+
 def compute_total_loss_weight(
     mb_list: MicroBatchList,
     loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor],
@@ -52,7 +70,9 @@ def compute_total_loss_weight(
         The total loss weight (scalar tensor) after all_reduce.
     """
     total_weight = (
-        torch.stack([loss_weight_fn(mb) for mb in mb_list.mbs])
+        torch.stack(
+            [compute_microbatch_loss_weight(mb, loss_weight_fn) for mb in mb_list.mbs]
+        )
         .sum()
         .detach()
         .clone()
@@ -138,7 +158,11 @@ def reorder_and_pad_outputs(
         The processed outputs, padded and stacked along batch dimension.
     """
     res = aggregate_fn(outputs)
+    semantic_batch_size = len(output_seqlens)
+    output_seqlens = [*output_seqlens, *([1] * mb_list.transport_dummy_count)]
     seqlens = [output_seqlens[i] for i in mb_list.forward_indices]
     unpacked = unpack_sequence(res, lens=seqlens, dim=0)
     reordered = reorder_list(unpacked, mb_list.backward_indices)
+    if mb_list.transport_dummy_count:
+        reordered = reordered[:semantic_batch_size]
     return pad_and_stack_tensors_along_first_dim(reordered)

areal/engine/fsdp_engine.py

@@ -60,6 +60,7 @@
 from areal.api.io_struct import DeviceRuntimeInfo
 from areal.engine.core import (
     aggregate_eval_losses,
+    compute_microbatch_loss_weight,
     compute_total_loss_weight,
     reorder_and_pad_outputs,
 )
@@ -582,13 +583,15 @@ def rollout_batch(
         group_size: int = 1,
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
+        min_usable_group_size: int = 1,
     ) -> list[dict[str, Any]]:
         self._check_rollout_engine_connected()
         return self.rollout_coordinator.rollout_batch(
             data,
             workflow=workflow,
             workflow_kwargs=workflow_kwargs,
             group_size=group_size,
+            min_usable_group_size=min_usable_group_size,
             reward_normalization=reward_normalization,
             drop_incomplete_group=drop_incomplete_group,
         )
@@ -603,6 +606,7 @@ def prepare_batch(
         dynamic_bs: bool = False,
         reward_normalization: bool = False,
         drop_incomplete_group: bool = False,
+        min_usable_group_size: int = 1,
     ) -> list[dict[str, Any]]:
         self._check_rollout_engine_connected()
         return self.rollout_coordinator.prepare_batch(
@@ -611,6 +615,7 @@ def prepare_batch(
             workflow_kwargs=workflow_kwargs,
             should_accept_fn=should_accept_fn,
             group_size=group_size,
+            min_usable_group_size=min_usable_group_size,
             dynamic_bs=dynamic_bs,
             reward_normalization=reward_normalization,
             drop_incomplete_group=drop_incomplete_group,
@@ -782,7 +787,9 @@ def train_batch(
         input_batched, _ = self._normalize_batch_input(input_)
 
         # Step 1: Prepare micro-batches
-        mb_list = self._prepare_mb_list(input_batched).to(self.device)
+        mb_list = self._prepare_mb_list(input_batched, allow_transport_padding=True).to(
+            self.device
+        )
 
         # Step 2: Compute total loss weight
         total_loss_weight = compute_total_loss_weight(
@@ -822,7 +829,9 @@ def eval_batch(
         input_batched, _ = self._normalize_batch_input(input_)
 
         # Step 1: Prepare micro-batches
-        mb_list = self._prepare_mb_list(input_batched).to(self.device)
+        mb_list = self._prepare_mb_list(input_batched, allow_transport_padding=True).to(
+            self.device
+        )
 
         # Step 2: Compute total loss weight
         total_loss_weight = compute_total_loss_weight(
@@ -880,7 +889,9 @@ def forward_batch(
         batch_size = len(output_seqlens)
 
         # Step 2: Prepare micro-batches
-        mb_list = self._prepare_mb_list(input_batched).to(self.device)
+        mb_list = self._prepare_mb_list(input_batched, allow_transport_padding=True).to(
+            self.device
+        )
 
         # Step 3: Forward using process_output_fn callback, collecting results
         outputs: list[torch.Tensor] = []
@@ -1878,7 +1889,12 @@ def _load_optimizer_state(self, path: str):
         self.optimizer.load_state_dict(optimizer_state_dict)
         dist.barrier(group=self.cpu_group)
 
-    def _prepare_mb_list(self, input_: dict[str, Any]) -> MicroBatchList:
+    def _prepare_mb_list(
+        self,
+        input_: dict[str, Any],
+        *,
+        allow_transport_padding: bool = False,
+    ) -> MicroBatchList:
         assert "attention_mask" in input_ and "input_ids" in input_
         input_ = input_.copy()
 
@@ -1941,7 +1957,12 @@ def _prepare_mb_list(self, input_: dict[str, Any]) -> MicroBatchList:
         else:
             input_ = amend_position_ids(input_)
 
-        mb_list = split_padded_tensor_dict_into_mb_list(input_, self.config.mb_spec)
+        mb_list = split_padded_tensor_dict_into_mb_list(
+            input_,
+            self.config.mb_spec,
+            group=self.data_parallel_group if allow_transport_padding else None,
+            allow_transport_padding=allow_transport_padding,
+        )
         mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs]
         mb_list = pad_mb_list(
             mb_list,
@@ -2150,7 +2171,7 @@ def _compute_logprobs_and_loss(
         loss_multiplier: float = 1.0,
     ) -> torch.Tensor:
         """Compute logprobs/entropy and return scaled loss."""
-        local_weight = loss_weight_fn(ctx.mb_input)
+        local_weight = compute_microbatch_loss_weight(ctx.mb_input, loss_weight_fn)
         if local_weight == 0:
             return logits.mean() * 0.0