From 43294f4a7e3d6e09fe510c0bcf4c33eaca4bc53a Mon Sep 17 00:00:00 2001 From: EazyReal <8047065+EazyReal@users.noreply.github.com> Date: Tue, 7 Jul 2026 11:16:03 -0700 Subject: [PATCH 1/4] feat(ppo): support actor loss aggregation modes Represent actor policy-gradient loss aggregation as an explicit distributed loss-reduction contract so token, sequence, prompt, and constant modes share the same engine boundary. Key changes: - Add LossReduction/LossTerm and migrate train/eval engine call sites - Add token_mean, seq_mean, prompt_mean, and constant actor reductions - Preserve variable prompt groups without rollout filtering changes - Reject unsupported teacher distillation with non-token aggregation - Fail fast on non-positive global loss normalizers --- areal/api/__init__.py | 8 + areal/api/cli_args.py | 74 +- areal/api/engine_api.py | 127 ++- areal/engine/core/__init__.py | 8 +- areal/engine/core/train_engine.py | 180 ++++- areal/engine/fsdp_engine.py | 65 +- areal/engine/megatron_engine.py | 64 +- areal/experimental/engine/archon_engine.py | 53 +- areal/infra/rpc/serialization.py | 6 + areal/trainer/dpo/dpo_engine.py | 20 +- areal/trainer/ppo/actor.py | 171 +++- areal/trainer/ppo/critic.py | 12 +- areal/trainer/rw/rw_engine.py | 20 +- areal/trainer/sft/lm_engine.py | 14 +- areal/utils/data.py | 61 +- areal/utils/functional/__init__.py | 16 + areal/utils/functional/functional.py | 731 +++++++++++++++++- .../training_service/controller/controller.py | 16 +- docs/en/best_practices/perf_profiling.md | 4 +- docs/en/cli_reference.md | 5 +- docs/generate_cli_docs.py | 10 +- docs/zh/best_practices/perf_profiling.md | 4 +- docs/zh/cli_reference.md | 5 +- .../archon/torchrun/run_archon_engine_pp.py | 41 +- tests/fp8/model_hooks.py | 37 +- tests/test_cispo_loss.py | 99 ++- tests/test_dpo.py | 8 +- tests/test_eval_dispatch.py | 51 +- tests/test_loss_reduction.py | 118 +++ tests/test_megatron_engine.py | 29 +- tests/test_ppo_stats.py | 31 + tests/test_prompt_mean_loss.py | 609 +++++++++++++++ tests/test_rpc_callable_serialization.py | 21 + tests/test_train_engine.py | 14 +- tests/test_tree_training.py | 19 +- tests/torchrun/run_fsdp_dcp_distributed.py | 15 +- .../torchrun/run_fsdp_ulysses_train_batch.py | 14 +- .../run_megatron_engine_distributed.py | 55 +- .../run_megatron_engine_vlm_distributed.py | 14 +- .../v2/training_service/fake_train_engine.py | 20 +- tests/v2/training_service/test_worker_unit.py | 2 +- 41 files changed, 2475 insertions(+), 396 deletions(-) create mode 100644 tests/test_loss_reduction.py create mode 100644 tests/test_prompt_mean_loss.py create mode 100644 tests/test_rpc_callable_serialization.py diff --git a/areal/api/__init__.py b/areal/api/__init__.py index 3401bf8b89..aac3612000 100644 --- a/areal/api/__init__.py +++ b/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.engine_api", + "LossTerm": "areal.api.engine_api", + "LOSS_TERM_REDUCTION_MEAN": "areal.api.engine_api", + "LOSS_TERM_REDUCTION_SUM": "areal.api.engine_api", "InferenceEngine": "areal.api.engine_api", "Scheduler": "areal.api.scheduler_api", "Worker": "areal.api.scheduler_api", diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index 50d05aa781..1a63da3a9c 100644 --- a/areal/api/cli_args.py +++ b/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, @@ -1757,6 +1793,32 @@ 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'." + ) + # Validate CISPO configuration if self.use_cispo_loss: if self.use_sapo_loss: @@ -3173,6 +3235,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__() diff --git a/areal/api/engine_api.py b/areal/api/engine_api.py index f3ca01187c..060f509649 100644 --- a/areal/api/engine_api.py +++ b/areal/api/engine_api.py @@ -3,8 +3,9 @@ from __future__ import annotations import abc -from collections.abc import Callable +from collections.abc import Callable, Mapping from concurrent.futures import Future +from dataclasses import dataclass from typing import TYPE_CHECKING, Any import torch @@ -29,6 +30,96 @@ from areal.utils.data import MicroBatchList +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, + ), + ), + ) + + class TrainEngine(abc.ABC): @abc.abstractmethod def create_process_group(self, parallel_strategy: ParallelStrategy | None = None): @@ -363,8 +454,7 @@ def forward_backward_batch( 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_reduction: LossReduction, ) -> dict[str, float]: """Update the model with a batch of data and a loss function. @@ -379,15 +469,11 @@ def train_batch( Preferred format is ``list[dict[str, Any]]`` (trajectory list). Backward compatibility: a pre-batched ``dict[str, Any]`` is also accepted. - loss_fn : Callable[..., torch.Tensor] - The loss function. For actor (is_critic=False), it receives - (logprobs, entropy, input_data). For critic (is_critic=True), - it receives (values, input_data). Returns a scalar normalized loss. - 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. + loss_reduction : LossReduction + Reduction contract with one or more terms. For ``reduction="sum"``, + the loss term is divided by the global normalizer. For + ``reduction="mean"``, the loss term is a local normalized scalar and + is reweighted by the local normalizer before global normalization. Returns ------- @@ -402,8 +488,7 @@ def train_batch( 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_reduction: LossReduction, ) -> torch.Tensor | None: """Evaluate the model using the forward pass and loss function. @@ -418,15 +503,11 @@ def eval_batch( Preferred format is ``list[dict[str, Any]]`` (trajectory list). Backward compatibility: a pre-batched ``dict[str, Any]`` is also accepted. - loss_fn : Callable[..., torch.Tensor] - The loss function. For actor (is_critic=False), it receives - (logprobs, entropy, input_data). For critic (is_critic=True), - it receives (values, input_data). Returns a scalar normalized loss. - 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. + loss_reduction : LossReduction + Reduction contract with one or more terms. For ``reduction="sum"``, + the loss term is divided by the global normalizer. For + ``reduction="mean"``, the loss term is a local normalized scalar and + is reweighted by the local normalizer before global normalization. Returns ------- diff --git a/areal/engine/core/__init__.py b/areal/engine/core/__init__.py index 499d02fd22..a64f4dfebf 100644 --- a/areal/engine/core/__init__.py +++ b/areal/engine/core/__init__.py @@ -4,12 +4,16 @@ from areal.engine.core.train_engine import ( aggregate_eval_losses, - compute_total_loss_weight, + compute_global_normalizers, + compute_local_normalizers, 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", ] diff --git a/areal/engine/core/train_engine.py b/areal/engine/core/train_engine.py index 9bc83c2b77..e1e4c8e4e6 100644 --- a/areal/engine/core/train_engine.py +++ b/areal/engine/core/train_engine.py @@ -6,12 +6,18 @@ different training engine implementations (FSDP, Megatron, etc.). """ -from collections.abc import Callable +from collections.abc import Callable, Iterable, Iterator, Mapping from typing import Any import torch import torch.distributed as dist +from areal.api.engine_api import ( + LOSS_TERM_REDUCTION_MEAN, + LossFnOutput, + LossReduction, + LossTerm, +) from areal.infra.platforms import current_platform from areal.utils.data import ( MicroBatchList, @@ -21,48 +27,162 @@ ) __all__ = [ - "compute_total_loss_weight", + "compute_global_normalizers", + "compute_local_normalizers", + "scale_loss_for_reduction", "aggregate_eval_losses", "reorder_and_pad_outputs", ] -def compute_total_loss_weight( - mb_list: MicroBatchList, - loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor], - dp_group: dist.ProcessGroup, -) -> torch.Tensor: - """Compute total loss weight and all_reduce across data parallel group. +def compute_local_normalizers( + input_: dict[str, Any], loss_reduction: LossReduction +) -> dict[str, torch.Tensor]: + return {term.name: term.normalizer_fn(input_) for term in loss_reduction.terms} - This aggregates the loss weights from all micro-batches and reduces - them across the data parallel group to get a global normalization factor. - Parameters - ---------- - mb_list : MicroBatchList - The list of micro-batches. - loss_weight_fn : Callable[[dict[str, Any]], torch.Tensor] - Function to compute loss weight for each micro-batch. - dp_group : dist.ProcessGroup - The data parallel process group for all_reduce. - - Returns - ------- - torch.Tensor - The total loss weight (scalar tensor) after all_reduce. - """ - total_weight = ( - torch.stack([loss_weight_fn(mb) for mb in mb_list.mbs]) +def _sum_term_normalizer(mb_list: MicroBatchList, term: LossTerm) -> torch.Tensor: + return ( + torch.stack([term.normalizer_fn(mb) for mb in mb_list.mbs]) .sum() .detach() .clone() .to(dtype=torch.float32) ) - dist.all_reduce(total_weight, group=dp_group) - assert total_weight > 0, ( - "Global total loss weight must be positive after all_reduce" + + +def compute_global_normalizers( + mb_list: MicroBatchList, + loss_reduction: LossReduction, + dp_group: dist.ProcessGroup, +) -> dict[str, torch.Tensor]: + normalizers = torch.stack( + [_sum_term_normalizer(mb_list, term) for term in loss_reduction.terms] + ) + dist.all_reduce(normalizers, group=dp_group) + if torch.any(normalizers <= 0).item(): + invalid_mask = (normalizers <= 0).detach().cpu().tolist() + invalid = [ + term.name + for term, is_invalid in zip(loss_reduction.terms, invalid_mask, strict=True) + if is_invalid + ] + raise RuntimeError( + "Global loss normalizers must be positive after all_reduce; " + f"got non-positive normalizer for terms {invalid}." + ) + return { + term.name: normalizer + for term, normalizer in zip( + loss_reduction.terms, normalizers.unbind(), strict=True + ) + } + + +def _get_loss_term_value( + loss: LossFnOutput, loss_reduction: LossReduction, term: LossTerm +) -> torch.Tensor: + if isinstance(loss, Mapping): + try: + return loss[term.name] + except KeyError as e: + raise KeyError( + f"loss output is missing term {term.name!r}; " + f"available terms are {tuple(loss)}." + ) from e + if len(loss_reduction.terms) == 1: + return loss + raise TypeError("loss_fn must return a mapping for multi-term LossReduction.") + + +def _zero_like_loss(loss: LossFnOutput) -> torch.Tensor: + if isinstance(loss, Mapping): + first = next(iter(loss.values())) + return first * 0.0 + return loss * 0.0 + + +def _apply_local_reduction( + term_loss: torch.Tensor, + term: LossTerm, + local_normalizer: torch.Tensor, +) -> torch.Tensor: + if term.reduction == LOSS_TERM_REDUCTION_MEAN: + reduced = term_loss * local_normalizer + else: + reduced = term_loss + return torch.where( + local_normalizer != 0, + reduced, + torch.zeros_like(term_loss), + ) + + +def _scale_loss_term( + loss: LossFnOutput, + loss_reduction: LossReduction, + term: LossTerm, + local_normalizers: dict[str, torch.Tensor], + global_normalizers: dict[str, torch.Tensor], + loss_multiplier: float, +) -> torch.Tensor: + global_normalizer = global_normalizers[term.name] + term_loss = _get_loss_term_value(loss, loss_reduction, term) + term_loss = _apply_local_reduction(term_loss, term, local_normalizers[term.name]) + return term_loss / global_normalizer * loss_multiplier + + +def _iter_scaled_terms( + loss: LossFnOutput, + loss_reduction: LossReduction, + local_normalizers: dict[str, torch.Tensor], + global_normalizers: dict[str, torch.Tensor], + loss_multiplier: float, +) -> Iterator[torch.Tensor]: + for term in loss_reduction.terms: + yield _scale_loss_term( + loss, + loss_reduction, + term, + local_normalizers, + global_normalizers, + loss_multiplier, + ) + + +def _sum_scaled_terms( + terms: Iterable[torch.Tensor], + loss: LossFnOutput, +) -> torch.Tensor: + iterator = iter(terms) + try: + total = next(iterator) + except StopIteration: + return _zero_like_loss(loss) + + for term in iterator: + total = total + term + return total + + +def scale_loss_for_reduction( + loss: LossFnOutput, + loss_reduction: LossReduction, + local_normalizers: dict[str, torch.Tensor], + global_normalizers: dict[str, torch.Tensor], + loss_multiplier: float, +) -> torch.Tensor: + """Scale local loss terms into a globally normalized engine loss.""" + return _sum_scaled_terms( + _iter_scaled_terms( + loss, + loss_reduction, + local_normalizers, + global_normalizers, + loss_multiplier, + ), + loss, ) - return total_weight def aggregate_eval_losses( diff --git a/areal/engine/fsdp_engine.py b/areal/engine/fsdp_engine.py index 27fc005f86..061e8d3e33 100644 --- a/areal/engine/fsdp_engine.py +++ b/areal/engine/fsdp_engine.py @@ -48,6 +48,7 @@ FinetuneSpec, FSDPParallelStrategy, InferenceEngine, + LossReduction, ModelAllocation, ParallelStrategy, ParamSpec, @@ -60,8 +61,10 @@ from areal.api.io_struct import DeviceRuntimeInfo from areal.engine.core import ( aggregate_eval_losses, - compute_total_loss_weight, + compute_global_normalizers, + compute_local_normalizers, reorder_and_pad_outputs, + scale_loss_for_reduction, ) from areal.engine.core.distributed import ( init_custom_process_group, @@ -760,23 +763,19 @@ def forward_backward_batch( 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_reduction: LossReduction, ) -> dict[str, float]: self._ensure_ready() self.optimizer_zero_grad() input_batched, _ = self._normalize_batch_input(input_) - # Step 1: Prepare micro-batches mb_list = self._prepare_mb_list(input_batched).to(self.device) - # Step 2: Compute total loss weight - total_loss_weight = compute_total_loss_weight( - mb_list, loss_weight_fn, self.dp_group + global_normalizers = compute_global_normalizers( + mb_list, loss_reduction, self.dp_group ) - # Step 3: Forward-backward using process_output_fn callback def process_output( logits: torch.Tensor, ctx_dict: dict[str, Any] ) -> torch.Tensor: @@ -784,15 +783,13 @@ def process_output( return self._compute_logprobs_and_loss( logits, ctx, - loss_fn, - loss_weight_fn, - total_loss_weight, + loss_reduction, + global_normalizers, loss_multiplier=self.parallel_helper.dp_size, ) self.forward_backward_batch(mb_list, process_output, forward_only=False) - # Step 4: Optimizer step stats = self.optimizer_step() stats["num_micro_batches"] = len(mb_list.mbs) return stats @@ -801,22 +798,18 @@ def process_output( 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_reduction: LossReduction, ) -> torch.Tensor | None: self._ensure_ready() input_batched, _ = self._normalize_batch_input(input_) - # Step 1: Prepare micro-batches mb_list = self._prepare_mb_list(input_batched).to(self.device) - # Step 2: Compute total loss weight - total_loss_weight = compute_total_loss_weight( - mb_list, loss_weight_fn, self.dp_group + global_normalizers = compute_global_normalizers( + mb_list, loss_reduction, self.dp_group ) - # Step 3: Forward using process_output_fn callback, collecting losses losses: list[torch.Tensor] = [] def process_output( @@ -826,16 +819,14 @@ def process_output( loss = self._compute_logprobs_and_loss( logits, ctx, - loss_fn, - loss_weight_fn, - total_loss_weight, + loss_reduction, + global_normalizers, ) losses.append(loss.detach()) return loss self.forward_backward_batch(mb_list, process_output, forward_only=True) - # Step 4: Aggregate losses return aggregate_eval_losses(losses, self.dp_group) @torch.no_grad() @@ -849,7 +840,6 @@ def forward_batch( input_batched, meta = self._normalize_batch_input(input_) - # Step 1: Prepare sequence lengths if meta is not None: assert isinstance(input_, list) inferred_seqlens = [d["attention_mask"].shape[-1] for d in input_] @@ -866,10 +856,8 @@ def forward_batch( assert output_seqlens is not None batch_size = len(output_seqlens) - # Step 2: Prepare micro-batches mb_list = self._prepare_mb_list(input_batched).to(self.device) - # Step 3: Forward using process_output_fn callback, collecting results outputs: list[torch.Tensor] = [] def process_output(logits: torch.Tensor, ctx_dict: dict[str, Any]) -> None: @@ -880,7 +868,6 @@ def process_output(logits: torch.Tensor, ctx_dict: dict[str, Any]) -> None: self.forward_backward_batch(mb_list, process_output, forward_only=True) - # Step 4: Aggregate and reorder outputs if self.enable_tree_training: result = merge_packed_tree_results(outputs, batch_size) else: @@ -2055,15 +2042,12 @@ def _compute_logprobs_and_loss( self, logits: torch.Tensor, ctx: FSDPTrainContext, - loss_fn: Callable[..., torch.Tensor], - loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor], - total_loss_weight: torch.Tensor, + loss_reduction: LossReduction, + global_normalizers: dict[str, torch.Tensor], loss_multiplier: float = 1.0, ) -> torch.Tensor: """Compute logprobs/entropy and return scaled loss.""" - local_weight = loss_weight_fn(ctx.mb_input) - if local_weight == 0: - return logits.mean() * 0.0 + local_normalizers = compute_local_normalizers(ctx.mb_input, loss_reduction) if self.config.is_critic and self.enable_tree_training: raise NotImplementedError( @@ -2106,7 +2090,7 @@ def _compute_logprobs_and_loss( entropy = entropy[: -ctx.pad_length] vocab_min_logits = vocab_min_logits[: -ctx.pad_length] vocab_max_logits = vocab_max_logits[: -ctx.pad_length] - loss = loss_fn( + loss = loss_reduction.loss_fn( logprobs, entropy, ctx.mb_input, @@ -2117,10 +2101,15 @@ def _compute_logprobs_and_loss( values = self._compute_values(logits.squeeze(-1), ctx.ulysses_pad_size) if ctx.pad_length > 0: values = values[: -ctx.pad_length] - loss = loss_fn(values, ctx.mb_input) - - loss_scale = local_weight / total_loss_weight * loss_multiplier - return loss * loss_scale + loss = loss_reduction.loss_fn(values, ctx.mb_input) + + return scale_loss_for_reduction( + loss, + loss_reduction, + local_normalizers, + global_normalizers, + loss_multiplier, + ) def _compute_forward_result( self, diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index 6e268451a8..6e3163f703 100644 --- a/areal/engine/megatron_engine.py +++ b/areal/engine/megatron_engine.py @@ -39,6 +39,7 @@ from areal.api import ( FinetuneSpec, InferenceEngine, + LossReduction, MegatronParallelStrategy, ParallelStrategy, ParamSpec, @@ -51,8 +52,10 @@ from areal.api.io_struct import DeviceRuntimeInfo from areal.engine.core import ( aggregate_eval_losses, - compute_total_loss_weight, + compute_global_normalizers, + compute_local_normalizers, reorder_and_pad_outputs, + scale_loss_for_reduction, ) from areal.engine.core.distributed import ( init_custom_process_group, @@ -990,29 +993,25 @@ def _process_output(input_, output_): 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_reduction: LossReduction, ) -> dict[str, float]: self._ensure_ready() self.optimizer_zero_grad() input_batched, _ = self._normalize_batch_input(input_) - # Step 1: Prepare micro-batches mb_list = self._prepare_mb_list(input_batched).to(self.device) - # Step 2: Compute total loss weight. # Use DP+CP group: after CP all-gather each rank computes the full-sequence # loss, so all_gather's backward (reduce_scatter) sums cp_size identical - # gradients, amplifying by cp_size. Including CP in the weight all-reduce + # gradients, amplifying by cp_size. Including CP in the normalizer all-reduce # introduces a matching cp_size factor in the denominator, cancelling out. - total_loss_weight = compute_total_loss_weight( + global_normalizers = compute_global_normalizers( mb_list, - loss_weight_fn, + loss_reduction, mpu.get_data_parallel_group(with_context_parallel=True), ) - # Step 3: Forward-backward using Megatron's pipeline function. # `len(mb_list)` compensates Megatron Core's `output_tensor /= num_microbatches` # applied in the 2-tuple `(loss, {})` branch of # `megatron.core.pipeline_parallel.schedules._forward_step_helper`. Our @@ -1031,9 +1030,8 @@ def process_output( return self._compute_logprobs_and_loss( output, inputs, - loss_fn, - loss_weight_fn, - total_loss_weight, + loss_reduction, + global_normalizers, loss_multiplier=loss_multiplier, ) @@ -1043,7 +1041,6 @@ def process_output( forward_only=False, ) - # Step 4: Optimizer step stats = self.optimizer_step() stats["num_micro_batches"] = len(mb_list.mbs) return stats @@ -1052,38 +1049,33 @@ def process_output( 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_reduction: LossReduction, ) -> torch.Tensor | None: self._ensure_ready() input_batched, _ = self._normalize_batch_input(input_) - # Step 1: Prepare micro-batches mb_list = self._prepare_mb_list(input_batched).to(self.device) - # Step 2: Compute total loss weight (DP+CP, see train_batch comment). - total_loss_weight = compute_total_loss_weight( + global_normalizers = compute_global_normalizers( mb_list, - loss_weight_fn, + loss_reduction, mpu.get_data_parallel_group(with_context_parallel=True), ) - # Step 3: Forward using Megatron's pipeline function, collecting losses losses: list[torch.Tensor] = [] def process_output( output: torch.Tensor, inputs: dict[str, Any] ) -> torch.Tensor: loss = self._compute_logprobs_and_loss( - output, inputs, loss_fn, loss_weight_fn, total_loss_weight + output, inputs, loss_reduction, global_normalizers ) losses.append(loss.detach()) return loss self.forward_backward_batch(mb_list, process_output, forward_only=True) - # Step 4: Aggregate losses if mpu.is_pipeline_last_stage(): return aggregate_eval_losses( losses, mpu.get_data_parallel_group(with_context_parallel=True) @@ -1101,7 +1093,6 @@ def forward_batch( input_batched, meta = self._normalize_batch_input(input_) - # Step 1: Prepare sequence lengths if meta is not None: assert isinstance(input_, list) inferred_seqlens = [d["attention_mask"].shape[-1] for d in input_] @@ -1118,10 +1109,8 @@ def forward_batch( assert output_seqlens is not None batch_size = len(output_seqlens) - # Step 2: Prepare micro-batches mb_list = self._prepare_mb_list(input_batched).to(self.device) - # Step 3: Forward using Megatron's pipeline function, collecting results outputs: list[torch.Tensor] = [] def process_output(output: torch.Tensor, inputs: dict[str, Any]) -> None: @@ -1133,7 +1122,6 @@ def process_output(output: torch.Tensor, inputs: dict[str, Any]) -> None: mb_list, process_output, forward_only=True, gather_cp_output=True ) - # Step 4: Aggregate, reorder, and broadcast outputs res = None if mpu.is_pipeline_last_stage(): if self.enable_tree_training: @@ -2296,14 +2284,11 @@ def _compute_logprobs_and_loss( self, output: torch.Tensor, inputs: dict[str, Any], - loss_fn: Callable[..., torch.Tensor], - loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor], - total_loss_weight: torch.Tensor, + loss_reduction: LossReduction, + global_normalizers: dict[str, torch.Tensor], loss_multiplier: float = 1.0, ) -> torch.Tensor: - local_weight = loss_weight_fn(inputs) - if local_weight == 0: - return output.mean() * 0.0 + local_normalizers = compute_local_normalizers(inputs, loss_reduction) if self.config.is_critic and self.enable_tree_training: raise NotImplementedError( @@ -2418,7 +2403,7 @@ def _compute_logprobs_and_loss( inputs = { k: v for k, v in inputs.items() if not k.startswith("_cp_") } - loss = loss_fn( + loss = loss_reduction.loss_fn( logprobs, entropy, inputs, @@ -2429,10 +2414,15 @@ def _compute_logprobs_and_loss( ) else: values = output.squeeze(-1) - loss = loss_fn(values, inputs) - - loss_scale = local_weight / total_loss_weight * loss_multiplier - return loss * loss_scale + loss = loss_reduction.loss_fn(values, inputs) + + return scale_loss_for_reduction( + loss, + loss_reduction, + local_normalizers, + global_normalizers, + loss_multiplier, + ) def _compute_forward_result( self, diff --git a/areal/experimental/engine/archon_engine.py b/areal/experimental/engine/archon_engine.py index cb8da760c7..5e56f81e7e 100644 --- a/areal/experimental/engine/archon_engine.py +++ b/areal/experimental/engine/archon_engine.py @@ -28,6 +28,7 @@ from areal.api import ( FinetuneSpec, + LossReduction, ParallelStrategy, SaveLoadMeta, TrainEngine, @@ -41,8 +42,10 @@ ) from areal.engine.core.train_engine import ( aggregate_eval_losses, - compute_total_loss_weight, + compute_global_normalizers, + compute_local_normalizers, reorder_and_pad_outputs, + scale_loss_for_reduction, ) from areal.engine.fsdp_utils.grad import fsdp2_clip_grad_norm from areal.experimental.engine.archon_checkpoint import ( @@ -525,8 +528,7 @@ def forward_backward_batch( 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_reduction: LossReduction, ) -> dict[str, float]: """Train on a batch of data.""" assert self._initialized @@ -536,8 +538,8 @@ def train_batch( mb_list = self._prepare_mb_list(input_batched).to(self.device) - total_loss_weight = compute_total_loss_weight( - mb_list, loss_weight_fn, self.data_parallel_group + global_normalizers = compute_global_normalizers( + mb_list, loss_reduction, self.data_parallel_group ) def process_output( @@ -547,9 +549,8 @@ def process_output( return self._compute_logprobs_and_loss( logits, ctx, - loss_fn, - loss_weight_fn, - total_loss_weight, + loss_reduction, + global_normalizers, loss_multiplier=self.data_parallel_world_size, ) @@ -563,8 +564,7 @@ def process_output( 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_reduction: LossReduction, ) -> torch.Tensor | None: """Evaluate on a batch of data.""" assert self._initialized @@ -573,8 +573,8 @@ def eval_batch( mb_list = self._prepare_mb_list(input_batched).to(self.device) - total_loss_weight = compute_total_loss_weight( - mb_list, loss_weight_fn, self.data_parallel_group + global_normalizers = compute_global_normalizers( + mb_list, loss_reduction, self.data_parallel_group ) def process_output( @@ -584,9 +584,8 @@ def process_output( return self._compute_logprobs_and_loss( logits, ctx, - loss_fn, - loss_weight_fn, - total_loss_weight, + loss_reduction, + global_normalizers, ) losses = self.forward_backward_batch(mb_list, process_output, forward_only=True) @@ -1261,22 +1260,19 @@ def _compute_logprobs_and_loss( self, logits: torch.Tensor, ctx: ArchonTrainContext, - loss_fn: Callable[..., torch.Tensor], - loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor], - total_loss_weight: torch.Tensor, + loss_reduction: LossReduction, + global_normalizers: dict[str, torch.Tensor], loss_multiplier: float = 1.0, ) -> torch.Tensor: """Compute logprobs/entropy and return scaled loss.""" - local_weight = loss_weight_fn(ctx.mb_input) - if local_weight == 0: - return logits.mean() * 0.0 + local_normalizers = compute_local_normalizers(ctx.mb_input, loss_reduction) if not self.config.is_critic: result = self._gather_actor_train_outputs(logits, ctx) if result is None: return logits.mean() * 0.0 logprobs, entropy, vocab_min_logits, vocab_max_logits = result - loss = loss_fn( + loss = loss_reduction.loss_fn( logprobs, entropy, ctx.mb_input, @@ -1285,10 +1281,15 @@ def _compute_logprobs_and_loss( ) else: values = self._gather_critic_output(logits, ctx) - loss = loss_fn(values, ctx.mb_input) - - loss_scale = local_weight / total_loss_weight * loss_multiplier - return loss * loss_scale + loss = loss_reduction.loss_fn(values, ctx.mb_input) + + return scale_loss_for_reduction( + loss, + loss_reduction, + local_normalizers, + global_normalizers, + loss_multiplier, + ) def _compute_forward_result( self, diff --git a/areal/infra/rpc/serialization.py b/areal/infra/rpc/serialization.py index 2ea864c838..c7c23d2a45 100644 --- a/areal/infra/rpc/serialization.py +++ b/areal/infra/rpc/serialization.py @@ -649,6 +649,9 @@ def serialize_value(value: Any) -> Any: "value": value.value, } + if callable(value) and not isinstance(value, type): + raise ValueError("RPC serialization does not support callable values.") + # Primitives (int, float, str, bool) pass through unchanged return value @@ -749,6 +752,9 @@ def deserialize_value(value: Any) -> Any: f"Failed to deserialize ray.ObjectRef, treating as regular dict: {e}" ) + if value.get("type") == "callable": + raise ValueError("RPC deserialization does not support callable values.") + # Check for SerializedTensor marker if value.get("type") == "tensor": try: diff --git a/areal/trainer/dpo/dpo_engine.py b/areal/trainer/dpo/dpo_engine.py index 8cd6736b0b..9a4c9da14a 100644 --- a/areal/trainer/dpo/dpo_engine.py +++ b/areal/trainer/dpo/dpo_engine.py @@ -5,7 +5,7 @@ import torch -from areal.api import TrainEngine +from areal.api import LossReduction, TrainEngine from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value from areal.utils import logging, stats_tracker @@ -24,7 +24,7 @@ def _dpo_valid_pairs(x: dict[str, Any]) -> torch.Tensor: return seqlens.view(-1, 2).ne(0).all(dim=1) -def _dpo_loss_weight(x: dict[str, Any]) -> torch.Tensor: +def _dpo_loss_normalizer(x: dict[str, Any]) -> torch.Tensor: return _dpo_valid_pairs(x).count_nonzero().float() @@ -56,10 +56,12 @@ def _train_dpo(self, data: dict[str, Any]) -> None: self.engine.train() stats = self.engine.train_batch( input_=data, - loss_fn=functools.partial( - compute_dpo_loss, beta=self.beta, loss_type=self.loss_type + loss_reduction=LossReduction.mean( + loss_fn=functools.partial( + compute_dpo_loss, beta=self.beta, loss_type=self.loss_type + ), + normalizer_fn=_dpo_loss_normalizer, ), - loss_weight_fn=_dpo_loss_weight, ) stats_tracker.scalar(**stats) @@ -72,10 +74,12 @@ def _evaluate_dpo(self, data: dict[str, Any]) -> None: self.engine.eval() self.engine.eval_batch( input_=data, - loss_fn=functools.partial( - compute_dpo_loss, beta=self.beta, loss_type=self.loss_type + loss_reduction=LossReduction.mean( + loss_fn=functools.partial( + compute_dpo_loss, beta=self.beta, loss_type=self.loss_type + ), + normalizer_fn=_dpo_loss_normalizer, ), - loss_weight_fn=_dpo_loss_weight, ) @trace_perf("dpo_engine.compute_logp", category="compute") diff --git a/areal/trainer/ppo/actor.py b/areal/trainer/ppo/actor.py index 0da045557f..598b0b0105 100644 --- a/areal/trainer/ppo/actor.py +++ b/areal/trainer/ppo/actor.py @@ -5,7 +5,7 @@ import torch -from areal.api import TrainEngine +from areal.api import LossReduction, TrainEngine from areal.api.cli_args import MicroBatchSpec, PPOActorConfig, RejectionSamplingConfig from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value @@ -26,10 +26,14 @@ Normalization, TrajBatchMeta, batched_call, + concat_batch, split_padded_tensor_dict_into_mb_list, ) from areal.utils.functional import ( + LOSS_AGGREGATION_PROMPT_MEAN, + LOSS_AGGREGATION_TOKEN_MEAN, cispo_loss_fn, + make_pg_loss_normalizer_fn, ppo_actor_loss_fn, reward_overlong_penalty, sapo_loss_fn, @@ -42,6 +46,47 @@ logger = logging.getLogger("PPOActor") +def _use_sum_pg_loss(data: dict[str, Any], m2_threshold: float | None) -> bool: + return "teacher_logp" not in data and m2_threshold is None + + +def _make_actor_loss_normalizer_fn( + loss_aggregation: str, + group_size: int, + m2_threshold: float | None, + prox_logp_method: str, + current_version: int | None, +): + normalizer_fn = make_pg_loss_normalizer_fn(loss_aggregation, group_size) + if m2_threshold is None: + return normalizer_fn + + def m2po_normalizer(data: dict[str, Any]) -> torch.Tensor: + prox_logp_gt = data.get("prox_logp") + if ( + prox_logp_gt is None + and ProxLogpMethod(prox_logp_method).skips_forward_pass() + ): + raise ValueError( + "m2_threshold requires prox_logp in the batch when " + "prox_logp_method skips the proximal forward pass." + ) + prox_logp = _resolve_proximal_logp( + prox_logp_gt=prox_logp_gt, + prox_logp_method=prox_logp_method, + old_logp=data["logprobs"], + logprobs=data["logprobs"].detach(), + versions=data.get("versions"), + current_version=current_version, + ) + loss_mask = _apply_m2po_masking( + data["logprobs"], prox_logp, data["loss_mask"].bool(), m2_threshold + ) + return normalizer_fn({**data, "loss_mask": loss_mask}) + + return m2po_normalizer + + class PPOActor: def __init__(self, config: PPOActorConfig, engine: TrainEngine): self.config = config @@ -116,9 +161,7 @@ def _log_configuration(self): # Log other critical config logger.info("=" * 70) logger.info("Training Parameters:") - logger.info( - f" importance_sampling_level: {getattr(config, 'importance_sampling_level', 'token')}" - ) + logger.info(f" importance_sampling_level: {config.importance_sampling_level}") logger.info( f" adv_norm: {config.adv_norm if config.adv_norm else 'DISABLED (None)'}" ) @@ -262,7 +305,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 == 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"] @@ -338,42 +384,71 @@ def _ppo_update(self, data: dict[str, Any]) -> None: ) ########## Logging code ends ########## - # Pop keys that are no longer needed after advantage computation - # Note: "versions" is kept if needed for approximation/metrics in loss function for key in ["rewards", "tot_rewards", "kl_rewards"]: data.pop(key, None) # NOTE: calling engine.train() is critical to enabling gradient checkpointing self.engine.train() + has_prompt_group_sizes = ( + self.config.loss_aggregation == LOSS_AGGREGATION_PROMPT_MEAN + and "group_sizes" in data + ) + outer_granularity = 1 + if ( + self.config.loss_aggregation == LOSS_AGGREGATION_PROMPT_MEAN + and not has_prompt_group_sizes + ): + outer_granularity = self.config.group_size mb_inputs = split_padded_tensor_dict_into_mb_list( data, - mb_spec=MicroBatchSpec(n_mbs=self.config.ppo_n_minibatches), + mb_spec=MicroBatchSpec( + n_mbs=self.config.ppo_n_minibatches, + granularity=outer_granularity, + ), ) with stats_tracker.scope("update"): - # Get current version for proximal approximation metrics current_version = self.engine.get_version() for mb in mb_inputs.mbs: - 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(), + use_sum_reduction = _use_sum_pg_loss(mb, self.m2_threshold) + 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_aggregation=self.config.loss_aggregation, + group_size=self.config.group_size, + loss_aggregation_divisor=self.config.loss_aggregation_divisor, + return_sum=use_sum_reduction, ) + normalizer_fn = _make_actor_loss_normalizer_fn( + self.config.loss_aggregation, + self.config.group_size, + self.m2_threshold, + self.config.prox_logp_method, + current_version, + ) + if use_sum_reduction: + loss_reduction = LossReduction.sum( + loss_fn=loss_fn, + normalizer_fn=normalizer_fn, + ) + else: + loss_reduction = LossReduction.mean( + loss_fn=loss_fn, + normalizer_fn=normalizer_fn, + ) + train_stat = self.engine.train_batch(mb, loss_reduction=loss_reduction) stats_tracker.scalar(**train_stat) @@ -434,6 +509,10 @@ def grpo_loss_fn( sapo_tau_neg: float = 1.05, use_cispo_loss: bool = False, use_decoupled_loss: bool = False, + loss_aggregation: str = "token_mean", + group_size: int = 1, + loss_aggregation_divisor: float | None = None, + return_sum: bool = False, vocab_min_logits: torch.Tensor | None = None, vocab_max_logits: torch.Tensor | None = None, vocab_mean_logits: torch.Tensor | None = None, @@ -461,11 +540,9 @@ def grpo_loss_fn( current_version=current_version, ) - # Apply M2PO masking if threshold is set if m2_threshold is not None: loss_mask = _apply_m2po_masking(old_logp, prox_logp, loss_mask, m2_threshold) - # Use CISPO, SAPO, or PPO loss if use_cispo_loss: if use_sapo_loss: raise ValueError( @@ -487,6 +564,11 @@ def grpo_loss_fn( old_logprobs=old_logp, rejection_sampling=rejection_sampling, cu_seqlens=input_data.get("cu_seqlens"), + loss_aggregation=loss_aggregation, + group_size=group_size, + loss_aggregation_divisor=loss_aggregation_divisor, + group_sizes=input_data.get("group_sizes"), + return_sum=return_sum, ) elif use_sapo_loss: if use_decoupled_loss: @@ -503,6 +585,11 @@ def grpo_loss_fn( loss_mask=loss_mask, importance_sampling_level=importance_sampling_level, cu_seqlens=input_data.get("cu_seqlens"), + loss_aggregation=loss_aggregation, + group_size=group_size, + loss_aggregation_divisor=loss_aggregation_divisor, + group_sizes=input_data.get("group_sizes"), + return_sum=return_sum, ) else: loss, stat = ppo_actor_loss_fn( @@ -517,22 +604,32 @@ def grpo_loss_fn( rejection_sampling=rejection_sampling, importance_sampling_level=importance_sampling_level, cu_seqlens=input_data.get("cu_seqlens"), + loss_aggregation=loss_aggregation, + group_size=group_size, + loss_aggregation_divisor=loss_aggregation_divisor, + group_sizes=input_data.get("group_sizes"), + return_sum=return_sum, ) - # Joint Distillation KL Loss teacher_logp = input_data.get("teacher_logp") rkl_stat = None if teacher_logp is not None: - # Coefficients for RL and Knowledge Distillation + if loss_aggregation != LOSS_AGGREGATION_TOKEN_MEAN: + raise ValueError( + "teacher_logp distillation is only supported with " + "loss_aggregation='token_mean'." + ) + if return_sum: + raise ValueError( + "return_sum=True is not supported with teacher_logp. Use the " + "mean reduction so distillation keeps its own local normalization." + ) rl_loss_weight = input_data.get("rl_loss_weight", 1.0) distill_loss_weight = input_data.get("distill_loss_weight", 0.005) - teacher_logp = ( - teacher_logp.detach() - ) # detach to prevent gradient backprop to teacher + teacher_logp = teacher_logp.detach() if rl_loss_weight == 0: - # Pure KD using reverse KL (importance-sampling) rkl_reward = teacher_logp - logprobs.detach() importance_weight = torch.exp(logprobs - old_logp) @@ -543,7 +640,6 @@ def grpo_loss_fn( rkl_stat = -1 * rkl_weighted_term else: - # KDRL: Knowledge Distillation + Reinforcement Learning (joint loss) rkl_penalty_per_token = (logprobs - teacher_logp) * loss_mask rkl_penalty = rkl_penalty_per_token.sum() / loss_mask.sum().clamp(min=1) @@ -551,7 +647,6 @@ def grpo_loss_fn( rkl_stat = rkl_penalty_per_token - # Log training statistics stats_tracker.denominator( n_tokens=infer_token_denominator(input_data, loss_mask), n_valid_tokens=loss_mask.bool(), diff --git a/areal/trainer/ppo/critic.py b/areal/trainer/ppo/critic.py index 05f83c7441..aa9f3b7409 100644 --- a/areal/trainer/ppo/critic.py +++ b/areal/trainer/ppo/critic.py @@ -5,7 +5,7 @@ import torch -from areal.api import TrainEngine +from areal.api import LossReduction, TrainEngine from areal.api.cli_args import MicroBatchSpec, PPOCriticConfig from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value @@ -65,11 +65,13 @@ def _ppo_update(self, data: dict[str, Any]) -> None: for mb in mb_inputs.mbs: train_stat = self.engine.train_batch( mb, - loss_fn=functools.partial( - ppo_loss_fn, - eps_clip=self.config.eps_clip, + loss_reduction=LossReduction.mean( + loss_fn=functools.partial( + ppo_loss_fn, + eps_clip=self.config.eps_clip, + ), + normalizer_fn=lambda x: x["loss_mask"].count_nonzero(), ), - loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), ) stats_tracker.scalar(**train_stat) diff --git a/areal/trainer/rw/rw_engine.py b/areal/trainer/rw/rw_engine.py index 1e5cbb0acf..1a1d17a873 100644 --- a/areal/trainer/rw/rw_engine.py +++ b/areal/trainer/rw/rw_engine.py @@ -4,7 +4,7 @@ import torch -from areal.api import TrainEngine +from areal.api import LossReduction, TrainEngine from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value from areal.utils import logging, stats_tracker @@ -22,7 +22,7 @@ def _rw_valid_pairs(x: dict[str, Any]) -> torch.Tensor: return seqlens.view(-1, 2).ne(0).all(dim=1) -def _rw_loss_weight(x: dict[str, Any]) -> torch.Tensor: +def _rw_loss_normalizer(x: dict[str, Any]) -> torch.Tensor: return _rw_valid_pairs(x).count_nonzero().float() @@ -49,13 +49,15 @@ def train_rw(self, data: list[dict[str, Any]]) -> None: def _train_rw(self, data: dict[str, Any]) -> None: """Train on a batch (reward model).""" - if _rw_loss_weight(data) == 0: + if _rw_loss_normalizer(data) == 0: _log_empty_rw_stats(data["cu_seqlens"].device) self.engine.train() stats = self.engine.train_batch( input_=data, - loss_fn=compute_rw_loss, - loss_weight_fn=_rw_loss_weight, + loss_reduction=LossReduction.mean( + loss_fn=compute_rw_loss, + normalizer_fn=_rw_loss_normalizer, + ), ) stats_tracker.scalar(**stats) @@ -65,13 +67,15 @@ def evaluate_rw(self, data: list[dict[str, Any]]) -> None: batched_call(self._evaluate_rw, data, unpack=False) def _evaluate_rw(self, data: dict[str, Any]) -> None: - if _rw_loss_weight(data) == 0: + if _rw_loss_normalizer(data) == 0: _log_empty_rw_stats(data["cu_seqlens"].device) self.engine.eval() self.engine.eval_batch( input_=data, - loss_fn=compute_rw_loss, - loss_weight_fn=_rw_loss_weight, + loss_reduction=LossReduction.mean( + loss_fn=compute_rw_loss, + normalizer_fn=_rw_loss_normalizer, + ), ) diff --git a/areal/trainer/sft/lm_engine.py b/areal/trainer/sft/lm_engine.py index 5e5343bb2d..4c5e50b793 100644 --- a/areal/trainer/sft/lm_engine.py +++ b/areal/trainer/sft/lm_engine.py @@ -4,7 +4,7 @@ import torch -from areal.api import TrainEngine +from areal.api import LossReduction, TrainEngine from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value from areal.utils import stats_tracker @@ -29,8 +29,10 @@ def _train_lm(self, data: dict[str, Any]) -> None: data["loss_mask"] = torch.roll(data["loss_mask"].bool(), shifts=-1, dims=-1) stats = self.engine.train_batch( input_=data, - loss_fn=compute_packed_sft_loss, - loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), + loss_reduction=LossReduction.mean( + loss_fn=compute_packed_sft_loss, + normalizer_fn=lambda x: x["loss_mask"].count_nonzero(), + ), ) stats_tracker.scalar(**stats) @@ -44,8 +46,10 @@ def _evaluate_lm(self, data: dict[str, Any]) -> None: data["loss_mask"] = torch.roll(data["loss_mask"].bool(), shifts=-1, dims=-1) self.engine.eval_batch( input_=data, - loss_fn=compute_packed_sft_loss, - loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), + loss_reduction=LossReduction.mean( + loss_fn=compute_packed_sft_loss, + normalizer_fn=lambda x: x["loss_mask"].count_nonzero(), + ), ) diff --git a/areal/utils/data.py b/areal/utils/data.py index 1560742a1a..35a8e932b5 100644 --- a/areal/utils/data.py +++ b/areal/utils/data.py @@ -591,7 +591,7 @@ class MicroBatchItem(NamedTuple): """A single micro-batch item from MicroBatchList iteration. Attributes: - orig_mb: Original micro-batch dict (for loss_weight_fn, context) + orig_mb: Original micro-batch dict (for normalizer_fn, context) padded_mb: Padded micro-batch dict (for model forward) padding_length: Batch-level padding added to this micro-batch old_cu_seqlens: Original cu_seqlens before sequence alignment (or None) @@ -642,7 +642,7 @@ def __iter__(self) -> Iterator[MicroBatchItem]: Yields: MicroBatchItem containing: - - orig_mb: Original micro-batch dict (for loss_weight_fn, context) + - orig_mb: Original micro-batch dict (for normalizer_fn, context) - padded_mb: Padded micro-batch dict (for model forward) - padding_length: Batch-level padding added to this micro-batch - old_cu_seqlens: Original cu_seqlens before sequence alignment (or None) @@ -697,6 +697,36 @@ def to(self, *args, **kwargs): DEFAULT_MAX_TOKENS_PER_MB = int(1e12) +def _resolve_microbatch_sequence_groups( + bs: int, granularity: int, group_sizes: Any | 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 +751,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 +765,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 +779,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 +833,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( diff --git a/areal/utils/functional/__init__.py b/areal/utils/functional/__init__.py index b09c6af5f2..2be8dd3cf0 100644 --- a/areal/utils/functional/__init__.py +++ b/areal/utils/functional/__init__.py @@ -1,11 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 from areal.utils.functional.functional import ( + LOSS_AGGREGATION_CONSTANT, + LOSS_AGGREGATION_PROMPT_MEAN, + LOSS_AGGREGATION_SEQ_MEAN, + LOSS_AGGREGATION_TOKEN_MEAN, + LOSS_AGGREGATIONS_ALL, RejectionSamplingResult, + aggregate_pg_loss, + aggregate_pg_loss_sum, apply_rejection_sampling, cispo_loss_fn, dpo_pair_logratios, dpo_preference_loss, + make_pg_loss_normalizer_fn, masked_normalization, ppo_actor_loss_fn, ppo_critic_loss_fn, @@ -19,11 +27,19 @@ __all__ = [ # functional.py + "LOSS_AGGREGATION_CONSTANT", + "LOSS_AGGREGATION_PROMPT_MEAN", + "LOSS_AGGREGATION_SEQ_MEAN", + "LOSS_AGGREGATION_TOKEN_MEAN", + "LOSS_AGGREGATIONS_ALL", "RejectionSamplingResult", + "aggregate_pg_loss", + "aggregate_pg_loss_sum", "apply_rejection_sampling", "cispo_loss_fn", "dpo_pair_logratios", "dpo_preference_loss", + "make_pg_loss_normalizer_fn", "masked_normalization", "ppo_actor_loss_fn", "ppo_critic_loss_fn", diff --git a/areal/utils/functional/functional.py b/areal/utils/functional/functional.py index 8aaed7730a..035c8b6940 100644 --- a/areal/utils/functional/functional.py +++ b/areal/utils/functional/functional.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 import functools +import math +from collections.abc import Callable from dataclasses import dataclass from typing import Any @@ -229,8 +231,11 @@ def apply_rejection_sampling( # Step 1: Compute log ratio = log(π_proximal / π_behave) # Upcast operands to fp32 before subtraction to avoid precision loss in bf16/fp16. log_ratio = proximal_logprobs.detach().float() - old_logprobs.detach().float() - # Sanitize non-finite values (e.g. -inf - (-inf) = NaN) to prevent NaN propagation. - log_ratio = torch.where(torch.isfinite(log_ratio), log_ratio, 0.0) + # Sanitize NaNs (e.g. -inf - (-inf)) while preserving true infinities so + # threshold checks still reject divergent ratios. + log_ratio = torch.where( + torch.isnan(log_ratio), torch.zeros_like(log_ratio), log_ratio + ) # Step 2: Compute metric value (reuse existing KLEstimator sign conventions) if config.metric == "ratio": @@ -449,6 +454,655 @@ def compute_binary_kl_divergence( return p * torch.log(p / q) + (1 - p) * torch.log((1 - p) / (1 - q)) +LOSS_AGGREGATION_TOKEN_MEAN = "token_mean" +LOSS_AGGREGATION_SEQ_MEAN = "seq_mean" +LOSS_AGGREGATION_PROMPT_MEAN = "prompt_mean" +LOSS_AGGREGATION_CONSTANT = "constant" + + +_PgLossReducer = Callable[ + [ + torch.Tensor, + torch.Tensor, + torch.Tensor, + int, + float | None, + torch.Tensor | None, + list[int] | None, + str, + ], + torch.Tensor, +] +_PgLossSumReducer = _PgLossReducer +_PgLossNormalizerFn = Callable[[dict[str, Any]], torch.Tensor] +_PgLossNormalizerFactory = Callable[[int], _PgLossNormalizerFn] + + +def _masked_pg_loss(pg_loss: torch.Tensor, num_mask: torch.Tensor) -> torch.Tensor: + return torch.where(num_mask, pg_loss, 0).to(torch.float32) + + +def _resolve_pg_masks( + pg_loss: torch.Tensor, + loss_mask: torch.Tensor, + denom_mask: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + if loss_mask.shape != pg_loss.shape: + raise ValueError( + f"loss_mask shape {tuple(loss_mask.shape)} must match " + f"pg_loss shape {tuple(pg_loss.shape)}." + ) + if denom_mask is not None and denom_mask.shape != pg_loss.shape: + raise ValueError( + f"denom_mask shape {tuple(denom_mask.shape)} must match " + f"pg_loss shape {tuple(pg_loss.shape)}." + ) + num_mask = loss_mask.bool() + den_mask = num_mask if denom_mask is None else denom_mask.bool() + return num_mask, den_mask + + +def _sum_active_unit_means(num: torch.Tensor, den: torch.Tensor) -> torch.Tensor: + active = den > 0 + return torch.where(active, num / den.clamp_min(1), torch.zeros_like(num)).sum() + + +def _mean_active_unit_means(num: torch.Tensor, den: torch.Tensor) -> torch.Tensor: + active = den > 0 + active_count = active.count_nonzero().clamp_min(1) + return _sum_active_unit_means(num, den) / active_count + + +def _count_active_packed_sequences( + den_mask: torch.Tensor, + cu_seqlens: torch.Tensor, + n_seqs: int, +) -> torch.Tensor: + seq_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to( + device=den_mask.device, dtype=torch.long + ) + seq_ids = torch.arange(n_seqs, device=den_mask.device).repeat_interleave(seq_lens) + den = torch.zeros(n_seqs, dtype=torch.float32, device=den_mask.device) + den.scatter_add_(0, seq_ids, den_mask.to(dtype=torch.float32)) + return den.count_nonzero() + + +def _validate_pg_group_sizes(n_seqs: int, group_sizes: list[int] | None) -> list[int]: + if group_sizes is None: + raise ValueError("group_sizes must be provided for explicit prompt groups.") + 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 != n_seqs: + raise ValueError(f"group_sizes sum to {total} but sequence count is {n_seqs}.") + return sizes + + +def _group_ids_from_sizes(group_sizes: list[int], device: torch.device) -> torch.Tensor: + sizes = torch.tensor(group_sizes, dtype=torch.long, device=device) + return torch.arange(len(group_sizes), device=device).repeat_interleave(sizes) + + +def _pg_token_normalizer(data: dict[str, Any]) -> torch.Tensor: + return data["loss_mask"].count_nonzero() + + +def _active_sequence_mask(data: dict[str, Any]) -> torch.Tensor: + loss_mask = data["loss_mask"].bool() + cu_seqlens = data.get("cu_seqlens") + if cu_seqlens is None: + return loss_mask.any(dim=-1) + + n_seqs = cu_seqlens.numel() - 1 + seq_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to( + device=loss_mask.device, dtype=torch.long + ) + seq_ids = torch.arange(n_seqs, device=loss_mask.device).repeat_interleave(seq_lens) + active_tokens = torch.zeros(n_seqs, dtype=torch.float32, device=loss_mask.device) + active_tokens.scatter_add_(0, seq_ids, loss_mask.to(dtype=torch.float32)) + return active_tokens > 0 + + +def _pg_unit_normalizer(data: dict[str, Any], unit: int) -> torch.Tensor: + active_seqs = _active_sequence_mask(data) + n_seqs = active_seqs.numel() + group_sizes = data.get("group_sizes") + if group_sizes is not None: + sizes = _validate_pg_group_sizes(n_seqs, group_sizes) + group_ids = _group_ids_from_sizes(sizes, active_seqs.device) + active_groups = torch.zeros( + len(sizes), dtype=torch.float32, device=active_seqs.device + ) + active_groups.scatter_add_(0, group_ids, active_seqs.to(dtype=torch.float32)) + return active_groups.count_nonzero().to(torch.float32) + if n_seqs % unit != 0: + raise ValueError( + f"microbatch sequence count {n_seqs} is not divisible by group_size {unit}." + ) + return active_seqs.view(-1, unit).any(dim=-1).count_nonzero().to(torch.float32) + + +def _make_pg_unit_normalizer_fn(unit: int) -> _PgLossNormalizerFn: + def pg_unit_normalizer(data: dict[str, Any]) -> torch.Tensor: + return _pg_unit_normalizer(data, unit) + + return pg_unit_normalizer + + +def _make_pg_token_normalizer_fn(_group_size: int) -> _PgLossNormalizerFn: + return _pg_token_normalizer + + +def _make_pg_sequence_normalizer_fn(_group_size: int) -> _PgLossNormalizerFn: + return _make_pg_unit_normalizer_fn(unit=1) + + +def _make_pg_prompt_normalizer_fn(group_size: int) -> _PgLossNormalizerFn: + return _make_pg_unit_normalizer_fn(unit=group_size) + + +def _aggregate_token_mean( + pg_loss: torch.Tensor, + num_mask: torch.Tensor, + den_mask: torch.Tensor, + _group_size: int, + _loss_aggregation_divisor: float | None, + _cu_seqlens: torch.Tensor | None, + _group_sizes: list[int] | None, + _mode: str, +) -> torch.Tensor: + denom = den_mask.count_nonzero().clamp_min(1) + return torch.where(num_mask, pg_loss, 0).sum() / denom + + +def _aggregate_token_sum( + pg_loss: torch.Tensor, + num_mask: torch.Tensor, + _den_mask: torch.Tensor, + _group_size: int, + _loss_aggregation_divisor: float | None, + _cu_seqlens: torch.Tensor | None, + _group_sizes: list[int] | None, + _mode: str, +) -> torch.Tensor: + return _masked_pg_loss(pg_loss, num_mask).sum() + + +def _aggregate_constant( + pg_loss: torch.Tensor, + num_mask: torch.Tensor, + den_mask: torch.Tensor, + _group_size: int, + loss_aggregation_divisor: float | None, + cu_seqlens: torch.Tensor | None, + _group_sizes: list[int] | None, + _mode: str, +) -> torch.Tensor: + divisor = loss_aggregation_divisor + assert divisor is not None + + if cu_seqlens is None: + if pg_loss.ndim != 2: + raise ValueError( + "constant expects 2D pg_loss when cu_seqlens is None, " + f"got shape {tuple(pg_loss.shape)}." + ) + active_seqs = den_mask.any(dim=-1).count_nonzero() + else: + if pg_loss.ndim != 1: + raise ValueError( + "constant expects 1D pg_loss when cu_seqlens is provided, " + f"got shape {tuple(pg_loss.shape)}." + ) + n_seqs = cu_seqlens.numel() - 1 + active_seqs = _count_active_packed_sequences(den_mask, cu_seqlens, n_seqs) + return _masked_pg_loss(pg_loss, num_mask).sum() / ( + active_seqs.clamp_min(1) * divisor + ) + + +def _aggregate_constant_sum( + pg_loss: torch.Tensor, + num_mask: torch.Tensor, + _den_mask: torch.Tensor, + _group_size: int, + loss_aggregation_divisor: float | None, + cu_seqlens: torch.Tensor | None, + _group_sizes: list[int] | None, + _mode: str, +) -> torch.Tensor: + divisor = loss_aggregation_divisor + assert divisor is not None + + if cu_seqlens is None: + if pg_loss.ndim != 2: + raise ValueError( + "constant expects 2D pg_loss when cu_seqlens is None, " + f"got shape {tuple(pg_loss.shape)}." + ) + n_seqs = pg_loss.shape[0] + else: + if pg_loss.ndim != 1: + raise ValueError( + "constant expects 1D pg_loss when cu_seqlens is provided, " + f"got shape {tuple(pg_loss.shape)}." + ) + n_seqs = cu_seqlens.numel() - 1 + if n_seqs == 0: + return pg_loss.sum() * 0.0 + return _masked_pg_loss(pg_loss, num_mask).sum() / divisor + + +def _aggregate_seq_mean( + pg_loss: torch.Tensor, + num_mask: torch.Tensor, + den_mask: torch.Tensor, + _group_size: int, + _loss_aggregation_divisor: float | None, + cu_seqlens: torch.Tensor | None, + _group_sizes: list[int] | None, + mode: str, +) -> torch.Tensor: + return _aggregate_unit_mean( + pg_loss, + num_mask, + den_mask, + unit=1, + cu_seqlens=cu_seqlens, + group_sizes=None, + mode=mode, + ) + + +def _aggregate_seq_sum( + pg_loss: torch.Tensor, + num_mask: torch.Tensor, + den_mask: torch.Tensor, + _group_size: int, + _loss_aggregation_divisor: float | None, + cu_seqlens: torch.Tensor | None, + _group_sizes: list[int] | None, + mode: str, +) -> torch.Tensor: + return _aggregate_unit_sum( + pg_loss, + num_mask, + den_mask, + unit=1, + cu_seqlens=cu_seqlens, + group_sizes=None, + mode=mode, + ) + + +def _aggregate_prompt_mean( + pg_loss: torch.Tensor, + num_mask: torch.Tensor, + den_mask: torch.Tensor, + group_size: int, + _loss_aggregation_divisor: float | None, + cu_seqlens: torch.Tensor | None, + group_sizes: list[int] | None, + mode: str, +) -> torch.Tensor: + return _aggregate_unit_mean( + pg_loss, + num_mask, + den_mask, + unit=group_size, + cu_seqlens=cu_seqlens, + group_sizes=group_sizes, + mode=mode, + ) + + +def _aggregate_prompt_sum( + pg_loss: torch.Tensor, + num_mask: torch.Tensor, + den_mask: torch.Tensor, + group_size: int, + _loss_aggregation_divisor: float | None, + cu_seqlens: torch.Tensor | None, + group_sizes: list[int] | None, + mode: str, +) -> torch.Tensor: + return _aggregate_unit_sum( + pg_loss, + num_mask, + den_mask, + unit=group_size, + cu_seqlens=cu_seqlens, + group_sizes=group_sizes, + mode=mode, + ) + + +def _aggregate_unit_mean( + pg_loss: torch.Tensor, + num_mask: torch.Tensor, + den_mask: torch.Tensor, + unit: int, + cu_seqlens: torch.Tensor | None, + group_sizes: list[int] | None, + mode: str, +) -> torch.Tensor: + masked_pg = _masked_pg_loss(pg_loss, num_mask) + den_f = den_mask.to(torch.float32) + + if cu_seqlens is None: + return _aggregate_padded_unit_mean( + pg_loss, unit, group_sizes, mode, masked_pg, den_f + ) + return _aggregate_packed_unit_mean( + pg_loss, unit, cu_seqlens, group_sizes, mode, masked_pg, den_f + ) + + +def _aggregate_unit_sum( + pg_loss: torch.Tensor, + num_mask: torch.Tensor, + den_mask: torch.Tensor, + unit: int, + cu_seqlens: torch.Tensor | None, + group_sizes: list[int] | None, + mode: str, +) -> torch.Tensor: + masked_pg = _masked_pg_loss(pg_loss, num_mask) + den_f = den_mask.to(torch.float32) + + if cu_seqlens is None: + return _aggregate_padded_unit_sum( + pg_loss, unit, group_sizes, mode, masked_pg, den_f + ) + return _aggregate_packed_unit_sum( + pg_loss, unit, cu_seqlens, group_sizes, mode, masked_pg, den_f + ) + + +def _aggregate_padded_unit_mean( + pg_loss: torch.Tensor, + unit: int, + group_sizes: list[int] | None, + mode: str, + masked_pg: torch.Tensor, + den_f: torch.Tensor, +) -> torch.Tensor: + if pg_loss.ndim != 2: + raise ValueError( + f"{mode} expects 2D pg_loss when cu_seqlens is None, " + f"got shape {tuple(pg_loss.shape)}." + ) + n, t = pg_loss.shape + if group_sizes is not None: + sizes = _validate_pg_group_sizes(n, group_sizes) + group_ids = _group_ids_from_sizes(sizes, pg_loss.device) + num = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) + den = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) + num.scatter_add_(0, group_ids, masked_pg.sum(dim=1)) + den.scatter_add_(0, group_ids, den_f.sum(dim=1)) + return _mean_active_unit_means(num, den) + if n % unit != 0: + raise ValueError( + f"{mode}: microbatch sequence count {n} is not divisible by " + f"group_size {unit}. Configure mb_spec.granularity >= group_size " + f"so prompt-groups stay contiguous within each microbatch." + ) + g = n // unit + if g == 0: + return pg_loss.sum() * 0.0 + num = masked_pg.view(g, unit, t).sum(dim=(1, 2)) + den = den_f.view(g, unit, t).sum(dim=(1, 2)) + return _mean_active_unit_means(num, den) + + +def _aggregate_padded_unit_sum( + pg_loss: torch.Tensor, + unit: int, + group_sizes: list[int] | None, + mode: str, + masked_pg: torch.Tensor, + den_f: torch.Tensor, +) -> torch.Tensor: + if pg_loss.ndim != 2: + raise ValueError( + f"{mode} expects 2D pg_loss when cu_seqlens is None, " + f"got shape {tuple(pg_loss.shape)}." + ) + n, t = pg_loss.shape + if group_sizes is not None: + sizes = _validate_pg_group_sizes(n, group_sizes) + group_ids = _group_ids_from_sizes(sizes, pg_loss.device) + num = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) + den = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) + num.scatter_add_(0, group_ids, masked_pg.sum(dim=1)) + den.scatter_add_(0, group_ids, den_f.sum(dim=1)) + return _sum_active_unit_means(num, den) + if n % unit != 0: + raise ValueError( + f"{mode}: microbatch sequence count {n} is not divisible by " + f"group_size {unit}. Configure mb_spec.granularity >= group_size " + f"so prompt-groups stay contiguous within each microbatch." + ) + g = n // unit + if g == 0: + return pg_loss.sum() * 0.0 + num = masked_pg.view(g, unit, t).sum(dim=(1, 2)) + den = den_f.view(g, unit, t).sum(dim=(1, 2)) + return _sum_active_unit_means(num, den) + + +def _aggregate_packed_unit_mean( + pg_loss: torch.Tensor, + unit: int, + cu_seqlens: torch.Tensor, + group_sizes: list[int] | None, + mode: str, + masked_pg: torch.Tensor, + den_f: torch.Tensor, +) -> torch.Tensor: + if pg_loss.ndim != 1: + raise ValueError( + f"{mode} expects 1D pg_loss when cu_seqlens is provided, " + f"got shape {tuple(pg_loss.shape)}." + ) + n_seqs = cu_seqlens.numel() - 1 + seq_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to( + device=pg_loss.device, dtype=torch.long + ) + if group_sizes is not None: + sizes = _validate_pg_group_sizes(n_seqs, group_sizes) + seq_group_ids = _group_ids_from_sizes(sizes, pg_loss.device) + group_ids = seq_group_ids.repeat_interleave(seq_lens) + num = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) + den = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) + num.scatter_add_(0, group_ids, masked_pg) + den.scatter_add_(0, group_ids, den_f) + return _mean_active_unit_means(num, den) + if n_seqs % unit != 0: + raise ValueError( + f"{mode}: packed microbatch has {n_seqs} sequences, not " + f"divisible by group_size {unit}. Configure " + f"mb_spec.granularity >= group_size." + ) + g = n_seqs // unit + if g == 0: + return pg_loss.sum() * 0.0 + seq_ids = torch.arange(n_seqs, device=pg_loss.device) + group_ids = (seq_ids // unit).repeat_interleave(seq_lens) + num = torch.zeros(g, dtype=torch.float32, device=pg_loss.device) + den = torch.zeros(g, dtype=torch.float32, device=pg_loss.device) + num.scatter_add_(0, group_ids, masked_pg) + den.scatter_add_(0, group_ids, den_f) + return _mean_active_unit_means(num, den) + + +def _aggregate_packed_unit_sum( + pg_loss: torch.Tensor, + unit: int, + cu_seqlens: torch.Tensor, + group_sizes: list[int] | None, + mode: str, + masked_pg: torch.Tensor, + den_f: torch.Tensor, +) -> torch.Tensor: + if pg_loss.ndim != 1: + raise ValueError( + f"{mode} expects 1D pg_loss when cu_seqlens is provided, " + f"got shape {tuple(pg_loss.shape)}." + ) + n_seqs = cu_seqlens.numel() - 1 + seq_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to( + device=pg_loss.device, dtype=torch.long + ) + if group_sizes is not None: + sizes = _validate_pg_group_sizes(n_seqs, group_sizes) + seq_group_ids = _group_ids_from_sizes(sizes, pg_loss.device) + group_ids = seq_group_ids.repeat_interleave(seq_lens) + num = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) + den = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) + num.scatter_add_(0, group_ids, masked_pg) + den.scatter_add_(0, group_ids, den_f) + return _sum_active_unit_means(num, den) + if n_seqs % unit != 0: + raise ValueError( + f"{mode}: packed microbatch has {n_seqs} sequences, not " + f"divisible by group_size {unit}. Configure " + f"mb_spec.granularity >= group_size." + ) + g = n_seqs // unit + if g == 0: + return pg_loss.sum() * 0.0 + seq_ids = torch.arange(n_seqs, device=pg_loss.device) + group_ids = (seq_ids // unit).repeat_interleave(seq_lens) + num = torch.zeros(g, dtype=torch.float32, device=pg_loss.device) + den = torch.zeros(g, dtype=torch.float32, device=pg_loss.device) + num.scatter_add_(0, group_ids, masked_pg) + den.scatter_add_(0, group_ids, den_f) + return _sum_active_unit_means(num, den) + + +_PG_LOSS_REDUCERS: dict[str, _PgLossReducer] = { + LOSS_AGGREGATION_TOKEN_MEAN: _aggregate_token_mean, + LOSS_AGGREGATION_SEQ_MEAN: _aggregate_seq_mean, + LOSS_AGGREGATION_PROMPT_MEAN: _aggregate_prompt_mean, + LOSS_AGGREGATION_CONSTANT: _aggregate_constant, +} +_PG_LOSS_SUM_REDUCERS: dict[str, _PgLossSumReducer] = { + LOSS_AGGREGATION_TOKEN_MEAN: _aggregate_token_sum, + LOSS_AGGREGATION_SEQ_MEAN: _aggregate_seq_sum, + LOSS_AGGREGATION_PROMPT_MEAN: _aggregate_prompt_sum, + LOSS_AGGREGATION_CONSTANT: _aggregate_constant_sum, +} +_PG_LOSS_NORMALIZER_FACTORIES: dict[str, _PgLossNormalizerFactory] = { + LOSS_AGGREGATION_TOKEN_MEAN: _make_pg_token_normalizer_fn, + LOSS_AGGREGATION_SEQ_MEAN: _make_pg_sequence_normalizer_fn, + LOSS_AGGREGATION_PROMPT_MEAN: _make_pg_prompt_normalizer_fn, + LOSS_AGGREGATION_CONSTANT: _make_pg_sequence_normalizer_fn, +} +LOSS_AGGREGATIONS_ALL = tuple(_PG_LOSS_REDUCERS) + + +def _validate_pg_loss_aggregation_config( + loss_aggregation: str, loss_aggregation_divisor: float | None +) -> None: + if loss_aggregation not in _PG_LOSS_REDUCERS: + raise ValueError( + f"loss_aggregation must be one of {LOSS_AGGREGATIONS_ALL}, " + f"got {loss_aggregation!r}." + ) + + if loss_aggregation == LOSS_AGGREGATION_CONSTANT: + if ( + loss_aggregation_divisor is None + or not math.isfinite(loss_aggregation_divisor) + or loss_aggregation_divisor <= 0 + ): + raise ValueError( + "loss_aggregation_divisor must be a positive finite value " + "when loss_aggregation='constant'." + ) + elif loss_aggregation_divisor is not None: + raise ValueError( + "loss_aggregation_divisor is only used when loss_aggregation='constant'." + ) + + +def make_pg_loss_normalizer_fn( + loss_aggregation: str, group_size: int +) -> _PgLossNormalizerFn: + """Return the denominator normalizer paired with policy-gradient aggregation.""" + try: + return _PG_LOSS_NORMALIZER_FACTORIES[loss_aggregation](group_size) + except KeyError as e: + raise ValueError( + f"loss_aggregation must be one of {LOSS_AGGREGATIONS_ALL}, " + f"got {loss_aggregation!r}." + ) from e + + +def aggregate_pg_loss( + pg_loss: torch.Tensor, + loss_mask: torch.Tensor, + loss_aggregation: str = LOSS_AGGREGATION_TOKEN_MEAN, + group_size: int = 1, + loss_aggregation_divisor: float | None = None, + cu_seqlens: torch.Tensor | None = None, + denom_mask: torch.Tensor | None = None, + group_sizes: list[int] | None = None, +) -> torch.Tensor: + """Reduce per-token policy-gradient loss to a microbatch scalar. + + ``token_mean`` averages valid tokens. ``seq_mean`` averages per-sequence + token means. ``prompt_mean`` averages prompt-group means, using + ``group_sizes`` for partial groups and ``group_size`` otherwise. + ``constant`` averages each sequence's masked token sum divided by + ``loss_aggregation_divisor``. Pair the returned scalar with + ``make_pg_loss_normalizer_fn`` for distributed reduction. + """ + _validate_pg_loss_aggregation_config(loss_aggregation, loss_aggregation_divisor) + + num_mask, den_mask = _resolve_pg_masks(pg_loss, loss_mask, denom_mask) + + return _PG_LOSS_REDUCERS[loss_aggregation]( + pg_loss, + num_mask, + den_mask, + group_size, + loss_aggregation_divisor, + cu_seqlens, + group_sizes, + loss_aggregation, + ) + + +def aggregate_pg_loss_sum( + pg_loss: torch.Tensor, + loss_mask: torch.Tensor, + loss_aggregation: str = LOSS_AGGREGATION_TOKEN_MEAN, + group_size: int = 1, + loss_aggregation_divisor: float | None = None, + cu_seqlens: torch.Tensor | None = None, + denom_mask: torch.Tensor | None = None, + group_sizes: list[int] | None = None, +) -> torch.Tensor: + """Reduce per-token policy-gradient loss to a local sum term.""" + _validate_pg_loss_aggregation_config(loss_aggregation, loss_aggregation_divisor) + + num_mask, den_mask = _resolve_pg_masks(pg_loss, loss_mask, denom_mask) + + return _PG_LOSS_SUM_REDUCERS[loss_aggregation]( + pg_loss, + num_mask, + den_mask, + group_size, + loss_aggregation_divisor, + cu_seqlens, + group_sizes, + loss_aggregation, + ) + + def ppo_actor_loss_fn( logprobs: torch.Tensor, proximal_logprobs: torch.Tensor, @@ -461,6 +1115,11 @@ def ppo_actor_loss_fn( rejection_sampling: RejectionSamplingConfig | None = None, importance_sampling_level: str = "token", cu_seqlens: torch.Tensor | None = None, + loss_aggregation: str = LOSS_AGGREGATION_TOKEN_MEAN, + group_size: int = 1, + loss_aggregation_divisor: float | None = None, + group_sizes: list[int] | None = None, + return_sum: bool = False, ) -> tuple[torch.Tensor, dict]: """PPO actor loss function with optional rejection sampling. @@ -499,11 +1158,8 @@ def ppo_actor_loss_fn( Shape: [batch_size + 1], where cu_seqlens[i] marks the start of sequence i. Not needed for 2D padded inputs (sequences identified by batch dimension). """ - # Save original count BEFORE rejection sampling may modify loss_mask. - # This keeps the denominator consistent with loss_weight_fn in actor.py, - # which always uses the original loss_mask from input_data. Without this, - # mask mode would inflate per-token gradients by N_original / N_kept. - loss_mask_count = loss_mask.count_nonzero() or 1 + # Rejection masking narrows the numerator but keeps the original denominator. + orig_loss_mask = loss_mask # === Apply rejection sampling (replaces old compute_behave_imp_weight) === if rejection_sampling is not None: @@ -562,7 +1218,17 @@ def ppo_actor_loss_fn( pg_loss = pg_loss * behave_imp_weight logging_loss = pg_loss.detach() - pg_loss = torch.where(loss_mask, pg_loss, 0).sum() / loss_mask_count + aggregate_fn = aggregate_pg_loss_sum if return_sum else aggregate_pg_loss + pg_loss = aggregate_fn( + pg_loss, + loss_mask, + loss_aggregation=loss_aggregation, + group_size=group_size, + loss_aggregation_divisor=loss_aggregation_divisor, + cu_seqlens=cu_seqlens, + denom_mask=orig_loss_mask, + group_sizes=group_sizes, + ) clip_mask.logical_and_(loss_mask) dual_clip_mask.logical_and_(loss_mask) stat = dict( @@ -592,6 +1258,11 @@ def sapo_loss_fn( loss_mask: torch.Tensor, importance_sampling_level: str = "token", cu_seqlens: torch.Tensor | None = None, + loss_aggregation: str = LOSS_AGGREGATION_TOKEN_MEAN, + group_size: int = 1, + loss_aggregation_divisor: float | None = None, + group_sizes: list[int] | None = None, + return_sum: bool = False, ) -> tuple[torch.Tensor, dict]: """SAPO (Soft Adaptive Policy Optimization) loss with asymmetric sigmoid gates. @@ -613,7 +1284,6 @@ def sapo_loss_fn( """ if tau_pos <= 0 or tau_neg <= 0: raise ValueError("SAPO temperatures (tau_pos, tau_neg) must be positive.") - loss_mask_count = loss_mask.count_nonzero() or 1 advantages = advantages.detach() log_ratio = logprobs - old_logprobs @@ -644,16 +1314,23 @@ def sapo_loss_fn( # Compute loss pg_loss = -soft_gate * advantages logging_loss = pg_loss.detach() - pg_loss = torch.where(loss_mask, pg_loss, 0).sum() / loss_mask_count + aggregate_fn = aggregate_pg_loss_sum if return_sum else aggregate_pg_loss + pg_loss = aggregate_fn( + pg_loss, + loss_mask, + loss_aggregation=loss_aggregation, + group_size=group_size, + loss_aggregation_divisor=loss_aggregation_divisor, + cu_seqlens=cu_seqlens, + group_sizes=group_sizes, + ) - # Return stat dict compatible with PPO (fake clip_mask for logging compatibility) stat = dict( loss=logging_loss, importance_weight=ratio.detach(), approx_kl=log_ratio.detach(), - clip_mask=torch.zeros_like(loss_mask, dtype=torch.bool), # SAPO doesn't clip + clip_mask=torch.zeros_like(loss_mask, dtype=torch.bool), dual_clip_mask=torch.zeros_like(loss_mask, dtype=torch.bool), - # SAPO-specific stats (scaled gates for consistency) sapo_soft_gate=soft_gate.detach(), sapo_scaled_gate_pos=scaled_gate_pos.detach(), sapo_scaled_gate_neg=scaled_gate_neg.detach(), @@ -672,6 +1349,11 @@ def cispo_loss_fn( old_logprobs: torch.Tensor | None = None, rejection_sampling: RejectionSamplingConfig | None = None, cu_seqlens: torch.Tensor | None = None, + loss_aggregation: str = LOSS_AGGREGATION_TOKEN_MEAN, + group_size: int = 1, + loss_aggregation_divisor: float | None = None, + group_sizes: list[int] | None = None, + return_sum: bool = False, ) -> tuple[torch.Tensor, dict]: """CISPO (Clipped IS-weight Policy Optimization) loss from MiniMax-M1. @@ -691,8 +1373,8 @@ def cispo_loss_fn( Advantages are never clipped. The clip bounds reuse the same delta-from-1 convention as :func:`ppo_actor_loss_fn`. CISPO is canonically single-sided: pass ``eps_clip=1.0`` (lower bound 0) with ``eps_clip_higher=4.0`` for the - wide MiniMax-M1 range. Token level only -- the geometric-mean sequence ratio - of GSPO is not part of the MiniMax-M1 surrogate. + wide MiniMax-M1 range. CISPO uses token-level importance weights; the + policy-gradient loss can still be reduced with any supported aggregation. Decoupled loss: when ``rejection_sampling`` is set, the detached ``pi_proximal/pi_behave`` weight (from ``old_logprobs``) rescales each token's @@ -727,10 +1409,8 @@ def cispo_loss_fn( "CISPO requires a positive eps_clip_higher; the asymmetric upper " f"clip is the defining knob (MiniMax-M1 Eq. 4-5). Got {eps_clip_higher!r}." ) - # Pre-rejection token count, so the denominator matches loss_weight_fn. - loss_mask_count = loss_mask.count_nonzero() or 1 + orig_loss_mask = loss_mask - # Decoupled off-policy correction: the pi_proximal/pi_behave weight. if rejection_sampling is not None: rs_result = apply_rejection_sampling( proximal_logprobs=proximal_logprobs, @@ -745,21 +1425,29 @@ def cispo_loss_fn( advantages = advantages.detach() - # Stop-gradient on the clipped IS weight; gradient flows through logprobs only. log_ratio = (logprobs - proximal_logprobs).detach() ratio = torch.exp(log_ratio) ratio_clipped = torch.clamp(ratio, 1.0 - eps_clip, 1.0 + eps_clip_higher).detach() pg_loss = -ratio_clipped * advantages * logprobs if rejection_sampling is not None: - # behave_imp_weight is detached at source -> still a valid policy gradient. behave_approx_kl = proximal_logprobs.detach() - old_logprobs.detach() behave_mask = (behave_imp_weight > 0).logical_and(loss_mask.bool()) behave_approx_kl = torch.where(behave_mask, behave_approx_kl, 0.0) pg_loss = pg_loss * behave_imp_weight logging_loss = pg_loss.detach() - pg_loss = torch.where(loss_mask, pg_loss, 0).sum() / loss_mask_count + aggregate_fn = aggregate_pg_loss_sum if return_sum else aggregate_pg_loss + pg_loss = aggregate_fn( + pg_loss, + loss_mask, + loss_aggregation=loss_aggregation, + group_size=group_size, + loss_aggregation_divisor=loss_aggregation_divisor, + cu_seqlens=cu_seqlens, + denom_mask=orig_loss_mask, + group_sizes=group_sizes, + ) clip_mask = (ratio_clipped != ratio).logical_and(loss_mask) stat = dict( @@ -767,7 +1455,6 @@ def cispo_loss_fn( importance_weight=ratio.detach(), approx_kl=log_ratio, clip_mask=clip_mask, - # CISPO has no dual clip; zeros keep the stat schema stable. dual_clip_mask=torch.zeros_like(loss_mask, dtype=torch.bool), ) if rejection_sampling is not None: diff --git a/areal/v2/training_service/controller/controller.py b/areal/v2/training_service/controller/controller.py index 16f136095e..dcc09e95e9 100644 --- a/areal/v2/training_service/controller/controller.py +++ b/areal/v2/training_service/controller/controller.py @@ -703,8 +703,8 @@ def _require_list_batch(input_: Any, method_name: str) -> list[dict[str, Any]]: def train_batch( self, input_: list[dict[str, Any]] | None = None, - loss_fn: Any = None, - loss_weight_fn: Any = None, + *, + loss_reduction: Any, ) -> Any: from areal.infra.rpc.serialization import serialize_value @@ -714,9 +714,7 @@ def train_batch( payload = { "args": serialize_value([batch]), - "kwargs": serialize_value( - {"loss_fn": loss_fn, "loss_weight_fn": loss_weight_fn} - ), + "kwargs": serialize_value({"loss_reduction": loss_reduction}), } return self._gateway_post_result("/train_batch", payload) @@ -740,8 +738,8 @@ def forward_batch( def eval_batch( self, input_: list[dict[str, Any]] | None = None, - loss_fn: Any = None, - loss_weight_fn: Any = None, + *, + loss_reduction: Any, ) -> Any: from areal.infra.rpc.serialization import serialize_value @@ -751,9 +749,7 @@ def eval_batch( payload = { "args": serialize_value([batch]), - "kwargs": serialize_value( - {"loss_fn": loss_fn, "loss_weight_fn": loss_weight_fn} - ), + "kwargs": serialize_value({"loss_reduction": loss_reduction}), } return self._gateway_post_result("/eval_batch", payload) diff --git a/docs/en/best_practices/perf_profiling.md b/docs/en/best_practices/perf_profiling.md index f25eb5afaf..1080449976 100644 --- a/docs/en/best_practices/perf_profiling.md +++ b/docs/en/best_practices/perf_profiling.md @@ -85,11 +85,11 @@ ppo_update, etc.). **Example** (from `areal/engine/fsdp_engine.py`): ```python +from areal.api import LossReduction from areal.utils.perf_tracer import trace_perf @trace_perf("fsdp_engine.train_batch") -def train_batch(self, input_: dict[str, Any], loss_fn, loss_weight_fn): - # Training logic here +def train_batch(self, input_: dict[str, Any], loss_reduction: LossReduction): ... ``` diff --git a/docs/en/cli_reference.md b/docs/en/cli_reference.md index 5e4b9ca5d6..2175205a08 100644 --- a/docs/en/cli_reference.md +++ b/docs/en/cli_reference.md @@ -406,7 +406,10 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `use_sapo_loss` | boolean | `False` | Use SAPO loss (mutually exclusive with PPO clipping) | | `sapo_tau_pos` | float | `1.0` | SAPO temperature for positive advantages | | `sapo_tau_neg` | float | `1.05` | SAPO temperature for negative advantages | -| `use_cispo_loss` | boolean | `False` | 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. | +| `use_cispo_loss` | boolean | `False` | 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. Requires eps_clip_higher > 0; recommended eps_clip=1.0 (single-sided, lower bound 0) with eps_clip_higher=4.0. | +| `loss_aggregation` | string | `"token_mean"` | 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. **Choices:** `token_mean`, `seq_mean`, `prompt_mean`, `constant` | +| `loss_aggregation_divisor` | float \| None | `None` | Positive fixed denominator L for loss_aggregation='constant'. Unused by other loss aggregation modes. | +| `group_size` | integer | `1` | Internal prompt-group size for prompt_mean; derived from gconfig.n_samples by PPOConfig.__post_init__. | | `recompute_logprob` | boolean | `False` | Recompute log probability and replace the log probability returned by inference. | | `use_decoupled_loss` | boolean | `False` | Use the decoupled loss. Implicitly enables recompute_logprob. | | `rejection_sampling` | [`RejectionSamplingConfig`](section-rejection-sampling) \| None | `None` | Rejection sampling configuration for filtering stale samples. None disables filtering (equivalent to old behave_imp_weight_mode='disabled'). Only effective when use_decoupled_loss=True. | diff --git a/docs/generate_cli_docs.py b/docs/generate_cli_docs.py index 270d1f7b05..993e9a75c2 100644 --- a/docs/generate_cli_docs.py +++ b/docs/generate_cli_docs.py @@ -219,6 +219,7 @@ def format_default_value(field_obj) -> str: def generate_config_section( config_class, all_dataclasses: dict[str, Any], + lang: str = "en", title: str = "", description: str = "", anchor: str = "", @@ -269,10 +270,9 @@ def generate_config_section( default_value = format_default_value(field) # Get help text from metadata - help_text = field.metadata.get( - "help", - "-", - ) + help_text = field.metadata.get(f"help_{lang}") + if help_text is None: + help_text = field.metadata.get("help", "-") # Get choices if available choices = field.metadata.get("choices") @@ -364,7 +364,7 @@ def generate_cli_documentation(lang: str = "en"): # Generate documentation sections automatically for category_name, class_list in categories.items(): for class_name, cls in class_list: - doc += generate_config_section(cls, all_dataclasses) + doc += generate_config_section(cls, all_dataclasses, lang=lang) return doc diff --git a/docs/zh/best_practices/perf_profiling.md b/docs/zh/best_practices/perf_profiling.md index 950366fadd..cd52f415bc 100644 --- a/docs/zh/best_practices/perf_profiling.md +++ b/docs/zh/best_practices/perf_profiling.md @@ -75,11 +75,11 @@ python -m areal.tools.perf_trace_converter logs/**/perf_tracer/traces-*.jsonl me **示例**（来自 `areal/engine/fsdp_engine.py`）： ```python +from areal.api import LossReduction from areal.utils.perf_tracer import trace_perf @trace_perf("fsdp_engine.train_batch") -def train_batch(self, input_: dict[str, Any], loss_fn, loss_weight_fn): - # Training logic here +def train_batch(self, input_: dict[str, Any], loss_reduction: LossReduction): ... ``` diff --git a/docs/zh/cli_reference.md b/docs/zh/cli_reference.md index 53a04daa75..96ba57d397 100644 --- a/docs/zh/cli_reference.md +++ b/docs/zh/cli_reference.md @@ -404,7 +404,10 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `use_sapo_loss` | boolean | `False` | Use SAPO loss (mutually exclusive with PPO clipping) | | `sapo_tau_pos` | float | `1.0` | SAPO temperature for positive advantages | | `sapo_tau_neg` | float | `1.05` | SAPO temperature for negative advantages | -| `use_cispo_loss` | boolean | `False` | 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. | +| `use_cispo_loss` | boolean | `False` | 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. Requires eps_clip_higher > 0; recommended eps_clip=1.0 (single-sided, lower bound 0) with eps_clip_higher=4.0. | +| `loss_aggregation` | string | `"token_mean"` | 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 | `None` | loss_aggregation='constant' 使用的正数固定分母 L。其他 loss aggregation 模式不使用。 | +| `group_size` | integer | `1` | prompt_mean 使用的内部 prompt group 大小；由 PPOConfig.__post_init__ 根据 gconfig.n_samples 推导。 | | `recompute_logprob` | boolean | `False` | Recompute log probability and replace the log probability returned by inference. | | `use_decoupled_loss` | boolean | `False` | Use the decoupled loss. Implicitly enables recompute_logprob. | | `rejection_sampling` | [`RejectionSamplingConfig`](section-rejection-sampling) \| None | `None` | Rejection sampling configuration for filtering stale samples. None disables filtering (equivalent to old behave_imp_weight_mode='disabled'). Only effective when use_decoupled_loss=True. | diff --git a/tests/experimental/archon/torchrun/run_archon_engine_pp.py b/tests/experimental/archon/torchrun/run_archon_engine_pp.py index d9d94d6b26..8ced33fd45 100644 --- a/tests/experimental/archon/torchrun/run_archon_engine_pp.py +++ b/tests/experimental/archon/torchrun/run_archon_engine_pp.py @@ -31,7 +31,7 @@ ) from tests.utils import get_model_path -from areal.api import FinetuneSpec, ParallelStrategy +from areal.api import FinetuneSpec, LossReduction, ParallelStrategy from areal.api.cli_args import MicroBatchSpec, OptimizerConfig, TrainEngineConfig from areal.experimental.engine.archon_engine import ArchonEngine @@ -95,8 +95,17 @@ def mock_loss_fn( return torch.mean(logprobs) -def mock_loss_weight_fn(input_data: dict) -> torch.Tensor: - """Mock loss weight function for testing.""" +def mock_loss_sum_fn( + logprobs: torch.Tensor, + entropy: torch.Tensor, + input_data: dict, + **kwargs, +) -> torch.Tensor: + """Mock local numerator loss for testing sum reduction.""" + return torch.sum(logprobs) + + +def mock_normalizer_fn(input_data: dict) -> torch.Tensor: return input_data["cu_seqlens"][-1].float() @@ -181,8 +190,10 @@ def test_eval_batch(engine: ArchonEngine, mock_input: dict) -> bool: try: loss = engine.eval_batch( mock_input, - loss_fn=mock_loss_fn, - loss_weight_fn=mock_loss_weight_fn, + loss_reduction=LossReduction.mean( + loss_fn=mock_loss_fn, + normalizer_fn=mock_normalizer_fn, + ), ) # In PP mode, eval_batch may return None (TODO in ArchonEngine) @@ -211,19 +222,35 @@ def test_train_batch(engine: ArchonEngine, mock_input: dict) -> bool: try: result = engine.train_batch( mock_input, - loss_fn=mock_loss_fn, - loss_weight_fn=mock_loss_weight_fn, + loss_reduction=LossReduction.mean( + loss_fn=mock_loss_fn, + normalizer_fn=mock_normalizer_fn, + ), + ) + sum_result = engine.train_batch( + mock_input, + loss_reduction=LossReduction.sum( + loss_fn=mock_loss_sum_fn, + normalizer_fn=mock_normalizer_fn, + ), ) print_rank0(" train_batch result:") print_rank0(f" update_successful = {result.get('update_successful')}") print_rank0(f" grad_norm = {result.get('grad_norm'):.6f}") print_rank0(f" lr = {result.get('lr'):.2e}") + print_rank0(" train_batch sum-reduction result:") + print_rank0(f" update_successful = {sum_result.get('update_successful')}") + print_rank0(f" grad_norm = {sum_result.get('grad_norm'):.6f}") + print_rank0(f" lr = {sum_result.get('lr'):.2e}") # Verify grad_norm is valid grad_norm = result.get("grad_norm") if grad_norm is None or not torch.isfinite(torch.tensor(grad_norm)): print_rank0(" WARNING: grad_norm is invalid") + sum_grad_norm = sum_result.get("grad_norm") + if sum_grad_norm is None or not torch.isfinite(torch.tensor(sum_grad_norm)): + print_rank0(" WARNING: sum-reduction grad_norm is invalid") print_rank0(" train_batch PASSED") return True diff --git a/tests/fp8/model_hooks.py b/tests/fp8/model_hooks.py index e95e5ae1ef..a0bfaa2c15 100644 --- a/tests/fp8/model_hooks.py +++ b/tests/fp8/model_hooks.py @@ -14,8 +14,9 @@ print_gemm_profile, ) +from areal.api import LossReduction from areal.engine import MegatronEngine -from areal.engine.core.train_engine import compute_total_loss_weight +from areal.engine.core.train_engine import compute_global_normalizers from areal.engine.megatron_utils.megatron import get_named_parameters from areal.utils import logging @@ -120,10 +121,8 @@ def collect_gradients_after_train_batch( for model in engine.model: model.zero_grad_buffer() - # Step 1: Prepare micro-batches mb_list = engine._prepare_mb_list(input_).to(engine.device) - # Step 2: Define loss functions def sft_loss_fn(logprobs, entropy, input_): """SFT loss function based on compute_packed_sft_loss.""" del entropy # SFT does not use entropy @@ -146,16 +145,21 @@ def sft_loss_fn(logprobs, entropy, input_): loss = -logprobs.sum() / num_valid return loss - def loss_weight_fn(mb): - """Loss weight function based on number of valid tokens.""" + def normalizer_fn(mb): + """Loss normalizer based on number of valid tokens.""" return mb["loss_mask"].count_nonzero() - # Step 3: Compute total loss weight - total_loss_weight = compute_total_loss_weight( - mb_list, loss_weight_fn, mpu.get_data_parallel_group() + loss_reduction = LossReduction.mean( + loss_fn=sft_loss_fn, + normalizer_fn=normalizer_fn, + ) + + global_normalizers = compute_global_normalizers( + mb_list, + loss_reduction, + mpu.get_data_parallel_group(with_context_parallel=True), ) - # Step 4: Forward-backward using Megatron's pipeline function loss_multiplier = ( mpu.get_data_parallel_world_size() * engine.optimizer.get_loss_scale().item() ) @@ -164,9 +168,8 @@ def process_output(output: torch.Tensor, inputs: dict[str, Any]) -> torch.Tensor return engine._compute_logprobs_and_loss( output, inputs, - loss_fn=sft_loss_fn, - loss_weight_fn=loss_weight_fn, - total_loss_weight=total_loss_weight, + loss_reduction=loss_reduction, + global_normalizers=global_normalizers, loss_multiplier=loss_multiplier, ) @@ -190,7 +193,6 @@ def process_output(output: torch.Tensor, inputs: dict[str, Any]) -> torch.Tensor else: engine.forward_backward_batch(mb_list, process_output, forward_only=False) - # Step 5: Collect gradients before optimizer.step() gradients = {} for name, param in get_named_parameters(engine.model, num_experts=None): if param.requires_grad: @@ -545,16 +547,21 @@ def sft_loss_fn(logprobs, entropy, input_): loss = -logprobs.sum() / num_valid return loss - def loss_weight_fn(mb): + def normalizer_fn(mb): return mb["loss_mask"].count_nonzero() + loss_reduction = LossReduction.mean( + loss_fn=sft_loss_fn, + normalizer_fn=normalizer_fn, + ) + # Use engine's train_batch but collect gradients before optimizer step engine.optimizer.zero_grad() for model_chunk in engine.model: model_chunk.zero_grad_buffer() # Forward and backward - engine.train_batch(input_, sft_loss_fn, loss_weight_fn) + engine.train_batch(input_, loss_reduction=loss_reduction) # Collect gradients from all components (focusing on the selected layers) model = get_model_from_engine(engine) diff --git a/tests/test_cispo_loss.py b/tests/test_cispo_loss.py index 44745a5a31..da666319cd 100644 --- a/tests/test_cispo_loss.py +++ b/tests/test_cispo_loss.py @@ -1,29 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 -"""CISPO surrogate (MiniMax-M1 Eq. 4-5) invariants. - -Two defining properties, each across a PPO-like band (0.2 / 0.28) and the wide -MiniMax-M1 band (1.0 / 4.0): - -1. closed-form surrogate value + which tokens the clip flags; -2. gradient routing -- ``logprobs.grad == -sg(clip(ratio)) * A / N`` exactly, - with zero gradient through the importance-ratio path (the test fails if the - stop-gradient detach is dropped). -""" - import pytest import torch from areal.api.cli_args import PPOActorConfig, RejectionSamplingConfig -from areal.utils.functional import cispo_loss_fn +from areal.utils.functional import apply_rejection_sampling, cispo_loss_fn -# (eps_clip, eps_clip_higher): a PPO-like band that clips most tokens, and the -# wide MiniMax-M1 band [0, 5] that clips none of the fixture ratios. BANDS = [(0.2, 0.28), (1.0, 4.0)] def _inputs(): - # log_ratio spans below / inside / above a tight band. The last token is - # masked out, so it must contribute neither loss nor gradient. log_ratio = torch.tensor([-2.0, -0.5, 0.0, 0.3, 1.5, 0.1]) logprobs = torch.tensor([-1.0, -2.0, -0.3, -1.2, -0.7, -2.5], requires_grad=True) proximal = logprobs.detach() - log_ratio # logprobs - proximal == log_ratio @@ -54,9 +39,7 @@ def test_cispo_closed_form_value_and_clip_mask(eps_clip, eps_clip_higher): ) torch.testing.assert_close(loss, expected) - # importance_weight logs the *unclipped* ratio for diagnostics. torch.testing.assert_close(stat["importance_weight"], ratio) - # clip_mask flags band-exit on either side, intersected with the loss mask. expected_clip = (ratio != ratio_clipped) & loss_mask assert torch.equal(stat["clip_mask"], expected_clip) assert not stat["dual_clip_mask"].any() @@ -65,8 +48,6 @@ def test_cispo_closed_form_value_and_clip_mask(eps_clip, eps_clip_higher): @pytest.mark.parametrize("eps_clip,eps_clip_higher", BANDS) def test_cispo_gradient_routes_through_logprobs_only(eps_clip, eps_clip_higher): logprobs, proximal, advantages, loss_mask, log_ratio = _inputs() - # Make the proximal logp a grad-tracking leaf: a correct stop-gradient must - # leave it with no gradient. proximal = proximal.clone().requires_grad_(True) loss, _ = cispo_loss_fn( @@ -82,14 +63,10 @@ def test_cispo_gradient_routes_through_logprobs_only(eps_clip, eps_clip_higher): ratio = torch.exp(log_ratio) ratio_clipped = ratio.clamp(1.0 - eps_clip, 1.0 + eps_clip_higher) n = loss_mask.count_nonzero() - # Gradient flows ONLY through the explicit `logprobs` factor; the clipped - # ratio is a stop-gradient constant. If the detach were dropped, ratio would - # depend on logprobs and add a term, breaking this exact equality. expected_grad = torch.where( loss_mask, -ratio_clipped * advantages / n, torch.zeros_like(ratio) ) torch.testing.assert_close(logprobs.grad, expected_grad) - # No gradient leaks to the importance-ratio (proximal) path. assert proximal.grad is None @@ -108,11 +85,8 @@ def test_cispo_rejects_nonpositive_eps_clip_higher(): def test_cispo_decoupled_applies_behave_imp_weight(): - # Decoupled loss: surrogate rescaled by detached pi_proximal/pi_behave; grad - # still routes only through logprobs. eps_clip, eps_clip_higher = 1.0, 4.0 logprobs, proximal, advantages, loss_mask, log_ratio = _inputs() - # upper is wide so clamp is a no-op and behave_imp_weight == exp(behave_log_ratio). behave_log_ratio = torch.tensor([0.2, -0.4, 0.0, 0.1, -0.3, 0.5]) old_logprobs = proximal - behave_log_ratio # proximal - old == behave_log_ratio rs = RejectionSamplingConfig( @@ -140,7 +114,6 @@ def test_cispo_decoupled_applies_behave_imp_weight(): per_token = -ratio_clipped * advantages * logprobs.detach() * behave_w expected = torch.where(loss_mask, per_token, torch.zeros_like(per_token)).sum() / n torch.testing.assert_close(loss, expected) - # behave_imp_weight is zeroed on masked tokens, so compare valid positions. torch.testing.assert_close( stat["behave_imp_weight"][loss_mask], behave_w[loss_mask] ) @@ -155,19 +128,81 @@ def test_cispo_decoupled_applies_behave_imp_weight(): assert old_leaf.grad is None +def test_cispo_respects_sequence_loss_aggregation(): + logprobs = torch.tensor([[-2.0, -2.0], [-4.0, -1.0]], requires_grad=True) + proximal = logprobs.detach() + advantages = torch.ones_like(logprobs) + loss_mask = torch.tensor([[1, 1], [1, 0]], dtype=torch.bool) + + loss, _ = cispo_loss_fn( + logprobs=logprobs, + proximal_logprobs=proximal, + advantages=advantages, + eps_clip=1.0, + eps_clip_higher=4.0, + loss_mask=loss_mask, + loss_aggregation="seq_mean", + ) + + torch.testing.assert_close(loss, torch.tensor(3.0)) + + +def test_cispo_prompt_mean_accepts_partial_group_sizes(): + logprobs = torch.tensor( + [[-2.0, -2.0], [-2.0, -2.0], [-10.0, -10.0]], requires_grad=True + ) + proximal = logprobs.detach() + advantages = torch.ones_like(logprobs) + loss_mask = torch.ones_like(logprobs, dtype=torch.bool) + + loss, _ = cispo_loss_fn( + logprobs=logprobs, + proximal_logprobs=proximal, + advantages=advantages, + eps_clip=1.0, + eps_clip_higher=4.0, + loss_mask=loss_mask, + loss_aggregation="prompt_mean", + group_size=2, + group_sizes=[2, 1], + ) + + torch.testing.assert_close(loss, torch.tensor(6.0)) + + +def test_rejection_sampling_preserves_infinite_log_ratios(): + rs = RejectionSamplingConfig( + level="token", action="mask", metric="ratio", upper=10.0 + ) + proximal = torch.tensor([0.0, -torch.inf, -torch.inf, torch.inf]) + old = torch.tensor([-torch.inf, -torch.inf, 0.0, -torch.inf]) + loss_mask = torch.ones(4, dtype=torch.bool) + + result = apply_rejection_sampling( + proximal_logprobs=proximal, + old_logprobs=old, + loss_mask=loss_mask, + cu_seqlens=None, + config=rs, + ) + + assert torch.equal(result.loss_mask, torch.tensor([False, True, True, False])) + + def test_cispo_config_validation(): - # Requires a positive upper clip. with pytest.raises(ValueError, match="eps_clip_higher"): PPOActorConfig(use_cispo_loss=True, eps_clip_higher=None) - # Mutually exclusive with SAPO. with pytest.raises(ValueError, match="mutually exclusive"): PPOActorConfig(use_cispo_loss=True, use_sapo_loss=True, eps_clip_higher=4.0) - # Token level only. with pytest.raises(ValueError, match="importance_sampling_level"): PPOActorConfig( use_cispo_loss=True, eps_clip_higher=4.0, importance_sampling_level="sequence", ) - # Valid configuration does not raise. + PPOActorConfig( + use_cispo_loss=True, + eps_clip_higher=4.0, + loss_aggregation="seq_mean", + ) PPOActorConfig(use_cispo_loss=True, eps_clip=1.0, eps_clip_higher=4.0) diff --git a/tests/test_dpo.py b/tests/test_dpo.py index f6b9445c58..738e2e37e4 100644 --- a/tests/test_dpo.py +++ b/tests/test_dpo.py @@ -4,7 +4,7 @@ import torch from areal.trainer.dpo.dpo_engine import ( - _dpo_loss_weight, + _dpo_loss_normalizer, _dpo_valid_pairs, compute_dpo_loss, ) @@ -219,11 +219,11 @@ def test_valid_pairs_with_empty(self): assert valid[0].item() is True assert valid[1].item() is False - def test_loss_weight(self): + def test_loss_normalizer(self): cu_seqlens = torch.tensor([0, 5, 10, 15, 20], dtype=torch.int32) input_ = {"cu_seqlens": cu_seqlens} - weight = _dpo_loss_weight(input_) - assert weight.item() == 2.0 + normalizer = _dpo_loss_normalizer(input_) + assert normalizer.item() == 2.0 class TestDPOLossIntraSequenceShift: diff --git a/tests/test_eval_dispatch.py b/tests/test_eval_dispatch.py index 30c40706fa..cd11b19c07 100644 --- a/tests/test_eval_dispatch.py +++ b/tests/test_eval_dispatch.py @@ -7,8 +7,9 @@ import torch.distributed as dist import torch.nn.functional as F +from areal.api import LossReduction from areal.api.cli_args import MicroBatchSpec -from areal.engine.core.train_engine import compute_total_loss_weight +from areal.engine.core.train_engine import compute_global_normalizers from areal.infra.controller.train_controller import ( _dispatch_tensors, _pad_eval_batch, @@ -16,7 +17,7 @@ from areal.trainer.rw.rw_engine import ( RWController, RWEngine, - _rw_loss_weight, + _rw_loss_normalizer, compute_rw_loss, ) from areal.trainer.sft.lm_engine import LMController @@ -203,7 +204,7 @@ def test_pad_eval_batch_default_group_size_unchanged(self): assert len(padded) == 8 assert _count_dummies(padded) == 1 - def test_compute_total_loss_weight_allows_local_zero( + def test_compute_global_normalizers_allows_local_zero( self, monkeypatch: pytest.MonkeyPatch ): mb_list = MicroBatchList( @@ -221,13 +222,45 @@ def _mock_all_reduce( monkeypatch.setattr(dist, "all_reduce", _mock_all_reduce) - total_weight = compute_total_loss_weight( + global_normalizers = compute_global_normalizers( mb_list=mb_list, - loss_weight_fn=lambda _mb: torch.tensor(0.0), + loss_reduction=LossReduction.mean( + loss_fn=lambda: torch.tensor(0.0), + normalizer_fn=lambda _mb: torch.tensor(0.0), + ), dp_group=cast(dist.ProcessGroup, object()), ) - torch.testing.assert_close(total_weight, torch.tensor(3.0), rtol=0.0, atol=0.0) + torch.testing.assert_close( + global_normalizers["loss"], torch.tensor(3.0), rtol=0.0, atol=0.0 + ) + + def test_compute_global_normalizers_rejects_global_zero( + self, monkeypatch: pytest.MonkeyPatch + ): + mb_list = MicroBatchList( + data={}, + mb_spec=MicroBatchSpec(), + mbs=[{"attention_mask": torch.zeros((1, 1), dtype=torch.bool)}], + group_lens=[1], + ) + + def _mock_all_reduce( + tensor: torch.Tensor, group: dist.ProcessGroup | None = None + ): + del tensor, group + + monkeypatch.setattr(dist, "all_reduce", _mock_all_reduce) + + with pytest.raises(RuntimeError, match="Global loss normalizers"): + compute_global_normalizers( + mb_list=mb_list, + loss_reduction=LossReduction.mean( + loss_fn=lambda: torch.tensor(0.0), + normalizer_fn=lambda _mb: torch.tensor(0.0), + ), + dp_group=cast(dist.ProcessGroup, object()), + ) class TestRWDispatchGrouping: @@ -323,12 +356,12 @@ def test_pad_aligns_non_divisible_input( class TestRWDummyPairSemantics: - def test_rw_loss_weight_counts_only_valid_pairs(self): + def test_rw_loss_normalizer_counts_only_valid_pairs(self): input_data = _make_rw_input([5, 4, 0, 0]) - loss_weight = _rw_loss_weight(input_data) + normalizer = _rw_loss_normalizer(input_data) - torch.testing.assert_close(loss_weight, torch.tensor(1.0), rtol=0.0, atol=0.0) + torch.testing.assert_close(normalizer, torch.tensor(1.0), rtol=0.0, atol=0.0) def test_compute_rw_loss_ignores_dummy_pairs_in_loss_and_metrics(self): tracker = DistributedStatsTracker() diff --git a/tests/test_loss_reduction.py b/tests/test_loss_reduction.py new file mode 100644 index 0000000000..268e85e391 --- /dev/null +++ b/tests/test_loss_reduction.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +import torch.distributed as dist + +from areal.api import ( + LOSS_TERM_REDUCTION_SUM, + LossReduction, + LossTerm, +) +from areal.engine.core.train_engine import ( + compute_global_normalizers, + scale_loss_for_reduction, +) +from areal.utils.data import MicroBatchList + + +def _normalizer(_data): + return torch.tensor(1.0) + + +def test_mean_scaling_preserves_local_mean_order(): + reduction = LossReduction.mean( + loss_fn=lambda: torch.tensor(2.0), normalizer_fn=_normalizer + ) + local_normalizer = torch.tensor(3.0) + global_normalizer = torch.tensor(12.0) + loss = torch.tensor(2.0) + + scaled = scale_loss_for_reduction( + loss, + reduction, + {"loss": local_normalizer}, + {"loss": global_normalizer}, + loss_multiplier=1.0, + ) + + torch.testing.assert_close( + scaled, loss * local_normalizer / global_normalizer, rtol=0, atol=0 + ) + + +def test_sum_scaling_uses_global_normalizer_directly(): + reduction = LossReduction.sum( + loss_fn=lambda: torch.tensor(6.0), normalizer_fn=_normalizer + ) + local_sum = torch.tensor(6.0) + + scaled = scale_loss_for_reduction( + local_sum, + reduction, + {"loss": torch.tensor(3.0)}, + {"loss": torch.tensor(12.0)}, + loss_multiplier=1.0, + ) + + torch.testing.assert_close(scaled, local_sum / 12.0, rtol=0, atol=0) + + +def test_local_zero_normalizer_masks_nan_loss_value(): + reduction = LossReduction.mean( + loss_fn=lambda: torch.tensor(float("nan")), normalizer_fn=_normalizer + ) + + scaled = scale_loss_for_reduction( + torch.tensor(float("nan")), + reduction, + {"loss": torch.tensor(0.0)}, + {"loss": torch.tensor(12.0)}, + loss_multiplier=1.0, + ) + + torch.testing.assert_close(scaled, torch.tensor(0.0), rtol=0, atol=0) + + +def test_multi_term_scaling_uses_each_terms_normalizer(): + reduction = LossReduction( + loss_fn=lambda: { + "pg": torch.tensor(6.0), + "kd": torch.tensor(2.0), + }, + terms=( + LossTerm( + "pg", normalizer_fn=_normalizer, reduction=LOSS_TERM_REDUCTION_SUM + ), + LossTerm( + "kd", normalizer_fn=_normalizer, reduction=LOSS_TERM_REDUCTION_SUM + ), + ), + ) + + scaled = scale_loss_for_reduction( + {"pg": torch.tensor(6.0), "kd": torch.tensor(2.0)}, + reduction, + {"pg": torch.tensor(3.0), "kd": torch.tensor(2.0)}, + {"pg": torch.tensor(12.0), "kd": torch.tensor(4.0)}, + loss_multiplier=1.0, + ) + + torch.testing.assert_close(scaled, torch.tensor(1.0), rtol=0, atol=0) + + +def test_global_normalizer_must_be_positive(monkeypatch): + reduction = LossReduction.sum( + loss_fn=lambda: torch.tensor(0.0), + normalizer_fn=lambda data: data["loss_mask"].count_nonzero(), + ) + mb_list = MicroBatchList( + data={}, + mb_spec=None, + mbs=[{"loss_mask": torch.zeros(2, dtype=torch.bool)}], + group_lens=[1], + ) + monkeypatch.setattr(dist, "all_reduce", lambda tensor, group=None: tensor) + + with pytest.raises(RuntimeError, match="Global loss normalizers"): + compute_global_normalizers(mb_list, reduction, dp_group=None) diff --git a/tests/test_megatron_engine.py b/tests/test_megatron_engine.py index 853f64f2dc..a48e01e134 100644 --- a/tests/test_megatron_engine.py +++ b/tests/test_megatron_engine.py @@ -10,7 +10,7 @@ from tests.utils import get_model_path -from areal.api import FinetuneSpec, SaveLoadMeta +from areal.api import FinetuneSpec, LossReduction, SaveLoadMeta from areal.api.alloc_mode import ModelAllocation from areal.api.cli_args import ( MegatronEngineConfig, @@ -70,6 +70,16 @@ def mock_loss_fn( return torch.mean(logprobs) +def mock_loss_sum_fn( + logprobs: torch.Tensor, + entropy: torch.Tensor, + input_data: dict, + **kwargs, +) -> torch.Tensor: + """Mock local numerator loss for testing sum reduction.""" + return torch.sum(logprobs) + + # Cannot use a "module" scope since process groups can only be initialized once. @pytest.fixture def engine(): @@ -115,11 +125,22 @@ def test_simple_train(engine, mock_input): engine.train() train_result = engine.train_batch( mock_input, - loss_fn=mock_loss_fn, - loss_weight_fn=lambda x: torch.tensor(1.0, device=engine.device), + loss_reduction=LossReduction.mean( + loss_fn=mock_loss_fn, + normalizer_fn=lambda x: torch.tensor(1.0, device=engine.device), + ), + ) + sum_train_result = engine.train_batch( + mock_input, + loss_reduction=LossReduction.sum( + loss_fn=mock_loss_sum_fn, + normalizer_fn=lambda x: x["loss_mask"].count_nonzero() + if "loss_mask" in x + else x["cu_seqlens"][-1], + ), ) engine.step_lr_scheduler() - logger.info(f"Train done, result={train_result}") + logger.info(f"Train done, result={train_result}, sum_result={sum_train_result}") @torch.no_grad() diff --git a/tests/test_ppo_stats.py b/tests/test_ppo_stats.py index f2a05cb2e2..7dd1afa8de 100644 --- a/tests/test_ppo_stats.py +++ b/tests/test_ppo_stats.py @@ -1,5 +1,6 @@ from unittest.mock import MagicMock, patch +import pytest import torch from areal.trainer.ppo.actor import grpo_loss_fn @@ -87,6 +88,36 @@ def test_grpo_loss_fn_uses_full_cu_seqlens_for_n_tokens(): assert torch.all(n_tokens) +def test_teacher_logp_requires_token_mean_loss_aggregation(): + input_data = { + "input_ids": torch.tensor([[11, 12]]), + "logprobs": torch.zeros(1, 2), + "advantages": torch.ones(1, 2), + "loss_mask": torch.ones(1, 2, dtype=torch.bool), + "prox_logp": torch.zeros(1, 2), + "versions": torch.zeros(1, 2, dtype=torch.int32), + "teacher_logp": torch.zeros(1, 2), + } + + with patch("areal.trainer.ppo.actor.stats_tracker") as mock_tracker: + mock_tracker.denominator = MagicMock() + mock_tracker.stat = MagicMock() + mock_tracker.scope = MagicMock() + mock_tracker.scope.return_value.__enter__ = MagicMock() + mock_tracker.scope.return_value.__exit__ = MagicMock() + + with pytest.raises(ValueError, match="teacher_logp distillation"): + grpo_loss_fn( + logprobs=torch.zeros(1, 2), + entropy=torch.zeros(1, 2), + input_data=input_data, + eps_clip=0.2, + eps_clip_higher=None, + c_clip=None, + loss_aggregation="seq_mean", + ) + + def test_critic_loss_fn_uses_full_cu_seqlens_for_n_tokens(): input_data = { "input_ids": torch.tensor([11, 12]), diff --git a/tests/test_prompt_mean_loss.py b/tests/test_prompt_mean_loss.py new file mode 100644 index 0000000000..a291b0584d --- /dev/null +++ b/tests/test_prompt_mean_loss.py @@ -0,0 +1,609 @@ +# SPDX-License-Identifier: Apache-2.0 +import pytest +import torch + +from areal.api.cli_args import ( + GenerationHyperparameters, + GRPOConfig, + MicroBatchSpec, + PPOActorConfig, +) +from areal.trainer.ppo.actor import _make_actor_loss_normalizer_fn +from areal.utils.constants import ( + PROX_LOGP_METHOD_LOGLINEAR, + PROX_LOGP_METHOD_RECOMPUTE, +) +from areal.utils.data import split_padded_tensor_dict_into_mb_list +from areal.utils.functional import ( + aggregate_pg_loss, + aggregate_pg_loss_sum, + make_pg_loss_normalizer_fn, +) + +PG = torch.tensor([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [3.0, 3.0, 3.0], [3.0, 3.0, 3.0]]) +MASK = torch.tensor( + [[1.0, 1.0, 1.0], [1.0, 0.0, 0.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]] +) +PROMPT_MEAN = 2.0 +TOKEN_MEAN = 2.2 +CONSTANT_DIVISOR = 10.0 +CONSTANT_MEAN = 0.55 + +SEQ_PG = torch.tensor([[2.0, 2.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0]]) +SEQ_MASK = torch.tensor([[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0]]) + + +def test_token_mean_is_global_token_average(): + loss = aggregate_pg_loss(PG, MASK, loss_aggregation="token_mean") + torch.testing.assert_close(loss, torch.tensor(TOKEN_MEAN)) + + +def test_seq_mean_weights_each_sequence_equally(): + loss = aggregate_pg_loss(SEQ_PG, SEQ_MASK, loss_aggregation="seq_mean") + torch.testing.assert_close(loss, torch.tensor(1.5)) + token = aggregate_pg_loss(SEQ_PG, SEQ_MASK, loss_aggregation="token_mean") + torch.testing.assert_close(token, torch.tensor(8.0 / 6.0)) + assert not torch.allclose(loss, token) + + +def test_seq_mean_packed_matches_padded(): + pg = torch.tensor([2.0, 2.0, 1.0, 1.0, 1.0, 1.0]) + mask = torch.ones(6) + cu_seqlens = torch.tensor([0, 2, 6], dtype=torch.int32) + loss = aggregate_pg_loss( + pg, mask, loss_aggregation="seq_mean", cu_seqlens=cu_seqlens + ) + torch.testing.assert_close(loss, torch.tensor(1.5)) + + +def test_prompt_mean_weights_each_group_equally_2d(): + loss = aggregate_pg_loss(PG, MASK, loss_aggregation="prompt_mean", group_size=2) + torch.testing.assert_close(loss, torch.tensor(PROMPT_MEAN)) + assert not torch.allclose(loss, torch.tensor(TOKEN_MEAN)) + + +def test_prompt_mean_packed_matches_padded(): + pg = PG.reshape(-1) + mask = MASK.reshape(-1) + cu_seqlens = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32) + loss = aggregate_pg_loss( + pg, mask, loss_aggregation="prompt_mean", group_size=2, cu_seqlens=cu_seqlens + ) + torch.testing.assert_close(loss, torch.tensor(PROMPT_MEAN)) + + +def test_prompt_mean_accepts_partial_group_sizes(): + pg = torch.tensor([[2.0, 2.0], [2.0, 2.0], [10.0, 10.0]]) + mask = torch.ones_like(pg) + + loss = aggregate_pg_loss( + pg, + mask, + loss_aggregation="prompt_mean", + group_size=2, + group_sizes=[2, 1], + ) + + torch.testing.assert_close(loss, torch.tensor(6.0)) + + +def test_prompt_mean_partial_group_packed_matches_padded(): + pg = torch.tensor([[2.0, 2.0], [2.0, 2.0], [10.0, 10.0]]) + mask = torch.ones_like(pg) + cu_seqlens = torch.tensor([0, 2, 4, 6], dtype=torch.int32) + + padded = aggregate_pg_loss( + pg, + mask, + loss_aggregation="prompt_mean", + group_size=2, + group_sizes=[2, 1], + ) + packed = aggregate_pg_loss( + pg.reshape(-1), + mask.reshape(-1), + loss_aggregation="prompt_mean", + group_size=2, + cu_seqlens=cu_seqlens, + group_sizes=[2, 1], + ) + + torch.testing.assert_close(packed, padded) + + +def test_constant_normalizes_token_sum_by_fixed_sequence_divisor(): + loss = aggregate_pg_loss( + PG, + MASK, + loss_aggregation="constant", + loss_aggregation_divisor=CONSTANT_DIVISOR, + ) + torch.testing.assert_close(loss, torch.tensor(CONSTANT_MEAN)) + + +def test_constant_packed_matches_padded(): + pg = PG.reshape(-1) + mask = MASK.reshape(-1) + cu_seqlens = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32) + loss = aggregate_pg_loss( + pg, + mask, + loss_aggregation="constant", + loss_aggregation_divisor=CONSTANT_DIVISOR, + cu_seqlens=cu_seqlens, + ) + torch.testing.assert_close(loss, torch.tensor(CONSTANT_MEAN)) + + +@pytest.mark.parametrize( + "aggregation", ["token_mean", "seq_mean", "prompt_mean", "constant"] +) +def test_loss_normalizer_pairing_realizes_global_mean(aggregation): + group_size = 2 if aggregation == "prompt_mean" else 1 + divisor = CONSTANT_DIVISOR if aggregation == "constant" else None + normalizer_fn = make_pg_loss_normalizer_fn(aggregation, group_size) + + full = aggregate_pg_loss( + PG, + MASK, + loss_aggregation=aggregation, + group_size=group_size, + loss_aggregation_divisor=divisor, + ) + + num = torch.tensor(0.0) + den = torch.tensor(0.0) + for s in (slice(0, 2), slice(2, 4)): + mb_pg, mb_mask = PG[s], MASK[s] + loss_mb = aggregate_pg_loss( + mb_pg, + mb_mask, + loss_aggregation=aggregation, + group_size=group_size, + loss_aggregation_divisor=divisor, + ) + normalizer = normalizer_fn({"loss_mask": mb_mask}) + num = num + loss_mb * normalizer + den = den + normalizer + torch.testing.assert_close(num / den, full) + + +@pytest.mark.parametrize( + "aggregation", ["token_mean", "seq_mean", "prompt_mean", "constant"] +) +def test_loss_sum_pairing_realizes_global_mean(aggregation): + group_size = 2 if aggregation == "prompt_mean" else 1 + divisor = CONSTANT_DIVISOR if aggregation == "constant" else None + normalizer_fn = make_pg_loss_normalizer_fn(aggregation, group_size) + + full = aggregate_pg_loss( + PG, + MASK, + loss_aggregation=aggregation, + group_size=group_size, + loss_aggregation_divisor=divisor, + ) + + num = torch.tensor(0.0) + den = torch.tensor(0.0) + for s in (slice(0, 2), slice(2, 4)): + mb_pg, mb_mask = PG[s], MASK[s] + num = num + aggregate_pg_loss_sum( + mb_pg, + mb_mask, + loss_aggregation=aggregation, + group_size=group_size, + loss_aggregation_divisor=divisor, + ) + den = den + normalizer_fn({"loss_mask": mb_mask}) + torch.testing.assert_close(num / den, full) + + +@pytest.mark.parametrize( + "aggregation", ["token_mean", "seq_mean", "prompt_mean", "constant"] +) +def test_loss_normalizer_pairing_realizes_global_mean_for_packed_inputs(aggregation): + group_size = 2 if aggregation == "prompt_mean" else 1 + divisor = CONSTANT_DIVISOR if aggregation == "constant" else None + normalizer_fn = make_pg_loss_normalizer_fn(aggregation, group_size) + + pg = PG.reshape(-1) + mask = MASK.reshape(-1) + cu_seqlens = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32) + full = aggregate_pg_loss( + pg, + mask, + loss_aggregation=aggregation, + group_size=group_size, + loss_aggregation_divisor=divisor, + cu_seqlens=cu_seqlens, + ) + + num = torch.tensor(0.0) + den = torch.tensor(0.0) + for s in (slice(0, 6), slice(6, 12)): + mb_pg, mb_mask = pg[s], mask[s] + mb_cu_seqlens = torch.tensor([0, 3, 6], dtype=torch.int32) + loss_mb = aggregate_pg_loss( + mb_pg, + mb_mask, + loss_aggregation=aggregation, + group_size=group_size, + loss_aggregation_divisor=divisor, + cu_seqlens=mb_cu_seqlens, + ) + normalizer = normalizer_fn({"loss_mask": mb_mask, "cu_seqlens": mb_cu_seqlens}) + num = num + loss_mb * normalizer + den = den + normalizer + torch.testing.assert_close(num / den, full) + + +@pytest.mark.parametrize( + "aggregation", ["token_mean", "seq_mean", "prompt_mean", "constant"] +) +def test_loss_sum_pairing_realizes_global_mean_for_packed_inputs(aggregation): + group_size = 2 if aggregation == "prompt_mean" else 1 + divisor = CONSTANT_DIVISOR if aggregation == "constant" else None + normalizer_fn = make_pg_loss_normalizer_fn(aggregation, group_size) + + pg = PG.reshape(-1) + mask = MASK.reshape(-1) + cu_seqlens = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32) + full = aggregate_pg_loss( + pg, + mask, + loss_aggregation=aggregation, + group_size=group_size, + loss_aggregation_divisor=divisor, + cu_seqlens=cu_seqlens, + ) + + num = torch.tensor(0.0) + den = torch.tensor(0.0) + for s in (slice(0, 6), slice(6, 12)): + mb_pg, mb_mask = pg[s], mask[s] + mb_cu_seqlens = torch.tensor([0, 3, 6], dtype=torch.int32) + num = num + aggregate_pg_loss_sum( + mb_pg, + mb_mask, + loss_aggregation=aggregation, + group_size=group_size, + loss_aggregation_divisor=divisor, + cu_seqlens=mb_cu_seqlens, + ) + den = den + normalizer_fn({"loss_mask": mb_mask, "cu_seqlens": mb_cu_seqlens}) + torch.testing.assert_close(num / den, full) + + +def test_denom_mask_uses_pre_rejection_count(): + pg = torch.tensor([[2.0, 2.0, 2.0, 2.0]]) + loss_mask = torch.tensor([[1.0, 1.0, 0.0, 0.0]]) + denom_mask = torch.tensor([[1.0, 1.0, 1.0, 1.0]]) + loss = aggregate_pg_loss( + pg, loss_mask, loss_aggregation="token_mean", denom_mask=denom_mask + ) + torch.testing.assert_close(loss, torch.tensor(1.0)) + without = aggregate_pg_loss(pg, loss_mask, loss_aggregation="token_mean") + torch.testing.assert_close(without, torch.tensor(2.0)) + + +@pytest.mark.parametrize( + ("aggregation", "group_size"), [("seq_mean", 1), ("prompt_mean", 2)] +) +def test_denom_mask_applies_to_unit_mean_denominator(aggregation, group_size): + pg = torch.tensor([[2.0, 2.0, 2.0, 2.0], [4.0, 4.0, 4.0, 4.0]]) + loss_mask = torch.tensor([[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0]]) + denom_mask = torch.ones_like(loss_mask) + + loss = aggregate_pg_loss( + pg, + loss_mask, + loss_aggregation=aggregation, + group_size=group_size, + denom_mask=denom_mask, + ) + torch.testing.assert_close(loss, torch.tensor(1.5)) + + without = aggregate_pg_loss( + pg, loss_mask, loss_aggregation=aggregation, group_size=group_size + ) + torch.testing.assert_close(without, torch.tensor(3.0)) + + +@pytest.mark.parametrize( + ("aggregation", "group_size"), [("seq_mean", 1), ("prompt_mean", 2)] +) +def test_unit_mean_skips_units_without_denominator(aggregation, group_size): + pg = torch.tensor([[2.0, 2.0], [2.0, 2.0], [100.0, 100.0], [100.0, 100.0]]) + loss_mask = torch.tensor([[1.0, 1.0], [1.0, 1.0], [0.0, 0.0], [0.0, 0.0]]) + + loss = aggregate_pg_loss( + pg, + loss_mask, + loss_aggregation=aggregation, + group_size=group_size, + ) + + torch.testing.assert_close(loss, torch.tensor(2.0)) + + +@pytest.mark.parametrize( + ("aggregation", "group_size"), [("seq_mean", 1), ("prompt_mean", 2)] +) +def test_packed_unit_mean_skips_units_without_denominator(aggregation, group_size): + pg = torch.tensor([2.0, 2.0, 2.0, 2.0, 100.0, 100.0, 100.0, 100.0]) + loss_mask = torch.tensor([1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0]) + cu_seqlens = torch.tensor([0, 2, 4, 6, 8], dtype=torch.int32) + + loss = aggregate_pg_loss( + pg, + loss_mask, + loss_aggregation=aggregation, + group_size=group_size, + cu_seqlens=cu_seqlens, + ) + + torch.testing.assert_close(loss, torch.tensor(2.0)) + + +def test_constant_skips_sequences_without_denominator(): + pg = torch.tensor([[2.0, 2.0], [100.0, 100.0]]) + loss_mask = torch.tensor([[1.0, 1.0], [0.0, 0.0]]) + + loss = aggregate_pg_loss( + pg, + loss_mask, + loss_aggregation="constant", + loss_aggregation_divisor=2.0, + ) + + torch.testing.assert_close(loss, torch.tensor(2.0)) + + +def test_packed_constant_skips_sequences_without_denominator(): + pg = torch.tensor([2.0, 2.0, 100.0, 100.0]) + loss_mask = torch.tensor([1.0, 1.0, 0.0, 0.0]) + cu_seqlens = torch.tensor([0, 2, 4], dtype=torch.int32) + + loss = aggregate_pg_loss( + pg, + loss_mask, + loss_aggregation="constant", + loss_aggregation_divisor=2.0, + cu_seqlens=cu_seqlens, + ) + + torch.testing.assert_close(loss, torch.tensor(2.0)) + + +@pytest.mark.parametrize( + ("aggregation", "group_size"), [("seq_mean", 1), ("prompt_mean", 2)] +) +def test_denom_mask_keeps_empty_numerator_units_in_denominator(aggregation, group_size): + pg = torch.tensor([[2.0, 2.0], [4.0, 4.0]]) + loss_mask = torch.tensor([[0.0, 0.0], [1.0, 1.0]]) + denom_mask = torch.ones_like(loss_mask) + + loss = aggregate_pg_loss( + pg, + loss_mask, + loss_aggregation=aggregation, + group_size=group_size, + denom_mask=denom_mask, + ) + without = aggregate_pg_loss( + pg, + loss_mask, + loss_aggregation=aggregation, + group_size=group_size, + ) + + torch.testing.assert_close(loss, torch.tensor(2.0)) + torch.testing.assert_close(without, torch.tensor(4.0)) + + +def test_pg_loss_rejects_broadcastable_loss_mask_shape(): + pg = torch.ones(2, 3) + loss_mask = torch.ones(2, 1) + + with pytest.raises(ValueError, match="loss_mask shape"): + aggregate_pg_loss(pg, loss_mask, loss_aggregation="token_mean") + + +def test_pg_loss_sum_rejects_broadcastable_denom_mask_shape(): + pg = torch.ones(2, 3) + loss_mask = torch.ones_like(pg) + denom_mask = torch.ones(1, 3) + + with pytest.raises(ValueError, match="denom_mask shape"): + aggregate_pg_loss_sum( + pg, + loss_mask, + loss_aggregation="prompt_mean", + group_size=2, + denom_mask=denom_mask, + ) + + +def test_loss_normalizer_counts_active_units(): + seq_normalizer = make_pg_loss_normalizer_fn("seq_mean", 1) + prompt_normalizer = make_pg_loss_normalizer_fn("prompt_mean", 2) + mask = torch.tensor([[1.0, 1.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]) + + torch.testing.assert_close(seq_normalizer({"loss_mask": mask}), torch.tensor(1.0)) + torch.testing.assert_close( + prompt_normalizer({"loss_mask": mask}), torch.tensor(1.0) + ) + + +def test_prompt_normalizer_counts_partial_groups(): + prompt_normalizer = make_pg_loss_normalizer_fn("prompt_mean", 2) + mask = torch.tensor([[1.0, 0.0], [0.0, 0.0], [1.0, 1.0]]) + + torch.testing.assert_close( + prompt_normalizer({"loss_mask": mask, "group_sizes": [2, 1]}), + torch.tensor(2.0), + ) + + +def test_m2po_normalizer_uses_post_filter_mask(): + normalizer = _make_actor_loss_normalizer_fn( + "token_mean", + group_size=1, + m2_threshold=0.1, + prox_logp_method=PROX_LOGP_METHOD_RECOMPUTE, + current_version=3, + ) + data = { + "logprobs": torch.tensor([[0.0, 0.2, 1.0]]), + "prox_logp": torch.zeros(1, 3), + "loss_mask": torch.ones(1, 3, dtype=torch.bool), + } + + torch.testing.assert_close(normalizer(data), torch.tensor(2)) + + +def test_m2po_normalizer_rejects_missing_prox_logp_for_loglinear(): + normalizer = _make_actor_loss_normalizer_fn( + "token_mean", + group_size=1, + m2_threshold=0.1, + prox_logp_method=PROX_LOGP_METHOD_LOGLINEAR, + current_version=3, + ) + data = { + "logprobs": torch.tensor([[0.0, 0.2, 1.0]]), + "loss_mask": torch.ones(1, 3, dtype=torch.bool), + "versions": torch.zeros(1, 3, dtype=torch.long), + } + + with pytest.raises(ValueError, match="m2_threshold requires prox_logp"): + normalizer(data) + + +def test_packed_loss_normalizer_counts_active_units(): + normalizer = make_pg_loss_normalizer_fn("seq_mean", 1) + mask = torch.tensor([1.0, 1.0, 0.0, 0.0]) + cu_seqlens = torch.tensor([0, 2, 4], dtype=torch.int32) + + torch.testing.assert_close( + normalizer({"loss_mask": mask, "cu_seqlens": cu_seqlens}), + torch.tensor(1.0), + ) + + +def test_prompt_mean_group_size_one_equals_seq_mean(): + a = aggregate_pg_loss(PG, MASK, loss_aggregation="prompt_mean", group_size=1) + b = aggregate_pg_loss(PG, MASK, loss_aggregation="seq_mean") + torch.testing.assert_close(a, b) + + +def test_prompt_mean_rejects_ragged_group_count(): + pg = torch.ones(3, 2) + mask = torch.ones(3, 2) + with pytest.raises(ValueError, match="not divisible by group_size"): + aggregate_pg_loss(pg, mask, loss_aggregation="prompt_mean", group_size=2) + + +def test_prompt_sum_pairing_realizes_global_mean_for_partial_groups(): + pg = torch.tensor([[2.0, 2.0], [2.0, 2.0], [10.0, 10.0]]) + mask = torch.ones_like(pg) + normalizer_fn = make_pg_loss_normalizer_fn("prompt_mean", 2) + full = aggregate_pg_loss( + pg, + mask, + loss_aggregation="prompt_mean", + group_size=2, + group_sizes=[2, 1], + ) + + num = torch.tensor(0.0) + den = torch.tensor(0.0) + for seq_slice, group_sizes in ((slice(0, 2), [2]), (slice(2, 3), [1])): + num = num + aggregate_pg_loss_sum( + pg[seq_slice], + mask[seq_slice], + loss_aggregation="prompt_mean", + group_size=2, + group_sizes=group_sizes, + ) + den = den + normalizer_fn( + {"loss_mask": mask[seq_slice], "group_sizes": group_sizes} + ) + + torch.testing.assert_close(num / den, full) + + +def test_split_padded_tensor_dict_preserves_partial_group_sizes(): + data = { + "attention_mask": torch.tensor( + [[1, 1, 0], [1, 1, 0], [1, 1, 1]], dtype=torch.bool + ), + "input_ids": torch.arange(9).view(3, 3), + "loss_mask": torch.tensor([[1, 1, 0], [1, 1, 0], [1, 1, 1]], dtype=torch.bool), + "group_sizes": [2, 1], + } + + mb_list = split_padded_tensor_dict_into_mb_list( + data, MicroBatchSpec(n_mbs=2, granularity=2) + ) + + assert [mb["group_sizes"] for mb in mb_list.mbs] == [[2], [1]] + + +def test_config_derives_group_size_from_n_samples(): + cfg = GRPOConfig( + gconfig=GenerationHyperparameters(n_samples=4), + actor=PPOActorConfig( + loss_aggregation="prompt_mean", mb_spec=MicroBatchSpec(granularity=4) + ), + ) + assert cfg.actor.group_size == 4 + + +def test_config_hand_set_group_size_cannot_silently_take_effect(): + cfg = GRPOConfig( + gconfig=GenerationHyperparameters(n_samples=4), + actor=PPOActorConfig( + loss_aggregation="prompt_mean", + group_size=99, + mb_spec=MicroBatchSpec(granularity=4), + ), + ) + assert cfg.actor.group_size == 4 + + +def test_config_granularity_is_not_mutated_for_prompt_mean(): + cfg = GRPOConfig( + gconfig=GenerationHyperparameters(n_samples=4), + actor=PPOActorConfig( + loss_aggregation="prompt_mean", mb_spec=MicroBatchSpec(granularity=2) + ), + ) + assert cfg.actor.mb_spec.granularity == 2 + + +def test_config_validation(): + with pytest.raises(ValueError, match="n_samples >= 2"): + GRPOConfig( + gconfig=GenerationHyperparameters(n_samples=1), + actor=PPOActorConfig(loss_aggregation="prompt_mean"), + ) + with pytest.raises(ValueError, match="loss_aggregation must be"): + PPOActorConfig(loss_aggregation="bogus") + with pytest.raises(ValueError, match="loss_aggregation_divisor"): + PPOActorConfig(loss_aggregation="constant") + with pytest.raises(ValueError, match="loss_aggregation_divisor"): + PPOActorConfig(loss_aggregation="constant", loss_aggregation_divisor=0) + with pytest.raises(ValueError, match="loss_aggregation_divisor"): + PPOActorConfig( + loss_aggregation="constant", loss_aggregation_divisor=float("inf") + ) + with pytest.raises(ValueError, match="only used"): + PPOActorConfig(loss_aggregation="seq_mean", loss_aggregation_divisor=10) + PPOActorConfig(loss_aggregation="constant", loss_aggregation_divisor=10) + GRPOConfig(gconfig=GenerationHyperparameters(n_samples=1)) + GRPOConfig( + gconfig=GenerationHyperparameters(n_samples=1), + actor=PPOActorConfig(loss_aggregation="seq_mean"), + ) diff --git a/tests/test_rpc_callable_serialization.py b/tests/test_rpc_callable_serialization.py new file mode 100644 index 0000000000..51f67e3d82 --- /dev/null +++ b/tests/test_rpc_callable_serialization.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +import pytest + +from areal.infra.rpc.serialization import ( + deserialize_value, + serialize_value, +) + + +def _sample_rpc_callable(value: int) -> int: + return value + 1 + + +def test_callable_serialization_rejected(): + with pytest.raises(ValueError, match="callable values"): + serialize_value(_sample_rpc_callable) + + +def test_callable_payload_rejected(): + with pytest.raises(ValueError, match="callable values"): + deserialize_value({"type": "callable", "key": "missing"}) diff --git a/tests/test_train_engine.py b/tests/test_train_engine.py index 3a72fe4f0d..937e35b430 100644 --- a/tests/test_train_engine.py +++ b/tests/test_train_engine.py @@ -12,7 +12,7 @@ from tests.utils import get_model_path -from areal.api import FinetuneSpec, SaveLoadMeta +from areal.api import FinetuneSpec, LossReduction, SaveLoadMeta from areal.api.cli_args import ( FSDPEngineConfig, MicroBatchSpec, @@ -177,8 +177,10 @@ def test_eval_batch(engine, mock_input): engine.config.mb_spec = MicroBatchSpec(n_mbs=2, max_tokens_per_mb=100) eval_result = engine.eval_batch( input_=mock_input, - loss_fn=mock_loss_fn, - loss_weight_fn=lambda x: x["cu_seqlens"][-1], + loss_reduction=LossReduction.mean( + loss_fn=mock_loss_fn, + normalizer_fn=lambda x: x["cu_seqlens"][-1], + ), ) assert isinstance(eval_result, torch.Tensor), "Evaluation should return a tensor" assert eval_result.is_cuda, "Evaluation tensor should be on CUDA device" @@ -191,8 +193,10 @@ def test_train_batch(engine, mock_input): engine.config.mb_spec = MicroBatchSpec(n_mbs=2, max_tokens_per_mb=100) train_result = engine.train_batch( input_=mock_input, - loss_fn=mock_loss_fn, - loss_weight_fn=lambda x: x["cu_seqlens"][-1], + loss_reduction=LossReduction.mean( + loss_fn=mock_loss_fn, + normalizer_fn=lambda x: x["cu_seqlens"][-1], + ), ) assert isinstance(train_result, dict), "Training should return a dictionary" assert train_result["grad_norm"] is not None diff --git a/tests/test_tree_training.py b/tests/test_tree_training.py index e88c358033..f20550be00 100644 --- a/tests/test_tree_training.py +++ b/tests/test_tree_training.py @@ -6,7 +6,7 @@ from tests.utils import get_model_path import areal -from areal.api import FinetuneSpec +from areal.api import FinetuneSpec, LossReduction from areal.api.alloc_mode import ModelAllocation from areal.api.cli_args import ( MicroBatchSpec, @@ -313,11 +313,18 @@ def test_tree_training_forward_backward(engine_type, tree_attn_backend): use_triton = tree_attn_backend == "triton" def loss_fn(logprobs, entropy, input_data, **kwargs): - return logprobs.mean() + loss_mask = input_data["loss_mask"].to(dtype=logprobs.dtype) + denom = loss_mask.count_nonzero().clamp_min(1) + return (logprobs * loss_mask).sum() / denom - def loss_weight_fn(input_data): + def normalizer_fn(input_data): return input_data["loss_mask"].count_nonzero() + loss_reduction = LossReduction.mean( + loss_fn=loss_fn, + normalizer_fn=normalizer_fn, + ) + areal.models.tree_attn.tree.USE_TRITON_TREE_ATTN = use_triton areal.models.tree_attn.module_fsdp.USE_TRITON_TREE_ATTN = use_triton areal.models.tree_attn.module_megatron.USE_TRITON_TREE_ATTN = use_triton @@ -328,8 +335,7 @@ def loss_weight_fn(input_data): baseline_engine.train() _ = baseline_engine.train_batch( inputs, - loss_fn=loss_fn, - loss_weight_fn=loss_weight_fn, + loss_reduction=loss_reduction, ) # Collect baseline gradients and parameters @@ -354,8 +360,7 @@ def loss_weight_fn(input_data): tree_engine.train() _ = tree_engine.train_batch( inputs, - loss_fn=loss_fn, - loss_weight_fn=loss_weight_fn, + loss_reduction=loss_reduction, ) if engine_type == "fsdp": diff --git a/tests/torchrun/run_fsdp_dcp_distributed.py b/tests/torchrun/run_fsdp_dcp_distributed.py index 3b5e0f9876..776f014180 100644 --- a/tests/torchrun/run_fsdp_dcp_distributed.py +++ b/tests/torchrun/run_fsdp_dcp_distributed.py @@ -11,7 +11,7 @@ from tests.utils import get_model_path -from areal.api import FinetuneSpec, SaveLoadMeta +from areal.api import FinetuneSpec, LossReduction, SaveLoadMeta from areal.api.alloc_mode import ModelAllocation from areal.api.cli_args import MicroBatchSpec, OptimizerConfig, TrainEngineConfig from areal.engine import FSDPEngine @@ -186,8 +186,13 @@ def test_train_dcp_save_load(alloc_mode: str, output: str | None = None): # Train step 1 engine.train() + loss_reduction = LossReduction.mean( + loss_fn=mock_loss_fn, + normalizer_fn=lambda x: x["cu_seqlens"][-1], + ) train_result = engine.train_batch( - input_data, loss_fn=mock_loss_fn, loss_weight_fn=lambda x: x["cu_seqlens"][-1] + input_data, + loss_reduction=loss_reduction, ) print(f"Rank {rank} train step 1 result: {train_result}") @@ -197,7 +202,8 @@ def test_train_dcp_save_load(alloc_mode: str, output: str | None = None): # Train step 2 train_result = engine.train_batch( - input_data, loss_fn=mock_loss_fn, loss_weight_fn=lambda x: x["cu_seqlens"][-1] + input_data, + loss_reduction=loss_reduction, ) print(f"Rank {rank} train step 2 result: {train_result}") @@ -215,7 +221,8 @@ def test_train_dcp_save_load(alloc_mode: str, output: str | None = None): # Train step 2 again after load engine.train() train_result = engine.train_batch( - input_data, loss_fn=mock_loss_fn, loss_weight_fn=lambda x: x["cu_seqlens"][-1] + input_data, + loss_reduction=loss_reduction, ) print(f"Rank {rank} train step 2 after load result: {train_result}") diff --git a/tests/torchrun/run_fsdp_ulysses_train_batch.py b/tests/torchrun/run_fsdp_ulysses_train_batch.py index 088f7c475d..9fc8dd6bb2 100644 --- a/tests/torchrun/run_fsdp_ulysses_train_batch.py +++ b/tests/torchrun/run_fsdp_ulysses_train_batch.py @@ -7,7 +7,7 @@ from tests.utils import get_model_path -from areal.api import FinetuneSpec, ParallelStrategy +from areal.api import FinetuneSpec, LossReduction, ParallelStrategy from areal.api.cli_args import ( MicroBatchSpec, OptimizerConfig, @@ -149,8 +149,10 @@ def test_ulysses(model_type: str): engine.train() engine.train_batch( input_=input, - loss_fn=mock_loss_fn, - loss_weight_fn=lambda x: x["cu_seqlens"][-1], + loss_reduction=LossReduction.mean( + loss_fn=mock_loss_fn, + normalizer_fn=lambda x: x["cu_seqlens"][-1], + ), ) engine.destroy() else: @@ -159,8 +161,10 @@ def test_ulysses(model_type: str): for input in input_chunks: engine_golden.train_batch( input_=input, - loss_fn=mock_loss_fn, - loss_weight_fn=lambda x: x["cu_seqlens"][-1], + loss_reduction=LossReduction.mean( + loss_fn=mock_loss_fn, + normalizer_fn=lambda x: x["cu_seqlens"][-1], + ), ) engine_golden.destroy() diff --git a/tests/torchrun/run_megatron_engine_distributed.py b/tests/torchrun/run_megatron_engine_distributed.py index 2384d50a90..cfb890986d 100644 --- a/tests/torchrun/run_megatron_engine_distributed.py +++ b/tests/torchrun/run_megatron_engine_distributed.py @@ -9,7 +9,7 @@ from megatron.core import parallel_state as mpu from transformers import AutoTokenizer -from areal.api import FinetuneSpec, SaveLoadMeta +from areal.api import FinetuneSpec, LossReduction, SaveLoadMeta from areal.api.alloc_mode import ModelAllocation from areal.api.cli_args import ( MegatronEngineConfig, @@ -239,6 +239,13 @@ def mock_loss_fn( return torch.mean(logprobs) +def mock_loss_sum_fn( + logprobs: torch.Tensor, entropy: torch.Tensor, input_data: dict, **kwargs +) -> torch.Tensor: + """Mock local numerator loss for testing sum reduction.""" + return torch.sum(logprobs) + + def test_train( model_type: str, alloc_mode: str, output: str | None = None, vpp_size: int = 1 ): @@ -261,11 +268,23 @@ def test_train( train_result = engine.train_batch( input_=bcasted_input, - loss_fn=mock_loss_fn, - loss_weight_fn=lambda x: x["cu_seqlens"][-1], + loss_reduction=LossReduction.mean( + loss_fn=mock_loss_fn, + normalizer_fn=lambda x: x["cu_seqlens"][-1], + ), ) print(f"final rank {rank} train_result: {train_result}") + + sum_train_result = engine.train_batch( + input_=bcasted_input, + loss_reduction=LossReduction.sum( + loss_fn=mock_loss_sum_fn, + normalizer_fn=lambda x: x["cu_seqlens"][-1], + ), + ) + + print(f"final rank {rank} sum_train_result: {sum_train_result}") current_platform.synchronize() dist.barrier() engine.destroy() @@ -325,8 +344,10 @@ def test_grad_norm_mb_invariance( result = engine.train_batch( input_=bcasted_input, - loss_fn=mock_loss_fn, - loss_weight_fn=lambda x: x["cu_seqlens"][-1], + loss_reduction=LossReduction.mean( + loss_fn=mock_loss_fn, + normalizer_fn=lambda x: x["cu_seqlens"][-1], + ), ) print( f"rank {rank} max_tokens_per_mb={max_tokens_per_mb} train_result={result}" @@ -400,8 +421,10 @@ def test_train_dcp_save_load( # train step 1 train_result = engine.train_batch( input_=bcasted_input, - loss_fn=mock_loss_fn, - loss_weight_fn=lambda x: x["cu_seqlens"][-1], + loss_reduction=LossReduction.mean( + loss_fn=mock_loss_fn, + normalizer_fn=lambda x: x["cu_seqlens"][-1], + ), ) print(f"final rank {rank} train_result: {train_result}") @@ -415,8 +438,10 @@ def test_train_dcp_save_load( # train step 2 engine.train_batch( input_=bcasted_input, - loss_fn=mock_loss_fn, - loss_weight_fn=lambda x: x["cu_seqlens"][-1], + loss_reduction=LossReduction.mean( + loss_fn=mock_loss_fn, + normalizer_fn=lambda x: x["cu_seqlens"][-1], + ), ) with torch.no_grad(): @@ -433,8 +458,10 @@ def test_train_dcp_save_load( # train step 2 after recover engine.train_batch( input_=bcasted_input, - loss_fn=mock_loss_fn, - loss_weight_fn=lambda x: x["cu_seqlens"][-1], + loss_reduction=LossReduction.mean( + loss_fn=mock_loss_fn, + normalizer_fn=lambda x: x["cu_seqlens"][-1], + ), ) current_platform.synchronize() @@ -576,8 +603,10 @@ def test_train_hf_save_load( ) train_result = engine.train_batch( input_=bcasted_input, - loss_fn=mock_loss_fn, - loss_weight_fn=lambda x: x["cu_seqlens"][-1], + loss_reduction=LossReduction.mean( + loss_fn=mock_loss_fn, + normalizer_fn=lambda x: x["cu_seqlens"][-1], + ), ) print(f"rank {rank} train_result: {train_result}") current_platform.synchronize() diff --git a/tests/torchrun/run_megatron_engine_vlm_distributed.py b/tests/torchrun/run_megatron_engine_vlm_distributed.py index 00c9faf8ec..904651e16a 100644 --- a/tests/torchrun/run_megatron_engine_vlm_distributed.py +++ b/tests/torchrun/run_megatron_engine_vlm_distributed.py @@ -22,7 +22,7 @@ import torch import torch.distributed as dist -from areal.api import FinetuneSpec, SaveLoadMeta +from areal.api import FinetuneSpec, LossReduction, SaveLoadMeta from areal.api.alloc_mode import ModelAllocation from areal.api.cli_args import ( MegatronEngineConfig, @@ -280,10 +280,16 @@ def test_vlm_train(backend: str, output: str | None = None): engine.train() train_result = engine.train_batch( input_=bcasted_input, - loss_fn=lambda logprobs, entropy, input_data, **kwargs: torch.mean( - logprobs + loss_reduction=LossReduction.mean( + loss_fn=lambda logprobs, entropy, input_data, **kwargs: ( + logprobs + * input_data["loss_mask"].to( + dtype=logprobs.dtype, device=logprobs.device + ) + ).sum() + / input_data["loss_mask"].count_nonzero().clamp_min(1), + normalizer_fn=lambda x: x["loss_mask"].count_nonzero(), ), - loss_weight_fn=lambda x: torch.tensor(1.0, device=engine.device), ) assert "grad_norm" in train_result, f"Missing grad_norm: {train_result}" diff --git a/tests/v2/training_service/fake_train_engine.py b/tests/v2/training_service/fake_train_engine.py index e886a0ccad..028a0c3433 100644 --- a/tests/v2/training_service/fake_train_engine.py +++ b/tests/v2/training_service/fake_train_engine.py @@ -140,21 +140,27 @@ def forward_backward_batch( def train_batch( self, input_: dict[str, Any], - loss_fn=None, - loss_weight_fn=None, + loss_reduction: Any, ) -> dict[str, float]: - return { + result = { "total": _sum_numbers(input_), "version": float(self._version), "train_mode": float(self._train_mode), } + if loss_reduction is not None: + normalizer = loss_reduction.terms[0].normalizer_fn(input_) + if hasattr(normalizer, "item"): + normalizer = normalizer.item() + result["loss_terms"] = float(len(loss_reduction.terms)) + result["normalizer"] = float(normalizer) + return result def eval_batch( self, input_: dict[str, Any], - loss_fn=None, - loss_weight_fn=None, + loss_reduction: Any, ) -> float: + _ = loss_reduction return _sum_numbers(input_) + self._version def forward_batch( @@ -172,11 +178,11 @@ def forward_batch( def train_lm(self, input_, **kwargs): _ = kwargs - return self.train_batch(input_) + return self.train_batch(input_, loss_reduction=None) def evaluate_lm(self, input_, **kwargs): _ = kwargs - return self.eval_batch(input_) + return self.eval_batch(input_, loss_reduction=None) def export_stats(self) -> dict[str, float]: return { diff --git a/tests/v2/training_service/test_worker_unit.py b/tests/v2/training_service/test_worker_unit.py index 69b080625c..407e5b9b1a 100644 --- a/tests/v2/training_service/test_worker_unit.py +++ b/tests/v2/training_service/test_worker_unit.py @@ -99,7 +99,7 @@ def test_train_batch_after_create_engine(self, client): "args": serialize_value( [{"token_ids": [1, 2, 3], "metadata": {"weight": 2.0}}] ), - "kwargs": serialize_value({}), + "kwargs": serialize_value({"loss_reduction": None}), }, ) assert train_resp.status_code == 200 From 924095ebbc188e5a8c2836ef8717d026517dd335 Mon Sep 17 00:00:00 2001 From: EazyReal <8047065+EazyReal@users.noreply.github.com> Date: Fri, 10 Jul 2026 20:27:10 -0700 Subject: [PATCH 2/4] refactor: consolidate loss reductions and preserve engine callbacks Centralize policy-gradient aggregation and route engine loss contracts through one validated reduction. Adapt the original callback-based engine API at the boundary while removing the PR-only helper signatures. Reuse rejection masks for distillation and fail fast for unsupported packed reductions. Key changes:\n- Move generic loss contracts into a dedicated API module\n- Replace duplicated actor aggregation helpers with PolicyGradientReduction\n- Preserve original callback calls across FSDP, Megatron, and Archon\n- Restore the original v2 RPC callback boundary --- areal/api/__init__.py | 8 +- areal/api/engine_api.py | 129 +-- areal/api/loss_api.py | 137 ++++ areal/engine/core/__init__.py | 2 + areal/engine/core/train_engine.py | 27 +- areal/engine/fsdp_engine.py | 19 +- areal/engine/megatron_engine.py | 19 +- areal/experimental/engine/archon_engine.py | 19 +- areal/infra/rpc/serialization.py | 6 - areal/trainer/ppo/actor.py | 89 +- areal/utils/functional/__init__.py | 20 +- areal/utils/functional/functional.py | 773 ++---------------- areal/utils/functional/loss_aggregation.py | 305 +++++++ .../training_service/controller/controller.py | 16 +- tests/test_cispo_loss.py | 11 +- tests/test_loss_reduction.py | 27 + tests/test_ppo_stats.py | 44 +- tests/test_prompt_mean_loss.py | 220 +++-- tests/test_rpc_callable_serialization.py | 21 - .../v2/training_service/fake_train_engine.py | 20 +- tests/v2/training_service/test_worker_unit.py | 2 +- 21 files changed, 955 insertions(+), 959 deletions(-) create mode 100644 areal/api/loss_api.py create mode 100644 areal/utils/functional/loss_aggregation.py delete mode 100644 tests/test_rpc_callable_serialization.py diff --git a/areal/api/__init__.py b/areal/api/__init__.py index aac3612000..a6b5a67f24 100644 --- a/areal/api/__init__.py +++ b/areal/api/__init__.py @@ -32,10 +32,10 @@ _LAZY_IMPORTS = { "TrainEngine": "areal.api.engine_api", - "LossReduction": "areal.api.engine_api", - "LossTerm": "areal.api.engine_api", - "LOSS_TERM_REDUCTION_MEAN": "areal.api.engine_api", - "LOSS_TERM_REDUCTION_SUM": "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", diff --git a/areal/api/engine_api.py b/areal/api/engine_api.py index 060f509649..705f3eaad4 100644 --- a/areal/api/engine_api.py +++ b/areal/api/engine_api.py @@ -3,9 +3,8 @@ from __future__ import annotations import abc -from collections.abc import Callable, Mapping +from collections.abc import Callable from concurrent.futures import Future -from dataclasses import dataclass from typing import TYPE_CHECKING, Any import torch @@ -23,6 +22,10 @@ SaveLoadMeta, WeightUpdateMeta, ) +from areal.api.loss_api import ( + LossReductionInput, + LossWeightFn, +) if TYPE_CHECKING: from areal.api.workflow_api import WorkflowLike @@ -30,96 +33,6 @@ from areal.utils.data import MicroBatchList -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, - ), - ), - ) - - class TrainEngine(abc.ABC): @abc.abstractmethod def create_process_group(self, parallel_strategy: ParallelStrategy | None = None): @@ -454,7 +367,10 @@ def forward_backward_batch( def train_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReduction, + loss_reduction: LossReductionInput | None = None, + loss_weight_fn: LossWeightFn | None = None, + *, + loss_fn: Callable[..., torch.Tensor] | None = None, ) -> dict[str, float]: """Update the model with a batch of data and a loss function. @@ -469,11 +385,12 @@ def train_batch( Preferred format is ``list[dict[str, Any]]`` (trajectory list). Backward compatibility: a pre-batched ``dict[str, Any]`` is also accepted. - loss_reduction : LossReduction - Reduction contract with one or more terms. For ``reduction="sum"``, - the loss term is divided by the global normalizer. For - ``reduction="mean"``, the loss term is a local normalized scalar and - is reweighted by the local normalizer before global normalization. + loss_reduction : LossReductionInput, optional + New reduction contract, or the original positional ``loss_fn``. + loss_weight_fn : Callable, optional + Original AReaL denominator callback, paired with ``loss_fn``. + loss_fn : Callable, optional + Original AReaL keyword loss callback. Returns ------- @@ -488,7 +405,10 @@ def train_batch( def eval_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReduction, + loss_reduction: LossReductionInput | None = None, + loss_weight_fn: LossWeightFn | None = None, + *, + loss_fn: Callable[..., torch.Tensor] | None = None, ) -> torch.Tensor | None: """Evaluate the model using the forward pass and loss function. @@ -503,11 +423,12 @@ def eval_batch( Preferred format is ``list[dict[str, Any]]`` (trajectory list). Backward compatibility: a pre-batched ``dict[str, Any]`` is also accepted. - loss_reduction : LossReduction - Reduction contract with one or more terms. For ``reduction="sum"``, - the loss term is divided by the global normalizer. For - ``reduction="mean"``, the loss term is a local normalized scalar and - is reweighted by the local normalizer before global normalization. + loss_reduction : LossReductionInput, optional + New reduction contract, or the original positional ``loss_fn``. + loss_weight_fn : Callable, optional + Original AReaL denominator callback, paired with ``loss_fn``. + loss_fn : Callable, optional + Original AReaL keyword loss callback. Returns ------- diff --git a/areal/api/loss_api.py b/areal/api/loss_api.py new file mode 100644 index 0000000000..95e088f6dd --- /dev/null +++ b/areal/api/loss_api.py @@ -0,0 +1,137 @@ +# 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, + ), + ), + ) + + +LossReductionInput = LossReduction | Callable[..., torch.Tensor] +LossWeightFn = Callable[[dict[str, Any]], torch.Tensor] + + +def coerce_loss_reduction( + loss_reduction: LossReductionInput | None = None, + loss_weight_fn: LossWeightFn | None = None, + *, + loss_fn: Callable[..., torch.Tensor] | None = None, +) -> LossReduction: + """Normalize the original loss callback API at the engine boundary. + + AReaL originally accepted ``loss_fn`` and ``loss_weight_fn`` separately. + The engine internals now consume ``LossReduction``; this adapter keeps the + original positional and keyword calls working without carrying two code + paths through every backend. + """ + if loss_fn is not None: + if loss_reduction is not None: + raise TypeError("pass either loss_fn or loss_reduction, not both") + loss_reduction = loss_fn + if isinstance(loss_reduction, LossReduction): + if loss_weight_fn is not None: + raise TypeError( + "loss_weight_fn is only valid with the original loss_fn API" + ) + return loss_reduction + if not callable(loss_reduction) or loss_weight_fn is None: + raise TypeError( + "train/eval batch requires LossReduction or both loss_fn and loss_weight_fn" + ) + return LossReduction.mean(loss_reduction, loss_weight_fn) diff --git a/areal/engine/core/__init__.py b/areal/engine/core/__init__.py index a64f4dfebf..13e4655f59 100644 --- a/areal/engine/core/__init__.py +++ b/areal/engine/core/__init__.py @@ -6,12 +6,14 @@ 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", diff --git a/areal/engine/core/train_engine.py b/areal/engine/core/train_engine.py index e1e4c8e4e6..4a87cc4fa3 100644 --- a/areal/engine/core/train_engine.py +++ b/areal/engine/core/train_engine.py @@ -12,7 +12,7 @@ import torch import torch.distributed as dist -from areal.api.engine_api import ( +from areal.api.loss_api import ( LOSS_TERM_REDUCTION_MEAN, LossFnOutput, LossReduction, @@ -27,6 +27,7 @@ ) __all__ = [ + "compute_total_loss_weight", "compute_global_normalizers", "compute_local_normalizers", "scale_loss_for_reduction", @@ -35,6 +36,30 @@ ] +def compute_total_loss_weight( + mb_list: MicroBatchList, + loss_weight_fn: Callable[[dict[str, Any]], torch.Tensor], + dp_group: dist.ProcessGroup, +) -> torch.Tensor: + """Compute total loss weight and all_reduce across data parallel group. + + This is the original pre-PR helper retained for callers that still use the + callback-based engine API. + """ + total_weight = ( + torch.stack([loss_weight_fn(mb) for mb in mb_list.mbs]) + .sum() + .detach() + .clone() + .to(dtype=torch.float32) + ) + dist.all_reduce(total_weight, group=dp_group) + assert total_weight > 0, ( + "Global total loss weight must be positive after all_reduce" + ) + return total_weight + + def compute_local_normalizers( input_: dict[str, Any], loss_reduction: LossReduction ) -> dict[str, torch.Tensor]: diff --git a/areal/engine/fsdp_engine.py b/areal/engine/fsdp_engine.py index 061e8d3e33..21b2b93e52 100644 --- a/areal/engine/fsdp_engine.py +++ b/areal/engine/fsdp_engine.py @@ -59,6 +59,7 @@ ) from areal.api.cli_args import OptimizerConfig, PerfTracerConfig, TrainEngineConfig from areal.api.io_struct import DeviceRuntimeInfo +from areal.api.loss_api import LossReductionInput, LossWeightFn, coerce_loss_reduction from areal.engine.core import ( aggregate_eval_losses, compute_global_normalizers, @@ -763,9 +764,15 @@ def forward_backward_batch( def train_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReduction, + loss_reduction: LossReductionInput | None = None, + loss_weight_fn: LossWeightFn | None = None, + *, + loss_fn: Callable[..., torch.Tensor] | None = None, ) -> dict[str, float]: self._ensure_ready() + loss_reduction = coerce_loss_reduction( + loss_reduction, loss_weight_fn, loss_fn=loss_fn + ) self.optimizer_zero_grad() input_batched, _ = self._normalize_batch_input(input_) @@ -798,9 +805,15 @@ def process_output( def eval_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReduction, + loss_reduction: LossReductionInput | None = None, + loss_weight_fn: LossWeightFn | None = None, + *, + loss_fn: Callable[..., torch.Tensor] | None = None, ) -> torch.Tensor | None: self._ensure_ready() + loss_reduction = coerce_loss_reduction( + loss_reduction, loss_weight_fn, loss_fn=loss_fn + ) input_batched, _ = self._normalize_batch_input(input_) @@ -2048,6 +2061,8 @@ def _compute_logprobs_and_loss( ) -> torch.Tensor: """Compute logprobs/entropy and return scaled loss.""" local_normalizers = compute_local_normalizers(ctx.mb_input, loss_reduction) + if all(normalizer == 0 for normalizer in local_normalizers.values()): + return logits.mean() * 0.0 if self.config.is_critic and self.enable_tree_training: raise NotImplementedError( diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index 6e3163f703..0c0ac9b3cd 100644 --- a/areal/engine/megatron_engine.py +++ b/areal/engine/megatron_engine.py @@ -50,6 +50,7 @@ ) from areal.api.cli_args import MicroBatchSpec, PerfTracerConfig, TrainEngineConfig from areal.api.io_struct import DeviceRuntimeInfo +from areal.api.loss_api import LossReductionInput, LossWeightFn, coerce_loss_reduction from areal.engine.core import ( aggregate_eval_losses, compute_global_normalizers, @@ -993,9 +994,15 @@ def _process_output(input_, output_): def train_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReduction, + loss_reduction: LossReductionInput | None = None, + loss_weight_fn: LossWeightFn | None = None, + *, + loss_fn: Callable[..., torch.Tensor] | None = None, ) -> dict[str, float]: self._ensure_ready() + loss_reduction = coerce_loss_reduction( + loss_reduction, loss_weight_fn, loss_fn=loss_fn + ) self.optimizer_zero_grad() input_batched, _ = self._normalize_batch_input(input_) @@ -1049,9 +1056,15 @@ def process_output( def eval_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReduction, + loss_reduction: LossReductionInput | None = None, + loss_weight_fn: LossWeightFn | None = None, + *, + loss_fn: Callable[..., torch.Tensor] | None = None, ) -> torch.Tensor | None: self._ensure_ready() + loss_reduction = coerce_loss_reduction( + loss_reduction, loss_weight_fn, loss_fn=loss_fn + ) input_batched, _ = self._normalize_batch_input(input_) @@ -2289,6 +2302,8 @@ def _compute_logprobs_and_loss( loss_multiplier: float = 1.0, ) -> torch.Tensor: local_normalizers = compute_local_normalizers(inputs, loss_reduction) + if all(normalizer == 0 for normalizer in local_normalizers.values()): + return output.mean() * 0.0 if self.config.is_critic and self.enable_tree_training: raise NotImplementedError( diff --git a/areal/experimental/engine/archon_engine.py b/areal/experimental/engine/archon_engine.py index 5e56f81e7e..7d7c73d5e6 100644 --- a/areal/experimental/engine/archon_engine.py +++ b/areal/experimental/engine/archon_engine.py @@ -36,6 +36,7 @@ ) from areal.api.cli_args import MicroBatchSpec from areal.api.io_struct import DeviceRuntimeInfo +from areal.api.loss_api import LossReductionInput, LossWeightFn, coerce_loss_reduction from areal.engine.core.distributed import ( patch_dist_group_timeout, warmup_process_groups, @@ -528,10 +529,16 @@ def forward_backward_batch( def train_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReduction, + loss_reduction: LossReductionInput | None = None, + loss_weight_fn: LossWeightFn | None = None, + *, + loss_fn: Callable[..., torch.Tensor] | None = None, ) -> dict[str, float]: """Train on a batch of data.""" assert self._initialized + loss_reduction = coerce_loss_reduction( + loss_reduction, loss_weight_fn, loss_fn=loss_fn + ) self.optimizer_zero_grad() input_batched, _ = self._normalize_batch_input(input_) @@ -564,10 +571,16 @@ def process_output( def eval_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReduction, + loss_reduction: LossReductionInput | None = None, + loss_weight_fn: LossWeightFn | None = None, + *, + loss_fn: Callable[..., torch.Tensor] | None = None, ) -> torch.Tensor | None: """Evaluate on a batch of data.""" assert self._initialized + loss_reduction = coerce_loss_reduction( + loss_reduction, loss_weight_fn, loss_fn=loss_fn + ) input_batched, _ = self._normalize_batch_input(input_) @@ -1266,6 +1279,8 @@ def _compute_logprobs_and_loss( ) -> torch.Tensor: """Compute logprobs/entropy and return scaled loss.""" local_normalizers = compute_local_normalizers(ctx.mb_input, loss_reduction) + if all(normalizer == 0 for normalizer in local_normalizers.values()): + return logits.mean() * 0.0 if not self.config.is_critic: result = self._gather_actor_train_outputs(logits, ctx) diff --git a/areal/infra/rpc/serialization.py b/areal/infra/rpc/serialization.py index c7c23d2a45..2ea864c838 100644 --- a/areal/infra/rpc/serialization.py +++ b/areal/infra/rpc/serialization.py @@ -649,9 +649,6 @@ def serialize_value(value: Any) -> Any: "value": value.value, } - if callable(value) and not isinstance(value, type): - raise ValueError("RPC serialization does not support callable values.") - # Primitives (int, float, str, bool) pass through unchanged return value @@ -752,9 +749,6 @@ def deserialize_value(value: Any) -> Any: f"Failed to deserialize ray.ObjectRef, treating as regular dict: {e}" ) - if value.get("type") == "callable": - raise ValueError("RPC deserialization does not support callable values.") - # Check for SerializedTensor marker if value.get("type") == "tensor": try: diff --git a/areal/trainer/ppo/actor.py b/areal/trainer/ppo/actor.py index 598b0b0105..db6f82c70e 100644 --- a/areal/trainer/ppo/actor.py +++ b/areal/trainer/ppo/actor.py @@ -32,8 +32,9 @@ from areal.utils.functional import ( LOSS_AGGREGATION_PROMPT_MEAN, LOSS_AGGREGATION_TOKEN_MEAN, + PolicyGradientReduction, + apply_rejection_sampling, cispo_loss_fn, - make_pg_loss_normalizer_fn, ppo_actor_loss_fn, reward_overlong_penalty, sapo_loss_fn, @@ -57,7 +58,10 @@ def _make_actor_loss_normalizer_fn( prox_logp_method: str, current_version: int | None, ): - normalizer_fn = make_pg_loss_normalizer_fn(loss_aggregation, group_size) + normalizer_fn = PolicyGradientReduction( + mode=loss_aggregation, + group_size=group_size, + ).normalizer_fn if m2_threshold is None: return normalizer_fn @@ -408,6 +412,11 @@ def _ppo_update(self, data: dict[str, Any]) -> None: with stats_tracker.scope("update"): current_version = self.engine.get_version() + pg_reduction = PolicyGradientReduction( + mode=self.config.loss_aggregation, + group_size=self.config.group_size, + divisor=self.config.loss_aggregation_divisor, + ) for mb in mb_inputs.mbs: use_sum_reduction = _use_sum_pg_loss(mb, self.m2_threshold) @@ -426,10 +435,8 @@ def _ppo_update(self, data: dict[str, Any]) -> None: 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_aggregation=self.config.loss_aggregation, - group_size=self.config.group_size, - loss_aggregation_divisor=self.config.loss_aggregation_divisor, - return_sum=use_sum_reduction, + pg_reduction=pg_reduction, + local_mean=not use_sum_reduction, ) normalizer_fn = _make_actor_loss_normalizer_fn( self.config.loss_aggregation, @@ -509,10 +516,8 @@ def grpo_loss_fn( sapo_tau_neg: float = 1.05, use_cispo_loss: bool = False, use_decoupled_loss: bool = False, - loss_aggregation: str = "token_mean", - group_size: int = 1, - loss_aggregation_divisor: float | None = None, - return_sum: bool = False, + pg_reduction: PolicyGradientReduction | None = None, + local_mean: bool = True, vocab_min_logits: torch.Tensor | None = None, vocab_max_logits: torch.Tensor | None = None, vocab_mean_logits: torch.Tensor | None = None, @@ -543,6 +548,17 @@ 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() + rejection_sampling_result = None + if rejection_sampling is not None and not use_sapo_loss: + rejection_sampling_result = apply_rejection_sampling( + proximal_logprobs=prox_logp, + old_logprobs=old_logp, + loss_mask=loss_mask, + cu_seqlens=input_data.get("cu_seqlens"), + config=rejection_sampling, + ) + if use_cispo_loss: if use_sapo_loss: raise ValueError( @@ -562,13 +578,11 @@ def grpo_loss_fn( eps_clip_higher=eps_clip_higher, loss_mask=loss_mask, old_logprobs=old_logp, - rejection_sampling=rejection_sampling, + rejection_sampling_result=rejection_sampling_result, cu_seqlens=input_data.get("cu_seqlens"), - loss_aggregation=loss_aggregation, - group_size=group_size, - loss_aggregation_divisor=loss_aggregation_divisor, group_sizes=input_data.get("group_sizes"), - return_sum=return_sum, + pg_reduction=pg_reduction, + local_mean=local_mean, ) elif use_sapo_loss: if use_decoupled_loss: @@ -585,11 +599,9 @@ def grpo_loss_fn( loss_mask=loss_mask, importance_sampling_level=importance_sampling_level, cu_seqlens=input_data.get("cu_seqlens"), - loss_aggregation=loss_aggregation, - group_size=group_size, - loss_aggregation_divisor=loss_aggregation_divisor, group_sizes=input_data.get("group_sizes"), - return_sum=return_sum, + pg_reduction=pg_reduction, + local_mean=local_mean, ) else: loss, stat = ppo_actor_loss_fn( @@ -601,25 +613,28 @@ def grpo_loss_fn( loss_mask=loss_mask, c_clip=c_clip, proximal_logprobs=prox_logp, - rejection_sampling=rejection_sampling, + rejection_sampling_result=rejection_sampling_result, importance_sampling_level=importance_sampling_level, cu_seqlens=input_data.get("cu_seqlens"), - loss_aggregation=loss_aggregation, - group_size=group_size, - loss_aggregation_divisor=loss_aggregation_divisor, group_sizes=input_data.get("group_sizes"), - return_sum=return_sum, + pg_reduction=pg_reduction, + local_mean=local_mean, ) teacher_logp = input_data.get("teacher_logp") + effective_loss_mask = ( + rejection_sampling_result.loss_mask + if rejection_sampling_result is not None + else loss_mask + ) rkl_stat = None if teacher_logp is not None: - if loss_aggregation != LOSS_AGGREGATION_TOKEN_MEAN: + if pg_reduction.mode != LOSS_AGGREGATION_TOKEN_MEAN: raise ValueError( "teacher_logp distillation is only supported with " "loss_aggregation='token_mean'." ) - if return_sum: + if not local_mean: raise ValueError( "return_sum=True is not supported with teacher_logp. Use the " "mean reduction so distillation keeps its own local normalization." @@ -633,23 +648,29 @@ def grpo_loss_fn( rkl_reward = teacher_logp - logprobs.detach() importance_weight = torch.exp(logprobs - old_logp) - rkl_weighted_term = importance_weight * rkl_reward * loss_mask + rkl_weighted_term = importance_weight * rkl_reward * effective_loss_mask kd_coef = -1 * distill_loss_weight - loss = kd_coef * rkl_weighted_term.sum() / loss_mask.sum().clamp(min=1) + loss = ( + kd_coef + * rkl_weighted_term.sum() + / effective_loss_mask.sum().clamp(min=1) + ) rkl_stat = -1 * rkl_weighted_term else: - rkl_penalty_per_token = (logprobs - teacher_logp) * loss_mask - rkl_penalty = rkl_penalty_per_token.sum() / loss_mask.sum().clamp(min=1) + rkl_penalty_per_token = (logprobs - teacher_logp) * effective_loss_mask + rkl_penalty = rkl_penalty_per_token.sum() / effective_loss_mask.sum().clamp( + min=1 + ) loss = rl_loss_weight * loss + distill_loss_weight * rkl_penalty rkl_stat = rkl_penalty_per_token stats_tracker.denominator( - n_tokens=infer_token_denominator(input_data, loss_mask), - n_valid_tokens=loss_mask.bool(), + n_tokens=infer_token_denominator(input_data, effective_loss_mask), + n_valid_tokens=effective_loss_mask.bool(), clipped_tokens=stat["clip_mask"], dual_clipped_tokens=stat["dual_clip_mask"], ) @@ -708,7 +729,7 @@ def grpo_loss_fn( ) # Log proximal approximation metrics - compute_logp_mask = stat.get("behave_mask", loss_mask) + compute_logp_mask = stat.get("behave_mask", effective_loss_mask) _log_proximal_approximation_stats( prox_logp_method=prox_logp_method, prox_logp_gt=prox_logp_gt, @@ -721,7 +742,7 @@ def grpo_loss_fn( # Log version staleness metrics if "versions" in input_data and current_version is not None: - version_metrics_mask = stat.get("behave_mask", loss_mask) + version_metrics_mask = stat.get("behave_mask", effective_loss_mask) _log_version_staleness_stats( versions=input_data["versions"], current_version=current_version, diff --git a/areal/utils/functional/__init__.py b/areal/utils/functional/__init__.py index 2be8dd3cf0..7df2108d98 100644 --- a/areal/utils/functional/__init__.py +++ b/areal/utils/functional/__init__.py @@ -1,25 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 from areal.utils.functional.functional import ( - LOSS_AGGREGATION_CONSTANT, - LOSS_AGGREGATION_PROMPT_MEAN, - LOSS_AGGREGATION_SEQ_MEAN, - LOSS_AGGREGATION_TOKEN_MEAN, - LOSS_AGGREGATIONS_ALL, RejectionSamplingResult, - aggregate_pg_loss, - aggregate_pg_loss_sum, apply_rejection_sampling, cispo_loss_fn, dpo_pair_logratios, dpo_preference_loss, - make_pg_loss_normalizer_fn, masked_normalization, ppo_actor_loss_fn, ppo_critic_loss_fn, reward_overlong_penalty, sapo_loss_fn, ) +from areal.utils.functional.loss_aggregation import ( + LOSS_AGGREGATION_CONSTANT, + LOSS_AGGREGATION_PROMPT_MEAN, + LOSS_AGGREGATION_SEQ_MEAN, + LOSS_AGGREGATION_TOKEN_MEAN, + LOSS_AGGREGATIONS_ALL, + PolicyGradientReduction, +) from areal.utils.functional.vocab_parallel import ( gather_logprobs, gather_logprobs_entropy, @@ -32,14 +32,12 @@ "LOSS_AGGREGATION_SEQ_MEAN", "LOSS_AGGREGATION_TOKEN_MEAN", "LOSS_AGGREGATIONS_ALL", + "PolicyGradientReduction", "RejectionSamplingResult", - "aggregate_pg_loss", - "aggregate_pg_loss_sum", "apply_rejection_sampling", "cispo_loss_fn", "dpo_pair_logratios", "dpo_preference_loss", - "make_pg_loss_normalizer_fn", "masked_normalization", "ppo_actor_loss_fn", "ppo_critic_loss_fn", diff --git a/areal/utils/functional/functional.py b/areal/utils/functional/functional.py index 035c8b6940..61c9de90f4 100644 --- a/areal/utils/functional/functional.py +++ b/areal/utils/functional/functional.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 import functools -import math -from collections.abc import Callable from dataclasses import dataclass from typing import Any @@ -12,6 +10,9 @@ from areal.api.cli_args import RejectionSamplingConfig from areal.utils.data import KLEstimator +from areal.utils.functional.loss_aggregation import ( + PolicyGradientReduction, +) @torch.no_grad() @@ -442,665 +443,51 @@ def apply_rejection_sampling( ) -def compute_binary_kl_divergence( - log_p: torch.Tensor, log_q: torch.Tensor, eps: float = 1e-8 -) -> torch.Tensor: - """KL(P||Q) for Bernoulli distributions parameterized by log-probabilities. - Treats each element as a Bernoulli: P = [p, 1-p], Q = [q, 1-q]. - KL(P||Q) = p*log(p/q) + (1-p)*log((1-p)/(1-q)) - """ - p = torch.clamp(torch.exp(log_p), eps, 1.0 - eps) - q = torch.clamp(torch.exp(log_q), eps, 1.0 - eps) - return p * torch.log(p / q) + (1 - p) * torch.log((1 - p) / (1 - q)) - - -LOSS_AGGREGATION_TOKEN_MEAN = "token_mean" -LOSS_AGGREGATION_SEQ_MEAN = "seq_mean" -LOSS_AGGREGATION_PROMPT_MEAN = "prompt_mean" -LOSS_AGGREGATION_CONSTANT = "constant" - - -_PgLossReducer = Callable[ - [ - torch.Tensor, - torch.Tensor, - torch.Tensor, - int, - float | None, - torch.Tensor | None, - list[int] | None, - str, - ], - torch.Tensor, -] -_PgLossSumReducer = _PgLossReducer -_PgLossNormalizerFn = Callable[[dict[str, Any]], torch.Tensor] -_PgLossNormalizerFactory = Callable[[int], _PgLossNormalizerFn] - - -def _masked_pg_loss(pg_loss: torch.Tensor, num_mask: torch.Tensor) -> torch.Tensor: - return torch.where(num_mask, pg_loss, 0).to(torch.float32) - - -def _resolve_pg_masks( - pg_loss: torch.Tensor, +def _resolve_rejection_sampling_result( + *, + proximal_logprobs: torch.Tensor, + old_logprobs: torch.Tensor | None, loss_mask: torch.Tensor, - denom_mask: torch.Tensor | None, -) -> tuple[torch.Tensor, torch.Tensor]: - if loss_mask.shape != pg_loss.shape: - raise ValueError( - f"loss_mask shape {tuple(loss_mask.shape)} must match " - f"pg_loss shape {tuple(pg_loss.shape)}." - ) - if denom_mask is not None and denom_mask.shape != pg_loss.shape: - raise ValueError( - f"denom_mask shape {tuple(denom_mask.shape)} must match " - f"pg_loss shape {tuple(pg_loss.shape)}." - ) - num_mask = loss_mask.bool() - den_mask = num_mask if denom_mask is None else denom_mask.bool() - return num_mask, den_mask - - -def _sum_active_unit_means(num: torch.Tensor, den: torch.Tensor) -> torch.Tensor: - active = den > 0 - return torch.where(active, num / den.clamp_min(1), torch.zeros_like(num)).sum() - - -def _mean_active_unit_means(num: torch.Tensor, den: torch.Tensor) -> torch.Tensor: - active = den > 0 - active_count = active.count_nonzero().clamp_min(1) - return _sum_active_unit_means(num, den) / active_count - - -def _count_active_packed_sequences( - den_mask: torch.Tensor, - cu_seqlens: torch.Tensor, - n_seqs: int, -) -> torch.Tensor: - seq_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to( - device=den_mask.device, dtype=torch.long - ) - seq_ids = torch.arange(n_seqs, device=den_mask.device).repeat_interleave(seq_lens) - den = torch.zeros(n_seqs, dtype=torch.float32, device=den_mask.device) - den.scatter_add_(0, seq_ids, den_mask.to(dtype=torch.float32)) - return den.count_nonzero() - - -def _validate_pg_group_sizes(n_seqs: int, group_sizes: list[int] | None) -> list[int]: - if group_sizes is None: - raise ValueError("group_sizes must be provided for explicit prompt groups.") - 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 != n_seqs: - raise ValueError(f"group_sizes sum to {total} but sequence count is {n_seqs}.") - return sizes - - -def _group_ids_from_sizes(group_sizes: list[int], device: torch.device) -> torch.Tensor: - sizes = torch.tensor(group_sizes, dtype=torch.long, device=device) - return torch.arange(len(group_sizes), device=device).repeat_interleave(sizes) - - -def _pg_token_normalizer(data: dict[str, Any]) -> torch.Tensor: - return data["loss_mask"].count_nonzero() - - -def _active_sequence_mask(data: dict[str, Any]) -> torch.Tensor: - loss_mask = data["loss_mask"].bool() - cu_seqlens = data.get("cu_seqlens") - if cu_seqlens is None: - return loss_mask.any(dim=-1) - - n_seqs = cu_seqlens.numel() - 1 - seq_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to( - device=loss_mask.device, dtype=torch.long - ) - seq_ids = torch.arange(n_seqs, device=loss_mask.device).repeat_interleave(seq_lens) - active_tokens = torch.zeros(n_seqs, dtype=torch.float32, device=loss_mask.device) - active_tokens.scatter_add_(0, seq_ids, loss_mask.to(dtype=torch.float32)) - return active_tokens > 0 - - -def _pg_unit_normalizer(data: dict[str, Any], unit: int) -> torch.Tensor: - active_seqs = _active_sequence_mask(data) - n_seqs = active_seqs.numel() - group_sizes = data.get("group_sizes") - if group_sizes is not None: - sizes = _validate_pg_group_sizes(n_seqs, group_sizes) - group_ids = _group_ids_from_sizes(sizes, active_seqs.device) - active_groups = torch.zeros( - len(sizes), dtype=torch.float32, device=active_seqs.device - ) - active_groups.scatter_add_(0, group_ids, active_seqs.to(dtype=torch.float32)) - return active_groups.count_nonzero().to(torch.float32) - if n_seqs % unit != 0: - raise ValueError( - f"microbatch sequence count {n_seqs} is not divisible by group_size {unit}." - ) - return active_seqs.view(-1, unit).any(dim=-1).count_nonzero().to(torch.float32) - - -def _make_pg_unit_normalizer_fn(unit: int) -> _PgLossNormalizerFn: - def pg_unit_normalizer(data: dict[str, Any]) -> torch.Tensor: - return _pg_unit_normalizer(data, unit) - - return pg_unit_normalizer - - -def _make_pg_token_normalizer_fn(_group_size: int) -> _PgLossNormalizerFn: - return _pg_token_normalizer - - -def _make_pg_sequence_normalizer_fn(_group_size: int) -> _PgLossNormalizerFn: - return _make_pg_unit_normalizer_fn(unit=1) - - -def _make_pg_prompt_normalizer_fn(group_size: int) -> _PgLossNormalizerFn: - return _make_pg_unit_normalizer_fn(unit=group_size) - - -def _aggregate_token_mean( - pg_loss: torch.Tensor, - num_mask: torch.Tensor, - den_mask: torch.Tensor, - _group_size: int, - _loss_aggregation_divisor: float | None, - _cu_seqlens: torch.Tensor | None, - _group_sizes: list[int] | None, - _mode: str, -) -> torch.Tensor: - denom = den_mask.count_nonzero().clamp_min(1) - return torch.where(num_mask, pg_loss, 0).sum() / denom - - -def _aggregate_token_sum( - pg_loss: torch.Tensor, - num_mask: torch.Tensor, - _den_mask: torch.Tensor, - _group_size: int, - _loss_aggregation_divisor: float | None, - _cu_seqlens: torch.Tensor | None, - _group_sizes: list[int] | None, - _mode: str, -) -> torch.Tensor: - return _masked_pg_loss(pg_loss, num_mask).sum() - - -def _aggregate_constant( - pg_loss: torch.Tensor, - num_mask: torch.Tensor, - den_mask: torch.Tensor, - _group_size: int, - loss_aggregation_divisor: float | None, - cu_seqlens: torch.Tensor | None, - _group_sizes: list[int] | None, - _mode: str, -) -> torch.Tensor: - divisor = loss_aggregation_divisor - assert divisor is not None - - if cu_seqlens is None: - if pg_loss.ndim != 2: - raise ValueError( - "constant expects 2D pg_loss when cu_seqlens is None, " - f"got shape {tuple(pg_loss.shape)}." - ) - active_seqs = den_mask.any(dim=-1).count_nonzero() - else: - if pg_loss.ndim != 1: - raise ValueError( - "constant expects 1D pg_loss when cu_seqlens is provided, " - f"got shape {tuple(pg_loss.shape)}." - ) - n_seqs = cu_seqlens.numel() - 1 - active_seqs = _count_active_packed_sequences(den_mask, cu_seqlens, n_seqs) - return _masked_pg_loss(pg_loss, num_mask).sum() / ( - active_seqs.clamp_min(1) * divisor - ) - - -def _aggregate_constant_sum( - pg_loss: torch.Tensor, - num_mask: torch.Tensor, - _den_mask: torch.Tensor, - _group_size: int, - loss_aggregation_divisor: float | None, cu_seqlens: torch.Tensor | None, - _group_sizes: list[int] | None, - _mode: str, -) -> torch.Tensor: - divisor = loss_aggregation_divisor - assert divisor is not None - - if cu_seqlens is None: - if pg_loss.ndim != 2: + config: RejectionSamplingConfig | None, + result: RejectionSamplingResult | None, +) -> RejectionSamplingResult | None: + """Resolve rejection sampling once when multiple loss terms share a mask.""" + if result is not None: + if config is not None: raise ValueError( - "constant expects 2D pg_loss when cu_seqlens is None, " - f"got shape {tuple(pg_loss.shape)}." + "pass either rejection_sampling or rejection_sampling_result, not both" ) - n_seqs = pg_loss.shape[0] - else: - if pg_loss.ndim != 1: + if old_logprobs is None: raise ValueError( - "constant expects 1D pg_loss when cu_seqlens is provided, " - f"got shape {tuple(pg_loss.shape)}." + "old_logprobs are required when rejection_sampling_result is enabled." ) - n_seqs = cu_seqlens.numel() - 1 - if n_seqs == 0: - return pg_loss.sum() * 0.0 - return _masked_pg_loss(pg_loss, num_mask).sum() / divisor - - -def _aggregate_seq_mean( - pg_loss: torch.Tensor, - num_mask: torch.Tensor, - den_mask: torch.Tensor, - _group_size: int, - _loss_aggregation_divisor: float | None, - cu_seqlens: torch.Tensor | None, - _group_sizes: list[int] | None, - mode: str, -) -> torch.Tensor: - return _aggregate_unit_mean( - pg_loss, - num_mask, - den_mask, - unit=1, - cu_seqlens=cu_seqlens, - group_sizes=None, - mode=mode, - ) - - -def _aggregate_seq_sum( - pg_loss: torch.Tensor, - num_mask: torch.Tensor, - den_mask: torch.Tensor, - _group_size: int, - _loss_aggregation_divisor: float | None, - cu_seqlens: torch.Tensor | None, - _group_sizes: list[int] | None, - mode: str, -) -> torch.Tensor: - return _aggregate_unit_sum( - pg_loss, - num_mask, - den_mask, - unit=1, - cu_seqlens=cu_seqlens, - group_sizes=None, - mode=mode, - ) - - -def _aggregate_prompt_mean( - pg_loss: torch.Tensor, - num_mask: torch.Tensor, - den_mask: torch.Tensor, - group_size: int, - _loss_aggregation_divisor: float | None, - cu_seqlens: torch.Tensor | None, - group_sizes: list[int] | None, - mode: str, -) -> torch.Tensor: - return _aggregate_unit_mean( - pg_loss, - num_mask, - den_mask, - unit=group_size, - cu_seqlens=cu_seqlens, - group_sizes=group_sizes, - mode=mode, - ) - - -def _aggregate_prompt_sum( - pg_loss: torch.Tensor, - num_mask: torch.Tensor, - den_mask: torch.Tensor, - group_size: int, - _loss_aggregation_divisor: float | None, - cu_seqlens: torch.Tensor | None, - group_sizes: list[int] | None, - mode: str, -) -> torch.Tensor: - return _aggregate_unit_sum( - pg_loss, - num_mask, - den_mask, - unit=group_size, - cu_seqlens=cu_seqlens, - group_sizes=group_sizes, - mode=mode, - ) - - -def _aggregate_unit_mean( - pg_loss: torch.Tensor, - num_mask: torch.Tensor, - den_mask: torch.Tensor, - unit: int, - cu_seqlens: torch.Tensor | None, - group_sizes: list[int] | None, - mode: str, -) -> torch.Tensor: - masked_pg = _masked_pg_loss(pg_loss, num_mask) - den_f = den_mask.to(torch.float32) - - if cu_seqlens is None: - return _aggregate_padded_unit_mean( - pg_loss, unit, group_sizes, mode, masked_pg, den_f - ) - return _aggregate_packed_unit_mean( - pg_loss, unit, cu_seqlens, group_sizes, mode, masked_pg, den_f - ) - - -def _aggregate_unit_sum( - pg_loss: torch.Tensor, - num_mask: torch.Tensor, - den_mask: torch.Tensor, - unit: int, - cu_seqlens: torch.Tensor | None, - group_sizes: list[int] | None, - mode: str, -) -> torch.Tensor: - masked_pg = _masked_pg_loss(pg_loss, num_mask) - den_f = den_mask.to(torch.float32) - - if cu_seqlens is None: - return _aggregate_padded_unit_sum( - pg_loss, unit, group_sizes, mode, masked_pg, den_f - ) - return _aggregate_packed_unit_sum( - pg_loss, unit, cu_seqlens, group_sizes, mode, masked_pg, den_f - ) - - -def _aggregate_padded_unit_mean( - pg_loss: torch.Tensor, - unit: int, - group_sizes: list[int] | None, - mode: str, - masked_pg: torch.Tensor, - den_f: torch.Tensor, -) -> torch.Tensor: - if pg_loss.ndim != 2: - raise ValueError( - f"{mode} expects 2D pg_loss when cu_seqlens is None, " - f"got shape {tuple(pg_loss.shape)}." - ) - n, t = pg_loss.shape - if group_sizes is not None: - sizes = _validate_pg_group_sizes(n, group_sizes) - group_ids = _group_ids_from_sizes(sizes, pg_loss.device) - num = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) - den = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) - num.scatter_add_(0, group_ids, masked_pg.sum(dim=1)) - den.scatter_add_(0, group_ids, den_f.sum(dim=1)) - return _mean_active_unit_means(num, den) - if n % unit != 0: - raise ValueError( - f"{mode}: microbatch sequence count {n} is not divisible by " - f"group_size {unit}. Configure mb_spec.granularity >= group_size " - f"so prompt-groups stay contiguous within each microbatch." - ) - g = n // unit - if g == 0: - return pg_loss.sum() * 0.0 - num = masked_pg.view(g, unit, t).sum(dim=(1, 2)) - den = den_f.view(g, unit, t).sum(dim=(1, 2)) - return _mean_active_unit_means(num, den) - - -def _aggregate_padded_unit_sum( - pg_loss: torch.Tensor, - unit: int, - group_sizes: list[int] | None, - mode: str, - masked_pg: torch.Tensor, - den_f: torch.Tensor, -) -> torch.Tensor: - if pg_loss.ndim != 2: - raise ValueError( - f"{mode} expects 2D pg_loss when cu_seqlens is None, " - f"got shape {tuple(pg_loss.shape)}." - ) - n, t = pg_loss.shape - if group_sizes is not None: - sizes = _validate_pg_group_sizes(n, group_sizes) - group_ids = _group_ids_from_sizes(sizes, pg_loss.device) - num = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) - den = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) - num.scatter_add_(0, group_ids, masked_pg.sum(dim=1)) - den.scatter_add_(0, group_ids, den_f.sum(dim=1)) - return _sum_active_unit_means(num, den) - if n % unit != 0: - raise ValueError( - f"{mode}: microbatch sequence count {n} is not divisible by " - f"group_size {unit}. Configure mb_spec.granularity >= group_size " - f"so prompt-groups stay contiguous within each microbatch." - ) - g = n // unit - if g == 0: - return pg_loss.sum() * 0.0 - num = masked_pg.view(g, unit, t).sum(dim=(1, 2)) - den = den_f.view(g, unit, t).sum(dim=(1, 2)) - return _sum_active_unit_means(num, den) - - -def _aggregate_packed_unit_mean( - pg_loss: torch.Tensor, - unit: int, - cu_seqlens: torch.Tensor, - group_sizes: list[int] | None, - mode: str, - masked_pg: torch.Tensor, - den_f: torch.Tensor, -) -> torch.Tensor: - if pg_loss.ndim != 1: - raise ValueError( - f"{mode} expects 1D pg_loss when cu_seqlens is provided, " - f"got shape {tuple(pg_loss.shape)}." - ) - n_seqs = cu_seqlens.numel() - 1 - seq_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to( - device=pg_loss.device, dtype=torch.long - ) - if group_sizes is not None: - sizes = _validate_pg_group_sizes(n_seqs, group_sizes) - seq_group_ids = _group_ids_from_sizes(sizes, pg_loss.device) - group_ids = seq_group_ids.repeat_interleave(seq_lens) - num = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) - den = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) - num.scatter_add_(0, group_ids, masked_pg) - den.scatter_add_(0, group_ids, den_f) - return _mean_active_unit_means(num, den) - if n_seqs % unit != 0: - raise ValueError( - f"{mode}: packed microbatch has {n_seqs} sequences, not " - f"divisible by group_size {unit}. Configure " - f"mb_spec.granularity >= group_size." - ) - g = n_seqs // unit - if g == 0: - return pg_loss.sum() * 0.0 - seq_ids = torch.arange(n_seqs, device=pg_loss.device) - group_ids = (seq_ids // unit).repeat_interleave(seq_lens) - num = torch.zeros(g, dtype=torch.float32, device=pg_loss.device) - den = torch.zeros(g, dtype=torch.float32, device=pg_loss.device) - num.scatter_add_(0, group_ids, masked_pg) - den.scatter_add_(0, group_ids, den_f) - return _mean_active_unit_means(num, den) - - -def _aggregate_packed_unit_sum( - pg_loss: torch.Tensor, - unit: int, - cu_seqlens: torch.Tensor, - group_sizes: list[int] | None, - mode: str, - masked_pg: torch.Tensor, - den_f: torch.Tensor, -) -> torch.Tensor: - if pg_loss.ndim != 1: + return result + if config is None: + return None + if old_logprobs is None: raise ValueError( - f"{mode} expects 1D pg_loss when cu_seqlens is provided, " - f"got shape {tuple(pg_loss.shape)}." + "old_logprobs are required when rejection_sampling is enabled." ) - n_seqs = cu_seqlens.numel() - 1 - seq_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to( - device=pg_loss.device, dtype=torch.long + return apply_rejection_sampling( + proximal_logprobs=proximal_logprobs, + old_logprobs=old_logprobs, + loss_mask=loss_mask, + cu_seqlens=cu_seqlens, + config=config, ) - if group_sizes is not None: - sizes = _validate_pg_group_sizes(n_seqs, group_sizes) - seq_group_ids = _group_ids_from_sizes(sizes, pg_loss.device) - group_ids = seq_group_ids.repeat_interleave(seq_lens) - num = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) - den = torch.zeros(len(sizes), dtype=torch.float32, device=pg_loss.device) - num.scatter_add_(0, group_ids, masked_pg) - den.scatter_add_(0, group_ids, den_f) - return _sum_active_unit_means(num, den) - if n_seqs % unit != 0: - raise ValueError( - f"{mode}: packed microbatch has {n_seqs} sequences, not " - f"divisible by group_size {unit}. Configure " - f"mb_spec.granularity >= group_size." - ) - g = n_seqs // unit - if g == 0: - return pg_loss.sum() * 0.0 - seq_ids = torch.arange(n_seqs, device=pg_loss.device) - group_ids = (seq_ids // unit).repeat_interleave(seq_lens) - num = torch.zeros(g, dtype=torch.float32, device=pg_loss.device) - den = torch.zeros(g, dtype=torch.float32, device=pg_loss.device) - num.scatter_add_(0, group_ids, masked_pg) - den.scatter_add_(0, group_ids, den_f) - return _sum_active_unit_means(num, den) - - -_PG_LOSS_REDUCERS: dict[str, _PgLossReducer] = { - LOSS_AGGREGATION_TOKEN_MEAN: _aggregate_token_mean, - LOSS_AGGREGATION_SEQ_MEAN: _aggregate_seq_mean, - LOSS_AGGREGATION_PROMPT_MEAN: _aggregate_prompt_mean, - LOSS_AGGREGATION_CONSTANT: _aggregate_constant, -} -_PG_LOSS_SUM_REDUCERS: dict[str, _PgLossSumReducer] = { - LOSS_AGGREGATION_TOKEN_MEAN: _aggregate_token_sum, - LOSS_AGGREGATION_SEQ_MEAN: _aggregate_seq_sum, - LOSS_AGGREGATION_PROMPT_MEAN: _aggregate_prompt_sum, - LOSS_AGGREGATION_CONSTANT: _aggregate_constant_sum, -} -_PG_LOSS_NORMALIZER_FACTORIES: dict[str, _PgLossNormalizerFactory] = { - LOSS_AGGREGATION_TOKEN_MEAN: _make_pg_token_normalizer_fn, - LOSS_AGGREGATION_SEQ_MEAN: _make_pg_sequence_normalizer_fn, - LOSS_AGGREGATION_PROMPT_MEAN: _make_pg_prompt_normalizer_fn, - LOSS_AGGREGATION_CONSTANT: _make_pg_sequence_normalizer_fn, -} -LOSS_AGGREGATIONS_ALL = tuple(_PG_LOSS_REDUCERS) - - -def _validate_pg_loss_aggregation_config( - loss_aggregation: str, loss_aggregation_divisor: float | None -) -> None: - if loss_aggregation not in _PG_LOSS_REDUCERS: - raise ValueError( - f"loss_aggregation must be one of {LOSS_AGGREGATIONS_ALL}, " - f"got {loss_aggregation!r}." - ) - - if loss_aggregation == LOSS_AGGREGATION_CONSTANT: - if ( - loss_aggregation_divisor is None - or not math.isfinite(loss_aggregation_divisor) - or loss_aggregation_divisor <= 0 - ): - raise ValueError( - "loss_aggregation_divisor must be a positive finite value " - "when loss_aggregation='constant'." - ) - elif loss_aggregation_divisor is not None: - raise ValueError( - "loss_aggregation_divisor is only used when loss_aggregation='constant'." - ) -def make_pg_loss_normalizer_fn( - loss_aggregation: str, group_size: int -) -> _PgLossNormalizerFn: - """Return the denominator normalizer paired with policy-gradient aggregation.""" - try: - return _PG_LOSS_NORMALIZER_FACTORIES[loss_aggregation](group_size) - except KeyError as e: - raise ValueError( - f"loss_aggregation must be one of {LOSS_AGGREGATIONS_ALL}, " - f"got {loss_aggregation!r}." - ) from e - - -def aggregate_pg_loss( - pg_loss: torch.Tensor, - loss_mask: torch.Tensor, - loss_aggregation: str = LOSS_AGGREGATION_TOKEN_MEAN, - group_size: int = 1, - loss_aggregation_divisor: float | None = None, - cu_seqlens: torch.Tensor | None = None, - denom_mask: torch.Tensor | None = None, - group_sizes: list[int] | None = None, +def compute_binary_kl_divergence( + log_p: torch.Tensor, log_q: torch.Tensor, eps: float = 1e-8 ) -> torch.Tensor: - """Reduce per-token policy-gradient loss to a microbatch scalar. - - ``token_mean`` averages valid tokens. ``seq_mean`` averages per-sequence - token means. ``prompt_mean`` averages prompt-group means, using - ``group_sizes`` for partial groups and ``group_size`` otherwise. - ``constant`` averages each sequence's masked token sum divided by - ``loss_aggregation_divisor``. Pair the returned scalar with - ``make_pg_loss_normalizer_fn`` for distributed reduction. + """KL(P||Q) for Bernoulli distributions parameterized by log-probabilities. + Treats each element as a Bernoulli: P = [p, 1-p], Q = [q, 1-q]. + KL(P||Q) = p*log(p/q) + (1-p)*log((1-p)/(1-q)) """ - _validate_pg_loss_aggregation_config(loss_aggregation, loss_aggregation_divisor) - - num_mask, den_mask = _resolve_pg_masks(pg_loss, loss_mask, denom_mask) - - return _PG_LOSS_REDUCERS[loss_aggregation]( - pg_loss, - num_mask, - den_mask, - group_size, - loss_aggregation_divisor, - cu_seqlens, - group_sizes, - loss_aggregation, - ) - - -def aggregate_pg_loss_sum( - pg_loss: torch.Tensor, - loss_mask: torch.Tensor, - loss_aggregation: str = LOSS_AGGREGATION_TOKEN_MEAN, - group_size: int = 1, - loss_aggregation_divisor: float | None = None, - cu_seqlens: torch.Tensor | None = None, - denom_mask: torch.Tensor | None = None, - group_sizes: list[int] | None = None, -) -> torch.Tensor: - """Reduce per-token policy-gradient loss to a local sum term.""" - _validate_pg_loss_aggregation_config(loss_aggregation, loss_aggregation_divisor) - - num_mask, den_mask = _resolve_pg_masks(pg_loss, loss_mask, denom_mask) - - return _PG_LOSS_SUM_REDUCERS[loss_aggregation]( - pg_loss, - num_mask, - den_mask, - group_size, - loss_aggregation_divisor, - cu_seqlens, - group_sizes, - loss_aggregation, - ) + p = torch.clamp(torch.exp(log_p), eps, 1.0 - eps) + q = torch.clamp(torch.exp(log_q), eps, 1.0 - eps) + return p * torch.log(p / q) + (1 - p) * torch.log((1 - p) / (1 - q)) def ppo_actor_loss_fn( @@ -1113,13 +500,12 @@ def ppo_actor_loss_fn( eps_clip_higher: float | None = None, c_clip: float | None = None, rejection_sampling: RejectionSamplingConfig | None = None, + rejection_sampling_result: RejectionSamplingResult | None = None, importance_sampling_level: str = "token", cu_seqlens: torch.Tensor | None = None, - loss_aggregation: str = LOSS_AGGREGATION_TOKEN_MEAN, - group_size: int = 1, - loss_aggregation_divisor: float | None = None, group_sizes: list[int] | None = None, - return_sum: bool = False, + pg_reduction: PolicyGradientReduction | None = None, + local_mean: bool = True, ) -> tuple[torch.Tensor, dict]: """PPO actor loss function with optional rejection sampling. @@ -1162,14 +548,15 @@ def ppo_actor_loss_fn( orig_loss_mask = loss_mask # === Apply rejection sampling (replaces old compute_behave_imp_weight) === - if rejection_sampling is not None: - rs_result = apply_rejection_sampling( - proximal_logprobs=proximal_logprobs, - old_logprobs=old_logprobs, - loss_mask=loss_mask, - cu_seqlens=cu_seqlens, - config=rejection_sampling, - ) + rs_result = _resolve_rejection_sampling_result( + proximal_logprobs=proximal_logprobs, + old_logprobs=old_logprobs, + loss_mask=loss_mask, + cu_seqlens=cu_seqlens, + config=rejection_sampling, + result=rejection_sampling_result, + ) + if rs_result is not None: # mask mode updates loss_mask; clamp mode keeps it unchanged loss_mask = rs_result.loss_mask behave_imp_weight = rs_result.behave_imp_weight @@ -1211,23 +598,21 @@ def ppo_actor_loss_fn( dual_clip_mask = torch.zeros_like(clip_mask) # Apply behavioural importance weight from rejection sampling - if rejection_sampling is not None: + if rs_result is not None: behave_approx_kl = proximal_logprobs.detach() - old_logprobs.detach() behave_mask = (behave_imp_weight > 0).logical_and(loss_mask.bool()) behave_approx_kl = torch.where(behave_mask, behave_approx_kl, 0.0) pg_loss = pg_loss * behave_imp_weight logging_loss = pg_loss.detach() - aggregate_fn = aggregate_pg_loss_sum if return_sum else aggregate_pg_loss - pg_loss = aggregate_fn( + reduction = pg_reduction or PolicyGradientReduction() + pg_loss = reduction.aggregate( pg_loss, loss_mask, - loss_aggregation=loss_aggregation, - group_size=group_size, - loss_aggregation_divisor=loss_aggregation_divisor, + denominator_mask=orig_loss_mask, cu_seqlens=cu_seqlens, - denom_mask=orig_loss_mask, group_sizes=group_sizes, + local_mean=local_mean, ) clip_mask.logical_and_(loss_mask) dual_clip_mask.logical_and_(loss_mask) @@ -1239,7 +624,7 @@ def ppo_actor_loss_fn( dual_clip_mask=dual_clip_mask, ) - if rejection_sampling is not None: + if rs_result is not None: stat.update( behave_approx_kl=behave_approx_kl.detach(), behave_imp_weight=behave_imp_weight.detach(), @@ -1258,11 +643,9 @@ def sapo_loss_fn( loss_mask: torch.Tensor, importance_sampling_level: str = "token", cu_seqlens: torch.Tensor | None = None, - loss_aggregation: str = LOSS_AGGREGATION_TOKEN_MEAN, - group_size: int = 1, - loss_aggregation_divisor: float | None = None, group_sizes: list[int] | None = None, - return_sum: bool = False, + pg_reduction: PolicyGradientReduction | None = None, + local_mean: bool = True, ) -> tuple[torch.Tensor, dict]: """SAPO (Soft Adaptive Policy Optimization) loss with asymmetric sigmoid gates. @@ -1314,15 +697,13 @@ def sapo_loss_fn( # Compute loss pg_loss = -soft_gate * advantages logging_loss = pg_loss.detach() - aggregate_fn = aggregate_pg_loss_sum if return_sum else aggregate_pg_loss - pg_loss = aggregate_fn( + reduction = pg_reduction or PolicyGradientReduction() + pg_loss = reduction.aggregate( pg_loss, loss_mask, - loss_aggregation=loss_aggregation, - group_size=group_size, - loss_aggregation_divisor=loss_aggregation_divisor, cu_seqlens=cu_seqlens, group_sizes=group_sizes, + local_mean=local_mean, ) stat = dict( @@ -1348,12 +729,11 @@ def cispo_loss_fn( eps_clip_higher: float | None = None, old_logprobs: torch.Tensor | None = None, rejection_sampling: RejectionSamplingConfig | None = None, + rejection_sampling_result: RejectionSamplingResult | None = None, cu_seqlens: torch.Tensor | None = None, - loss_aggregation: str = LOSS_AGGREGATION_TOKEN_MEAN, - group_size: int = 1, - loss_aggregation_divisor: float | None = None, group_sizes: list[int] | None = None, - return_sum: bool = False, + pg_reduction: PolicyGradientReduction | None = None, + local_mean: bool = True, ) -> tuple[torch.Tensor, dict]: """CISPO (Clipped IS-weight Policy Optimization) loss from MiniMax-M1. @@ -1411,14 +791,15 @@ def cispo_loss_fn( ) orig_loss_mask = loss_mask - if rejection_sampling is not None: - rs_result = apply_rejection_sampling( - proximal_logprobs=proximal_logprobs, - old_logprobs=old_logprobs, - loss_mask=loss_mask, - cu_seqlens=cu_seqlens, - config=rejection_sampling, - ) + rs_result = _resolve_rejection_sampling_result( + proximal_logprobs=proximal_logprobs, + old_logprobs=old_logprobs, + loss_mask=loss_mask, + cu_seqlens=cu_seqlens, + config=rejection_sampling, + result=rejection_sampling_result, + ) + if rs_result is not None: loss_mask = rs_result.loss_mask behave_imp_weight = rs_result.behave_imp_weight filtered_fraction = rs_result.filtered_fraction @@ -1430,23 +811,21 @@ def cispo_loss_fn( ratio_clipped = torch.clamp(ratio, 1.0 - eps_clip, 1.0 + eps_clip_higher).detach() pg_loss = -ratio_clipped * advantages * logprobs - if rejection_sampling is not None: + if rs_result is not None: behave_approx_kl = proximal_logprobs.detach() - old_logprobs.detach() behave_mask = (behave_imp_weight > 0).logical_and(loss_mask.bool()) behave_approx_kl = torch.where(behave_mask, behave_approx_kl, 0.0) pg_loss = pg_loss * behave_imp_weight logging_loss = pg_loss.detach() - aggregate_fn = aggregate_pg_loss_sum if return_sum else aggregate_pg_loss - pg_loss = aggregate_fn( + reduction = pg_reduction or PolicyGradientReduction() + pg_loss = reduction.aggregate( pg_loss, loss_mask, - loss_aggregation=loss_aggregation, - group_size=group_size, - loss_aggregation_divisor=loss_aggregation_divisor, + denominator_mask=orig_loss_mask, cu_seqlens=cu_seqlens, - denom_mask=orig_loss_mask, group_sizes=group_sizes, + local_mean=local_mean, ) clip_mask = (ratio_clipped != ratio).logical_and(loss_mask) @@ -1457,7 +836,7 @@ def cispo_loss_fn( clip_mask=clip_mask, dual_clip_mask=torch.zeros_like(loss_mask, dtype=torch.bool), ) - if rejection_sampling is not None: + if rs_result is not None: stat.update( behave_approx_kl=behave_approx_kl.detach(), behave_imp_weight=behave_imp_weight.detach(), diff --git a/areal/utils/functional/loss_aggregation.py b/areal/utils/functional/loss_aggregation.py new file mode 100644 index 0000000000..ecbe4d388c --- /dev/null +++ b/areal/utils/functional/loss_aggregation.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Policy-gradient loss aggregation and distributed normalizer contracts.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any, Literal + +import torch + +LOSS_AGGREGATION_TOKEN_MEAN = "token_mean" +LOSS_AGGREGATION_SEQ_MEAN = "seq_mean" +LOSS_AGGREGATION_PROMPT_MEAN = "prompt_mean" +LOSS_AGGREGATION_CONSTANT = "constant" + +LossAggregationMode = Literal["token_mean", "seq_mean", "prompt_mean", "constant"] +LOSS_AGGREGATIONS_ALL = ( + LOSS_AGGREGATION_TOKEN_MEAN, + LOSS_AGGREGATION_SEQ_MEAN, + LOSS_AGGREGATION_PROMPT_MEAN, + LOSS_AGGREGATION_CONSTANT, +) + +GroupSizes = Sequence[int] | torch.Tensor + + +def _masked_loss(loss: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + return torch.where(mask, loss, 0).to(torch.float32) + + +def _resolve_masks( + loss: torch.Tensor, + loss_mask: torch.Tensor, + denominator_mask: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + if loss.shape != loss_mask.shape: + raise ValueError( + f"loss_mask shape {tuple(loss_mask.shape)} must match " + f"loss shape {tuple(loss.shape)}." + ) + if denominator_mask is not None and loss.shape != denominator_mask.shape: + raise ValueError( + f"denom_mask shape {tuple(denominator_mask.shape)} must match " + f"loss shape {tuple(loss.shape)}." + ) + numerator_mask = loss_mask.bool() + return numerator_mask, ( + numerator_mask if denominator_mask is None else denominator_mask.bool() + ) + + +def _sequence_sums( + values: torch.Tensor, cu_seqlens: torch.Tensor | None +) -> torch.Tensor: + """Sum token values per sequence for padded or packed inputs.""" + if cu_seqlens is None: + if values.ndim != 2: + raise ValueError( + "padded policy-gradient inputs must be 2D, " + f"got shape {tuple(values.shape)}." + ) + return values.sum(dim=-1) + + if values.ndim != 1: + raise ValueError( + "packed policy-gradient inputs must be 1D, " + f"got shape {tuple(values.shape)}." + ) + if cu_seqlens.ndim != 1: + raise ValueError(f"cu_seqlens must be 1D, got shape {tuple(cu_seqlens.shape)}.") + + n_sequences = cu_seqlens.numel() - 1 + sequence_lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).to( + device=values.device, dtype=torch.long + ) + if torch.any(sequence_lengths < 0): + raise ValueError("cu_seqlens must be non-decreasing.") + sequence_ids = torch.arange(n_sequences, device=values.device).repeat_interleave( + sequence_lengths + ) + if sequence_ids.numel() != values.numel(): + raise ValueError( + "cu_seqlens does not describe the packed loss: " + f"expected {sequence_ids.numel()} tokens, got {values.numel()}." + ) + result = torch.zeros(n_sequences, dtype=values.dtype, device=values.device) + return result.scatter_add_(0, sequence_ids, values) + + +def _validate_group_sizes( + n_sequences: int, group_sizes: GroupSizes | None +) -> list[int]: + if group_sizes is None: + raise ValueError("group_sizes are required for explicit prompt groups.") + if torch.is_tensor(group_sizes): + raw_sizes = group_sizes.detach().cpu().tolist() + else: + raw_sizes = list(group_sizes) + sizes = [int(size) for size in raw_sizes] + if any(size <= 0 for size in sizes): + raise ValueError(f"group_sizes must be positive, got {sizes}.") + if sum(sizes) != n_sequences: + raise ValueError( + f"group_sizes sum to {sum(sizes)} but sequence count is {n_sequences}." + ) + return sizes + + +def _group_ids( + n_sequences: int, + unit_size: int, + group_sizes: GroupSizes | None, + device: torch.device, +) -> tuple[torch.Tensor, int]: + if group_sizes is not None: + sizes = _validate_group_sizes(n_sequences, group_sizes) + return ( + torch.arange(len(sizes), device=device).repeat_interleave( + torch.tensor(sizes, device=device) + ), + len(sizes), + ) + if unit_size <= 0: + raise ValueError(f"group_size must be positive, got {unit_size}.") + if n_sequences % unit_size != 0: + raise ValueError( + f"sequence count {n_sequences} is not divisible by group_size {unit_size}." + ) + return ( + torch.arange(n_sequences, device=device) // unit_size, + n_sequences // unit_size, + ) + + +def _reduce_unit_means( + numerator: torch.Tensor, + denominator: torch.Tensor, + *, + local_mean: bool, +) -> torch.Tensor: + active = denominator > 0 + unit_means = torch.where( + active, + numerator / denominator.clamp_min(1), + torch.zeros_like(numerator), + ) + if not local_mean: + return unit_means.sum() + return unit_means.sum() / active.count_nonzero().clamp_min(1) + + +def _aggregate_units( + loss: torch.Tensor, + numerator_mask: torch.Tensor, + denominator_mask: torch.Tensor, + *, + unit_size: int, + group_sizes: GroupSizes | None, + cu_seqlens: torch.Tensor | None, + local_mean: bool, +) -> torch.Tensor: + masked_loss = _masked_loss(loss, numerator_mask) + sequence_numerators = _sequence_sums(masked_loss, cu_seqlens) + sequence_denominators = _sequence_sums( + denominator_mask.to(torch.float32), cu_seqlens + ) + ids, n_units = _group_ids( + sequence_numerators.numel(), + unit_size, + group_sizes, + loss.device, + ) + unit_numerators = torch.zeros(n_units, dtype=torch.float32, device=loss.device) + unit_denominators = torch.zeros(n_units, dtype=torch.float32, device=loss.device) + unit_numerators.scatter_add_(0, ids, sequence_numerators) + unit_denominators.scatter_add_(0, ids, sequence_denominators) + return _reduce_unit_means( + unit_numerators, + unit_denominators, + local_mean=local_mean, + ) + + +@dataclass(frozen=True, slots=True) +class PolicyGradientReduction: + """Validated local and distributed reduction policy for actor losses. + + ``aggregate(..., local_mean=True)`` returns a microbatch mean. With + ``local_mean=False`` it returns the corresponding numerator term expected by + the training engine's global normalizer contract. ``normalizer_fn`` returns + the matching number of active tokens, sequences, or prompt groups. + """ + + mode: str = LOSS_AGGREGATION_TOKEN_MEAN + group_size: int = 1 + divisor: float | None = None + + def __post_init__(self) -> None: + if self.mode not in LOSS_AGGREGATIONS_ALL: + raise ValueError( + f"loss_aggregation must be one of {LOSS_AGGREGATIONS_ALL}, " + f"got {self.mode!r}." + ) + if self.group_size <= 0: + raise ValueError(f"group_size must be positive, got {self.group_size}.") + if self.divisor is not None: + if ( + self.mode != LOSS_AGGREGATION_CONSTANT + or not math.isfinite(self.divisor) + or self.divisor <= 0 + ): + raise ValueError( + "divisor must be a positive finite value and is only valid " + "for loss_aggregation='constant'." + ) + + def _require_divisor(self) -> float: + if self.divisor is None: + raise ValueError( + "a positive divisor is required for loss_aggregation='constant'." + ) + return self.divisor + + def _require_sequence_boundaries( + self, loss_mask: torch.Tensor, cu_seqlens: torch.Tensor | None + ) -> None: + if cu_seqlens is None and loss_mask.ndim == 1: + raise ValueError( + f"loss_aggregation='{self.mode}' requires cu_seqlens for packed " + "inputs; tree-packed training currently supports only token_mean." + ) + + def normalizer_fn(self, data: dict[str, Any]) -> torch.Tensor: + """Return this reduction's active local denominator.""" + loss_mask = data["loss_mask"].bool() + if self.mode == LOSS_AGGREGATION_TOKEN_MEAN: + return loss_mask.count_nonzero() + + self._require_sequence_boundaries(loss_mask, data.get("cu_seqlens")) + sequence_denominators = _sequence_sums( + loss_mask.to(torch.float32), data.get("cu_seqlens") + ) + group_sizes = ( + data.get("group_sizes") + if self.mode == LOSS_AGGREGATION_PROMPT_MEAN + else None + ) + ids, n_units = _group_ids( + sequence_denominators.numel(), + self.group_size, + group_sizes, + loss_mask.device, + ) + unit_denominators = torch.zeros( + n_units, dtype=torch.float32, device=loss_mask.device + ) + unit_denominators.scatter_add_(0, ids, sequence_denominators) + return unit_denominators.count_nonzero().to(torch.float32) + + def aggregate( + self, + loss: torch.Tensor, + loss_mask: torch.Tensor, + *, + denominator_mask: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + group_sizes: GroupSizes | None = None, + local_mean: bool = True, + ) -> torch.Tensor: + """Aggregate a token-shaped policy-gradient loss.""" + numerator_mask, denominator_mask = _resolve_masks( + loss, loss_mask, denominator_mask + ) + if self.mode == LOSS_AGGREGATION_TOKEN_MEAN: + numerator = _masked_loss(loss, numerator_mask).sum() + if not local_mean: + return numerator + return numerator / denominator_mask.count_nonzero().clamp_min(1) + + self._require_sequence_boundaries(loss, cu_seqlens) + if self.mode == LOSS_AGGREGATION_CONSTANT: + divisor = self._require_divisor() + numerator = _masked_loss(loss, numerator_mask).sum() + if not local_mean: + return numerator / divisor + active_sequences = _sequence_sums( + denominator_mask.to(torch.float32), cu_seqlens + ).count_nonzero() + return numerator / (active_sequences.clamp_min(1) * divisor) + + unit_size = self.group_size if self.mode == LOSS_AGGREGATION_PROMPT_MEAN else 1 + return _aggregate_units( + loss, + numerator_mask, + denominator_mask, + unit_size=unit_size, + group_sizes=( + group_sizes if self.mode == LOSS_AGGREGATION_PROMPT_MEAN else None + ), + cu_seqlens=cu_seqlens, + local_mean=local_mean, + ) diff --git a/areal/v2/training_service/controller/controller.py b/areal/v2/training_service/controller/controller.py index dcc09e95e9..16f136095e 100644 --- a/areal/v2/training_service/controller/controller.py +++ b/areal/v2/training_service/controller/controller.py @@ -703,8 +703,8 @@ def _require_list_batch(input_: Any, method_name: str) -> list[dict[str, Any]]: def train_batch( self, input_: list[dict[str, Any]] | None = None, - *, - loss_reduction: Any, + loss_fn: Any = None, + loss_weight_fn: Any = None, ) -> Any: from areal.infra.rpc.serialization import serialize_value @@ -714,7 +714,9 @@ def train_batch( payload = { "args": serialize_value([batch]), - "kwargs": serialize_value({"loss_reduction": loss_reduction}), + "kwargs": serialize_value( + {"loss_fn": loss_fn, "loss_weight_fn": loss_weight_fn} + ), } return self._gateway_post_result("/train_batch", payload) @@ -738,8 +740,8 @@ def forward_batch( def eval_batch( self, input_: list[dict[str, Any]] | None = None, - *, - loss_reduction: Any, + loss_fn: Any = None, + loss_weight_fn: Any = None, ) -> Any: from areal.infra.rpc.serialization import serialize_value @@ -749,7 +751,9 @@ def eval_batch( payload = { "args": serialize_value([batch]), - "kwargs": serialize_value({"loss_reduction": loss_reduction}), + "kwargs": serialize_value( + {"loss_fn": loss_fn, "loss_weight_fn": loss_weight_fn} + ), } return self._gateway_post_result("/eval_batch", payload) diff --git a/tests/test_cispo_loss.py b/tests/test_cispo_loss.py index da666319cd..6d076139a6 100644 --- a/tests/test_cispo_loss.py +++ b/tests/test_cispo_loss.py @@ -3,7 +3,11 @@ import torch from areal.api.cli_args import PPOActorConfig, RejectionSamplingConfig -from areal.utils.functional import apply_rejection_sampling, cispo_loss_fn +from areal.utils.functional import ( + PolicyGradientReduction, + apply_rejection_sampling, + cispo_loss_fn, +) BANDS = [(0.2, 0.28), (1.0, 4.0)] @@ -141,7 +145,7 @@ def test_cispo_respects_sequence_loss_aggregation(): eps_clip=1.0, eps_clip_higher=4.0, loss_mask=loss_mask, - loss_aggregation="seq_mean", + pg_reduction=PolicyGradientReduction(mode="seq_mean"), ) torch.testing.assert_close(loss, torch.tensor(3.0)) @@ -162,8 +166,7 @@ def test_cispo_prompt_mean_accepts_partial_group_sizes(): eps_clip=1.0, eps_clip_higher=4.0, loss_mask=loss_mask, - loss_aggregation="prompt_mean", - group_size=2, + pg_reduction=PolicyGradientReduction(mode="prompt_mean", group_size=2), group_sizes=[2, 1], ) diff --git a/tests/test_loss_reduction.py b/tests/test_loss_reduction.py index 268e85e391..c3a21f958f 100644 --- a/tests/test_loss_reduction.py +++ b/tests/test_loss_reduction.py @@ -9,6 +9,7 @@ LossReduction, LossTerm, ) +from areal.api.loss_api import coerce_loss_reduction from areal.engine.core.train_engine import ( compute_global_normalizers, scale_loss_for_reduction, @@ -20,6 +21,32 @@ def _normalizer(_data): return torch.tensor(1.0) +def test_coerce_loss_reduction_preserves_original_callback_api(): + def loss_fn(*_args): + return torch.tensor(2.0) + + def loss_weight_fn(_data): + return torch.tensor(3.0) + + positional = coerce_loss_reduction(loss_fn, loss_weight_fn) + keyword = coerce_loss_reduction( + loss_fn=loss_fn, + loss_weight_fn=loss_weight_fn, + ) + + for reduction in (positional, keyword): + assert reduction.loss_fn is loss_fn + assert reduction.terms[0].normalizer_fn is loss_weight_fn + assert reduction.terms[0].reduction == "mean" + + +def test_coerce_loss_reduction_rejects_mixed_contracts(): + reduction = LossReduction.mean(lambda: torch.tensor(1.0), _normalizer) + + with pytest.raises(TypeError, match="only valid with the original loss_fn API"): + coerce_loss_reduction(reduction, _normalizer) + + def test_mean_scaling_preserves_local_mean_order(): reduction = LossReduction.mean( loss_fn=lambda: torch.tensor(2.0), normalizer_fn=_normalizer diff --git a/tests/test_ppo_stats.py b/tests/test_ppo_stats.py index 7dd1afa8de..f93bba41d7 100644 --- a/tests/test_ppo_stats.py +++ b/tests/test_ppo_stats.py @@ -3,9 +3,11 @@ import pytest import torch +from areal.api.cli_args import RejectionSamplingConfig from areal.trainer.ppo.actor import grpo_loss_fn from areal.trainer.ppo.critic import ppo_loss_fn from areal.trainer.ppo.stats import infer_token_denominator +from areal.utils.functional import PolicyGradientReduction from areal.utils.stats_tracker import DistributedStatsTracker @@ -114,10 +116,50 @@ def test_teacher_logp_requires_token_mean_loss_aggregation(): eps_clip=0.2, eps_clip_higher=None, c_clip=None, - loss_aggregation="seq_mean", + pg_reduction=PolicyGradientReduction(mode="seq_mean"), ) +def test_teacher_distillation_excludes_rejected_tokens_from_gradient(): + logprobs = torch.zeros(1, 2, requires_grad=True) + input_data = { + "input_ids": torch.tensor([[11, 12]]), + "logprobs": torch.zeros(1, 2), + "advantages": torch.zeros(1, 2), + "loss_mask": torch.ones(1, 2, dtype=torch.bool), + "prox_logp": torch.tensor([[0.0, torch.log(torch.tensor(2.0))]]), + "versions": torch.zeros(1, 2, dtype=torch.int32), + "teacher_logp": torch.ones(1, 2), + "rl_loss_weight": 0.0, + "distill_loss_weight": 1.0, + } + rejection_sampling = RejectionSamplingConfig( + level="token", action="mask", metric="ratio", upper=1.5 + ) + + with patch("areal.trainer.ppo.actor.stats_tracker") as mock_tracker: + mock_tracker.denominator = MagicMock() + mock_tracker.stat = MagicMock() + mock_tracker.scope = MagicMock() + mock_tracker.scope.return_value.__enter__ = MagicMock() + mock_tracker.scope.return_value.__exit__ = MagicMock() + + loss = grpo_loss_fn( + logprobs=logprobs, + entropy=torch.zeros(1, 2), + input_data=input_data, + eps_clip=0.2, + eps_clip_higher=None, + c_clip=None, + rejection_sampling=rejection_sampling, + ) + + loss.backward() + assert logprobs.grad is not None + assert logprobs.grad[0, 0] != 0 + assert logprobs.grad[0, 1] == 0 + + def test_critic_loss_fn_uses_full_cu_seqlens_for_n_tokens(): input_data = { "input_ids": torch.tensor([11, 12]), diff --git a/tests/test_prompt_mean_loss.py b/tests/test_prompt_mean_loss.py index a291b0584d..4ddf1b8b66 100644 --- a/tests/test_prompt_mean_loss.py +++ b/tests/test_prompt_mean_loss.py @@ -12,12 +12,11 @@ from areal.utils.constants import ( PROX_LOGP_METHOD_LOGLINEAR, PROX_LOGP_METHOD_RECOMPUTE, + PROX_LOGP_METHOD_REUSE_TRAIN_LOGP, ) from areal.utils.data import split_padded_tensor_dict_into_mb_list from areal.utils.functional import ( - aggregate_pg_loss, - aggregate_pg_loss_sum, - make_pg_loss_normalizer_fn, + PolicyGradientReduction, ) PG = torch.tensor([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [3.0, 3.0, 3.0], [3.0, 3.0, 3.0]]) @@ -33,31 +32,124 @@ SEQ_MASK = torch.tensor([[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0]]) +def reduce_pg_loss( + pg_loss, + loss_mask, + loss_aggregation="token_mean", + group_size=1, + loss_aggregation_divisor=None, + cu_seqlens=None, + denom_mask=None, + group_sizes=None, +): + return PolicyGradientReduction( + mode=loss_aggregation, + group_size=group_size, + divisor=loss_aggregation_divisor, + ).aggregate( + pg_loss, + loss_mask, + denominator_mask=denom_mask, + cu_seqlens=cu_seqlens, + group_sizes=group_sizes, + ) + + +def reduce_pg_loss_sum( + pg_loss, + loss_mask, + loss_aggregation="token_mean", + group_size=1, + loss_aggregation_divisor=None, + cu_seqlens=None, + denom_mask=None, + group_sizes=None, +): + return PolicyGradientReduction( + mode=loss_aggregation, + group_size=group_size, + divisor=loss_aggregation_divisor, + ).aggregate( + pg_loss, + loss_mask, + denominator_mask=denom_mask, + cu_seqlens=cu_seqlens, + group_sizes=group_sizes, + local_mean=False, + ) + + +def make_normalizer(loss_aggregation, group_size): + return PolicyGradientReduction( + mode=loss_aggregation, + group_size=group_size, + ).normalizer_fn + + def test_token_mean_is_global_token_average(): - loss = aggregate_pg_loss(PG, MASK, loss_aggregation="token_mean") + loss = reduce_pg_loss(PG, MASK, loss_aggregation="token_mean") torch.testing.assert_close(loss, torch.tensor(TOKEN_MEAN)) def test_seq_mean_weights_each_sequence_equally(): - loss = aggregate_pg_loss(SEQ_PG, SEQ_MASK, loss_aggregation="seq_mean") + loss = reduce_pg_loss(SEQ_PG, SEQ_MASK, loss_aggregation="seq_mean") torch.testing.assert_close(loss, torch.tensor(1.5)) - token = aggregate_pg_loss(SEQ_PG, SEQ_MASK, loss_aggregation="token_mean") + token = reduce_pg_loss(SEQ_PG, SEQ_MASK, loss_aggregation="token_mean") torch.testing.assert_close(token, torch.tensor(8.0 / 6.0)) assert not torch.allclose(loss, token) +@pytest.mark.parametrize("aggregation", ["seq_mean", "prompt_mean"]) +def test_inactive_units_do_not_deflate_local_mean(aggregation): + pg = torch.tensor([[2.0, 2.0], [100.0, 100.0], [100.0, 100.0]]) + mask = torch.tensor([[1, 1], [0, 0], [0, 0]], dtype=torch.bool) + kwargs = { + "loss_aggregation": aggregation, + "group_size": 2, + "group_sizes": [2, 1], + } + + loss = reduce_pg_loss(pg, mask, **kwargs) + normalizer = make_normalizer( + aggregation, group_size=2 if aggregation == "prompt_mean" else 1 + )( + { + "loss_mask": mask, + **({"group_sizes": [2, 1]} if aggregation == "prompt_mean" else {}), + } + ) + + torch.testing.assert_close(loss, torch.tensor(2.0)) + torch.testing.assert_close(normalizer, torch.tensor(1.0)) + + +def test_inactive_units_do_not_deflate_packed_prompt_mean(): + pg = torch.tensor([2.0, 2.0, 100.0, 100.0, 100.0, 100.0]) + mask = torch.tensor([1, 1, 0, 0, 0, 0], dtype=torch.bool) + cu_seqlens = torch.tensor([0, 2, 4, 6], dtype=torch.int32) + + loss = reduce_pg_loss( + pg, + mask, + loss_aggregation="prompt_mean", + group_size=2, + group_sizes=[2, 1], + cu_seqlens=cu_seqlens, + ) + + torch.testing.assert_close(loss, torch.tensor(2.0)) + + def test_seq_mean_packed_matches_padded(): pg = torch.tensor([2.0, 2.0, 1.0, 1.0, 1.0, 1.0]) mask = torch.ones(6) cu_seqlens = torch.tensor([0, 2, 6], dtype=torch.int32) - loss = aggregate_pg_loss( - pg, mask, loss_aggregation="seq_mean", cu_seqlens=cu_seqlens - ) + loss = reduce_pg_loss(pg, mask, loss_aggregation="seq_mean", cu_seqlens=cu_seqlens) torch.testing.assert_close(loss, torch.tensor(1.5)) def test_prompt_mean_weights_each_group_equally_2d(): - loss = aggregate_pg_loss(PG, MASK, loss_aggregation="prompt_mean", group_size=2) + loss = reduce_pg_loss(PG, MASK, loss_aggregation="prompt_mean", group_size=2) torch.testing.assert_close(loss, torch.tensor(PROMPT_MEAN)) assert not torch.allclose(loss, torch.tensor(TOKEN_MEAN)) @@ -66,7 +158,7 @@ def test_prompt_mean_packed_matches_padded(): pg = PG.reshape(-1) mask = MASK.reshape(-1) cu_seqlens = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32) - loss = aggregate_pg_loss( + loss = reduce_pg_loss( pg, mask, loss_aggregation="prompt_mean", group_size=2, cu_seqlens=cu_seqlens ) torch.testing.assert_close(loss, torch.tensor(PROMPT_MEAN)) @@ -76,7 +168,7 @@ def test_prompt_mean_accepts_partial_group_sizes(): pg = torch.tensor([[2.0, 2.0], [2.0, 2.0], [10.0, 10.0]]) mask = torch.ones_like(pg) - loss = aggregate_pg_loss( + loss = reduce_pg_loss( pg, mask, loss_aggregation="prompt_mean", @@ -92,14 +184,14 @@ def test_prompt_mean_partial_group_packed_matches_padded(): mask = torch.ones_like(pg) cu_seqlens = torch.tensor([0, 2, 4, 6], dtype=torch.int32) - padded = aggregate_pg_loss( + padded = reduce_pg_loss( pg, mask, loss_aggregation="prompt_mean", group_size=2, group_sizes=[2, 1], ) - packed = aggregate_pg_loss( + packed = reduce_pg_loss( pg.reshape(-1), mask.reshape(-1), loss_aggregation="prompt_mean", @@ -112,7 +204,7 @@ def test_prompt_mean_partial_group_packed_matches_padded(): def test_constant_normalizes_token_sum_by_fixed_sequence_divisor(): - loss = aggregate_pg_loss( + loss = reduce_pg_loss( PG, MASK, loss_aggregation="constant", @@ -125,7 +217,7 @@ def test_constant_packed_matches_padded(): pg = PG.reshape(-1) mask = MASK.reshape(-1) cu_seqlens = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32) - loss = aggregate_pg_loss( + loss = reduce_pg_loss( pg, mask, loss_aggregation="constant", @@ -141,9 +233,9 @@ def test_constant_packed_matches_padded(): def test_loss_normalizer_pairing_realizes_global_mean(aggregation): group_size = 2 if aggregation == "prompt_mean" else 1 divisor = CONSTANT_DIVISOR if aggregation == "constant" else None - normalizer_fn = make_pg_loss_normalizer_fn(aggregation, group_size) + normalizer_fn = make_normalizer(aggregation, group_size) - full = aggregate_pg_loss( + full = reduce_pg_loss( PG, MASK, loss_aggregation=aggregation, @@ -155,7 +247,7 @@ def test_loss_normalizer_pairing_realizes_global_mean(aggregation): den = torch.tensor(0.0) for s in (slice(0, 2), slice(2, 4)): mb_pg, mb_mask = PG[s], MASK[s] - loss_mb = aggregate_pg_loss( + loss_mb = reduce_pg_loss( mb_pg, mb_mask, loss_aggregation=aggregation, @@ -174,9 +266,9 @@ def test_loss_normalizer_pairing_realizes_global_mean(aggregation): def test_loss_sum_pairing_realizes_global_mean(aggregation): group_size = 2 if aggregation == "prompt_mean" else 1 divisor = CONSTANT_DIVISOR if aggregation == "constant" else None - normalizer_fn = make_pg_loss_normalizer_fn(aggregation, group_size) + normalizer_fn = make_normalizer(aggregation, group_size) - full = aggregate_pg_loss( + full = reduce_pg_loss( PG, MASK, loss_aggregation=aggregation, @@ -188,7 +280,7 @@ def test_loss_sum_pairing_realizes_global_mean(aggregation): den = torch.tensor(0.0) for s in (slice(0, 2), slice(2, 4)): mb_pg, mb_mask = PG[s], MASK[s] - num = num + aggregate_pg_loss_sum( + num = num + reduce_pg_loss_sum( mb_pg, mb_mask, loss_aggregation=aggregation, @@ -205,12 +297,12 @@ def test_loss_sum_pairing_realizes_global_mean(aggregation): def test_loss_normalizer_pairing_realizes_global_mean_for_packed_inputs(aggregation): group_size = 2 if aggregation == "prompt_mean" else 1 divisor = CONSTANT_DIVISOR if aggregation == "constant" else None - normalizer_fn = make_pg_loss_normalizer_fn(aggregation, group_size) + normalizer_fn = make_normalizer(aggregation, group_size) pg = PG.reshape(-1) mask = MASK.reshape(-1) cu_seqlens = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32) - full = aggregate_pg_loss( + full = reduce_pg_loss( pg, mask, loss_aggregation=aggregation, @@ -224,7 +316,7 @@ def test_loss_normalizer_pairing_realizes_global_mean_for_packed_inputs(aggregat for s in (slice(0, 6), slice(6, 12)): mb_pg, mb_mask = pg[s], mask[s] mb_cu_seqlens = torch.tensor([0, 3, 6], dtype=torch.int32) - loss_mb = aggregate_pg_loss( + loss_mb = reduce_pg_loss( mb_pg, mb_mask, loss_aggregation=aggregation, @@ -244,12 +336,12 @@ def test_loss_normalizer_pairing_realizes_global_mean_for_packed_inputs(aggregat def test_loss_sum_pairing_realizes_global_mean_for_packed_inputs(aggregation): group_size = 2 if aggregation == "prompt_mean" else 1 divisor = CONSTANT_DIVISOR if aggregation == "constant" else None - normalizer_fn = make_pg_loss_normalizer_fn(aggregation, group_size) + normalizer_fn = make_normalizer(aggregation, group_size) pg = PG.reshape(-1) mask = MASK.reshape(-1) cu_seqlens = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32) - full = aggregate_pg_loss( + full = reduce_pg_loss( pg, mask, loss_aggregation=aggregation, @@ -263,7 +355,7 @@ def test_loss_sum_pairing_realizes_global_mean_for_packed_inputs(aggregation): for s in (slice(0, 6), slice(6, 12)): mb_pg, mb_mask = pg[s], mask[s] mb_cu_seqlens = torch.tensor([0, 3, 6], dtype=torch.int32) - num = num + aggregate_pg_loss_sum( + num = num + reduce_pg_loss_sum( mb_pg, mb_mask, loss_aggregation=aggregation, @@ -279,11 +371,11 @@ def test_denom_mask_uses_pre_rejection_count(): pg = torch.tensor([[2.0, 2.0, 2.0, 2.0]]) loss_mask = torch.tensor([[1.0, 1.0, 0.0, 0.0]]) denom_mask = torch.tensor([[1.0, 1.0, 1.0, 1.0]]) - loss = aggregate_pg_loss( + loss = reduce_pg_loss( pg, loss_mask, loss_aggregation="token_mean", denom_mask=denom_mask ) torch.testing.assert_close(loss, torch.tensor(1.0)) - without = aggregate_pg_loss(pg, loss_mask, loss_aggregation="token_mean") + without = reduce_pg_loss(pg, loss_mask, loss_aggregation="token_mean") torch.testing.assert_close(without, torch.tensor(2.0)) @@ -295,7 +387,7 @@ def test_denom_mask_applies_to_unit_mean_denominator(aggregation, group_size): loss_mask = torch.tensor([[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0]]) denom_mask = torch.ones_like(loss_mask) - loss = aggregate_pg_loss( + loss = reduce_pg_loss( pg, loss_mask, loss_aggregation=aggregation, @@ -304,7 +396,7 @@ def test_denom_mask_applies_to_unit_mean_denominator(aggregation, group_size): ) torch.testing.assert_close(loss, torch.tensor(1.5)) - without = aggregate_pg_loss( + without = reduce_pg_loss( pg, loss_mask, loss_aggregation=aggregation, group_size=group_size ) torch.testing.assert_close(without, torch.tensor(3.0)) @@ -317,7 +409,7 @@ def test_unit_mean_skips_units_without_denominator(aggregation, group_size): pg = torch.tensor([[2.0, 2.0], [2.0, 2.0], [100.0, 100.0], [100.0, 100.0]]) loss_mask = torch.tensor([[1.0, 1.0], [1.0, 1.0], [0.0, 0.0], [0.0, 0.0]]) - loss = aggregate_pg_loss( + loss = reduce_pg_loss( pg, loss_mask, loss_aggregation=aggregation, @@ -335,7 +427,7 @@ def test_packed_unit_mean_skips_units_without_denominator(aggregation, group_siz loss_mask = torch.tensor([1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0]) cu_seqlens = torch.tensor([0, 2, 4, 6, 8], dtype=torch.int32) - loss = aggregate_pg_loss( + loss = reduce_pg_loss( pg, loss_mask, loss_aggregation=aggregation, @@ -350,7 +442,7 @@ def test_constant_skips_sequences_without_denominator(): pg = torch.tensor([[2.0, 2.0], [100.0, 100.0]]) loss_mask = torch.tensor([[1.0, 1.0], [0.0, 0.0]]) - loss = aggregate_pg_loss( + loss = reduce_pg_loss( pg, loss_mask, loss_aggregation="constant", @@ -365,7 +457,7 @@ def test_packed_constant_skips_sequences_without_denominator(): loss_mask = torch.tensor([1.0, 1.0, 0.0, 0.0]) cu_seqlens = torch.tensor([0, 2, 4], dtype=torch.int32) - loss = aggregate_pg_loss( + loss = reduce_pg_loss( pg, loss_mask, loss_aggregation="constant", @@ -384,14 +476,14 @@ def test_denom_mask_keeps_empty_numerator_units_in_denominator(aggregation, grou loss_mask = torch.tensor([[0.0, 0.0], [1.0, 1.0]]) denom_mask = torch.ones_like(loss_mask) - loss = aggregate_pg_loss( + loss = reduce_pg_loss( pg, loss_mask, loss_aggregation=aggregation, group_size=group_size, denom_mask=denom_mask, ) - without = aggregate_pg_loss( + without = reduce_pg_loss( pg, loss_mask, loss_aggregation=aggregation, @@ -407,7 +499,7 @@ def test_pg_loss_rejects_broadcastable_loss_mask_shape(): loss_mask = torch.ones(2, 1) with pytest.raises(ValueError, match="loss_mask shape"): - aggregate_pg_loss(pg, loss_mask, loss_aggregation="token_mean") + reduce_pg_loss(pg, loss_mask, loss_aggregation="token_mean") def test_pg_loss_sum_rejects_broadcastable_denom_mask_shape(): @@ -416,7 +508,7 @@ def test_pg_loss_sum_rejects_broadcastable_denom_mask_shape(): denom_mask = torch.ones(1, 3) with pytest.raises(ValueError, match="denom_mask shape"): - aggregate_pg_loss_sum( + reduce_pg_loss_sum( pg, loss_mask, loss_aggregation="prompt_mean", @@ -426,8 +518,8 @@ def test_pg_loss_sum_rejects_broadcastable_denom_mask_shape(): def test_loss_normalizer_counts_active_units(): - seq_normalizer = make_pg_loss_normalizer_fn("seq_mean", 1) - prompt_normalizer = make_pg_loss_normalizer_fn("prompt_mean", 2) + seq_normalizer = make_normalizer("seq_mean", 1) + prompt_normalizer = make_normalizer("prompt_mean", 2) mask = torch.tensor([[1.0, 1.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]) torch.testing.assert_close(seq_normalizer({"loss_mask": mask}), torch.tensor(1.0)) @@ -437,7 +529,7 @@ def test_loss_normalizer_counts_active_units(): def test_prompt_normalizer_counts_partial_groups(): - prompt_normalizer = make_pg_loss_normalizer_fn("prompt_mean", 2) + prompt_normalizer = make_normalizer("prompt_mean", 2) mask = torch.tensor([[1.0, 0.0], [0.0, 0.0], [1.0, 1.0]]) torch.testing.assert_close( @@ -481,8 +573,25 @@ def test_m2po_normalizer_rejects_missing_prox_logp_for_loglinear(): normalizer(data) +def test_m2po_normalizer_rejects_missing_prox_logp_for_reused_train_logp(): + normalizer = _make_actor_loss_normalizer_fn( + "token_mean", + group_size=1, + m2_threshold=0.1, + prox_logp_method=PROX_LOGP_METHOD_REUSE_TRAIN_LOGP, + current_version=3, + ) + data = { + "logprobs": torch.tensor([[0.0, 0.2, 1.0]]), + "loss_mask": torch.ones(1, 3, dtype=torch.bool), + } + + with pytest.raises(ValueError, match="m2_threshold requires prox_logp"): + normalizer(data) + + def test_packed_loss_normalizer_counts_active_units(): - normalizer = make_pg_loss_normalizer_fn("seq_mean", 1) + normalizer = make_normalizer("seq_mean", 1) mask = torch.tensor([1.0, 1.0, 0.0, 0.0]) cu_seqlens = torch.tensor([0, 2, 4], dtype=torch.int32) @@ -492,9 +601,20 @@ def test_packed_loss_normalizer_counts_active_units(): ) +def test_non_token_packed_reduction_requires_sequence_boundaries(): + reduction = PolicyGradientReduction(mode="seq_mean") + loss = torch.ones(4) + mask = torch.ones(4, dtype=torch.bool) + + with pytest.raises(ValueError, match="requires cu_seqlens"): + reduction.aggregate(loss, mask) + with pytest.raises(ValueError, match="requires cu_seqlens"): + reduction.normalizer_fn({"loss_mask": mask}) + + def test_prompt_mean_group_size_one_equals_seq_mean(): - a = aggregate_pg_loss(PG, MASK, loss_aggregation="prompt_mean", group_size=1) - b = aggregate_pg_loss(PG, MASK, loss_aggregation="seq_mean") + a = reduce_pg_loss(PG, MASK, loss_aggregation="prompt_mean", group_size=1) + b = reduce_pg_loss(PG, MASK, loss_aggregation="seq_mean") torch.testing.assert_close(a, b) @@ -502,14 +622,14 @@ def test_prompt_mean_rejects_ragged_group_count(): pg = torch.ones(3, 2) mask = torch.ones(3, 2) with pytest.raises(ValueError, match="not divisible by group_size"): - aggregate_pg_loss(pg, mask, loss_aggregation="prompt_mean", group_size=2) + reduce_pg_loss(pg, mask, loss_aggregation="prompt_mean", group_size=2) def test_prompt_sum_pairing_realizes_global_mean_for_partial_groups(): pg = torch.tensor([[2.0, 2.0], [2.0, 2.0], [10.0, 10.0]]) mask = torch.ones_like(pg) - normalizer_fn = make_pg_loss_normalizer_fn("prompt_mean", 2) - full = aggregate_pg_loss( + normalizer_fn = make_normalizer("prompt_mean", 2) + full = reduce_pg_loss( pg, mask, loss_aggregation="prompt_mean", @@ -520,7 +640,7 @@ def test_prompt_sum_pairing_realizes_global_mean_for_partial_groups(): num = torch.tensor(0.0) den = torch.tensor(0.0) for seq_slice, group_sizes in ((slice(0, 2), [2]), (slice(2, 3), [1])): - num = num + aggregate_pg_loss_sum( + num = num + reduce_pg_loss_sum( pg[seq_slice], mask[seq_slice], loss_aggregation="prompt_mean", diff --git a/tests/test_rpc_callable_serialization.py b/tests/test_rpc_callable_serialization.py deleted file mode 100644 index 51f67e3d82..0000000000 --- a/tests/test_rpc_callable_serialization.py +++ /dev/null @@ -1,21 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -import pytest - -from areal.infra.rpc.serialization import ( - deserialize_value, - serialize_value, -) - - -def _sample_rpc_callable(value: int) -> int: - return value + 1 - - -def test_callable_serialization_rejected(): - with pytest.raises(ValueError, match="callable values"): - serialize_value(_sample_rpc_callable) - - -def test_callable_payload_rejected(): - with pytest.raises(ValueError, match="callable values"): - deserialize_value({"type": "callable", "key": "missing"}) diff --git a/tests/v2/training_service/fake_train_engine.py b/tests/v2/training_service/fake_train_engine.py index 028a0c3433..e886a0ccad 100644 --- a/tests/v2/training_service/fake_train_engine.py +++ b/tests/v2/training_service/fake_train_engine.py @@ -140,27 +140,21 @@ def forward_backward_batch( def train_batch( self, input_: dict[str, Any], - loss_reduction: Any, + loss_fn=None, + loss_weight_fn=None, ) -> dict[str, float]: - result = { + return { "total": _sum_numbers(input_), "version": float(self._version), "train_mode": float(self._train_mode), } - if loss_reduction is not None: - normalizer = loss_reduction.terms[0].normalizer_fn(input_) - if hasattr(normalizer, "item"): - normalizer = normalizer.item() - result["loss_terms"] = float(len(loss_reduction.terms)) - result["normalizer"] = float(normalizer) - return result def eval_batch( self, input_: dict[str, Any], - loss_reduction: Any, + loss_fn=None, + loss_weight_fn=None, ) -> float: - _ = loss_reduction return _sum_numbers(input_) + self._version def forward_batch( @@ -178,11 +172,11 @@ def forward_batch( def train_lm(self, input_, **kwargs): _ = kwargs - return self.train_batch(input_, loss_reduction=None) + return self.train_batch(input_) def evaluate_lm(self, input_, **kwargs): _ = kwargs - return self.eval_batch(input_, loss_reduction=None) + return self.eval_batch(input_) def export_stats(self) -> dict[str, float]: return { diff --git a/tests/v2/training_service/test_worker_unit.py b/tests/v2/training_service/test_worker_unit.py index 407e5b9b1a..69b080625c 100644 --- a/tests/v2/training_service/test_worker_unit.py +++ b/tests/v2/training_service/test_worker_unit.py @@ -99,7 +99,7 @@ def test_train_batch_after_create_engine(self, client): "args": serialize_value( [{"token_ids": [1, 2, 3], "metadata": {"weight": 2.0}}] ), - "kwargs": serialize_value({"loss_reduction": None}), + "kwargs": serialize_value({}), }, ) assert train_resp.status_code == 200 From 9ddd7fd9f2ea79281d231f275782829bb0f46723 Mon Sep 17 00:00:00 2001 From: EazyReal <8047065+EazyReal@users.noreply.github.com> Date: Sat, 11 Jul 2026 17:31:45 -0700 Subject: [PATCH 3/4] fix: preserve original engine callbacks in loss reductions Keep the pre-PR TrainEngine callback contract as the stable extension boundary while exposing advanced reductions through explicit methods. Key changes: - Restore original trainer and abstract engine callback signatures - Harden distributed normalizers and policy reduction invariants - Keep KD rejection normalization partition invariant - Reject impossible M2PO configurations before training - Add compatibility and regression coverage across all backends Refs: #1443 --- areal/api/cli_args.py | 20 ++++- areal/api/engine_api.py | 82 ++++++++++++++----- areal/api/loss_api.py | 31 ------- areal/engine/core/train_engine.py | 9 +- areal/engine/fsdp_engine.py | 39 +++++---- areal/engine/megatron_engine.py | 39 +++++---- areal/experimental/engine/archon_engine.py | 41 ++++++---- areal/trainer/dpo/dpo_engine.py | 18 ++-- areal/trainer/ppo/actor.py | 39 +++------ areal/trainer/ppo/critic.py | 12 ++- areal/trainer/rw/rw_engine.py | 14 ++-- areal/trainer/sft/lm_engine.py | 14 ++-- areal/utils/functional/loss_aggregation.py | 18 ++-- docs/en/best_practices/perf_profiling.md | 4 +- docs/zh/best_practices/perf_profiling.md | 4 +- .../archon/torchrun/run_archon_engine_pp.py | 14 ++-- tests/fp8/model_hooks.py | 2 +- tests/test_loss_reduction.py | 64 ++++++++++----- tests/test_megatron_engine.py | 8 +- tests/test_ppo_stats.py | 40 +++++++++ tests/test_prompt_mean_loss.py | 56 +++++++++++-- tests/test_train_engine.py | 14 ++-- tests/test_tree_training.py | 4 +- tests/torchrun/run_fsdp_dcp_distributed.py | 6 +- .../torchrun/run_fsdp_ulysses_train_batch.py | 4 +- .../run_megatron_engine_distributed.py | 14 ++-- .../run_megatron_engine_vlm_distributed.py | 2 +- 27 files changed, 372 insertions(+), 240 deletions(-) diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index 1a63da3a9c..70d88d26f6 100644 --- a/areal/api/cli_args.py +++ b/areal/api/cli_args.py @@ -1744,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( @@ -1818,6 +1832,10 @@ def __post_init__(self): "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: diff --git a/areal/api/engine_api.py b/areal/api/engine_api.py index 705f3eaad4..fa90fe90c0 100644 --- a/areal/api/engine_api.py +++ b/areal/api/engine_api.py @@ -23,7 +23,9 @@ WeightUpdateMeta, ) from areal.api.loss_api import ( - LossReductionInput, + LOSS_TERM_REDUCTION_MEAN, + LossReduction, + LossTerm, LossWeightFn, ) @@ -367,10 +369,8 @@ def forward_backward_batch( def train_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReductionInput | None = None, - loss_weight_fn: LossWeightFn | None = None, - *, - loss_fn: Callable[..., torch.Tensor] | None = None, + loss_fn: Callable[..., torch.Tensor], + loss_weight_fn: LossWeightFn, ) -> dict[str, float]: """Update the model with a batch of data and a loss function. @@ -385,12 +385,14 @@ def train_batch( Preferred format is ``list[dict[str, Any]]`` (trajectory list). Backward compatibility: a pre-batched ``dict[str, Any]`` is also accepted. - loss_reduction : LossReductionInput, optional - New reduction contract, or the original positional ``loss_fn``. - loss_weight_fn : Callable, optional - Original AReaL denominator callback, paired with ``loss_fn``. - loss_fn : Callable, optional - Original AReaL keyword loss callback. + loss_fn : Callable[..., torch.Tensor] + The loss function. For actor (is_critic=False), it receives + (logprobs, entropy, input_data). For critic (is_critic=True), + it receives (values, input_data). Returns a scalar normalized loss. + 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. Returns ------- @@ -400,15 +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_reduction: LossReductionInput | None = None, - loss_weight_fn: LossWeightFn | None = None, - *, - loss_fn: Callable[..., torch.Tensor] | None = None, + loss_fn: Callable[..., torch.Tensor], + loss_weight_fn: LossWeightFn, ) -> torch.Tensor | None: """Evaluate the model using the forward pass and loss function. @@ -423,12 +437,14 @@ def eval_batch( Preferred format is ``list[dict[str, Any]]`` (trajectory list). Backward compatibility: a pre-batched ``dict[str, Any]`` is also accepted. - loss_reduction : LossReductionInput, optional - New reduction contract, or the original positional ``loss_fn``. - loss_weight_fn : Callable, optional - Original AReaL denominator callback, paired with ``loss_fn``. - loss_fn : Callable, optional - Original AReaL keyword loss callback. + loss_fn : Callable[..., torch.Tensor] + The loss function. For actor (is_critic=False), it receives + (logprobs, entropy, input_data). For critic (is_critic=True), + it receives (values, input_data). Returns a scalar normalized loss. + 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. Returns ------- @@ -438,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( diff --git a/areal/api/loss_api.py b/areal/api/loss_api.py index 95e088f6dd..ee55e915df 100644 --- a/areal/api/loss_api.py +++ b/areal/api/loss_api.py @@ -103,35 +103,4 @@ def sum( ) -LossReductionInput = LossReduction | Callable[..., torch.Tensor] LossWeightFn = Callable[[dict[str, Any]], torch.Tensor] - - -def coerce_loss_reduction( - loss_reduction: LossReductionInput | None = None, - loss_weight_fn: LossWeightFn | None = None, - *, - loss_fn: Callable[..., torch.Tensor] | None = None, -) -> LossReduction: - """Normalize the original loss callback API at the engine boundary. - - AReaL originally accepted ``loss_fn`` and ``loss_weight_fn`` separately. - The engine internals now consume ``LossReduction``; this adapter keeps the - original positional and keyword calls working without carrying two code - paths through every backend. - """ - if loss_fn is not None: - if loss_reduction is not None: - raise TypeError("pass either loss_fn or loss_reduction, not both") - loss_reduction = loss_fn - if isinstance(loss_reduction, LossReduction): - if loss_weight_fn is not None: - raise TypeError( - "loss_weight_fn is only valid with the original loss_fn API" - ) - return loss_reduction - if not callable(loss_reduction) or loss_weight_fn is None: - raise TypeError( - "train/eval batch requires LossReduction or both loss_fn and loss_weight_fn" - ) - return LossReduction.mean(loss_reduction, loss_weight_fn) diff --git a/areal/engine/core/train_engine.py b/areal/engine/core/train_engine.py index 4a87cc4fa3..27a6a3335c 100644 --- a/areal/engine/core/train_engine.py +++ b/areal/engine/core/train_engine.py @@ -85,16 +85,17 @@ def compute_global_normalizers( [_sum_term_normalizer(mb_list, term) for term in loss_reduction.terms] ) dist.all_reduce(normalizers, group=dp_group) - if torch.any(normalizers <= 0).item(): - invalid_mask = (normalizers <= 0).detach().cpu().tolist() + invalid_normalizers = ~torch.isfinite(normalizers) | (normalizers <= 0) + if torch.any(invalid_normalizers).item(): + invalid_mask = invalid_normalizers.detach().cpu().tolist() invalid = [ term.name for term, is_invalid in zip(loss_reduction.terms, invalid_mask, strict=True) if is_invalid ] raise RuntimeError( - "Global loss normalizers must be positive after all_reduce; " - f"got non-positive normalizer for terms {invalid}." + "Global loss normalizers must be finite and positive after all_reduce; " + f"got invalid normalizer for terms {invalid}." ) return { term.name: normalizer diff --git a/areal/engine/fsdp_engine.py b/areal/engine/fsdp_engine.py index 21b2b93e52..18b7a7a12c 100644 --- a/areal/engine/fsdp_engine.py +++ b/areal/engine/fsdp_engine.py @@ -59,7 +59,7 @@ ) from areal.api.cli_args import OptimizerConfig, PerfTracerConfig, TrainEngineConfig from areal.api.io_struct import DeviceRuntimeInfo -from areal.api.loss_api import LossReductionInput, LossWeightFn, coerce_loss_reduction +from areal.api.loss_api import LossWeightFn from areal.engine.core import ( aggregate_eval_losses, compute_global_normalizers, @@ -764,15 +764,19 @@ def forward_backward_batch( def train_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReductionInput | None = None, - loss_weight_fn: LossWeightFn | None = None, - *, - loss_fn: Callable[..., torch.Tensor] | None = None, + loss_fn: Callable[..., torch.Tensor], + loss_weight_fn: LossWeightFn, ) -> dict[str, float]: - self._ensure_ready() - loss_reduction = coerce_loss_reduction( - loss_reduction, loss_weight_fn, loss_fn=loss_fn + return self.train_batch_with_reduction( + input_, LossReduction.mean(loss_fn, loss_weight_fn) ) + + def train_batch_with_reduction( + self, + input_: list[dict[str, Any]] | dict[str, Any], + loss_reduction: LossReduction, + ) -> dict[str, float]: + self._ensure_ready() self.optimizer_zero_grad() input_batched, _ = self._normalize_batch_input(input_) @@ -805,16 +809,21 @@ def process_output( def eval_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReductionInput | None = None, - loss_weight_fn: LossWeightFn | None = None, - *, - loss_fn: Callable[..., torch.Tensor] | None = None, + loss_fn: Callable[..., torch.Tensor], + loss_weight_fn: LossWeightFn, ) -> torch.Tensor | None: - self._ensure_ready() - loss_reduction = coerce_loss_reduction( - loss_reduction, loss_weight_fn, loss_fn=loss_fn + return self.eval_batch_with_reduction( + input_, LossReduction.mean(loss_fn, loss_weight_fn) ) + @torch.no_grad() + def eval_batch_with_reduction( + self, + input_: list[dict[str, Any]] | dict[str, Any], + loss_reduction: LossReduction, + ) -> torch.Tensor | None: + self._ensure_ready() + input_batched, _ = self._normalize_batch_input(input_) mb_list = self._prepare_mb_list(input_batched).to(self.device) diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index 0c0ac9b3cd..74d65bcafb 100644 --- a/areal/engine/megatron_engine.py +++ b/areal/engine/megatron_engine.py @@ -50,7 +50,7 @@ ) from areal.api.cli_args import MicroBatchSpec, PerfTracerConfig, TrainEngineConfig from areal.api.io_struct import DeviceRuntimeInfo -from areal.api.loss_api import LossReductionInput, LossWeightFn, coerce_loss_reduction +from areal.api.loss_api import LossWeightFn from areal.engine.core import ( aggregate_eval_losses, compute_global_normalizers, @@ -994,15 +994,19 @@ def _process_output(input_, output_): def train_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReductionInput | None = None, - loss_weight_fn: LossWeightFn | None = None, - *, - loss_fn: Callable[..., torch.Tensor] | None = None, + loss_fn: Callable[..., torch.Tensor], + loss_weight_fn: LossWeightFn, ) -> dict[str, float]: - self._ensure_ready() - loss_reduction = coerce_loss_reduction( - loss_reduction, loss_weight_fn, loss_fn=loss_fn + return self.train_batch_with_reduction( + input_, LossReduction.mean(loss_fn, loss_weight_fn) ) + + def train_batch_with_reduction( + self, + input_: list[dict[str, Any]] | dict[str, Any], + loss_reduction: LossReduction, + ) -> dict[str, float]: + self._ensure_ready() self.optimizer_zero_grad() input_batched, _ = self._normalize_batch_input(input_) @@ -1056,16 +1060,21 @@ def process_output( def eval_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReductionInput | None = None, - loss_weight_fn: LossWeightFn | None = None, - *, - loss_fn: Callable[..., torch.Tensor] | None = None, + loss_fn: Callable[..., torch.Tensor], + loss_weight_fn: LossWeightFn, ) -> torch.Tensor | None: - self._ensure_ready() - loss_reduction = coerce_loss_reduction( - loss_reduction, loss_weight_fn, loss_fn=loss_fn + return self.eval_batch_with_reduction( + input_, LossReduction.mean(loss_fn, loss_weight_fn) ) + @torch.no_grad() + def eval_batch_with_reduction( + self, + input_: list[dict[str, Any]] | dict[str, Any], + loss_reduction: LossReduction, + ) -> torch.Tensor | None: + self._ensure_ready() + input_batched, _ = self._normalize_batch_input(input_) mb_list = self._prepare_mb_list(input_batched).to(self.device) diff --git a/areal/experimental/engine/archon_engine.py b/areal/experimental/engine/archon_engine.py index 7d7c73d5e6..136567b5b1 100644 --- a/areal/experimental/engine/archon_engine.py +++ b/areal/experimental/engine/archon_engine.py @@ -36,7 +36,7 @@ ) from areal.api.cli_args import MicroBatchSpec from areal.api.io_struct import DeviceRuntimeInfo -from areal.api.loss_api import LossReductionInput, LossWeightFn, coerce_loss_reduction +from areal.api.loss_api import LossWeightFn from areal.engine.core.distributed import ( patch_dist_group_timeout, warmup_process_groups, @@ -529,16 +529,21 @@ def forward_backward_batch( def train_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReductionInput | None = None, - loss_weight_fn: LossWeightFn | None = None, - *, - loss_fn: Callable[..., torch.Tensor] | None = None, + loss_fn: Callable[..., torch.Tensor], + loss_weight_fn: LossWeightFn, ) -> dict[str, float]: """Train on a batch of data.""" - assert self._initialized - loss_reduction = coerce_loss_reduction( - loss_reduction, loss_weight_fn, loss_fn=loss_fn + return self.train_batch_with_reduction( + input_, LossReduction.mean(loss_fn, loss_weight_fn) ) + + def train_batch_with_reduction( + self, + input_: list[dict[str, Any]] | dict[str, Any], + loss_reduction: LossReduction, + ) -> dict[str, float]: + """Train on a batch with an explicit reduction contract.""" + assert self._initialized self.optimizer_zero_grad() input_batched, _ = self._normalize_batch_input(input_) @@ -571,17 +576,23 @@ def process_output( def eval_batch( self, input_: list[dict[str, Any]] | dict[str, Any], - loss_reduction: LossReductionInput | None = None, - loss_weight_fn: LossWeightFn | None = None, - *, - loss_fn: Callable[..., torch.Tensor] | None = None, + loss_fn: Callable[..., torch.Tensor], + loss_weight_fn: LossWeightFn, ) -> torch.Tensor | None: """Evaluate on a batch of data.""" - assert self._initialized - loss_reduction = coerce_loss_reduction( - loss_reduction, loss_weight_fn, loss_fn=loss_fn + return self.eval_batch_with_reduction( + input_, LossReduction.mean(loss_fn, loss_weight_fn) ) + @torch.no_grad() + def eval_batch_with_reduction( + self, + input_: list[dict[str, Any]] | dict[str, Any], + loss_reduction: LossReduction, + ) -> torch.Tensor | None: + """Evaluate on a batch with an explicit reduction contract.""" + assert self._initialized + input_batched, _ = self._normalize_batch_input(input_) mb_list = self._prepare_mb_list(input_batched).to(self.device) diff --git a/areal/trainer/dpo/dpo_engine.py b/areal/trainer/dpo/dpo_engine.py index 9a4c9da14a..e9c80668ec 100644 --- a/areal/trainer/dpo/dpo_engine.py +++ b/areal/trainer/dpo/dpo_engine.py @@ -5,7 +5,7 @@ import torch -from areal.api import LossReduction, TrainEngine +from areal.api import TrainEngine from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value from areal.utils import logging, stats_tracker @@ -56,12 +56,10 @@ def _train_dpo(self, data: dict[str, Any]) -> None: self.engine.train() stats = self.engine.train_batch( input_=data, - loss_reduction=LossReduction.mean( - loss_fn=functools.partial( - compute_dpo_loss, beta=self.beta, loss_type=self.loss_type - ), - normalizer_fn=_dpo_loss_normalizer, + loss_fn=functools.partial( + compute_dpo_loss, beta=self.beta, loss_type=self.loss_type ), + loss_weight_fn=_dpo_loss_normalizer, ) stats_tracker.scalar(**stats) @@ -74,12 +72,10 @@ def _evaluate_dpo(self, data: dict[str, Any]) -> None: self.engine.eval() self.engine.eval_batch( input_=data, - loss_reduction=LossReduction.mean( - loss_fn=functools.partial( - compute_dpo_loss, beta=self.beta, loss_type=self.loss_type - ), - normalizer_fn=_dpo_loss_normalizer, + loss_fn=functools.partial( + compute_dpo_loss, beta=self.beta, loss_type=self.loss_type ), + loss_weight_fn=_dpo_loss_normalizer, ) @trace_perf("dpo_engine.compute_logp", category="compute") diff --git a/areal/trainer/ppo/actor.py b/areal/trainer/ppo/actor.py index db6f82c70e..99a0587419 100644 --- a/areal/trainer/ppo/actor.py +++ b/areal/trainer/ppo/actor.py @@ -5,7 +5,7 @@ import torch -from areal.api import LossReduction, TrainEngine +from areal.api import TrainEngine from areal.api.cli_args import MicroBatchSpec, PPOActorConfig, RejectionSamplingConfig from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value @@ -47,13 +47,10 @@ logger = logging.getLogger("PPOActor") -def _use_sum_pg_loss(data: dict[str, Any], m2_threshold: float | None) -> bool: - return "teacher_logp" not in data and m2_threshold is None - - def _make_actor_loss_normalizer_fn( loss_aggregation: str, group_size: int, + loss_aggregation_divisor: float | None, m2_threshold: float | None, prox_logp_method: str, current_version: int | None, @@ -61,6 +58,7 @@ def _make_actor_loss_normalizer_fn( normalizer_fn = PolicyGradientReduction( mode=loss_aggregation, group_size=group_size, + divisor=loss_aggregation_divisor, ).normalizer_fn if m2_threshold is None: return normalizer_fn @@ -419,7 +417,6 @@ def _ppo_update(self, data: dict[str, Any]) -> None: ) for mb in mb_inputs.mbs: - use_sum_reduction = _use_sum_pg_loss(mb, self.m2_threshold) loss_fn = functools.partial( grpo_loss_fn, eps_clip=self.config.eps_clip, @@ -436,26 +433,21 @@ def _ppo_update(self, data: dict[str, Any]) -> None: use_cispo_loss=self.config.use_cispo_loss, use_decoupled_loss=self.config.use_decoupled_loss, pg_reduction=pg_reduction, - local_mean=not use_sum_reduction, + local_mean=True, ) normalizer_fn = _make_actor_loss_normalizer_fn( self.config.loss_aggregation, self.config.group_size, + self.config.loss_aggregation_divisor, self.m2_threshold, self.config.prox_logp_method, current_version, ) - if use_sum_reduction: - loss_reduction = LossReduction.sum( - loss_fn=loss_fn, - normalizer_fn=normalizer_fn, - ) - else: - loss_reduction = LossReduction.mean( - loss_fn=loss_fn, - normalizer_fn=normalizer_fn, - ) - train_stat = self.engine.train_batch(mb, loss_reduction=loss_reduction) + train_stat = self.engine.train_batch( + mb, + loss_fn=loss_fn, + loss_weight_fn=normalizer_fn, + ) stats_tracker.scalar(**train_stat) @@ -641,6 +633,7 @@ def grpo_loss_fn( ) rl_loss_weight = input_data.get("rl_loss_weight", 1.0) distill_loss_weight = input_data.get("distill_loss_weight", 0.005) + loss_normalizer = loss_mask.count_nonzero().clamp_min(1) teacher_logp = teacher_logp.detach() @@ -651,18 +644,12 @@ def grpo_loss_fn( rkl_weighted_term = importance_weight * rkl_reward * effective_loss_mask kd_coef = -1 * distill_loss_weight - loss = ( - kd_coef - * rkl_weighted_term.sum() - / effective_loss_mask.sum().clamp(min=1) - ) + loss = kd_coef * rkl_weighted_term.sum() / loss_normalizer rkl_stat = -1 * rkl_weighted_term else: rkl_penalty_per_token = (logprobs - teacher_logp) * effective_loss_mask - rkl_penalty = rkl_penalty_per_token.sum() / effective_loss_mask.sum().clamp( - min=1 - ) + rkl_penalty = rkl_penalty_per_token.sum() / loss_normalizer loss = rl_loss_weight * loss + distill_loss_weight * rkl_penalty diff --git a/areal/trainer/ppo/critic.py b/areal/trainer/ppo/critic.py index aa9f3b7409..05f83c7441 100644 --- a/areal/trainer/ppo/critic.py +++ b/areal/trainer/ppo/critic.py @@ -5,7 +5,7 @@ import torch -from areal.api import LossReduction, TrainEngine +from areal.api import TrainEngine from areal.api.cli_args import MicroBatchSpec, PPOCriticConfig from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value @@ -65,13 +65,11 @@ def _ppo_update(self, data: dict[str, Any]) -> None: for mb in mb_inputs.mbs: train_stat = self.engine.train_batch( mb, - loss_reduction=LossReduction.mean( - loss_fn=functools.partial( - ppo_loss_fn, - eps_clip=self.config.eps_clip, - ), - normalizer_fn=lambda x: x["loss_mask"].count_nonzero(), + loss_fn=functools.partial( + ppo_loss_fn, + eps_clip=self.config.eps_clip, ), + loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), ) stats_tracker.scalar(**train_stat) diff --git a/areal/trainer/rw/rw_engine.py b/areal/trainer/rw/rw_engine.py index 1a1d17a873..37ccc909cf 100644 --- a/areal/trainer/rw/rw_engine.py +++ b/areal/trainer/rw/rw_engine.py @@ -4,7 +4,7 @@ import torch -from areal.api import LossReduction, TrainEngine +from areal.api import TrainEngine from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value from areal.utils import logging, stats_tracker @@ -54,10 +54,8 @@ def _train_rw(self, data: dict[str, Any]) -> None: self.engine.train() stats = self.engine.train_batch( input_=data, - loss_reduction=LossReduction.mean( - loss_fn=compute_rw_loss, - normalizer_fn=_rw_loss_normalizer, - ), + loss_fn=compute_rw_loss, + loss_weight_fn=_rw_loss_normalizer, ) stats_tracker.scalar(**stats) @@ -72,10 +70,8 @@ def _evaluate_rw(self, data: dict[str, Any]) -> None: self.engine.eval() self.engine.eval_batch( input_=data, - loss_reduction=LossReduction.mean( - loss_fn=compute_rw_loss, - normalizer_fn=_rw_loss_normalizer, - ), + loss_fn=compute_rw_loss, + loss_weight_fn=_rw_loss_normalizer, ) diff --git a/areal/trainer/sft/lm_engine.py b/areal/trainer/sft/lm_engine.py index 4c5e50b793..5e5343bb2d 100644 --- a/areal/trainer/sft/lm_engine.py +++ b/areal/trainer/sft/lm_engine.py @@ -4,7 +4,7 @@ import torch -from areal.api import LossReduction, TrainEngine +from areal.api import TrainEngine from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value from areal.utils import stats_tracker @@ -29,10 +29,8 @@ def _train_lm(self, data: dict[str, Any]) -> None: data["loss_mask"] = torch.roll(data["loss_mask"].bool(), shifts=-1, dims=-1) stats = self.engine.train_batch( input_=data, - loss_reduction=LossReduction.mean( - loss_fn=compute_packed_sft_loss, - normalizer_fn=lambda x: x["loss_mask"].count_nonzero(), - ), + loss_fn=compute_packed_sft_loss, + loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), ) stats_tracker.scalar(**stats) @@ -46,10 +44,8 @@ def _evaluate_lm(self, data: dict[str, Any]) -> None: data["loss_mask"] = torch.roll(data["loss_mask"].bool(), shifts=-1, dims=-1) self.engine.eval_batch( input_=data, - loss_reduction=LossReduction.mean( - loss_fn=compute_packed_sft_loss, - normalizer_fn=lambda x: x["loss_mask"].count_nonzero(), - ), + loss_fn=compute_packed_sft_loss, + loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), ) diff --git a/areal/utils/functional/loss_aggregation.py b/areal/utils/functional/loss_aggregation.py index ecbe4d388c..102989bbc4 100644 --- a/areal/utils/functional/loss_aggregation.py +++ b/areal/utils/functional/loss_aggregation.py @@ -194,7 +194,7 @@ class PolicyGradientReduction: the matching number of active tokens, sequences, or prompt groups. """ - mode: str = LOSS_AGGREGATION_TOKEN_MEAN + mode: LossAggregationMode = LOSS_AGGREGATION_TOKEN_MEAN group_size: int = 1 divisor: float | None = None @@ -206,16 +206,22 @@ def __post_init__(self) -> None: ) if self.group_size <= 0: raise ValueError(f"group_size must be positive, got {self.group_size}.") - if self.divisor is not None: + if self.mode != LOSS_AGGREGATION_PROMPT_MEAN and self.group_size != 1: + raise ValueError( + "group_size is only valid for loss_aggregation='prompt_mean'." + ) + if self.mode == LOSS_AGGREGATION_CONSTANT: if ( - self.mode != LOSS_AGGREGATION_CONSTANT + self.divisor is None or not math.isfinite(self.divisor) or self.divisor <= 0 ): raise ValueError( - "divisor must be a positive finite value and is only valid " - "for loss_aggregation='constant'." + "divisor must be a positive finite value for " + "loss_aggregation='constant'." ) + elif self.divisor is not None: + raise ValueError("divisor is only valid for loss_aggregation='constant'.") def _require_divisor(self) -> float: if self.divisor is None: @@ -250,7 +256,7 @@ def normalizer_fn(self, data: dict[str, Any]) -> torch.Tensor: ) ids, n_units = _group_ids( sequence_denominators.numel(), - self.group_size, + self.group_size if self.mode == LOSS_AGGREGATION_PROMPT_MEAN else 1, group_sizes, loss_mask.device, ) diff --git a/docs/en/best_practices/perf_profiling.md b/docs/en/best_practices/perf_profiling.md index 1080449976..f25eb5afaf 100644 --- a/docs/en/best_practices/perf_profiling.md +++ b/docs/en/best_practices/perf_profiling.md @@ -85,11 +85,11 @@ ppo_update, etc.). **Example** (from `areal/engine/fsdp_engine.py`): ```python -from areal.api import LossReduction from areal.utils.perf_tracer import trace_perf @trace_perf("fsdp_engine.train_batch") -def train_batch(self, input_: dict[str, Any], loss_reduction: LossReduction): +def train_batch(self, input_: dict[str, Any], loss_fn, loss_weight_fn): + # Training logic here ... ``` diff --git a/docs/zh/best_practices/perf_profiling.md b/docs/zh/best_practices/perf_profiling.md index cd52f415bc..950366fadd 100644 --- a/docs/zh/best_practices/perf_profiling.md +++ b/docs/zh/best_practices/perf_profiling.md @@ -75,11 +75,11 @@ python -m areal.tools.perf_trace_converter logs/**/perf_tracer/traces-*.jsonl me **示例**（来自 `areal/engine/fsdp_engine.py`）： ```python -from areal.api import LossReduction from areal.utils.perf_tracer import trace_perf @trace_perf("fsdp_engine.train_batch") -def train_batch(self, input_: dict[str, Any], loss_reduction: LossReduction): +def train_batch(self, input_: dict[str, Any], loss_fn, loss_weight_fn): + # Training logic here ... ``` diff --git a/tests/experimental/archon/torchrun/run_archon_engine_pp.py b/tests/experimental/archon/torchrun/run_archon_engine_pp.py index 8ced33fd45..4c200a15e4 100644 --- a/tests/experimental/archon/torchrun/run_archon_engine_pp.py +++ b/tests/experimental/archon/torchrun/run_archon_engine_pp.py @@ -190,10 +190,8 @@ def test_eval_batch(engine: ArchonEngine, mock_input: dict) -> bool: try: loss = engine.eval_batch( mock_input, - loss_reduction=LossReduction.mean( - loss_fn=mock_loss_fn, - normalizer_fn=mock_normalizer_fn, - ), + loss_fn=mock_loss_fn, + loss_weight_fn=mock_normalizer_fn, ) # In PP mode, eval_batch may return None (TODO in ArchonEngine) @@ -222,12 +220,10 @@ def test_train_batch(engine: ArchonEngine, mock_input: dict) -> bool: try: result = engine.train_batch( mock_input, - loss_reduction=LossReduction.mean( - loss_fn=mock_loss_fn, - normalizer_fn=mock_normalizer_fn, - ), + loss_fn=mock_loss_fn, + loss_weight_fn=mock_normalizer_fn, ) - sum_result = engine.train_batch( + sum_result = engine.train_batch_with_reduction( mock_input, loss_reduction=LossReduction.sum( loss_fn=mock_loss_sum_fn, diff --git a/tests/fp8/model_hooks.py b/tests/fp8/model_hooks.py index a0bfaa2c15..d8d2a9d39a 100644 --- a/tests/fp8/model_hooks.py +++ b/tests/fp8/model_hooks.py @@ -561,7 +561,7 @@ def normalizer_fn(mb): model_chunk.zero_grad_buffer() # Forward and backward - engine.train_batch(input_, loss_reduction=loss_reduction) + engine.train_batch_with_reduction(input_, loss_reduction=loss_reduction) # Collect gradients from all components (focusing on the selected layers) model = get_model_from_engine(engine) diff --git a/tests/test_loss_reduction.py b/tests/test_loss_reduction.py index c3a21f958f..3eb7048bf9 100644 --- a/tests/test_loss_reduction.py +++ b/tests/test_loss_reduction.py @@ -8,8 +8,8 @@ LOSS_TERM_REDUCTION_SUM, LossReduction, LossTerm, + TrainEngine, ) -from areal.api.loss_api import coerce_loss_reduction from areal.engine.core.train_engine import ( compute_global_normalizers, scale_loss_for_reduction, @@ -21,30 +21,51 @@ def _normalizer(_data): return torch.tensor(1.0) -def test_coerce_loss_reduction_preserves_original_callback_api(): - def loss_fn(*_args): +def _make_original_callback_engine(): + calls = [] + + def stub(*_args, **_kwargs): + return None + + def train_batch(self, input_, loss_fn, loss_weight_fn): + calls.append(("train", input_, loss_fn, loss_weight_fn)) + return {"lr": 1.0} + + def eval_batch(self, input_, loss_fn, loss_weight_fn): + calls.append(("eval", input_, loss_fn, loss_weight_fn)) return torch.tensor(2.0) - def loss_weight_fn(_data): - return torch.tensor(3.0) + implementations = {name: stub for name in TrainEngine.__abstractmethods__} + implementations.update(train_batch=train_batch, eval_batch=eval_batch) + engine_type = type("OriginalCallbackEngine", (TrainEngine,), implementations) + return engine_type(), calls - positional = coerce_loss_reduction(loss_fn, loss_weight_fn) - keyword = coerce_loss_reduction( - loss_fn=loss_fn, - loss_weight_fn=loss_weight_fn, - ) - for reduction in (positional, keyword): - assert reduction.loss_fn is loss_fn - assert reduction.terms[0].normalizer_fn is loss_weight_fn - assert reduction.terms[0].reduction == "mean" +def test_original_train_engine_subclass_supports_mean_reduction_adapter(): + engine, calls = _make_original_callback_engine() + + def loss_fn(): + return torch.tensor(2.0) + + reduction = LossReduction.mean(loss_fn, _normalizer) + + train_result = engine.train_batch_with_reduction({"x": 1}, reduction) + eval_result = engine.eval_batch_with_reduction({"x": 2}, reduction) + + assert train_result == {"lr": 1.0} + torch.testing.assert_close(eval_result, torch.tensor(2.0)) + assert calls == [ + ("train", {"x": 1}, loss_fn, _normalizer), + ("eval", {"x": 2}, loss_fn, _normalizer), + ] -def test_coerce_loss_reduction_rejects_mixed_contracts(): - reduction = LossReduction.mean(lambda: torch.tensor(1.0), _normalizer) +def test_original_train_engine_subclass_rejects_advanced_reduction(): + engine, _ = _make_original_callback_engine() + reduction = LossReduction.sum(lambda: torch.tensor(1.0), _normalizer) - with pytest.raises(TypeError, match="only valid with the original loss_fn API"): - coerce_loss_reduction(reduction, _normalizer) + with pytest.raises(NotImplementedError, match="original callback API"): + engine.train_batch_with_reduction({}, reduction) def test_mean_scaling_preserves_local_mean_order(): @@ -128,10 +149,11 @@ def test_multi_term_scaling_uses_each_terms_normalizer(): torch.testing.assert_close(scaled, torch.tensor(1.0), rtol=0, atol=0) -def test_global_normalizer_must_be_positive(monkeypatch): +@pytest.mark.parametrize("normalizer", [0.0, -1.0, float("nan"), float("inf")]) +def test_global_normalizer_must_be_finite_and_positive(monkeypatch, normalizer): reduction = LossReduction.sum( loss_fn=lambda: torch.tensor(0.0), - normalizer_fn=lambda data: data["loss_mask"].count_nonzero(), + normalizer_fn=lambda _data: torch.tensor(normalizer), ) mb_list = MicroBatchList( data={}, @@ -141,5 +163,5 @@ def test_global_normalizer_must_be_positive(monkeypatch): ) monkeypatch.setattr(dist, "all_reduce", lambda tensor, group=None: tensor) - with pytest.raises(RuntimeError, match="Global loss normalizers"): + with pytest.raises(RuntimeError, match="finite and positive"): compute_global_normalizers(mb_list, reduction, dp_group=None) diff --git a/tests/test_megatron_engine.py b/tests/test_megatron_engine.py index a48e01e134..b96b9fa090 100644 --- a/tests/test_megatron_engine.py +++ b/tests/test_megatron_engine.py @@ -125,12 +125,10 @@ def test_simple_train(engine, mock_input): engine.train() train_result = engine.train_batch( mock_input, - loss_reduction=LossReduction.mean( - loss_fn=mock_loss_fn, - normalizer_fn=lambda x: torch.tensor(1.0, device=engine.device), - ), + loss_fn=mock_loss_fn, + loss_weight_fn=lambda x: torch.tensor(1.0, device=engine.device), ) - sum_train_result = engine.train_batch( + sum_train_result = engine.train_batch_with_reduction( mock_input, loss_reduction=LossReduction.sum( loss_fn=mock_loss_sum_fn, diff --git a/tests/test_ppo_stats.py b/tests/test_ppo_stats.py index f93bba41d7..0a49a90e03 100644 --- a/tests/test_ppo_stats.py +++ b/tests/test_ppo_stats.py @@ -160,6 +160,46 @@ def test_teacher_distillation_excludes_rejected_tokens_from_gradient(): assert logprobs.grad[0, 1] == 0 +def test_teacher_distillation_rejection_uses_partition_invariant_normalizer(): + rejection_sampling = RejectionSamplingConfig( + level="token", action="mask", metric="ratio", upper=1.5 + ) + + def compute_loss(prox_logp): + n_tokens = prox_logp.numel() + input_data = { + "input_ids": torch.arange(n_tokens).unsqueeze(0), + "logprobs": torch.zeros_like(prox_logp), + "advantages": torch.zeros_like(prox_logp), + "loss_mask": torch.ones_like(prox_logp, dtype=torch.bool), + "prox_logp": prox_logp, + "versions": torch.zeros_like(prox_logp, dtype=torch.int32), + "teacher_logp": torch.ones_like(prox_logp), + "rl_loss_weight": 0.0, + "distill_loss_weight": 1.0, + } + return grpo_loss_fn( + logprobs=torch.zeros_like(prox_logp), + entropy=torch.zeros_like(prox_logp), + input_data=input_data, + eps_clip=0.2, + eps_clip_higher=None, + c_clip=None, + rejection_sampling=rejection_sampling, + ) + + rejected = torch.log(torch.tensor(2.0)) + prox_logp = torch.tensor([[0.0, rejected, rejected, rejected]]) + with patch("areal.trainer.ppo.actor.stats_tracker"): + full_loss = compute_loss(prox_logp) + partitioned_loss = ( + compute_loss(prox_logp[:, :2]) + compute_loss(prox_logp[:, 2:]) + ) / 2 + + torch.testing.assert_close(full_loss, torch.tensor(-0.25)) + torch.testing.assert_close(partitioned_loss, full_loss) + + def test_critic_loss_fn_uses_full_cu_seqlens_for_n_tokens(): input_data = { "input_ids": torch.tensor([11, 12]), diff --git a/tests/test_prompt_mean_loss.py b/tests/test_prompt_mean_loss.py index 4ddf1b8b66..45a53b5ad8 100644 --- a/tests/test_prompt_mean_loss.py +++ b/tests/test_prompt_mean_loss.py @@ -79,10 +79,11 @@ def reduce_pg_loss_sum( ) -def make_normalizer(loss_aggregation, group_size): +def make_normalizer(loss_aggregation, group_size, loss_aggregation_divisor=None): return PolicyGradientReduction( mode=loss_aggregation, group_size=group_size, + divisor=loss_aggregation_divisor, ).normalizer_fn @@ -105,7 +106,7 @@ def test_inactive_units_do_not_deflate_local_mean(aggregation): mask = torch.tensor([[1, 1], [0, 0], [0, 0]], dtype=torch.bool) kwargs = { "loss_aggregation": aggregation, - "group_size": 2, + "group_size": 2 if aggregation == "prompt_mean" else 1, "group_sizes": [2, 1], } @@ -233,7 +234,7 @@ def test_constant_packed_matches_padded(): def test_loss_normalizer_pairing_realizes_global_mean(aggregation): group_size = 2 if aggregation == "prompt_mean" else 1 divisor = CONSTANT_DIVISOR if aggregation == "constant" else None - normalizer_fn = make_normalizer(aggregation, group_size) + normalizer_fn = make_normalizer(aggregation, group_size, divisor) full = reduce_pg_loss( PG, @@ -266,7 +267,7 @@ def test_loss_normalizer_pairing_realizes_global_mean(aggregation): def test_loss_sum_pairing_realizes_global_mean(aggregation): group_size = 2 if aggregation == "prompt_mean" else 1 divisor = CONSTANT_DIVISOR if aggregation == "constant" else None - normalizer_fn = make_normalizer(aggregation, group_size) + normalizer_fn = make_normalizer(aggregation, group_size, divisor) full = reduce_pg_loss( PG, @@ -297,7 +298,7 @@ def test_loss_sum_pairing_realizes_global_mean(aggregation): def test_loss_normalizer_pairing_realizes_global_mean_for_packed_inputs(aggregation): group_size = 2 if aggregation == "prompt_mean" else 1 divisor = CONSTANT_DIVISOR if aggregation == "constant" else None - normalizer_fn = make_normalizer(aggregation, group_size) + normalizer_fn = make_normalizer(aggregation, group_size, divisor) pg = PG.reshape(-1) mask = MASK.reshape(-1) @@ -336,7 +337,7 @@ def test_loss_normalizer_pairing_realizes_global_mean_for_packed_inputs(aggregat def test_loss_sum_pairing_realizes_global_mean_for_packed_inputs(aggregation): group_size = 2 if aggregation == "prompt_mean" else 1 divisor = CONSTANT_DIVISOR if aggregation == "constant" else None - normalizer_fn = make_normalizer(aggregation, group_size) + normalizer_fn = make_normalizer(aggregation, group_size, divisor) pg = PG.reshape(-1) mask = MASK.reshape(-1) @@ -542,6 +543,7 @@ def test_m2po_normalizer_uses_post_filter_mask(): normalizer = _make_actor_loss_normalizer_fn( "token_mean", group_size=1, + loss_aggregation_divisor=None, m2_threshold=0.1, prox_logp_method=PROX_LOGP_METHOD_RECOMPUTE, current_version=3, @@ -555,10 +557,25 @@ def test_m2po_normalizer_uses_post_filter_mask(): torch.testing.assert_close(normalizer(data), torch.tensor(2)) +def test_constant_actor_normalizer_carries_valid_reduction_config(): + normalizer = _make_actor_loss_normalizer_fn( + "constant", + group_size=1, + loss_aggregation_divisor=10.0, + m2_threshold=None, + prox_logp_method=PROX_LOGP_METHOD_RECOMPUTE, + current_version=3, + ) + data = {"loss_mask": torch.tensor([[1, 0], [1, 1]], dtype=torch.bool)} + + torch.testing.assert_close(normalizer(data), torch.tensor(2.0)) + + def test_m2po_normalizer_rejects_missing_prox_logp_for_loglinear(): normalizer = _make_actor_loss_normalizer_fn( "token_mean", group_size=1, + loss_aggregation_divisor=None, m2_threshold=0.1, prox_logp_method=PROX_LOGP_METHOD_LOGLINEAR, current_version=3, @@ -577,6 +594,7 @@ def test_m2po_normalizer_rejects_missing_prox_logp_for_reused_train_logp(): normalizer = _make_actor_loss_normalizer_fn( "token_mean", group_size=1, + loss_aggregation_divisor=None, m2_threshold=0.1, prox_logp_method=PROX_LOGP_METHOD_REUSE_TRAIN_LOGP, current_version=3, @@ -612,6 +630,21 @@ def test_non_token_packed_reduction_requires_sequence_boundaries(): reduction.normalizer_fn({"loss_mask": mask}) +@pytest.mark.parametrize("mode", ["token_mean", "seq_mean", "constant"]) +def test_group_size_is_only_valid_for_prompt_mean(mode): + with pytest.raises(ValueError, match="group_size is only valid"): + PolicyGradientReduction( + mode=mode, + group_size=2, + divisor=10.0 if mode == "constant" else None, + ) + + +def test_constant_reduction_requires_divisor_at_construction(): + with pytest.raises(ValueError, match="positive finite"): + PolicyGradientReduction(mode="constant") + + def test_prompt_mean_group_size_one_equals_seq_mean(): a = reduce_pg_loss(PG, MASK, loss_aggregation="prompt_mean", group_size=1) b = reduce_pg_loss(PG, MASK, loss_aggregation="seq_mean") @@ -721,6 +754,17 @@ def test_config_validation(): ) with pytest.raises(ValueError, match="only used"): PPOActorConfig(loss_aggregation="seq_mean", loss_aggregation_divisor=10) + with pytest.raises(ValueError, match="group_size is only valid"): + PPOActorConfig(loss_aggregation="seq_mean", group_size=2) + with pytest.raises(ValueError, match="positive finite"): + PPOActorConfig(m2_threshold=0) + with pytest.raises(ValueError, match="incompatible"): + PPOActorConfig(m2_threshold=0.1, prox_logp_method=PROX_LOGP_METHOD_LOGLINEAR) + with pytest.raises(ValueError, match="incompatible"): + PPOActorConfig( + m2_threshold=0.1, + prox_logp_method=PROX_LOGP_METHOD_REUSE_TRAIN_LOGP, + ) PPOActorConfig(loss_aggregation="constant", loss_aggregation_divisor=10) GRPOConfig(gconfig=GenerationHyperparameters(n_samples=1)) GRPOConfig( diff --git a/tests/test_train_engine.py b/tests/test_train_engine.py index 937e35b430..3a72fe4f0d 100644 --- a/tests/test_train_engine.py +++ b/tests/test_train_engine.py @@ -12,7 +12,7 @@ from tests.utils import get_model_path -from areal.api import FinetuneSpec, LossReduction, SaveLoadMeta +from areal.api import FinetuneSpec, SaveLoadMeta from areal.api.cli_args import ( FSDPEngineConfig, MicroBatchSpec, @@ -177,10 +177,8 @@ def test_eval_batch(engine, mock_input): engine.config.mb_spec = MicroBatchSpec(n_mbs=2, max_tokens_per_mb=100) eval_result = engine.eval_batch( input_=mock_input, - loss_reduction=LossReduction.mean( - loss_fn=mock_loss_fn, - normalizer_fn=lambda x: x["cu_seqlens"][-1], - ), + loss_fn=mock_loss_fn, + loss_weight_fn=lambda x: x["cu_seqlens"][-1], ) assert isinstance(eval_result, torch.Tensor), "Evaluation should return a tensor" assert eval_result.is_cuda, "Evaluation tensor should be on CUDA device" @@ -193,10 +191,8 @@ def test_train_batch(engine, mock_input): engine.config.mb_spec = MicroBatchSpec(n_mbs=2, max_tokens_per_mb=100) train_result = engine.train_batch( input_=mock_input, - loss_reduction=LossReduction.mean( - loss_fn=mock_loss_fn, - normalizer_fn=lambda x: x["cu_seqlens"][-1], - ), + loss_fn=mock_loss_fn, + loss_weight_fn=lambda x: x["cu_seqlens"][-1], ) assert isinstance(train_result, dict), "Training should return a dictionary" assert train_result["grad_norm"] is not None diff --git a/tests/test_tree_training.py b/tests/test_tree_training.py index f20550be00..ea47011645 100644 --- a/tests/test_tree_training.py +++ b/tests/test_tree_training.py @@ -333,7 +333,7 @@ def normalizer_fn(input_data): # Create baseline engine baseline_engine = _create_engine(engine_type, port="7777") baseline_engine.train() - _ = baseline_engine.train_batch( + _ = baseline_engine.train_batch_with_reduction( inputs, loss_reduction=loss_reduction, ) @@ -358,7 +358,7 @@ def normalizer_fn(input_data): experiment_name="test_tree", ) tree_engine.train() - _ = tree_engine.train_batch( + _ = tree_engine.train_batch_with_reduction( inputs, loss_reduction=loss_reduction, ) diff --git a/tests/torchrun/run_fsdp_dcp_distributed.py b/tests/torchrun/run_fsdp_dcp_distributed.py index 776f014180..b6ce198f81 100644 --- a/tests/torchrun/run_fsdp_dcp_distributed.py +++ b/tests/torchrun/run_fsdp_dcp_distributed.py @@ -190,7 +190,7 @@ def test_train_dcp_save_load(alloc_mode: str, output: str | None = None): loss_fn=mock_loss_fn, normalizer_fn=lambda x: x["cu_seqlens"][-1], ) - train_result = engine.train_batch( + train_result = engine.train_batch_with_reduction( input_data, loss_reduction=loss_reduction, ) @@ -201,7 +201,7 @@ def test_train_dcp_save_load(alloc_mode: str, output: str | None = None): print(f"Rank {rank} checkpoint saved") # Train step 2 - train_result = engine.train_batch( + train_result = engine.train_batch_with_reduction( input_data, loss_reduction=loss_reduction, ) @@ -220,7 +220,7 @@ def test_train_dcp_save_load(alloc_mode: str, output: str | None = None): # Train step 2 again after load engine.train() - train_result = engine.train_batch( + train_result = engine.train_batch_with_reduction( input_data, loss_reduction=loss_reduction, ) diff --git a/tests/torchrun/run_fsdp_ulysses_train_batch.py b/tests/torchrun/run_fsdp_ulysses_train_batch.py index 9fc8dd6bb2..f5137cff8e 100644 --- a/tests/torchrun/run_fsdp_ulysses_train_batch.py +++ b/tests/torchrun/run_fsdp_ulysses_train_batch.py @@ -147,7 +147,7 @@ def test_ulysses(model_type: str): model_type, mb_spec, ulysses_sp_size=2, init_optimizer=True ) engine.train() - engine.train_batch( + engine.train_batch_with_reduction( input_=input, loss_reduction=LossReduction.mean( loss_fn=mock_loss_fn, @@ -159,7 +159,7 @@ def test_ulysses(model_type: str): engine_golden = make_engine(model_type, mb_spec, init_optimizer=True) engine_golden.train() for input in input_chunks: - engine_golden.train_batch( + engine_golden.train_batch_with_reduction( input_=input, loss_reduction=LossReduction.mean( loss_fn=mock_loss_fn, diff --git a/tests/torchrun/run_megatron_engine_distributed.py b/tests/torchrun/run_megatron_engine_distributed.py index cfb890986d..a5b5d4073d 100644 --- a/tests/torchrun/run_megatron_engine_distributed.py +++ b/tests/torchrun/run_megatron_engine_distributed.py @@ -266,7 +266,7 @@ def test_train( group=engine.context_and_model_parallel_group, ) - train_result = engine.train_batch( + train_result = engine.train_batch_with_reduction( input_=bcasted_input, loss_reduction=LossReduction.mean( loss_fn=mock_loss_fn, @@ -276,7 +276,7 @@ def test_train( print(f"final rank {rank} train_result: {train_result}") - sum_train_result = engine.train_batch( + sum_train_result = engine.train_batch_with_reduction( input_=bcasted_input, loss_reduction=LossReduction.sum( loss_fn=mock_loss_sum_fn, @@ -342,7 +342,7 @@ def test_grad_norm_mb_invariance( group=engine.context_and_model_parallel_group, ) - result = engine.train_batch( + result = engine.train_batch_with_reduction( input_=bcasted_input, loss_reduction=LossReduction.mean( loss_fn=mock_loss_fn, @@ -419,7 +419,7 @@ def test_train_dcp_save_load( ) # train step 1 - train_result = engine.train_batch( + train_result = engine.train_batch_with_reduction( input_=bcasted_input, loss_reduction=LossReduction.mean( loss_fn=mock_loss_fn, @@ -436,7 +436,7 @@ def test_train_dcp_save_load( engine.save(save_load_meta) # train step 2 - engine.train_batch( + engine.train_batch_with_reduction( input_=bcasted_input, loss_reduction=LossReduction.mean( loss_fn=mock_loss_fn, @@ -456,7 +456,7 @@ def test_train_dcp_save_load( engine.train() # train step 2 after recover - engine.train_batch( + engine.train_batch_with_reduction( input_=bcasted_input, loss_reduction=LossReduction.mean( loss_fn=mock_loss_fn, @@ -601,7 +601,7 @@ def test_train_hf_save_load( src_rank=engine.current_data_parallel_head(), group=engine.context_and_model_parallel_group, ) - train_result = engine.train_batch( + train_result = engine.train_batch_with_reduction( input_=bcasted_input, loss_reduction=LossReduction.mean( loss_fn=mock_loss_fn, diff --git a/tests/torchrun/run_megatron_engine_vlm_distributed.py b/tests/torchrun/run_megatron_engine_vlm_distributed.py index 904651e16a..be7f1a023b 100644 --- a/tests/torchrun/run_megatron_engine_vlm_distributed.py +++ b/tests/torchrun/run_megatron_engine_vlm_distributed.py @@ -278,7 +278,7 @@ def test_vlm_train(backend: str, output: str | None = None): bcasted_input = _make_input(engine) engine.train() - train_result = engine.train_batch( + train_result = engine.train_batch_with_reduction( input_=bcasted_input, loss_reduction=LossReduction.mean( loss_fn=lambda logprobs, entropy, input_data, **kwargs: ( From 2496822a9342fd4d4194d1dbd9262d796402fd0e Mon Sep 17 00:00:00 2001 From: EazyReal <8047065+EazyReal@users.noreply.github.com> Date: Sat, 11 Jul 2026 17:48:22 -0700 Subject: [PATCH 4/4] docs(engine): describe loss weight aggregation Keep function documentation focused on behavior and contract rather than refactor history or compatibility rationale. --- areal/engine/core/train_engine.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/areal/engine/core/train_engine.py b/areal/engine/core/train_engine.py index 27a6a3335c..564f774fe3 100644 --- a/areal/engine/core/train_engine.py +++ b/areal/engine/core/train_engine.py @@ -43,8 +43,22 @@ def compute_total_loss_weight( ) -> torch.Tensor: """Compute total loss weight and all_reduce across data parallel group. - This is the original pre-PR helper retained for callers that still use the - callback-based engine API. + This aggregates the loss weights from all micro-batches and reduces + them across the data parallel group to get a global normalization factor. + + Parameters + ---------- + mb_list : MicroBatchList + The list of micro-batches. + loss_weight_fn : Callable[[dict[str, Any]], torch.Tensor] + Function to compute loss weight for each micro-batch. + dp_group : dist.ProcessGroup + The data parallel process group for all_reduce. + + Returns + ------- + torch.Tensor + The total loss weight (scalar tensor) after all_reduce. """ total_weight = ( torch.stack([loss_weight_fn(mb) for mb in mb_list.mbs])