draft(engine): preserve explicit loss reduction prototype

areal/api/__init__.py

@@ -4,6 +4,10 @@
     "RolloutWorkflow",
     "AsyncRewardWrapper",
     "TrainEngine",
+    "LossReduction",
+    "LossTerm",
+    "LOSS_TERM_REDUCTION_MEAN",
+    "LOSS_TERM_REDUCTION_SUM",
     "InferenceEngine",
     "Scheduler",
     "Worker",
@@ -28,6 +32,10 @@
 
 _LAZY_IMPORTS = {
     "TrainEngine": "areal.api.engine_api",
+    "LossReduction": "areal.api.loss_api",
+    "LossTerm": "areal.api.loss_api",
+    "LOSS_TERM_REDUCTION_MEAN": "areal.api.loss_api",
+    "LOSS_TERM_REDUCTION_SUM": "areal.api.loss_api",
     "InferenceEngine": "areal.api.engine_api",
     "Scheduler": "areal.api.scheduler_api",
     "Worker": "areal.api.scheduler_api",

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
@@ -1610,12 +1611,47 @@ 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 "
+            "Eq. 4-5). Mutually exclusive with SAPO. 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' averages valid "
+            "tokens, 'seq_mean' averages sequence means, and 'prompt_mean' averages "
+            "prompt-group means over gconfig.n_samples responses. 'constant' "
+            "averages each response's masked token sum divided by "
+            "loss_aggregation_divisor.",
+            "help_zh": "Policy-gradient loss 的归约方式。'token_mean' 对有效 token "
+            "求平均，'seq_mean' 对每条序列的均值求平均，'prompt_mean' 对 "
+            "gconfig.n_samples 个 response 组成的 prompt group 均值求平均。"
+            "'constant' 将每条 response 的 masked token loss 之和除以 "
+            "loss_aggregation_divisor 后再求平均。",
+            "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.",
+            "help_zh": "loss_aggregation='constant' 使用的正数固定分母 L。其他 "
+            "loss aggregation 模式不使用。",
+        },
+    )
+    group_size: int = field(
+        default=1,
+        metadata={
+            "help": "Internal prompt-group size for prompt_mean; derived from "
+            "gconfig.n_samples by PPOConfig.__post_init__.",
+            "help_zh": "prompt_mean 使用的内部 prompt group 大小；由 "
+            "PPOConfig.__post_init__ 根据 gconfig.n_samples 推导。",
+        },
+    )
+
     # Asynchronous RL
     recompute_logprob: bool = field(
         default=False,
@@ -1708,8 +1744,22 @@ def __post_init__(self):
 
         from areal.utils.constants import ProxLogpMethod
 
+        prox_logp_method = ProxLogpMethod(self.prox_logp_method)
+        if self.m2_threshold is not None:
+            if not math.isfinite(self.m2_threshold) or self.m2_threshold <= 0:
+                raise ValueError(
+                    "m2_threshold must be a positive finite value, "
+                    f"got {self.m2_threshold!r}."
+                )
+            if prox_logp_method.skips_forward_pass():
+                raise ValueError(
+                    "m2_threshold requires proximal log probabilities before the "
+                    "training forward pass and is incompatible with "
+                    f"prox_logp_method={self.prox_logp_method!r}."
+                )
+
         if (
-            ProxLogpMethod(self.prox_logp_method) == ProxLogpMethod.REUSE_TRAIN_LOGP
+            prox_logp_method == ProxLogpMethod.REUSE_TRAIN_LOGP
             and self.ppo_n_minibatches > 1
         ):
             logger.warning(
@@ -1757,6 +1807,36 @@ 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.loss_aggregation != "prompt_mean" and self.group_size != 1:
+            raise ValueError(
+                "group_size is only valid for loss_aggregation='prompt_mean'."
+            )
+
         # Validate CISPO configuration
         if self.use_cispo_loss:
             if self.use_sapo_loss:
@@ -3173,6 +3253,16 @@ def __post_init__(self):
         # the engine config. Single source of truth: gconfig.lora_name.
         if self.rollout.use_lora and not self.rollout.lora_name:
             self.rollout.lora_name = self.gconfig.lora_name
+
+        if self.actor.loss_aggregation == "prompt_mean":
+            self.actor.group_size = self.gconfig.n_samples
+            if self.actor.group_size < 2:
+                raise ValueError(
+                    "loss_aggregation='prompt_mean' requires gconfig.n_samples "
+                    ">= 2 (a single sample per prompt collapses to seq_mean), "
+                    f"got n_samples={self.gconfig.n_samples}."
+                )
+
         super().__post_init__()
 
 

areal/api/engine_api.py

@@ -22,6 +22,12 @@
     SaveLoadMeta,
     WeightUpdateMeta,
 )
+from areal.api.loss_api import (
+    LOSS_TERM_REDUCTION_MEAN,
+    LossReduction,
+    LossTerm,
+    LossWeightFn,
+)
 
 if TYPE_CHECKING:
     from areal.api.workflow_api import WorkflowLike
@@ -364,7 +370,7 @@ def train_batch(
         self,
         input_: list[dict[str, Any]] | dict[str, Any],
         loss_fn: Callable[..., torch.Tensor],
-        loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor],
+        loss_weight_fn: LossWeightFn,
     ) -> dict[str, float]:
         """Update the model with a batch of data and a loss function.
 
@@ -386,8 +392,7 @@ def train_batch(
         loss_weight_fn : Callable[[dict[str, Any]], torch.Tensor]
             A function used to calculate the weight of each micro-batch. Since
             loss_fn normalizes the loss for a micro-batch, we need a corresponding
-            weight for each micro-batch to normalize the loss globally. The weight
-            is usually the number of response tokens in the batch.
+            weight for each micro-batch to normalize the loss globally.
 
         Returns
         -------
@@ -397,13 +402,27 @@ def train_batch(
         """
         raise NotImplementedError()
 
