fix(engine): make ragged transport padding objective-safe

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,
 )
@@ -782,7 +783,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 +825,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 +885,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 +1879,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 +1947,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 +2159,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
 

areal/engine/megatron_engine.py

@@ -51,6 +51,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,
 )
@@ -1008,7 +1009,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.
         # Use DP+CP group: after CP all-gather each rank computes the full-sequence
@@ -1069,7 +1072,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 (DP+CP, see train_batch comment).
         total_loss_weight = compute_total_loss_weight(
@@ -1128,7 +1133,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 Megatron's pipeline function, collecting results
         outputs: list[torch.Tensor] = []
@@ -2250,7 +2257,12 @@ def _load_model_from_hf(self, path: str) -> None:
                 fp8_direct_convert=self.fp8_direct_convert,
             )
 
-    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_
         # Parallel sizes
         pp_size = self.parallel_strategy.pipeline_parallel_size
@@ -2301,6 +2313,7 @@ def _prepare_mb_list(self, input_: dict[str, Any]) -> MicroBatchList:
             input_,
             mb_spec,
             group=mpu.get_data_parallel_group(),
+            allow_transport_padding=allow_transport_padding,
         )
         mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs]
         # NOTE: Pad micro-batches to:
@@ -2362,7 +2375,7 @@ def _compute_logprobs_and_loss(
         total_loss_weight: torch.Tensor,
         loss_multiplier: float = 1.0,
     ) -> torch.Tensor:
-        local_weight = loss_weight_fn(inputs)
+        local_weight = compute_microbatch_loss_weight(inputs, loss_weight_fn)
         if local_weight == 0:
             return output.mean() * 0.0
 

areal/experimental/engine/archon_engine.py

@@ -41,6 +41,7 @@
 )
 from areal.engine.core.train_engine import (
     aggregate_eval_losses,
+    compute_microbatch_loss_weight,
     compute_total_loss_weight,
     reorder_and_pad_outputs,
 )
@@ -534,7 +535,9 @@ def train_batch(
 
         input_batched, _ = self._normalize_batch_input(input_)
 
-        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
+        )
 
         total_loss_weight = compute_total_loss_weight(
             mb_list, loss_weight_fn, self.data_parallel_group
@@ -571,7 +574,9 @@ def eval_batch(
 
         input_batched, _ = self._normalize_batch_input(input_)
 
-        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
+        )
 
         total_loss_weight = compute_total_loss_weight(
             mb_list, loss_weight_fn, self.data_parallel_group
@@ -629,7 +634,9 @@ def forward_batch(
         assert output_seqlens is not None
         batch_size = len(output_seqlens)
 
-        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
+        )
 
         def process_output(
             logits: torch.Tensor, ctx_dict: dict[str, Any]
@@ -1183,7 +1190,12 @@ def _normalize_batch_input(
             return concat_batch(input_)
         return input_, None
 
-    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()
 
@@ -1211,7 +1223,7 @@ def _prepare_mb_list(self, input_: dict[str, Any]) -> MicroBatchList:
             stages_per_rank = len(self.pp_stages)
             num_total_stages = pp_size * stages_per_rank
             n_seqs = input_["attention_mask"].shape[0]
-            if n_seqs < num_total_stages:
+            if n_seqs < num_total_stages and not allow_transport_padding:
                 raise RuntimeError(
                     f"Pipeline parallelism requires at least {num_total_stages} "
                     f"sequences (pp_size={pp_size} * stages_per_rank="
@@ -1227,7 +1239,12 @@ def _prepare_mb_list(self, input_: dict[str, Any]) -> MicroBatchList:
         else:
             mb_spec = self.config.mb_spec
 
-        mb_list = split_padded_tensor_dict_into_mb_list(input_, mb_spec)
+        mb_list = split_padded_tensor_dict_into_mb_list(
+            input_,
+            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]
 
         # LCM ensures page-aligned memory and exact CP slicing without extra padding.
@@ -1275,7 +1292,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
 

areal/models/tree_attn/tree.py

@@ -26,7 +26,7 @@
     precompute_tree_attention_data,
 )
 from areal.utils import logging, stats_tracker
-from areal.utils.data import MicroBatchList
+from areal.utils.data import TRANSPORT_DUMMY_KEY, MicroBatchList
 from areal.utils.perf_tracer import trace_perf, trace_scope
 
 logger = logging.getLogger("TreeAttentionCore")
@@ -404,6 +404,7 @@ def build_packed_tree_batch(
 
     # Build packed outputs for each tree
     mbs: list[dict[str, Any]] = []
+    padded_mbs: list[dict[str, Any]] = []
     padding_lengths: list[int] = []
     padded_to_lengths: list[int] = []
 
@@ -448,13 +449,17 @@ def build_packed_tree_batch(
                 non_packable_keys,
             )
 
-        mb = {
+        padded_mb = {
             "input_ids": input_ids,
             "position_ids": position_ids,
             "trie_node": trie,
             **extra_data,
         }
+        mb = dict(padded_mb)
+        if not trie.all_sequence_ids:
+            mb[TRANSPORT_DUMMY_KEY] = True
         mbs.append(mb)
+        padded_mbs.append(padded_mb)
         padding_lengths.append(padded_size - num_tokens)
         padded_to_lengths.append(padded_size)
 
@@ -465,7 +470,7 @@ def build_packed_tree_batch(
         mb_spec=mb_spec,
         mbs=mbs,
         group_lens=[num for num in num_tokens_list],
-        padded_mbs=mbs,
+        padded_mbs=padded_mbs,
         padding_lengths=padding_lengths,
         padded_to_lengths=padded_to_lengths,
         _max_seqlen=max(padded_to_lengths),