feat(ppo): make loss aggregation configurable - #1546
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces configurable policy-gradient loss aggregation modes (token_mean, seq_mean, prompt_mean, and constant) for PPO actor training, along with corresponding CLI configuration options, documentation updates, and comprehensive unit tests. The review feedback identifies a performance concern in the newly added _sequence_sums function in areal/utils/functional/loss_aggregation.py, where checking torch.any(sequence_lengths < 0) inside an if statement triggers a GPU-to-CPU synchronization on every microbatch. It is recommended to remove this redundant check from the hot path to avoid degrading training throughput.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Trust packer-owned cumulative sequence lengths and provide exact repeat output sizes so aggregation does not synchronize to discover device-side values or shapes. Refs: areal-project#1546
Allow PPO-family actor losses to weight tokens, sequences, prompt groups, or fixed-length response sums without changing the engine contract. Key changes: - Pair each local policy-gradient mean with its matching engine weight - Preserve explicit prompt-group boundaries across microbatch splitting - Keep token_mean and existing backend behavior unchanged - Document and regression-test padded, packed, ragged, and filtered inputs Refs: areal-project#1423 Signed-off-by: EazyReal <8047065+EazyReal@users.noreply.github.com>
Trust packer-owned cumulative sequence lengths and provide exact repeat output sizes so aggregation does not synchronize to discover device-side values or shapes. Refs: areal-project#1546 Signed-off-by: EazyReal <8047065+EazyReal@users.noreply.github.com>
d4997ad to
aeb6391
Compare
| loss_aggregation: str = field( | ||
| default="token_mean", | ||
| metadata={ | ||
| "help": "Policy-gradient loss reduction. 'token_mean' averages valid " |
There was a problem hiding this comment.
Please note the format of the docstring.
Present each supported loss aggregation choice as a name-description pair so generated references follow the existing config documentation style. Refs: areal-project#1546 Signed-off-by: EazyReal <8047065+EazyReal@users.noreply.github.com>
| "loss_aggregation_divisor. Non-token modes require sequence " | ||
| "boundaries; tree-packed actor training currently supports only " | ||
| "'token_mean'.", | ||
| "help_zh": "Policy-gradient loss 的归约方式。" |
There was a problem hiding this comment.
Please change the description to English.
There was a problem hiding this comment.
Removed — dropped the help_zh strings and reverted the doc generator's language plumbing, so config help stays English-only like every other field; the zh CLI reference is regenerated accordingly.
| return torch.where(mask, loss, 0).to(torch.float32) | ||
|
|
||
|
|
||
| def _resolve_masks( |
There was a problem hiding this comment.
These free functions are best added to PolicyGradientReduction.
There was a problem hiding this comment.
Moved them into PolicyGradientReduction as private methods — they had no consumer outside the class.
…uction Give the reduction contract a single home by moving the module-level aggregation helpers into the dataclass they serve; no consumer other than PolicyGradientReduction ever used them. Refs: areal-project#1546 Signed-off-by: EazyReal <8047065+EazyReal@users.noreply.github.com>
Drop the help_zh metadata and the language plumbing in the CLI docs generator; config help stays English everywhere, matching every other field, and the zh reference is regenerated accordingly. Refs: areal-project#1546 Signed-off-by: EazyReal <8047065+EazyReal@users.noreply.github.com>
…n-modes-final # Conflicts: # areal/utils/functional/functional.py
The help text already states tree-packed actor training supports only token_mean, but the combination passed config validation and died on the first train_batch (tree microbatches carry no cu_seqlens), after launch and a full rollout. tau2's shipped configs enable tree training, so any of their users opting into seq_mean/prompt_mean/constant would hit it. Signed-off-by: Max Wang <maxwill@vmax.ai>
group_sizes, pg_reduction, and denominator_mask were missing from the otherwise-complete Args sections of ppo_actor_loss_fn, sapo_loss_fn, and cispo_loss_fn; also drop a stray trailing comment in the CISPO tests. Signed-off-by: Max Wang <maxwill@vmax.ai>
2 participants
Description
AReaL currently hard-codes the actor policy-gradient objective as a global token
mean. This PR makes the averaging unit configurable without changing the
training-engine API or any backend implementation.
actor.loss_aggregationtoken_mean(default)seq_meanprompt_meanconstantloss_aggregation_divisorEach mode supplies a local mean and its matching weight to AReaL's existing
TrainEngine.train_batch(loss_fn, loss_weight_fn)contract. The engine alreadycombines microbatches as
sum(local_mean * local_weight) / sum(local_weight), producing the intendedglobal objective without a second distributed-reduction path.
The implementation also:
token_meandtype and reduction path;ragged prompt groups remain explicit and atomic;
numerator;
fails fast when a non-token packed input has no sequence boundaries;
term has a separate token-normalized objective; and
enable_tree_training, becausetree-packed batches carry no per-sequence boundaries to average over.
The diff deliberately contains no FSDP, Megatron, Archon, or public engine-API
changes. It replaces the earlier review surfaces in #1417 and #1443 with a
single clean implementation.
Related Issue
Fixes #1423.
Type of Change
Validation
python -m pytest -q tests/test_prompt_mean_loss.py tests/test_cispo_loss.py tests/test_ppo_stats.py tests/test_functional.py tests/test_rejection_sampling.py tests/test_seqpack.py— 251 passedpre-commit run --all-files— passed every hookgit diff --check— passedThe focused suite covers all four formulas, default token-mean parity, padded and
packed parity, ragged groups, nested group-preserving microbatch splits, filtered
numerators, and partition-invariant callback pairing. Multi-GPU backend tests
were not rerun because the backend implementations are unchanged.
Checklist
Contributing Guide
pre-commit run --all-files)current
main/review-prcommand/create-prBreaking Change Details (if applicable): No API or configuration migration is
needed, and the default
token_meanpath is unchanged for default configs. Onenumeric change does reach existing users: with
actor.m2_thresholdset, thepolicy-gradient denominator is now the valid-token count before M2PO narrows
the mask, which is what the engine's per-microbatch weight already used. The two
previously disagreed — the loss divided by the post-M2PO count while the weight
was the pre-M2PO count — so each microbatch's gradient was scaled by
n_before / n_after. M2PO runs will therefore see different loss values andgradient magnitudes on the same data; that is the point of the fix, and no
configuration change is required to adopt it.
Additional Context
This is the sole public review surface for the loss-aggregation feature. The
broader experimental
LossReduction/LossTermengine abstraction has beenpreserved separately for future design work rather than mixed into this feature
PR.
Archived design prototype: EazyReal#4