+    def train_batch_with_reduction(
+        self,
+        input_: list[dict[str, Any]] | dict[str, Any],
+        loss_reduction: LossReduction,
+    ) -> dict[str, float]:
+        """Update the model using an explicit distributed reduction contract.
+
+        Engines written against the original callback API automatically support
+        a single mean term. Engines that support sum or multi-term reductions
+        should override this method.
+        """
+        term = self._original_loss_term(loss_reduction)
+        return self.train_batch(input_, loss_reduction.loss_fn, term.normalizer_fn)
+
     @torch.no_grad()
     @abc.abstractmethod
     def eval_batch(
         self,
         input_: list[dict[str, Any]] | dict[str, Any],
         loss_fn: Callable[..., torch.Tensor],
-        loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor],
+        loss_weight_fn: LossWeightFn,
     ) -> torch.Tensor | None:
         """Evaluate the model using the forward pass and loss function.
 
@@ -425,8 +444,7 @@ def eval_batch(
         loss_weight_fn : Callable[[dict[str, Any]], torch.Tensor]
             A function used to calculate the weight of each micro-batch. Since
             loss_fn normalizes the loss for a micro-batch, we need a corresponding
-            weight for each micro-batch to normalize the loss globally. The weight
-            is usually the number of response tokens in the batch.
+            weight for each micro-batch to normalize the loss globally.
 
         Returns
         -------
@@ -436,6 +454,30 @@ def eval_batch(
         """
         raise NotImplementedError()
 
