Skip to content

fix(rollout): train safely on incomplete groups - #1563

Open
EazyReal wants to merge 10 commits into
areal-project:mainfrom
EazyReal:vmax/mask-incomplete-rollout-groups
Open

fix(rollout): train safely on incomplete groups#1563
EazyReal wants to merge 10 commits into
areal-project:mainfrom
EazyReal:vmax/mask-incomplete-rollout-groups

Conversation

@EazyReal

@EazyReal EazyReal commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

Stacked on #1566. The first commit of this branch is #1566 (b0dbd4c4), so the
diff below carries #1566's nine-file transport-padding change on top of main.
Merging #1566 first shrinks this PR to its rollout-side core; reviewing both together
here also works. #1566 owns areal/engine/core/*, areal/models/tree_attn/tree.py,
and tests/test_tree_transport.py outright; this PR builds further on it in
areal/utils/data.py, the three train engines, and tests/test_utils.py.

Problem

Grouped rollout workflows already return None for unusable slots, and since #1454
main normalizes the survivors against their actual group sizes. Everything downstream
of the group still assumes the configured group size:

  • balanced_greedy_partition rejects a dispatch whose group count is not divisible by
    dp_size, so one incomplete group aborts the step that partial-group support exists
    to keep alive.
  • Ragged shards let data-parallel ranks disagree about how many optimizer minibatches
    exist, and a rank-local split can be dropped.
  • A group reduced to a single survivor has no peer to normalize against. Under group
    statistics it is centered against itself and trains as a zero-advantage row whose
    tokens only dilute the token-mean denominator.
  • Normalization._build_group_slices silently truncates the tail when a direct caller
    passes no group_sizes and the batch is not a multiple of group_size.
  • Dynamic collection retries None results forever, so a reward function that rejects
    everything spins the run behind a warning instead of failing.

Solution

  • retain every usable rollout slot exactly once and never duplicate a slot
  • derive the minimum usable group size from group-relative reward or advantage
    normalization: min(2, n_samples) when either uses group statistics, 1 otherwise,
    so batch-relative estimators and n_samples: 1 keep a usable singleton
  • add actor.min_usable_group_size as an explicit override of that derivation
    (None derives as above). PPOActorConfig.__post_init__ rejects non-positive values,
    and rejects anything below 2 while reward_norm/adv_norm uses group statistics.
    The v2 rollout path does not consume the option and now fails fast if it is set,
    matching how RolloutControllerV2 handles reward_normalization and
    drop_incomplete_group.
  • preserve actual ragged group boundaries through redistribution and sequence packing:
    balanced_greedy_partition accepts a non-divisible group count (partition
    cardinalities then differ by at most one), and _dispatch_tensors enforces the
    invariant that actually matters, n_groups >= dp_size
  • backfill DP-dispatchable batches and synchronize optimizer minibatches without
    dropping a rank-local split
  • abort a sustained collection stall in both collection loops (the trainer-local loop in
    rl_trainer.py and the DP-coordinated one in dist_rollout.py): at least 8
    consecutive rounds that add no trainable group and at least 30 minutes without
    progress. Requiring both keeps a legitimate fast all-reject burst alive, such as a
    staleness flush right after a weight update.
  • propagate deterministic workflow-contract violations terminally across local and
    controller/worker execution while ordinary objective rejection remains retryable
    None
  • add partial-group observability: target_slot_count, usable_slot_count,
    trainable_slot_count, fully_masked_group, singleton_slot_group,
    pre_filter_usable_slot_yield, and pre_filter_trainable_slot_yield per grouped
    rollout; usable_group_size, group_loss_weight, and per-size n_groups_size_<N> /
    group_loss_weight_size_<N> during PPO training
  • keep the existing globally token-weighted PPO actor objective, and state it in the
    grouped-rollout reference: a partial group with more valid response tokens carries
    more loss weight, which is the existing backward-compatible estimator rather than an
    implicit claim of equal prompt weighting
  • examples/openclaw/config.yaml moves reward_norm to mean_level: batch /
    std_level: batch. The demo runs n_samples: 1, so group centering erased the task
    reward outright, and its OpenAI-proxy agent exports one row per interaction, so group
    std collapsed a k-turn episode's identical rewards to sign(r - mean) * sqrt((k-1)/k)
    and discarded reward magnitude. adv_norm was already batch-level.

