fix(rollout): train safely on incomplete groups

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/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] = []
@@ -1872,7 +1883,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()
 
@@ -1935,7 +1951,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,
@@ -2142,7 +2163,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