+    @torch.no_grad()
+    def eval_batch_with_reduction(
+        self,
+        input_: list[dict[str, Any]] | dict[str, Any],
+        loss_reduction: LossReduction,
+    ) -> torch.Tensor | None:
+        """Evaluate the model using an explicit distributed reduction contract."""
+        term = self._original_loss_term(loss_reduction)
+        return self.eval_batch(input_, loss_reduction.loss_fn, term.normalizer_fn)
+
+    @staticmethod
+    def _original_loss_term(loss_reduction: LossReduction) -> LossTerm:
+        if (
+            len(loss_reduction.terms) != 1
+            or loss_reduction.terms[0].reduction != LOSS_TERM_REDUCTION_MEAN
+        ):
+            raise NotImplementedError(
+                "This TrainEngine implements the original callback API and only "
+                "supports a single mean loss term. Override "
+                "train_batch_with_reduction/eval_batch_with_reduction to support "
+                "sum or multi-term reductions."
+            )
+        return loss_reduction.terms[0]
+
     @torch.no_grad()
     @abc.abstractmethod
     def forward_batch(

areal/api/loss_api.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Contracts for loss values and their distributed normalizers."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+LOSS_TERM_REDUCTION_MEAN = "mean"
+LOSS_TERM_REDUCTION_SUM = "sum"
+LOSS_TERM_REDUCTIONS_ALL = (
+    LOSS_TERM_REDUCTION_MEAN,
+    LOSS_TERM_REDUCTION_SUM,
+)
+
+LossFnOutput = torch.Tensor | Mapping[str, torch.Tensor]
+
+
+@dataclass(frozen=True, slots=True)
+class LossTerm:
+    """One named term in a distributed loss reduction.
+
+    ``normalizer_fn`` returns this rank's scalar contribution to the global
+    normalizer for this term. ``reduction="mean"`` means the loss value is
+    already divided by that local normalizer. ``reduction="sum"`` means the
+    loss value is the local numerator term.
+    """
+
+    name: str
+    normalizer_fn: Callable[[dict[str, Any]], torch.Tensor]
+    reduction: str
+
+    def __post_init__(self) -> None:
+        if not self.name:
+            raise ValueError("LossTerm.name must be non-empty.")
+        if self.reduction not in LOSS_TERM_REDUCTIONS_ALL:
+            raise ValueError(
+                f"reduction must be one of {LOSS_TERM_REDUCTIONS_ALL}, "
+                f"got {self.reduction!r}."
+            )
+
+
+@dataclass(frozen=True, slots=True)
+class LossReduction:
+    """Loss function plus the distributed reduction contract for its outputs.
+
+    For a ``mean`` term, the engine computes
+    ``local_mean * local_normalizer / global_normalizer``. For a ``sum`` term,
+    the engine computes ``local_sum / global_normalizer``. If ``loss_fn``
+    returns a mapping, each term consumes the value under its own name.
+    """
+
+    loss_fn: Callable[..., LossFnOutput]
+    terms: tuple[LossTerm, ...]
+
+    def __post_init__(self) -> None:
+        if not self.terms:
+            raise ValueError("LossReduction requires at least one term.")
+        names = [term.name for term in self.terms]
+        if len(names) != len(set(names)):
+            raise ValueError(f"LossReduction term names must be unique, got {names}.")
+
+    @classmethod
+    def mean(
+        cls,
+        loss_fn: Callable[..., torch.Tensor],
+        normalizer_fn: Callable[[dict[str, Any]], torch.Tensor],
+        name: str = "loss",
+    ) -> LossReduction:
+        """Build a reduction for a loss value normalized within each microbatch."""
+        return cls(
+            loss_fn=loss_fn,
+            terms=(
+                LossTerm(
+                    name=name,
+                    normalizer_fn=normalizer_fn,
+                    reduction=LOSS_TERM_REDUCTION_MEAN,
+                ),
+            ),
+        )
+
+    @classmethod
+    def sum(
+        cls,
+        loss_fn: Callable[..., LossFnOutput],
+        normalizer_fn: Callable[[dict[str, Any]], torch.Tensor],
+        name: str = "loss",
+    ) -> LossReduction:
+        """Build a reduction for a local numerator term."""
+        return cls(
+            loss_fn=loss_fn,
+            terms=(
+                LossTerm(
+                    name=name,
+                    normalizer_fn=normalizer_fn,
+                    reduction=LOSS_TERM_REDUCTION_SUM,
+                ),
+            ),
+        )
+
+
+LossWeightFn = Callable[[dict[str, Any]], torch.Tensor]

areal/engine/core/__init__.py

@@ -4,12 +4,18 @@
 
 from areal.engine.core.train_engine import (
     aggregate_eval_losses,
+    compute_global_normalizers,
+    compute_local_normalizers,
     compute_total_loss_weight,
     reorder_and_pad_outputs,
+    scale_loss_for_reduction,
 )
 
 __all__ = [
     "aggregate_eval_losses",
     "compute_total_loss_weight",
+    "compute_global_normalizers",
+    "compute_local_normalizers",
     "reorder_and_pad_outputs",
+    "scale_loss_for_reduction",
 ]