There is deliberately no requires_peer protocol noun: the estimator owns its minimum
usable sample count, and the rollout layer receives only that normalized integer
contract.

Beyond the new failure paths above, default-behavior change against current main is
limited to the singleton case: under group statistics, a group reduced to one survivor
is dropped and the collector refills the slot with a trainable prompt, where main
trains it as a zero-advantage row.
gconfig.drop_incomplete_group=True still drops any group with a failed slot before the
minimum applies, unchanged.

Two notes for reviewers:

Review stack

GitHub cannot base an upstream PR on a contributor-fork branch, so this PR is
temporarily cumulative against main. The v2 follow-up will be opened upstream only
after this PR lands.

Part of #1559.

Verification

At head 3a297376:

  • affected CPU test surface —
    pytest tests/test_data_redistribution.py tests/test_eval_dispatch.py tests/test_grouped_rollout_workflow.py tests/test_incomplete_rollout_groups.py tests/test_reward_norm_variable_group.py tests/test_seqpack.py tests/test_train_controller.py tests/test_tree_transport.py tests/test_utils.py tests/v2/training_service/test_data_proxy_unit.py
    — 318 passed, 3 skipped. The skips are the 2/4/8-GPU redistribution cases;
    tests/test_data_redistribution.py still exercises the ragged all-gather and
    rank-disagreement paths over Gloo.
  • tests/test_rollout_controller.py and tests/test_serialization.py also change here
    but are not part of that run: importing areal.utils.testing_utils eagerly resolves
    every registered test model, including Qwen/Qwen3-30B-A3B, so collecting them needs
    a machine with those weights already cached.
  • hosted CI on this head — pre-commit and the Ubuntu, macOS, SGLang-extra, and
    vLLM-extra installation checks pass; Docker is skipped by workflow policy.
  • CodeRabbit reviewed the focused core mirror
    (EazyReal/AReaL#6); all four of its
    threads are resolved.

No multi-GPU Megatron or Archon pipeline run was performed. This is static, CPU/Gloo
test, and hosted-install readiness rather than hardware certification.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@EazyReal
EazyReal marked this pull request as draft July 28, 2026 01:02
@EazyReal
EazyReal force-pushed the vmax/mask-incomplete-rollout-groups branch from 8850bb4 to 18b64a8 Compare July 29, 2026 05:57
@EazyReal EazyReal changed the title feat(rollout): mask untrainable incomplete groups fix(rollout): train safely on incomplete groups Jul 29, 2026
@EazyReal
EazyReal force-pushed the vmax/mask-incomplete-rollout-groups branch from 18b64a8 to 6b11d22 Compare July 29, 2026 19:21
@EazyReal
EazyReal marked this pull request as ready for review July 31, 2026 02:33
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

Comment thread areal/trainer/rl_trainer.py Outdated
std_level = getattr(normalization, "std_level", None)
uses_group_statistics |= mean_level == "group" or std_level == "group"

if uses_group_statistics and target_group_size < 2:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The n_samples=1 setting in examples/openclaw/config.yaml conflicts with this, causing the demo to fail to run properly.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed — the minimum now relaxes for singleton target groups: minimum_usable_group_size returns min(2, n_samples), so n_samples=1 keeps its pre-PR behavior. A singleton group is complete by definition, so there is no partial-group hazard to guard (and the v2 rollout path never armed this check). The openclaw demo starts again; the config-level degeneracy (group centering with group_size=1 zeroes the normalized reward) remains the existing PPOActorConfig.__post_init__ UserWarning, as on main.

Side note: that warning does fire for openclaw's reward_norm (group/group with n_samples: 1). Switching the demo to mean_level: batch would give it a live reward signal — I left the example untouched here since that would be a semantic change to the demo.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  • The minimum_usable_group_size parameter would be better as a configurable option.
  • As for the demo, you're right. When this PR is implemented, mean_level should be changed from group to batch.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both done — 677c92a, 6aa093d:

  • actor.min_usable_group_size is now a config field. Default None keeps the derived value (2 under group statistics, 1 for a singleton target group, else 1); an explicit value replaces it and is still validated against group_size at workflow setup.
  • The demo's reward_norm.mean_level is now batch. I left std_level: group untouched: with n_samples: 1 every group is a singleton, whose std is pinned to 1, so it is a no-op scale. Happy to flip it to batch too if you prefer the demo read coherently.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correction to my earlier claim here: std_level: group was not a no-op for this demo. The OpenAI-proxy agent exports one row per interaction, so compute_advantages passes per-episode row counts as group sizes — a k-turn episode's identical rewards collapse to sign(r - mean) * sqrt((k-1)/k), discarding reward magnitude and scaling by turn count. Fixed in a6b354a: reward_norm is now batch/batch (matching adv_norm), and the unused group_size line is gone.

world_size = dist.get_world_size(group)
local_group_rank = dist.get_rank(group=group)
lengths: list[int | None] = [None] * world_size
dist.all_gather_object(lengths, len(trajectories), group=group)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the number of groups for any rank differs from others, take the slow path: perform serial broadcast_tensor_container $sum(lengths)$ times. Each broadcast includes object serialization and tensor transfer. It is recommended to first all-gather the metadata (shape/length) from each rank, then perform a padded all-gather (or all_gather_object to pack the entire list) per rank, rather than broadcasting each item individually. The fast path already exists, and the slow path requires equivalent batch processing.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reworked as you suggested: all_gather_ragged_tensor_container (areal/utils/data.py) first all-gathers per-item skeletons — lengths, structure, shapes, dtypes, and non-tensor leaves in a single all_gather_object — then moves all tensor payloads with one padded all_gather per (dtype, device-type) bucket derived from the gathered metadata. Every rank, including ranks with zero trajectories, issues the same collectives, so no per-item serial broadcasts remain; gathered tensors are cloned out of the padded buffers so a skewed distribution does not pin world_size × max-payload memory. The equal-length fast path is unchanged. The gloo torchrun test covers the ragged path end-to-end, including a new zero-trajectory-rank case.

raise ValueError(f"min_batch_size must be positive, got {min_batch_size}")

batch: list[dict[str, Any]] = []
while len(batch) < min_batch_size:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When dynamic_bs=True and the workflow fails to produce enough usable groups continuously (e.g., due to a bug in the reward function causing all instances to be rejected, or when min_usable_group_size=2 and the model only produces singleton groups), both loops spin indefinitely, only emitting warnings. Additionally:

  • The inner loop within each iteration performs 2 all_gather_object calls, continuously applying pressure on the DP head group's collective communication;
  • Non-head ranks are blocked on the broadcast_tensor_container in _broadcast_and_redistribute_trajectories (due to different groups, there's no deadlock, but it hangs indefinitely);
  • Warnings overwhelm the logs. This turns "training progress stagnation" from an observable failure into a silent hang.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed — both collection loops now abort on a sustained stall: at least 8 consecutive rounds that add no trainable group AND 30 minutes without progress (both constants next to the other liveness bounds in workflow_executor.py). Requiring both keeps legitimate fast all-reject bursts alive — e.g. a staleness flush after a weight update drains a deep queue of already-generated rejections in near-zero time and must not kill the run — while a reward bug or an unattainable min_usable_group_size trips the bound within one generation cycle past the timeout.

In the SPMD loop both signals ride the existing per-round all-gather, so every DP head raises on the same iteration and the coordinated error path delivers a terminal RuntimeError to all ranks — non-head ranks no longer hang in the broadcast, and the message names the likely causes plus the usable_slot_count/fully_masked_group stats to check. Warnings are rate-limited to once per minute. I kept these as module constants rather than config; happy to expose a knob if you'd prefer.

Comment thread areal/utils/seqpack.py
n = len(nums)
if n < K:
raise ValueError(f"Number of items ({n}) must be >= K ({K}).")
if n % K != 0:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

areal/infra/controller/train_controller.py:103 still uses balanced_greedy_partition(group_weights, K=dp_size), but the docstring explicitly depends on this check. Before dispatch, this is called to ensure balanced_greedy_partition receives a divisible input.

After removal: If n_groups % dp_size is not zero, the old code would raise a clear ValueError, but the new code silently produces partitions of unequal sizes (each DP rank receives a different number of groups). _pad_eval_batch still fills to dp_size * group_size, and its docstring has not been updated—there is a discrepancy between documentation and implementation. Additionally, areal/v2/training_service/data_proxy/dispatcher.py:10 also references this tool (tests/v2/training_service/test_data_proxy_unit.py is also modified +8/−2), and the v2 path is similarly affected but has not been analyzed.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Partly intentional, partly fixed. Ragged partitions are the point of the change: incomplete groups make training batches non-divisible by dp_size, and the old ValueError would abort exactly the runs this PR keeps alive. The safety moved rather than disappeared — _dispatch_tensors now rejects n_groups < dp_size (train_controller.py), engines absorb ragged shards via the transport padding from the prerequisite commit, and config-level non-divisibility is still rejected at dataloader construction. On v2: the data-proxy dispatcher imports the same _dispatch_tensors/_pad_eval_batch, so it inherits both the ragged partitioning and the new guard; its unit test was updated in this PR to lock the ragged-shard contract (shards [1, 2], order-preserving merge), and v2 eval routes still pad.

What was genuinely stale is now corrected: _pad_eval_batch's docstring no longer cites the removed divisibility guarantee (it states the eval-pads / training-ragged split and the invariants that replaced it), the cardinality note in balanced_greedy_partition moved into the description proper, and the v2 dispatcher module docstring is qualified to eval routes.

pre_filter_trainable_slot_yield=trainable_slot_count / self.group_size,
)

def _validate_slot_cardinality(self, slot_sizes: list[int]) -> None:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WorkflowContractError is terminal (from WorkflowContractFailure → unwrap_workflow_result → throw → kill the entire run without retrying).

That is: any custom workflow that produces multiple lines of training samples per episode, with group-level normalization enabled (the default configuration of GRPO), will terminate the entire training.

The built-in workflows in the repository have all been verified to be compliant (multi_turn.py, rlvr.py, vision_rlvr.py all return {k: v.unsqueeze(0)}, exactly 1 line), and there is no usage of concat_padded_tensors in the examples/. Therefore, this does not constitute a regression. However:

  • This is a tightening of the public contract for RolloutWorkflow, and yet this PR has rewritten docs/en|zh/reference/rollout_workflow.md extensively but completely omitted this constraint;

  • It is destructive for third-party/user-defined workflows (multi-turn per turn, tree search branches, etc.), and the failure mode is termination rather than degradation.

It is recommended to write it into rollout_workflow.md (EN+ZH) and declare it in the docstring of the RolloutWorkflow base class; consider providing more informative error messages (indicating the use of mean_level: batch or having the workflow return only one line each time).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Documented in all three places: the grouped-rollout reference (EN + ZH) has an explicit contract bullet, the RolloutWorkflow.arun_episode docstring states it, and the error message now names both remedies (return one sample per episode — for agent workflows agent.export_style: 'concat' — or switch mean_level/std_level to batch).

Two sharpenings from checking the blast radius: (1) the constraint also binds the first-party OpenAI-proxy agent path under the default export_style: individual when group normalization is on — the docs call that combination out explicitly; (2) the docs scope the promise honestly to grouped rollouts (n_samples >= 2): n_samples=1 installs no group wrapper, so a multi-sample episode there is normalized as its own group rather than rejected, and the docs say so.

The check itself stays fail-fast: training-side group boundaries are physical row counts, so a multi-row slot would silently treat per-turn rows as independent group members — termination beats silent corruption. Relaxing this via per-slot logical group ids would be a real feature for a separate PR.

batch = redist.data
error = self._synchronize_head_error(preparation_error)
if trajectories is not None and error is None:
try:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It appears that if any rank throws an error due to rank-local reasons during the middle of redistribute_trajectories (after entering a collective communication), other ranks will be permanently blocked on that collective communication — the error synchronization occurs outside of the try block, and there's no time to recover.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're right — the recoverable domain was wider than what is soundly recoverable. Restructured: redistribute_trajectories is now gather (collective phase) + _pack_gathered_trajectories (pure-local, taking explicit world_size/rank so the signature proves it cannot communicate), and only the local phase sits inside the recoverable try. After the gather every head holds identical data, so local failures are symmetric and always reach the error sync. A failure inside the gather collectives now propagates immediately instead of posting an error all-gather that would pair with peers' pending operations — that window is unrecoverable in userspace (there is no portable way to recover a half-posted collective); peers are released by the process-group timeout/NCCL watchdog, and torchrun tears down co-located ranks.

Comment thread areal/trainer/rl_trainer.py Outdated
logger = logging.getLogger("RLTrainer")


def _minimum_usable_group_size(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These free functions are better as methods of PPOActorConfig or attributes of NormConfig.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Moved: NormConfig gains a uses_group_statistics property and PPOActorConfig owns minimum_usable_group_size, using the same isinstance(dict, DictConfig) branch as the existing __post_init__ instead of the ad-hoc duck test. The two neighbors stayed as trainer-local helpers deliberately — _collect_trainable_rollout_batch takes a prepare callable and _minimum_consumer_batch_size takes train engines; neither reads config, so the config classes can't own them without pulling trainer/engine concepts into cli_args.

EazyReal added a commit to EazyReal/AReaL that referenced this pull request Aug 9, 2026
Address review on the incomplete-group transport and collection paths:

- Replace the serial per-item broadcast fallback for ragged trajectory
  lists with one metadata all-gather plus one padded all-gather per
  dtype/device bucket, so unequal group counts cost a handful of
  collectives instead of sum(lengths) container broadcasts. Ranks with
  zero trajectories join the same collectives via the gathered metadata,
  and gathered tensors are cloned out of the padded buffers so a skewed
  distribution does not pin world_size x max-payload memory.
- Split redistribute_trajectories into its gather phase and a pure-local
  packing phase, and recover only the latter: every head holds the same
  gathered data there, so failures are symmetric and always reach the
  error sync; a failure inside a collective now propagates instead of
  posting mismatched operations at peers.
- Abort dynamic preparation only on a sustained stall: at least eight
  consecutive rounds that add no trainable group AND thirty minutes
  without progress. Both signals ride the existing per-round all-gather,
  so every head raises on the same iteration and the coordinated error
  path turns a silent SPMD spin into a terminal error on all ranks,
  while fast legitimate all-reject bursts (e.g. a staleness flush after
  a weight update) stay alive.

Refs: areal-project#1563
Signed-off-by: EazyReal <8047065+EazyReal@users.noreply.github.com>
EazyReal added a commit to EazyReal/AReaL that referenced this pull request Aug 9, 2026
Address review on estimator-owned group minimums and the tightened
workflow contract:

- NormConfig gains uses_group_statistics and PPOActorConfig owns
  minimum_usable_group_size, replacing the trainer-side free function
  and its ad-hoc shape duck-typing.
- A singleton target group is complete by definition, so group
  statistics no longer hard-fail n_samples=1 configs; the previous
  ValueError broke examples/openclaw on startup while the v2 rollout
  path skipped the check entirely. Config-level degeneracy stays a
  UserWarning, as on main.
- The trainer collection loop shares the empty-round streak plus
  wall-clock stall bound and rate-limits its progress warnings.
- WorkflowContractError now names the two remedies (one sample per
  episode or batch-level normalization), and the one-sample-per-slot
  contract is documented in the RolloutWorkflow docstring and the
  grouped-rollout reference (EN/ZH), scoped honestly to grouped
  rollouts: n_samples=1 installs no group wrapper and is not checked.
- Correct docstrings that still described divisibility-based dispatch:
  balanced_greedy_partition, _pad_eval_batch, and the v2 data-proxy
  dispatcher module.

Refs: areal-project#1563
Signed-off-by: EazyReal <8047065+EazyReal@users.noreply.github.com>
@EazyReal

EazyReal commented Aug 9, 2026

Copy link
Copy Markdown
Contributor Author

Pushed two commits addressing all seven review threads: batched ragged gathers (one metadata all-gather + one padded all-gather per dtype bucket, zero-trajectory ranks included), a collective-free recoverable phase in redistribution, stall-bounded dynamic collection (8 empty rounds + 30 min without progress → coordinated terminal error on all ranks), singleton target groups keep minimum 1 (openclaw runs again), config-owned minimum_usable_group_size, an actionable WorkflowContractError, and the one-sample-per-slot contract documented in the EN/ZH references and the base-class docstring. Per-thread details in the inline replies.

Address review on the incomplete-group transport and collection paths:

- Replace the serial per-item broadcast fallback for ragged trajectory
  lists with one metadata all-gather plus one padded all-gather per
  dtype/device bucket, so unequal group counts cost a handful of
  collectives instead of sum(lengths) container broadcasts. Ranks with
  zero trajectories join the same collectives via the gathered metadata,
  and gathered tensors are cloned out of the padded buffers so a skewed
  distribution does not pin world_size x max-payload memory.
- Split redistribute_trajectories into its gather phase and a pure-local
  packing phase, and recover only the latter: every head holds the same
  gathered data there, so failures are symmetric and always reach the
  error sync; a failure inside a collective now propagates instead of
  posting mismatched operations at peers.
- Abort dynamic preparation only on a sustained stall: at least eight
  consecutive rounds that add no trainable group AND thirty minutes
  without progress. Both signals ride the existing per-round all-gather,
  so every head raises on the same iteration and the coordinated error
  path turns a silent SPMD spin into a terminal error on all ranks,
  while fast legitimate all-reject bursts (e.g. a staleness flush after
  a weight update) stay alive.

Refs: areal-project#1563
Signed-off-by: EazyReal <8047065+EazyReal@users.noreply.github.com>
Address review on estimator-owned group minimums and the tightened
workflow contract:

- NormConfig gains uses_group_statistics and PPOActorConfig owns
  minimum_usable_group_size, replacing the trainer-side free function
  and its ad-hoc shape duck-typing.
- A singleton target group is complete by definition, so group
  statistics no longer hard-fail n_samples=1 configs; the previous
  ValueError broke examples/openclaw on startup while the v2 rollout
  path skipped the check entirely. Config-level degeneracy stays a
  UserWarning, as on main.
- The trainer collection loop shares the empty-round streak plus
  wall-clock stall bound and rate-limits its progress warnings.
- WorkflowContractError now names the two remedies (one sample per
  episode or batch-level normalization), and the one-sample-per-slot
  contract is documented in the RolloutWorkflow docstring and the
  grouped-rollout reference (EN/ZH), scoped honestly to grouped
  rollouts: n_samples=1 installs no group wrapper and is not checked.
- Correct docstrings that still described divisibility-based dispatch:
  balanced_greedy_partition, _pad_eval_batch, and the v2 data-proxy
  dispatcher module.

Refs: areal-project#1563
Signed-off-by: EazyReal <8047065+EazyReal@users.noreply.github.com>
@EazyReal
EazyReal force-pushed the vmax/mask-incomplete-rollout-groups branch from 3e7fa87 to 9392e20 Compare August 9, 2026 23:04
Reviewer request on areal-project#1563: expose the partial-group trainability
threshold instead of only deriving it. PPOActorConfig gains
min_usable_group_size (default None); None keeps the derived value
(2 when reward_norm/adv_norm uses group statistics, 1 for singleton
target groups, else 1), an explicit value replaces it and is still
validated against group_size at workflow setup.

Signed-off-by: Max Wang <maxwill@vmax.ai>
With n_samples=1 every prompt group is a singleton, so group mean
centering erases the task reward. Batch centering keeps a live signal;
singleton group std is already pinned to 1 and stays a no-op.

Signed-off-by: Max Wang <maxwill@vmax.ai>
@EazyReal
EazyReal requested a review from CormickKneey as a code owner August 10, 2026 06:46
@sitabulaixizawaluduo

Copy link
Copy Markdown
Collaborator

Please resolve the conflict.

…e-rollout-groups

# Conflicts:
#	areal/trainer/ppo/actor.py
@EazyReal

Copy link
Copy Markdown
Contributor Author

Merged main in e8a483c — the PR is mergeable again. The only textual conflict was in areal/trainer/ppo/actor.py (this PR's _group_training_metrics vs #1572's _infer_prompt_lens added at the same spot; kept both). Unit and torchrun suites pass on the merged head.

Comment thread areal/api/cli_args.py
become a hazard; a singleton target group is complete by definition,
so it keeps the minimum of one.
"""
if self.min_usable_group_size is not None:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we add an assert to check that if group normalization is used, min_usable_group_size must be greater than 1 ?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added in f7ec6b3PPOActorConfig.__post_init__ now raises when an explicit min_usable_group_size is below 2 while reward_norm/adv_norm uses group statistics (and rejects non-positive values generally). The derived default is unchanged.

While re-checking the field's blast radius I also tightened two adjacent spots in 3a29737: the v2 rollout path never consumes the option, so an explicit setting there now fails fast instead of being silently ignored (matching RolloutControllerV2's handling of reward_normalization/drop_incomplete_group), and the slot-cardinality error plus docs now name the resolved minimum — not group normalization per se — as the trigger, since the two can diverge once the field is set explicitly.

Reviewer request on areal-project#1563: an explicit min_usable_group_size below 2
combined with group-relative normalization would train lone survivors
that have no group peers to normalize against, so PPOActorConfig now
rejects it at construction. The derived default is unaffected.

Signed-off-by: Max Wang <maxwill@vmax.ai>
Comment thread areal/api/cli_args.py Outdated
)
else:
uses_group_statistics = normalization.uses_group_statistics
if uses_group_statistics:

@sitabulaixizawaluduo sitabulaixizawaluduo Aug 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a forced modification of the user's settings. If the user does not want to use this feature, they must obtain n_samples trajectories; otherwise, the behavior will be subtly altered unless it explicitly sets min_usable_group_size=n_samples.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The strict behavior keeps its existing knob: drop_incomplete_group=True still drops any group with a failed slot before the minimum applies, unchanged from main.

On the default path, main today does not obtain n_samples trajectories either — a partial group is kept with a warning ("using remaining results" in GroupedRolloutWorkflow.arun_episode) and then normalized with fixed positional group_size slices, so its group statistics silently straddle neighboring groups. That mis-normalization is the bug this PR fixes: the default here trains the same partial groups main already trains, just with statistics computed over the actual members. The one real default change is that a group reduced to a single member under group statistics is now dropped (it has no peers to normalize against) instead of mis-normalized.

If you would rather have the default strict — only complete groups train unless the user opts in — I am happy to flip the derivation to min_usable_group_size = n_samples; it is a small change. I kept main's default group yield while fixing its statistics as the least surprising option, but the default is your call.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A more precise ledger, with one correction to my previous reply: the statistics fix for partial groups is not this PR — it already landed on main via #1454, which passes actual survivor counts into group normalization through batched_call metadata. Against current main, the only default-behavior change left in this PR is the singleton case: a group reduced to one survivor under group statistics trains as a zero-advantage row on main (it is centered against itself), while this PR drops it and lets the collector refill the slot with a trainable prompt.

So the choice is narrower than it looked:

  1. Keep the derived minimum of 2 plus the new assert (current state). Default batch composition changes only in the singleton case, where the dropped row carried zero advantage on main anyway — its tokens merely diluted the token-mean denominator.
  2. Derive 1 instead. Byte-for-byte default preservation, but singletons keep entering batches as dead rows, and the assert from your earlier comment becomes incoherent — explicitly writing the derived default would be rejected.

I recommend 1 and have left the PR in that state; happy to switch to 2 if you prefer strict preservation.

std_level: group is not the no-op I claimed on the review thread: the
OpenAI-proxy agent exports one row per interaction, so compute_advantages
passes per-episode row counts as group sizes and a k-turn episode's
identical rewards collapse to sign(r - mean) * sqrt((k-1)/k), discarding
reward magnitude. Batch std matches adv_norm and keeps the signal; the
group_size line had no remaining consumer.

Signed-off-by: Max Wang <maxwill@vmax.ai>
Two exactness gaps around the new config field. The v2 rollout path never
consumes it, so an explicit setting now fails fast there (matching
RolloutControllerV2's rejection of reward_normalization and
drop_incomplete_group) and the help text names the v1 scope. The
one-sample-per-slot contract is armed by the resolved minimum, not by
group normalization itself, so the error message, arun_episode docstring,
and reference docs now name the real trigger and include unsetting
actor.min_usable_group_size among the remedies; the agent tutorial and
export-style reference gain the same caveat where they teach
export_style: individual.

Signed-off-by: Max Wang <maxwill@vmax.ai>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants