feat(ppo): make loss aggregation configurable

areal/api/cli_args.py

@@ -2,6 +2,7 @@
 
 import argparse
 import json
+import math
 import os
 import warnings
 from dataclasses import MISSING as dataclass_missing
@@ -1702,12 +1703,33 @@ class PPOActorConfig(TrainEngineConfig):
         metadata={
             "help": "Use CISPO loss: clip the importance-sampling weight under "
             "stop-gradient and keep gradient on every token's log pi (MiniMax-M1 "
-            "Eq. 4-5). Mutually exclusive with SAPO. Token-level only. Requires "
-            "eps_clip_higher > 0; recommended eps_clip=1.0 (single-sided, lower "
-            "bound 0) with eps_clip_higher=4.0."
+            "Eq. 4-5). Mutually exclusive with SAPO. Uses token-level "
+            "importance-sampling ratios. Requires eps_clip_higher > 0; recommended "
+            "eps_clip=1.0 (single-sided, lower bound 0) with eps_clip_higher=4.0."
         },
     )
 
+    loss_aggregation: str = field(
+        default="token_mean",
+        metadata={
+            "help": "Policy-gradient loss reduction. "
+            "'token_mean': average over valid tokens. "
+            "'seq_mean': average per-response token means. "
+            "'prompt_mean': average per-prompt-group token means. "
+            "'constant': average each response's masked token sum divided by "
+            "loss_aggregation_divisor. Non-token modes require sequence "
+            "boundaries; tree-packed actor training currently supports only "
+            "'token_mean'.",
+            "choices": ["token_mean", "seq_mean", "prompt_mean", "constant"],
+        },
+    )
+    loss_aggregation_divisor: float | None = field(
+        default=None,
+        metadata={
+            "help": "Positive fixed denominator L for loss_aggregation='constant'. "
+            "Unused by other loss aggregation modes.",
+        },
+    )
     # Asynchronous RL
     recompute_logprob: bool = field(
         default=False,
@@ -1849,6 +1871,37 @@ def __post_init__(self):
                     "Please set `actor.use_decoupled_loss=false` in your configuration."
                 )
 
+        if self.loss_aggregation not in (
+            "token_mean",
+            "seq_mean",
+            "prompt_mean",
+            "constant",
+        ):
+            raise ValueError(
+                "loss_aggregation must be 'token_mean', 'seq_mean', "
+                f"'prompt_mean', or 'constant', got {self.loss_aggregation!r}."
+            )
+        if self.loss_aggregation == "constant":
+            if (
+                self.loss_aggregation_divisor is None
+                or not math.isfinite(self.loss_aggregation_divisor)
+                or self.loss_aggregation_divisor <= 0
+            ):
+                raise ValueError(
+                    "loss_aggregation_divisor must be a positive finite value "
+                    "when loss_aggregation='constant'."
+                )
+        elif self.loss_aggregation_divisor is not None:
+            raise ValueError(
+                "loss_aggregation_divisor is only used when "
+                "loss_aggregation='constant'."
+            )
+        if self.enable_tree_training and self.loss_aggregation != "token_mean":
+            raise ValueError(
+                f"loss_aggregation={self.loss_aggregation!r} needs per-sequence "
+                "boundaries, which tree-packed batches do not carry; tree "
+                "training supports only 'token_mean'."
+            )
         # Validate CISPO configuration
         if self.use_cispo_loss:
             if self.use_sapo_loss:

areal/trainer/ppo/actor.py

@@ -26,6 +26,7 @@
     Normalization,
     TrajBatchMeta,
     batched_call,
+    concat_batch,
     split_padded_tensor_dict_into_mb_list,
 )
 from areal.utils.functional import (
@@ -34,6 +35,7 @@
     reward_overlong_penalty,
     sapo_loss_fn,
 )
+from areal.utils.functional.loss_aggregation import PolicyGradientReduction
 from areal.utils.perf_tracer import trace_perf
 from areal.v2.training_service.controller.controller import (
     GatewayTrainController,
@@ -277,7 +279,10 @@ def _compute_advantages(
     @trace_perf("ppo_actor.ppo_update", category="compute")
     @stats_tracker.scope_func_wrapper("ppo_actor")
     def ppo_update(self, data: list[dict[str, Any]]) -> None:
-        batched_call(self._ppo_update, data, unpack=False)
+        batched, meta = concat_batch(data)
+        if self.config.loss_aggregation == "prompt_mean":
+            batched["group_sizes"] = meta.traj_group_sizes
+        self._ppo_update(batched)
 
     def _ppo_update(self, data: dict[str, Any]) -> None:
         attn_mask = data["attention_mask"]
@@ -372,27 +377,33 @@ def _ppo_update(self, data: dict[str, Any]) -> None:
         with stats_tracker.scope("update"):
             # Get current version for proximal approximation metrics
             current_version = self.engine.get_version()
+            pg_reduction = PolicyGradientReduction(
+                mode=self.config.loss_aggregation,
+                divisor=self.config.loss_aggregation_divisor,
+            )
 
             for mb in mb_inputs.mbs:
+                loss_fn = functools.partial(
+                    grpo_loss_fn,
+                    eps_clip=self.config.eps_clip,
+                    eps_clip_higher=self.config.eps_clip_higher,
+                    c_clip=self.config.c_clip,
+                    rejection_sampling=self.config.rejection_sampling,
+                    m2_threshold=self.m2_threshold,
+                    importance_sampling_level=self.config.importance_sampling_level,
+                    current_version=current_version,
+                    prox_logp_method=self.config.prox_logp_method,
+                    use_sapo_loss=self.config.use_sapo_loss,
+                    sapo_tau_pos=self.config.sapo_tau_pos,
+                    sapo_tau_neg=self.config.sapo_tau_neg,
+                    use_cispo_loss=self.config.use_cispo_loss,
+                    use_decoupled_loss=self.config.use_decoupled_loss,
+                    pg_reduction=pg_reduction,
+                )
                 train_stat = self.engine.train_batch(
                     mb,
-                    loss_fn=functools.partial(
-                        grpo_loss_fn,
-                        eps_clip=self.config.eps_clip,
-                        eps_clip_higher=self.config.eps_clip_higher,
-                        c_clip=self.config.c_clip,
-                        rejection_sampling=self.config.rejection_sampling,
-                        m2_threshold=self.m2_threshold,
-                        importance_sampling_level=self.config.importance_sampling_level,
-                        current_version=current_version,
-                        prox_logp_method=self.config.prox_logp_method,
-                        use_sapo_loss=self.config.use_sapo_loss,
-                        sapo_tau_pos=self.config.sapo_tau_pos,
-                        sapo_tau_neg=self.config.sapo_tau_neg,
-                        use_cispo_loss=self.config.use_cispo_loss,
-                        use_decoupled_loss=self.config.use_decoupled_loss,
-                    ),
-                    loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(),
+                    loss_fn=loss_fn,
+                    loss_weight_fn=pg_reduction.normalizer_fn,
                 )
                 stats_tracker.scalar(**train_stat)
 
@@ -454,6 +465,7 @@ def grpo_loss_fn(
     sapo_tau_neg: float = 1.05,
     use_cispo_loss: bool = False,
     use_decoupled_loss: bool = False,
+    pg_reduction: PolicyGradientReduction | None = None,
     vocab_min_logits: torch.Tensor | None = None,
     vocab_max_logits: torch.Tensor | None = None,
     vocab_mean_logits: torch.Tensor | None = None,
@@ -463,7 +475,8 @@ def grpo_loss_fn(
     pipeline micro batches, returns loss and logging stats."""
     old_logp = input_data["logprobs"]
     advantages = input_data["advantages"]
-    loss_mask = input_data["loss_mask"].bool()
+    denominator_mask = input_data["loss_mask"].bool()
+    loss_mask = denominator_mask
     prox_logp_gt = input_data.get("prox_logp")  # Could be None if skipped
 
     entropy = entropy.detach()
@@ -485,6 +498,8 @@ def grpo_loss_fn(
     if m2_threshold is not None:
         loss_mask = _apply_m2po_masking(old_logp, prox_logp, loss_mask, m2_threshold)
 
+    pg_reduction = pg_reduction or PolicyGradientReduction()
+
     # Use CISPO, SAPO, or PPO loss
     if use_cispo_loss:
         if use_sapo_loss:
@@ -507,6 +522,9 @@ def grpo_loss_fn(
             old_logprobs=old_logp,
             rejection_sampling=rejection_sampling,
             cu_seqlens=input_data.get("cu_seqlens"),
+            group_sizes=input_data.get("group_sizes"),
+            pg_reduction=pg_reduction,
+            denominator_mask=denominator_mask,
         )
     elif use_sapo_loss:
         if use_decoupled_loss:
@@ -523,6 +541,9 @@ def grpo_loss_fn(
             loss_mask=loss_mask,
             importance_sampling_level=importance_sampling_level,
             cu_seqlens=input_data.get("cu_seqlens"),
+            group_sizes=input_data.get("group_sizes"),
+            pg_reduction=pg_reduction,
+            denominator_mask=denominator_mask,
         )
     else:
         loss, stat = ppo_actor_loss_fn(
@@ -537,12 +558,20 @@ def grpo_loss_fn(
             rejection_sampling=rejection_sampling,
             importance_sampling_level=importance_sampling_level,
             cu_seqlens=input_data.get("cu_seqlens"),
+            group_sizes=input_data.get("group_sizes"),
+            pg_reduction=pg_reduction,
+            denominator_mask=denominator_mask,
         )
 
     # Joint Distillation KL Loss
     teacher_logp = input_data.get("teacher_logp")
     rkl_stat = None
     if teacher_logp is not None:
+        if pg_reduction.mode != "token_mean":
+            raise ValueError(
+                "teacher_logp distillation is only supported with "
+                "loss_aggregation='token_mean'."
+            )
         # Coefficients for RL and Knowledge Distillation
         rl_loss_weight = input_data.get("rl_loss_weight", 1.0)
         distill_loss_weight = input_data.get("distill_loss_weight", 0.005)

areal/utils/data.py

@@ -697,6 +697,38 @@ def to(self, *args, **kwargs):
 DEFAULT_MAX_TOKENS_PER_MB = int(1e12)
 
 
+def _resolve_microbatch_sequence_groups(
+    bs: int,
+    granularity: int,
+    group_sizes: Sequence[int] | torch.Tensor | None,
+) -> tuple[list[list[int]], list[int] | None]:
+    if group_sizes is None:
+        if bs % granularity != 0:
+            raise RuntimeError(
+                f"Batch size {bs} cannot divide granularity {granularity}."
+            )
+        return [
+            list(range(i * granularity, (i + 1) * granularity))
+            for i in range(bs // granularity)
+        ], None
+
+    if torch.is_tensor(group_sizes):
+        group_sizes = group_sizes.detach().cpu().tolist()
+    sizes = [int(size) for size in group_sizes]
+    if any(size <= 0 for size in sizes):
+        raise ValueError(f"group_sizes must be positive, got {sizes}.")
+    total = sum(sizes)
+    if total != bs:
+        raise ValueError(f"group_sizes sum to {total} but batch size is {bs}.")
+
+    groups = []
+    offset = 0
+    for size in sizes:
+        groups.append(list(range(offset, offset + size)))
+        offset += size
+    return groups, sizes
+
+
 def split_padded_tensor_dict_into_mb_list(
     data: dict[str, Any],
     mb_spec: MicroBatchSpec,
@@ -721,18 +753,12 @@ def split_padded_tensor_dict_into_mb_list(
         )
     granularity = mb_spec.granularity
     bs = data["attention_mask"].shape[0]
-    if bs % granularity != 0:
-        raise RuntimeError(f"Batch size {bs} cannot divide granularity {granularity}.")
+    seq_groups, explicit_group_sizes = _resolve_microbatch_sequence_groups(
+        bs, granularity, data.get("group_sizes")
+    )
     max_seqlen = data["attention_mask"].shape[1]
     seq_lens = data["attention_mask"].sum(1).long().cpu().numpy().tolist()
-    input_lens = (
-        data["attention_mask"]
-        .view(bs // granularity, granularity, -1)
-        .sum(dim=(1, 2))
-        .long()
-        .cpu()
-        .numpy()
-    )
+    input_lens = [sum(seq_lens[i] for i in group) for group in seq_groups]
 
     # check for multimodal input data
     multimodal_keys = {key for key in data if is_multi_modal_key(key)}
@@ -741,6 +767,8 @@ def split_padded_tensor_dict_into_mb_list(
     to_split = {}
     not_to_split = {}
     for key, value in data.items():
+        if key == "group_sizes":
+            continue
         if key in multimodal_keys:
             continue
         if key == "position_ids" or (
@@ -753,10 +781,13 @@ def split_padded_tensor_dict_into_mb_list(
 
     # split
     group_indices = allocate_balanced_mbs_synced(mb_spec, input_lens, group=group)
+    mb_group_sizes = (
+        [[len(seq_groups[i]) for i in group_index] for group_index in group_indices]
+        if explicit_group_sizes is not None
+        else None
+    )
     group_indices = [
-        seqpack.flat2d(
-            [list(range(i * granularity, (i + 1) * granularity)) for i in group_index]
-        )
+        seqpack.flat2d([seq_groups[i] for i in group_index])
         for group_index in group_indices
     ]
     splitted_lens = [
@@ -804,6 +835,8 @@ def _split(tensor):
     # organize splitted micro batches
     assert len(mbs) == len(splitted_lens), (len(mbs), len(splitted_lens))
     for i, (mb, lens) in enumerate(zip(mbs, splitted_lens)):
+        if mb_group_sizes is not None:
+            mb["group_sizes"] = mb_group_sizes[i]
         results.append({**mb, **not_to_split})
 
     return MicroBatchList(