diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index 21e51e255b..b0750719a8 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 @@ -1577,6 +1578,41 @@ class PPOActorConfig(TrainEngineConfig): }, ) + 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, @@ -1684,6 +1720,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: @@ -1702,6 +1764,11 @@ def __post_init__(self): "CISPO only supports importance_sampling_level='token'. " "Sequence-level (GSPO-style) CISPO has no published surrogate." ) + if self.loss_aggregation != "token_mean": + raise ValueError( + "CISPO only supports loss_aggregation='token_mean'. " + "Non-token-mean CISPO has no published surrogate." + ) super().__post_init__() @@ -2167,6 +2234,17 @@ class InferenceEngineConfig: default=1, metadata={"help": "Batch size for consuming rollouts from the queue."}, ) + min_valid_group_size: int = field( + default=1, + metadata={ + "help": "Minimum non-None trajectories required to keep a rollout " + "group. Default 1 keeps non-empty partial groups; set to " + "gconfig.n_samples to require full groups. Must be in [1, group_size].", + "help_zh": "保留一个 rollout group 所需的最小非 None trajectory 数。" + "默认值 1 会保留非空的 partial group；设为 gconfig.n_samples " + "则要求完整 group。必须在 [1, group_size] 范围内。", + }, + ) max_head_offpolicyness: int = field( default=0, metadata={ @@ -2325,6 +2403,10 @@ def __post_init__(self): ) if not self.admin_api_key or not self.admin_api_key.strip(): raise ValueError("admin_api_key must not be empty or whitespace-only") + if self.min_valid_group_size < 1: + raise ValueError( + f"min_valid_group_size must be >= 1, got {self.min_valid_group_size}" + ) if ( self._version == "v2" and self.agent is not None @@ -3070,6 +3152,50 @@ 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}." + ) + g = self.actor.mb_spec.granularity + if g % self.actor.group_size != 0: + new_g = -(-g // self.actor.group_size) * self.actor.group_size + logger.warning( + "loss_aggregation='prompt_mean': bumping " + "actor.mb_spec.granularity %d -> %d (a multiple of " + "gconfig.n_samples=%d) so each microbatch holds whole " + "prompt-groups.", + g, + new_g, + self.actor.group_size, + ) + self.actor.mb_spec.granularity = new_g + if ( + self.rollout is not None + and self.rollout.min_valid_group_size < self.actor.group_size + ): + logger.warning( + "loss_aggregation='prompt_mean' needs whole prompt-groups; " + "raising rollout.min_valid_group_size %d -> %d so " + "under-filled rollout groups are dropped, not mis-grouped.", + self.rollout.min_valid_group_size, + self.actor.group_size, + ) + self.rollout.min_valid_group_size = self.actor.group_size + + if ( + self.rollout is not None + and self.rollout.min_valid_group_size > self.gconfig.n_samples + ): + raise ValueError( + f"rollout.min_valid_group_size ({self.rollout.min_valid_group_size}) " + f"cannot exceed gconfig.n_samples ({self.gconfig.n_samples}); no " + "group could ever reach the threshold." + ) super().__post_init__() diff --git a/areal/infra/remote_inf_engine.py b/areal/infra/remote_inf_engine.py index fb91af095b..156214384a 100644 --- a/areal/infra/remote_inf_engine.py +++ b/areal/infra/remote_inf_engine.py @@ -74,12 +74,19 @@ def __init__( workflow: RolloutWorkflow, group_size: int, logger: Logger, + min_valid_group_size: int = 1, ): if group_size < 1: raise ValueError(f"group_size must be >= 1, got {group_size}") + if not 1 <= min_valid_group_size <= group_size: + raise ValueError( + f"min_valid_group_size must be in [1, group_size={group_size}], " + f"got {min_valid_group_size}" + ) self.workflow = workflow self.group_size = group_size self.logger = logger + self.min_valid_group_size = min_valid_group_size async def arun_episode( self, engine: InferenceEngine, data: dict[str, Any] @@ -92,11 +99,17 @@ async def arun_episode( valid_results = [r for r in results if r is not None] - # All results None -> return None if not valid_results: return None - # Some results None -> warn and continue with valid ones + if len(valid_results) < self.min_valid_group_size: + self.logger.warning( + f"GroupedRolloutWorkflow: only {len(valid_results)}/{len(results)} " + f"trajectories valid (< min_valid_group_size=" + f"{self.min_valid_group_size}); dropping group" + ) + return None + if len(valid_results) < len(results): self.logger.warning( f"GroupedRolloutWorkflow: {len(results) - len(valid_results)}/{len(results)} " @@ -621,7 +634,12 @@ def _resolve_workflow( raise ValueError("proxy_addr is required for online mode") resolved = self._wrap_openai_agent(None, proxy_addr=proxy_addr) if group_size > 1: - resolved = GroupedRolloutWorkflow(resolved, group_size, self.logger) + resolved = GroupedRolloutWorkflow( + resolved, + group_size, + self.logger, + min_valid_group_size=self.config.min_valid_group_size, + ) return resolved # 1. Already a RolloutWorkflow instance @@ -712,7 +730,12 @@ def _resolve_workflow( # Wrap with GroupedRolloutWorkflow if group_size > 1 if group_size > 1: - resolved = GroupedRolloutWorkflow(resolved, group_size, self.logger) + resolved = GroupedRolloutWorkflow( + resolved, + group_size, + self.logger, + min_valid_group_size=self.config.min_valid_group_size, + ) return resolved diff --git a/areal/trainer/ppo/actor.py b/areal/trainer/ppo/actor.py index 2d7972a450..84e7946fbc 100644 --- a/areal/trainer/ppo/actor.py +++ b/areal/trainer/ppo/actor.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 import functools +from collections.abc import Callable from typing import Any import torch @@ -28,9 +29,16 @@ KLEstimator, Normalization, batched_call, + concat_batch, + split_batch, split_padded_tensor_dict_into_mb_list, ) from areal.utils.functional import ( + LOSS_AGGREGATION_CONSTANT, + LOSS_AGGREGATION_PROMPT_MEAN, + LOSS_AGGREGATION_SEQ_MEAN, + LOSS_AGGREGATION_TOKEN_MEAN, + LOSS_AGGREGATIONS_ALL, cispo_loss_fn, ppo_actor_loss_fn, reward_overlong_penalty, @@ -40,6 +48,61 @@ logger = logging.getLogger("PPOActor") +_LossWeightFn = Callable[[dict[str, Any]], torch.Tensor] +_LossWeightFactory = Callable[[int], _LossWeightFn] + + +def _ppo_loss_weight(data: dict[str, Any]) -> torch.Tensor: + return data["loss_mask"].count_nonzero() + + +def _unit_count_weight(data: dict[str, Any], unit: int) -> torch.Tensor: + cu_seqlens = data.get("cu_seqlens") + if cu_seqlens is not None: + n_seqs = cu_seqlens.numel() - 1 + else: + n_seqs = data["loss_mask"].shape[0] + n_units = n_seqs // unit + return torch.tensor(n_units, dtype=torch.float32, device=data["loss_mask"].device) + + +def _make_unit_count_weight_fn(unit: int) -> _LossWeightFn: + def unit_count_weight(data: dict[str, Any]) -> torch.Tensor: + return _unit_count_weight(data, unit) + + return unit_count_weight + + +def _token_loss_weight_fn(_group_size: int) -> _LossWeightFn: + return _ppo_loss_weight + + +def _sequence_loss_weight_fn(_group_size: int) -> _LossWeightFn: + return _make_unit_count_weight_fn(unit=1) + + +def _prompt_loss_weight_fn(group_size: int) -> _LossWeightFn: + return _make_unit_count_weight_fn(unit=group_size) + + +_LOSS_WEIGHT_FACTORIES: dict[str, _LossWeightFactory] = { + LOSS_AGGREGATION_TOKEN_MEAN: _token_loss_weight_fn, + LOSS_AGGREGATION_SEQ_MEAN: _sequence_loss_weight_fn, + LOSS_AGGREGATION_PROMPT_MEAN: _prompt_loss_weight_fn, + LOSS_AGGREGATION_CONSTANT: _sequence_loss_weight_fn, +} + + +def _make_loss_weight_fn(loss_aggregation: str, group_size: int) -> _LossWeightFn: + """Return the microbatch weight paired with ``aggregate_pg_loss``.""" + try: + return _LOSS_WEIGHT_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 + class PPOActor: def __init__(self, config: PPOActorConfig, engine: TrainEngine): @@ -141,9 +204,13 @@ def _compute_logp(self, data: dict[str, Any]) -> torch.Tensor | None: @trace_perf("ppo_actor.compute_advantages", category="compute") def compute_advantages(self, data: list[dict[str, Any]]) -> list[dict[str, Any]]: - return batched_call(self._compute_advantages, data) + batched, meta = concat_batch(data) + result = self._compute_advantages(batched, group_sizes=meta.traj_group_sizes) + return split_batch(result, meta) - def _compute_advantages(self, data: dict[str, Any]) -> dict[str, Any]: + def _compute_advantages( + self, data: dict[str, Any], group_sizes: list[int] | None = None + ) -> dict[str, Any]: bs = data["input_ids"].shape[0] max_seqlen = data["input_ids"].shape[1] batch_indices = torch.arange( @@ -171,7 +238,7 @@ def _compute_advantages(self, data: dict[str, Any]) -> dict[str, Any]: reward_score, max=self.reward_clip, min=-self.reward_clip ) if self.reward_norm: - reward_score = self.reward_norm(reward_score) + reward_score = self.reward_norm(reward_score, group_sizes=group_sizes) loss_mask = data["loss_mask"].float() loss_mask = torch.roll(loss_mask, shifts=-1, dims=-1) @@ -237,7 +304,7 @@ def _compute_advantages(self, data: dict[str, Any]) -> dict[str, Any]: # Optionally perform advantage normalization. if self.adv_norm is not None: - advantages = self.adv_norm(advantages, loss_mask) + advantages = self.adv_norm(advantages, loss_mask, group_sizes=group_sizes) # Store data in the dict. data["advantages"] = advantages @@ -334,9 +401,17 @@ def _ppo_update(self, data: dict[str, Any]) -> None: data.pop(key, None) # NOTE: calling engine.train() is critical to enabling gradient checkpointing self.engine.train() + outer_granularity = ( + self.config.group_size + if self.config.loss_aggregation == "prompt_mean" + else 1 + ) 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"): @@ -361,8 +436,13 @@ 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, + ), + loss_weight_fn=_make_loss_weight_fn( + self.config.loss_aggregation, self.config.group_size ), - loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), ) stats_tracker.scalar(**train_stat) @@ -424,6 +504,9 @@ 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, vocab_min_logits: torch.Tensor | None = None, vocab_max_logits: torch.Tensor | None = None, ): @@ -488,6 +571,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, ) else: loss, stat = ppo_actor_loss_fn( @@ -502,6 +588,9 @@ 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, ) # Joint Distillation KL Loss diff --git a/areal/utils/data.py b/areal/utils/data.py index 09368e4ef2..fd9d7c1816 100644 --- a/areal/utils/data.py +++ b/areal/utils/data.py @@ -1400,6 +1400,7 @@ def __call__( loss_mask: torch.Tensor | None = None, high_precision: bool = True, reduce_group=None, + group_sizes: list[int] | None = None, ) -> torch.Tensor: bs = x.size(0) eps = self.eps @@ -1408,6 +1409,11 @@ def __call__( if loss_mask is not None and loss_mask.sum().item() == 0: return x.float() + # Per-group row boundaries, only needed for group-level normalization. + group_bounds = None + if self.mean_level == "group" or self.std_level == "group": + group_bounds = self._resolve_group_bounds(bs, group_sizes) + # Step 1: Compute mean if self.mean_level == "batch": mean = self._compute_mean( @@ -1421,13 +1427,12 @@ def __call__( mean = mean.expand_as(x) elif self.mean_level == "group": mean = torch.zeros_like(x) - for i in range(0, bs // self.group_size): - s = slice(i * self.group_size, (i + 1) * self.group_size) + for s in group_bounds: xx = x[s] m = loss_mask[s] if loss_mask is not None else None - # Special case: with group_size=1 and leave_one_out=True, mean should be 0 - if self.group_size == 1 and self.mean_leave1out: + # A group of one has no leave-one-out baseline -> mean 0. + if xx.size(0) == 1 and self.mean_leave1out: dtype = torch.float64 if high_precision else torch.float32 group_mean = torch.zeros( (1, *xx.shape[1:]), dtype=dtype, device=xx.device @@ -1465,14 +1470,13 @@ def __call__( std = std.expand_as(x) elif self.std_level == "group": std = torch.zeros_like(x) - for i in range(0, bs // self.group_size): - s = slice(i * self.group_size, (i + 1) * self.group_size) + for s in group_bounds: xx = x[s] m = loss_mask[s] if loss_mask is not None else None group_mean_slice = mean[s] # already computed and expanded - # Special case: with group_size=1 and std_unbiased=True, std should be 1 for numerical stability - if self.group_size == 1 and self.std_unbiased: + # A group of one has zero variance -> std 1 for stability. + if xx.size(0) == 1 and self.std_unbiased: dtype = torch.float64 if high_precision else torch.float32 group_std = torch.ones( (1, *xx.shape[1:]), dtype=dtype, device=xx.device @@ -1495,6 +1499,37 @@ def __call__( # Normalize return (x_centered / (std + eps)).float() + def _resolve_group_bounds( + self, bs: int, group_sizes: list[int] | None + ) -> list[slice]: + """Partition rows ``[0, bs)`` into per-group slices for group-level norm. + + ``group_sizes`` gives the actual per-group row counts, which may be + unequal (e.g. partial groups). When omitted, fall back to fixed-size + positional groups of ``self.group_size`` and require exact divisibility. + """ + if group_sizes is not None: + if any(k < 1 for k in group_sizes): + raise ValueError(f"group_sizes must be positive, got {group_sizes}") + total = sum(group_sizes) + if total != bs: + raise ValueError(f"group_sizes sum to {total} but batch size is {bs}") + bounds: list[slice] = [] + offset = 0 + for k in group_sizes: + bounds.append(slice(offset, offset + k)) + offset += k + return bounds + if bs % self.group_size != 0: + raise ValueError( + f"batch size {bs} is not divisible by group_size " + f"{self.group_size}; pass group_sizes for partial/unequal groups" + ) + return [ + slice(i * self.group_size, (i + 1) * self.group_size) + for i in range(bs // self.group_size) + ] + @staticmethod def _compute_mean( x: torch.Tensor, diff --git a/areal/utils/functional/__init__.py b/areal/utils/functional/__init__.py index b09c6af5f2..e4b69aafe4 100644 --- a/areal/utils/functional/__init__.py +++ b/areal/utils/functional/__init__.py @@ -1,7 +1,13 @@ # 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, apply_rejection_sampling, cispo_loss_fn, dpo_pair_logratios, @@ -19,7 +25,13 @@ __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", "apply_rejection_sampling", "cispo_loss_fn", "dpo_pair_logratios", diff --git a/areal/utils/functional/functional.py b/areal/utils/functional/functional.py index 8aaed7730a..691b283beb 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 @@ -449,6 +451,256 @@ 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, + str, + ], + torch.Tensor, +] + + +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 _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, + _mode: str, +) -> torch.Tensor: + denom = den_mask.count_nonzero().clamp_min(1) + return torch.where(num_mask, pg_loss, 0).sum() / denom + + +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, + _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() / (n_seqs * 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, + mode: str, +) -> torch.Tensor: + return _aggregate_unit_mean( + pg_loss, num_mask, den_mask, unit=1, cu_seqlens=cu_seqlens, 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, + mode: str, +) -> torch.Tensor: + return _aggregate_unit_mean( + pg_loss, + num_mask, + den_mask, + unit=group_size, + cu_seqlens=cu_seqlens, + 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, + 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, mode, masked_pg, den_f) + return _aggregate_packed_unit_mean( + pg_loss, unit, cu_seqlens, mode, masked_pg, den_f + ) + + +def _aggregate_padded_unit_mean( + pg_loss: torch.Tensor, + unit: int, + 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 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)).clamp_min(1) + return (num / den).mean() + + +def _aggregate_packed_unit_mean( + pg_loss: torch.Tensor, + unit: int, + cu_seqlens: torch.Tensor, + 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 + 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_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to( + device=pg_loss.device, dtype=torch.long + ) + 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 (num / den.clamp_min(1)).mean() + + +_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, +} +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 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, +) -> 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 means over ``group_size`` consecutive + sequences. ``constant`` averages each sequence's masked token sum divided + by ``loss_aggregation_divisor``. The returned microbatch mean must be paired + with ``_make_loss_weight_fn`` so distributed reduction computes the matching + global mean. ``denom_mask`` overrides only the denominator for data-dependent + reductions. + """ + _validate_pg_loss_aggregation_config(loss_aggregation, loss_aggregation_divisor) + + num_mask = loss_mask.bool() + den_mask = num_mask if denom_mask is None else denom_mask.bool() + + return _PG_LOSS_REDUCERS[loss_aggregation]( + pg_loss, + num_mask, + den_mask, + group_size, + loss_aggregation_divisor, + cu_seqlens, + loss_aggregation, + ) + + def ppo_actor_loss_fn( logprobs: torch.Tensor, proximal_logprobs: torch.Tensor, @@ -461,6 +713,9 @@ 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, ) -> tuple[torch.Tensor, dict]: """PPO actor loss function with optional rejection sampling. @@ -499,11 +754,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 +814,15 @@ 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 + pg_loss = aggregate_pg_loss( + 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, + ) clip_mask.logical_and_(loss_mask) dual_clip_mask.logical_and_(loss_mask) stat = dict( @@ -592,6 +852,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, ) -> tuple[torch.Tensor, dict]: """SAPO (Soft Adaptive Policy Optimization) loss with asymmetric sigmoid gates. @@ -613,7 +876,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,7 +906,14 @@ 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 + pg_loss = aggregate_pg_loss( + pg_loss, + loss_mask, + loss_aggregation=loss_aggregation, + group_size=group_size, + loss_aggregation_divisor=loss_aggregation_divisor, + cu_seqlens=cu_seqlens, + ) # Return stat dict compatible with PPO (fake clip_mask for logging compatibility) stat = dict( diff --git a/docs/en/cli_reference.md b/docs/en/cli_reference.md index 78d9515cde..fef8af680e 100644 --- a/docs/en/cli_reference.md +++ b/docs/en/cli_reference.md @@ -402,6 +402,9 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `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. | +| `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. | @@ -542,6 +545,7 @@ Configuration for inference servers, including offpolicyness control. | `max_concurrent_rollouts` | integer \| None | `None` | Maximum number of concurrent rollouts to the inference engine. Defaults to consumer_batch_size. | | `queue_size` | integer \| None | `None` | Input/Output queue size for async rollout. | | `consumer_batch_size` | integer | `1` | Batch size for consuming rollouts from the queue. | +| `min_valid_group_size` | integer | `1` | Minimum non-None trajectories required to keep a rollout group. Default 1 keeps non-empty partial groups; set to gconfig.n_samples to require full groups. Must be in \[1, group_size\]. | | `max_head_offpolicyness` | integer | `0` | Maximum off-policyness for the head. If the current version is more than this many versions behind, the request will not be accepted. | | `enable_rollout_tracing` | boolean | `False` | Whether to output verbose tracing messages for each generation request. | | `check_trajectory_format` | boolean | `False` | Whether to check the format of produced trajectories of a customized workflow. Useful when debugging the workflow in isolation. Should be False during RL training. | 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/cli_reference.md b/docs/zh/cli_reference.md index 455339316f..93fff2bf2b 100644 --- a/docs/zh/cli_reference.md +++ b/docs/zh/cli_reference.md @@ -400,6 +400,9 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `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. | +| `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. | @@ -540,6 +543,7 @@ Configuration for inference servers, including offpolicyness control. | `max_concurrent_rollouts` | integer \| None | `None` | Maximum number of concurrent rollouts to the inference engine. Defaults to consumer_batch_size. | | `queue_size` | integer \| None | `None` | Input/Output queue size for async rollout. | | `consumer_batch_size` | integer | `1` | Batch size for consuming rollouts from the queue. | +| `min_valid_group_size` | integer | `1` | 保留一个 rollout group 所需的最小非 None trajectory 数。默认值 1 会保留非空的 partial group；设为 gconfig.n_samples 则要求完整 group。必须在 \[1, group_size\] 范围内。 | | `max_head_offpolicyness` | integer | `0` | Maximum off-policyness for the head. If the current version is more than this many versions behind, the request will not be accepted. | | `enable_rollout_tracing` | boolean | `False` | Whether to output verbose tracing messages for each generation request. | | `check_trajectory_format` | boolean | `False` | Whether to check the format of produced trajectories of a customized workflow. Useful when debugging the workflow in isolation. Should be False during RL training. | diff --git a/tests/test_cispo_loss.py b/tests/test_cispo_loss.py index 44745a5a31..49b7cd5ae2 100644 --- a/tests/test_cispo_loss.py +++ b/tests/test_cispo_loss.py @@ -169,5 +169,18 @@ def test_cispo_config_validation(): eps_clip_higher=4.0, importance_sampling_level="sequence", ) + with pytest.raises(ValueError, match="loss_aggregation"): + PPOActorConfig( + use_cispo_loss=True, + eps_clip_higher=4.0, + loss_aggregation="seq_mean", + ) + with pytest.raises(ValueError, match="loss_aggregation"): + PPOActorConfig( + use_cispo_loss=True, + eps_clip_higher=4.0, + loss_aggregation="constant", + loss_aggregation_divisor=10, + ) # Valid configuration does not raise. PPOActorConfig(use_cispo_loss=True, eps_clip=1.0, eps_clip_higher=4.0) diff --git a/tests/test_grouped_rollout_min_valid.py b/tests/test_grouped_rollout_min_valid.py new file mode 100644 index 0000000000..3100996a2d --- /dev/null +++ b/tests/test_grouped_rollout_min_valid.py @@ -0,0 +1,63 @@ +import asyncio +import logging + +import pytest +import torch + +from areal.infra.remote_inf_engine import GroupedRolloutWorkflow + + +class _FakeWorkflow: + def __init__(self, n_none: int): + self._n_none = n_none + self._calls = 0 + + async def arun_episode(self, engine, data): + self._calls += 1 + if self._calls <= self._n_none: + return None + return { + "input_ids": torch.tensor([[1, 2, 3]]), + "attention_mask": torch.tensor([[1, 1, 1]]), + } + + +def _run(group_size, n_none, min_valid_group_size): + wf = GroupedRolloutWorkflow( + _FakeWorkflow(n_none), + group_size=group_size, + logger=logging.getLogger("test"), + min_valid_group_size=min_valid_group_size, + ) + return asyncio.run(wf.arun_episode(engine=None, data={})) + + +@pytest.mark.parametrize( + "min_valid, n_none, kept", + [ + (1, 1, True), + (2, 2, True), + (2, 3, False), + (4, 1, False), + ], +) +def test_min_valid_group_size_threshold(min_valid, n_none, kept): + out = _run(group_size=4, n_none=n_none, min_valid_group_size=min_valid) + assert (out is not None) == kept + if kept: + assert out["input_ids"].shape[0] == 4 - n_none + + +def test_all_none_returns_none(): + assert _run(group_size=4, n_none=4, min_valid_group_size=1) is None + + +@pytest.mark.parametrize("bad", [0, 5]) +def test_threshold_outside_valid_range_raises(bad): + with pytest.raises(ValueError, match="min_valid_group_size must be in"): + GroupedRolloutWorkflow( + _FakeWorkflow(0), + group_size=4, + logger=logging.getLogger("test"), + min_valid_group_size=bad, + ) diff --git a/tests/test_partial_group_norm.py b/tests/test_partial_group_norm.py new file mode 100644 index 0000000000..0ca099fc6e --- /dev/null +++ b/tests/test_partial_group_norm.py @@ -0,0 +1,98 @@ +import pytest +import torch + +from areal.api.cli_args import NormConfig +from areal.utils.data import Normalization, concat_batch + + +def _group_norm(mean_level="group", std_level="group", group_size=4, **kw): + return Normalization( + NormConfig( + mean_level=mean_level, std_level=std_level, group_size=group_size, **kw + ) + ) + + +def test_group_sizes_matches_positional_for_full_groups_reward_path(): + torch.manual_seed(0) + x = torch.randn(8) + norm = _group_norm() + torch.testing.assert_close(norm(x), norm(x, group_sizes=[4, 4])) + + +def test_group_sizes_matches_positional_for_full_groups_adv_path(): + torch.manual_seed(1) + x = torch.randn(8, 5) + mask = torch.ones(8, 5) + mask[:, 4] = 0.0 + norm = _group_norm() + torch.testing.assert_close(norm(x, mask), norm(x, mask, group_sizes=[4, 4])) + + +def test_trailing_partial_group_does_not_blow_up(): + torch.manual_seed(2) + x = torch.randn(11) * 3.0 + out = _group_norm()(x, group_sizes=[4, 4, 3]) + assert torch.isfinite(out).all() + for s in (slice(0, 4), slice(4, 8), slice(8, 11)): + torch.testing.assert_close(out[s].mean(), torch.tensor(0.0), atol=1e-5, rtol=0) + + +def test_mid_batch_partial_group_avoids_cross_prompt_sign_flip(): + x = torch.tensor([0.5, -0.5, 1.0, -1.0, 2.0, -2.0, 0.0, 3.0, 0.0, 1.0, 2.0, 100.0]) + norm = _group_norm() + + out_positional = norm(x) + out_grouped = norm(x, group_sizes=[4, 4, 3, 1]) + + torch.testing.assert_close(out_positional[:8], out_grouped[:8]) + assert out_grouped[10] > 0 and out_positional[10] < 0 + torch.testing.assert_close(out_grouped[11], torch.tensor(0.0), atol=1e-6, rtol=0) + + +@pytest.mark.parametrize("std_unbiased", [True, False]) +def test_singleton_group_is_finite_and_zero(std_unbiased): + torch.manual_seed(3) + x = torch.randn(5) * 10.0 + out = _group_norm(std_unbiased=std_unbiased)(x, group_sizes=[4, 1]) + assert torch.isfinite(out).all() + torch.testing.assert_close(out[4], torch.tensor(0.0), atol=1e-6, rtol=0) + + +@pytest.mark.parametrize("std_unbiased", [True, False]) +def test_singleton_group_masked_path_is_finite(std_unbiased): + torch.manual_seed(4) + x = torch.randn(5, 6) * 10.0 + mask = torch.ones(5, 6) + mask[:, 5] = 0.0 + out = _group_norm(std_unbiased=std_unbiased)(x, mask, group_sizes=[4, 1]) + assert torch.isfinite(out).all() + torch.testing.assert_close(out[:, 5], torch.zeros(5), atol=1e-6, rtol=0) + + +def test_non_divisible_batch_without_group_sizes_raises(): + with pytest.raises(ValueError, match="not divisible"): + _group_norm()(torch.randn(11)) + + +def test_group_sizes_must_sum_to_batch_size(): + with pytest.raises(ValueError, match="group_sizes sum"): + _group_norm()(torch.randn(8), group_sizes=[4, 3]) + + +def test_group_sizes_must_be_positive(): + with pytest.raises(ValueError, match="group_sizes must be positive"): + _group_norm()(torch.randn(8), group_sizes=[8, 0]) + + +def test_batch_level_ignores_group_sizes(): + out = _group_norm(mean_level="batch", std_level="batch")(torch.randn(11)) + assert torch.isfinite(out).all() + + +def test_concat_batch_records_partial_group_sizes(): + def traj(k): + return {"rewards": torch.randn(k), "attention_mask": torch.ones(k, 3)} + + _, meta = concat_batch([traj(4), traj(4), traj(3)]) + assert meta.traj_group_sizes == [4, 4, 3] diff --git a/tests/test_prompt_mean_loss.py b/tests/test_prompt_mean_loss.py new file mode 100644 index 0000000000..91aeeeed37 --- /dev/null +++ b/tests/test_prompt_mean_loss.py @@ -0,0 +1,289 @@ +# SPDX-License-Identifier: Apache-2.0 +import pytest +import torch + +from areal.api.cli_args import ( + GenerationHyperparameters, + GRPOConfig, + InferenceEngineConfig, + MicroBatchSpec, + PPOActorConfig, +) +from areal.trainer.ppo.actor import _make_loss_weight_fn +from areal.utils.functional import aggregate_pg_loss + +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_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_weight_pairing_realizes_global_mean(aggregation): + group_size = 2 if aggregation == "prompt_mean" else 1 + divisor = CONSTANT_DIVISOR if aggregation == "constant" else None + weight_fn = _make_loss_weight_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, + ) + w = weight_fn({"loss_mask": mb_mask}) + num = num + loss_mb * w + den = den + w + torch.testing.assert_close(num / den, full) + + +@pytest.mark.parametrize( + "aggregation", ["token_mean", "seq_mean", "prompt_mean", "constant"] +) +def test_loss_weight_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 + weight_fn = _make_loss_weight_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, + ) + w = weight_fn({"loss_mask": mb_mask, "cu_seqlens": mb_cu_seqlens}) + num = num + loss_mb * w + den = den + w + 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)) + + +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_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_auto_bumped_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 == 4 + cfg = GRPOConfig( + gconfig=GenerationHyperparameters(n_samples=8), + actor=PPOActorConfig(loss_aggregation="prompt_mean"), + ) + assert cfg.actor.mb_spec.granularity == 8 + + +def test_prompt_mean_drops_under_filled_groups(): + cfg = GRPOConfig( + gconfig=GenerationHyperparameters(n_samples=4), + actor=PPOActorConfig( + loss_aggregation="prompt_mean", mb_spec=MicroBatchSpec(granularity=4) + ), + ) + assert cfg.rollout.min_valid_group_size == 4 + + +def test_min_valid_group_size_cannot_exceed_n_samples(): + with pytest.raises(ValueError, match="cannot exceed gconfig.n_samples"): + GRPOConfig( + gconfig=GenerationHyperparameters(n_samples=4), + rollout=InferenceEngineConfig(min_valid_group_size=5), + ) + + +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"), + )