fix(rollout): train safely on incomplete groups - #6
Open
EazyReal wants to merge 21 commits into
Open
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe PR adds ChangesRollout API and propagation
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
EazyReal
force-pushed
the
vmax/mask-incomplete-rollout-groups
branch
6 times, most recently
from
July 24, 2026 11:06
50b6ffa to
682003a
Compare
EazyReal
force-pushed
the
vmax/mask-incomplete-rollout-groups
branch
from
July 29, 2026 05:57
8850bb4 to
18b64a8
Compare
EazyReal
changed the title
EazyReal
changed the base branch from
review/issue-1559-base
to
vmax/ragged-transport-safety
July 29, 2026 10:18
Owner
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
Owner
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
areal/api/engine_api.py (1)
205-226: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winBound
min_usable_group_sizebefore rollout retries.A grouped rollout can provide at most
group_sizeusable slots. Ifmin_usable_group_size > group_size,arun_episoderejects every group, and dynamic preparation can retry forever. Validate1 <= min_usable_group_size <= group_sizeat the shared rollout boundary and document the constraint.Also applies to: 247-269, 750-784, 858-887, 908-954
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@areal/api/engine_api.py` around lines 205 - 226, Validate min_usable_group_size at the shared rollout boundary before any rollout retries, requiring 1 <= min_usable_group_size <= group_size; reject invalid values early so arun_episode and dynamic preparation cannot retry indefinitely. Update the parameter documentation and all corresponding rollout entry points to state this constraint, including the methods covering the referenced grouped rollout paths.
🧹 Nitpick comments (3)
tests/test_train_controller.py (1)
579-608: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover non-default forwarding.
Both tests exercise only the default value. Call the controller methods with
group_size=2andmin_usable_group_size=2, then assert that2reaches the rollout mock.Also applies to: 610-638
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_train_controller.py` around lines 579 - 608, Update the prepare_batch delegation tests, including the additional test covering the same path, to pass group_size=2 and min_usable_group_size=2 to the controller and assert mock_rollout.prepare_batch receives both values as 2 instead of the defaults.areal/infra/remote_inf_engine.py (1)
724-730: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider documenting
min_usable_group_sizeonrollout_batch/prepare_batchtoo.Only
submitgained the parameter docs (Lines 1250-1251); the other two public entrypoints now accept the same argument with no docstring entry.Also applies to: 751-751, 849-849, 1233-1233, 1250-1251, 1275-1275, 1328-1328, 1361-1361, 1415-1415
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@areal/infra/remote_inf_engine.py` around lines 724 - 730, Document the min_usable_group_size parameter in the docstrings for the public rollout_batch and prepare_batch entrypoints, matching the existing description and behavior documented for submit. Ensure each entrypoint’s parameter list explains the valid range and purpose without changing the validation or implementation.tests/torchrun/redistribute.py (1)
14-24: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_HybridTrainEngineomitscpu_group/config, which the coordinator touches on the success path.Both assertions in
_test_hybrid_errorraise before_broadcast_and_redistribute_trajectoriesreachesdist.barrier(group=self.train_engine.cpu_group), so this passes today. Adding acpu_group(and aconfigstub) makes the fake engine resilient if the error path ever shifts, instead of failing with a confusingAttributeError.♻️ Suggested hardening
class _HybridTrainEngine: def __init__(self, rank, dp_group, model_group): self.rank = rank self.data_parallel_group = dp_group self.context_and_model_parallel_group = model_group + self.cpu_group = model_group + self.config = None🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/torchrun/redistribute.py` around lines 14 - 24, Update _HybridTrainEngine to define the cpu_group attribute and a config stub expected by _broadcast_and_redistribute_trajectories, initializing them in __init__ alongside the existing rank and process-group fields. Preserve the current test behavior while ensuring the fake engine can safely reach the coordinator’s success path without raising AttributeError.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@areal/infra/dist_rollout.py`:
- Around line 346-390: Bound the retry behavior in _prepare_until_dispatchable
when dynamic_bs is enabled: track attempts, emit progress diagnostics comparable
to _collect_trainable_rollout_batch, and apply the established retry/backoff
approach if available. After the configured limit is reached without a
dispatchable batch, raise a clear RuntimeError instead of looping indefinitely;
preserve the existing synchronization and fixed-batch behavior.
In `@areal/trainer/rl_trainer.py`:
- Around line 102-128: Bound retries in _collect_trainable_rollout_batch when
dynamic_bs is enabled by enforcing a finite attempt limit (or equivalent
timeout) while collecting undersized batches. After the limit is exceeded, raise
a clear actionable error that includes the collected and required group counts;
preserve the current successful return and fixed-batch RuntimeError behavior.
In `@areal/utils/data.py`:
- Around line 934-989: Update split_training_batch_into_microbatches so
effective_n_mbs provides at least enough slots for the largest local microbatch
contribution, while preserving the requested n_mbs limit when it is sufficient.
Replace the scheduled-slot assert with explicit collision handling that fails
loudly rather than overwriting data. Add coverage for a rank whose split count
exceeds n_mbs, such as by forcing a small max_tokens_per_mb, including the
synchronized schedule path used by PPO actor and critic updates.
In `@tests/test_data_redistribution.py`:
- Around line 52-78: Add timeout=60 to the subprocess.run invocation in
test_redistribute_ragged_cpu, matching the existing hybrid redistribution test
while preserving the current command and output assertions.
---
Outside diff comments:
In `@areal/api/engine_api.py`:
- Around line 205-226: Validate min_usable_group_size at the shared rollout
boundary before any rollout retries, requiring 1 <= min_usable_group_size <=
group_size; reject invalid values early so arun_episode and dynamic preparation
cannot retry indefinitely. Update the parameter documentation and all
corresponding rollout entry points to state this constraint, including the
methods covering the referenced grouped rollout paths.
---
Nitpick comments:
In `@areal/infra/remote_inf_engine.py`:
- Around line 724-730: Document the min_usable_group_size parameter in the
docstrings for the public rollout_batch and prepare_batch entrypoints, matching
the existing description and behavior documented for submit. Ensure each
entrypoint’s parameter list explains the valid range and purpose without
changing the validation or implementation.
In `@tests/test_train_controller.py`:
- Around line 579-608: Update the prepare_batch delegation tests, including the
additional test covering the same path, to pass group_size=2 and
min_usable_group_size=2 to the controller and assert mock_rollout.prepare_batch
receives both values as 2 instead of the defaults.
In `@tests/torchrun/redistribute.py`:
- Around line 14-24: Update _HybridTrainEngine to define the cpu_group attribute
and a config stub expected by _broadcast_and_redistribute_trajectories,
initializing them in __init__ alongside the existing rank and process-group
fields. Preserve the current test behavior while ensuring the fake engine can
safely reach the coordinator’s success path without raising AttributeError.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 24659cf9-e44a-4e85-bf21-a9c4bf2c4632
📒 Files selected for processing (27)
areal/api/engine_api.pyareal/engine/fsdp_engine.pyareal/engine/megatron_engine.pyareal/engine/sglang_remote.pyareal/engine/vllm_remote.pyareal/experimental/engine/archon_engine.pyareal/infra/controller/rollout_controller.pyareal/infra/controller/train_controller.pyareal/infra/dist_rollout.pyareal/infra/remote_inf_engine.pyareal/trainer/ppo/actor.pyareal/trainer/ppo/critic.pyareal/trainer/rl_trainer.pyareal/utils/data.pyareal/utils/seqpack.pydocs/en/reference/rollout_workflow.mddocs/zh/reference/rollout_workflow.mdtests/test_data_redistribution.pytests/test_eval_dispatch.pytests/test_grouped_rollout_workflow.pytests/test_incomplete_rollout_groups.pytests/test_reward_norm_variable_group.pytests/test_seqpack.pytests/test_train_controller.pytests/test_utils.pytests/torchrun/redistribute.pytests/v2/training_service/test_data_proxy_unit.py
EazyReal
force-pushed
the
vmax/mask-incomplete-rollout-groups
branch
from
July 29, 2026 19:21
18b64a8 to
6b11d22
Compare
Owner
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
Owner
Author
|
CodeRabbit disposition on current head
Current evidence: 310 affected tests passed with 3 environment skips, distributed ragged tests passed, and full pre-commit passed. |
…roject#1569) agenerate read the weight version twice: once when building the request and again when recording output_versions. A weight update landing in between labelled tokens with a version that did not generate them, which skews staleness checks such as the decoupled loss and rejection sampling. Pin the version before the request goes out and reuse it when recording, so each segment of a trajectory carries the version that actually served it.
…real-project#1544) use_deterministic_algorithms previously set deterministic_mode only on the built model's config. Several Megatron-Core and TransformerEngine code paths consume determinism settings at module construction and cache them in instance state: - VocabParallelEmbedding copies config.deterministic_mode at __init__; without it the nondeterministic F.embedding backward is used. - TEDotProductAttention validates NVTE_ALLOW_NONDETERMINISTIC_ALGO against deterministic_mode only at __init__, and TE snapshots its deterministic flag from that env var and the global torch switch at construction. Setting the flag after the model was built therefore engaged only runtime consumers (loss fusions, schedules) and silently left the layer-level kernels nondeterministic. Changes: - Set deterministic_mode on the TransformerConfig before the model is built, keeping the post-build call for runtime consumers. - Select AttnBackend.flash under deterministic mode: Megatron-Core owns the NVTE_*_ATTN selection env vars and asserts they match the config, and the cuDNN fused-attention deterministic backward needs workspaces that grow prohibitively with context length. - Export NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 to trainer processes from the launchers and to megatron worker specs in single-controller mode, so it is in place before TransformerEngine reads it regardless of TE version; warn when the setting may have come too late. With this, repeated runs of the same batch produce bitwise-identical training stats, including grad_norm.
…(incl. areal-project#1577) (areal-project#1579) * fix: emit PEFT-standard disk LoRA adapter keys so vLLM can load them, and align sglang best-effort unload test assertion (areal-project#1577) * test(inference-service): raise VLM controller init timeout to 600s to avoid vLLM startup timeout
…real-project#1573) Upstream SGLang does not enable the mamba cache path for the Bailing hybrid MoE family, so radix cache stays off and prefix reuse is lost. Add an opt-in script that applies the same eight edits the Ling hybrid fork carries, covering SGLang 0.5.9 and 0.5.10.post1, which moved the DP-attention padding helper behind a dp_size prologue. The script is reversible and supports a dry run, so the target tree can be checked before anything is written.
…ndering (areal-project#1499) * fix(openai): align proxy tool schemas with sglang chat-completions rendering The proxy path renders the request `tools` block through LiteLLM-style dicts, while sglang's native /v1/chat/completions route round-trips tools through its pydantic `Tool` model before applying the chat template. The two renderings differ in field order and default-field presence (e.g. sglang's `Function` dumps `strict: false` while the proxy omits it), so the same trajectory produces different prompt token ids on the two paths and training rewards drift from evaluation results. Round-trip tools through sglang's `Tool` model on the proxy path so the dicts fed to `apply_chat_template` byte-match sglang's own rendering. Flat Responses `FunctionToolParam` entries (top-level `name`/`parameters`, no `function` key) are first normalized to the nested Chat shape that both sglang's model and the chat template expect. The alignment only runs when the engine class name identifies an sglang backend; other backends are left unaligned, and per-tool validation failures fall back to the original dict. Ported from a verified internal fix. * test(openai): cover tool schema alignment against sglang rendering Assert the aligned dicts match what sglang's own Tool model produces, so a regression that reintroduces the prompt-token drift fails here rather than showing up as a train/eval reward mismatch. Also pin the documented fallbacks: flat Responses tools normalize to the nested Chat shape, a single invalid tool does not fail the batch, and unsupported entries pass through unchanged.
) * feat(mcore): apply fp32 lm head forward when enabled enable_fp32_lm_head only reached the model through mbridge extra args, so it was dropped for configs whose TransformerConfig rejects the field and the flag silently had no effect. Patch the lm_head/output_layer forward after the model is built instead, so the existing flag applies on every Megatron construction path. The patch is skipped when the flag is off, when Megatron already provides a native fp32 column-parallel head, and when a module was already patched, and it never runs for critic models, which replace the output layer with a value head. * fix(mcore): forward the tensor-parallel group from the fp32 lm head The patched forward mirrors ColumnParallelLinear, which passes self.tp_group to copy_to_tensor_model_parallel_region and gather_from_tensor_model_parallel_region. Both calls omitted it, so they fell back to the default tensor-model-parallel group and ran the collectives on the wrong ranks whenever the head was built with a non-default group. The sequence parallel path already forwarded the group, so the two behaved inconsistently. Take the group once and route all three calls through a helper that drops the keyword on megatron-core releases that do not accept it. Runs on the default group are unaffected: the wrappers resolve None through get_tensor_model_parallel_group_if_none, so passing it explicitly is identical to omitting it. Document the patch on _fp32_lm_head_forward_impl: what it fixes, the megatron-core version it was verified against, and how it drifts if upstream changes the forward or the collective signatures.
…l-project#1572) Rejection sampling narrows the loss mask, so the existing token counters no longer describe what the update actually trained on. Report the total, valid and masked token counts alongside the masked ratio, and split the log-prob drift into signed and absolute forms so a run's staleness is visible per step. prompt_len is derived from the first trained position rather than the difference of the two mask sums, which stops reporting the prompt as longer than it is once rejection removes generated tokens from the loss mask. The mask reaching _ppo_update is rolled left by one, so the roll is undone before the lookup.
areal-project#1545) * fix(infra): fail fast when a local inference server dies during launch _wait_for_server already receives the server process handle but never inspects it: when the server crashes at startup (port conflict, scheduler init error, OOM kill), the launcher keeps polling the health endpoint until setup_timeout expires and then raises a bare TimeoutError. Poll the process each iteration and raise immediately with the exit code and a pointer to the server traceback instead, which also makes the scheduler-level launch retries practical since a failed attempt now surfaces in seconds rather than after the full timeout. Log a heartbeat every 60s while waiting and include the address and timeout in the TimeoutError message. * test(infra): cover fail-fast on dead inference server process Requested in review: mock process.poll() returning an exit code and assert that _wait_for_server raises RuntimeError carrying the returncode instead of waiting for the health-check timeout. * fix(infra): shut the server down on any launch failure The fail-fast path raises RuntimeError, which escaped the TimeoutError handler in launch_server and skipped _shutdown_one_server. The skip happens to be harmless today because the address is not registered yet and the dead process needs no kill, but the two failure modes should not diverge. Catch both and report the underlying reason instead of assuming a timeout.
* perf: reduce Megatron training memory peaks Add an SFT profiling workflow and use its memory snapshots to remove full-sequence vocabulary and optimizer gradient peaks from Megatron training. Key changes: - Add rank-aware kernel and memory profiling for packed SFT workloads - Fuse FP32 vocab-parallel logprob storage with LM head backward - Add optional true chunked LM head loss with recomputed backward - Configure precision-aware optimizer fields before Megatron validation - Cover BF16/FP32 numerical parity and distributed TP/SP behavior * fix(models): avoid private storage identity checks Track the LM head output tensor weakly and compare storage through the public data_ptr API. This preserves allocator-address reuse protection without depending on PyTorch's private storage _cdata field. * test: compare parameter storage without object identity Parameter.data may return a fresh Tensor wrapper on each access. Verify that replicated parameters retain their data pointer and storage offset instead of comparing transient Python objects. * test: make recycled CUDA storage check deterministic Construct the replacement tensor from the original storage instead of relying on the caching allocator to immediately reuse a freed address after the full CI suite. * fix(engine): guard AReaL LM Head storage reuse Keep entropy differentiable unless its gradients are disabled, and make the optimized LM Head path opt-in. Warn when destructive storage reuse makes entropy non-differentiable, while rejecting unsupported NPU and tree-training combinations. * fix(engine): export standard FSDP LoRA adapter keys PEFT keeps the adapter name in live parameter FQNs, while serving engines expect serialized LoRA keys without it. Normalize per-parameter FSDP exports and align the SGLang best-effort assertion with its load behavior. * feat(engine): support chunked logits for padded models Enable chunked LM Head loss for text-only padded BSHD models such as Qwen3.5 and rename the public toggle to enable_chunked_logits so the configuration reflects its behavior. Key changes: - Add padded label construction and output repacking - Add Qwen3.5 and updated Qwen3 MoE profile recipes - Update CLI docs, validation, and regression coverage * fix(engine): configure logprob chunking explicitly Replace the profile-only environment override with a validated train-engine option so FSDP, Megatron, Archon, and tree paths use the same explicit value. Key changes: - add and document TrainEngineConfig.logprobs_chunk_size - pass the setting through every engine logprob path - translate the profile guide and remove out-of-scope FSDP LoRA changes - add config, launcher, and explicit chunk-size tests Refs: areal-project#1555
…project#1500) * feat(colocate): support AWEX colocated actor-rollout training Add shared-GPU colocated training: the Megatron actor and the SGLang rollout engine time-share the same devices, coordinated through AWEX weight transfer. - AWEX weight-sync writer with tag-based offload/onload and an SGLang engine plugin implementing the colocate reader protocol. - Colocation scheduling support in the Slurm scheduler and controllers (pause/resume generation around the train phase, engine-level offload orchestration, recover handling). - Megatron model registration and fast HF checkpoint loading for the Bailing hybrid MoE family. Both sides key the CUDA IPC handoff on physical GPU ids, mapping through the device mask. Deriving them differently only agrees when the mask starts at zero, so a process pinned to any other GPU would await a key its peer never publishes. Pausing for a weight update keeps SGLang's default mode, which aborts in-flight requests and returns their partial output so the client resumes them by extending the prompt. A second in-place pause raises the scheduler's paused flag so the colocate loop services awex work. Splitting it this way leaves the scheduler fully idle, which SGLang requires before releasing memory, and other backends keep their single-request pause. Supports SGLang 0.5.9 and 0.5.10, whose removed decode-stat hooks are now optional. Verified on multi-node colocated RL runs: training statistics align with the separated-GPU baseline, and single-batch replays produce bitwise-identical training stats. A single-node example test covers the colocate path end to end. * fix(colocate): gate colocation-only steps on the v1 AWEX colocate setup weight_update_mode alone does not identify the colocated v1 run. Controller v2 selects AWEX from use_lora and never reads that field, so a v2 separation run may legitimately carry weight_update_mode="awex" and would then take the v1 colocation handover: its offload was disabled, a stray meta server started, awex_colocate_mode was forced onto its SGLang servers, and both checkpoint branches were skipped so it wrote nothing at all. Gate on _is_v1_awex_colocate, which also requires controller v1 and an actor-rollout colocation strategy. The weight-update meta dispatch keeps its comparison: it sits in an elif chain that v2 short-circuits earlier. Behaviour is unchanged for every configuration that exists today; the three colocated example configs still select the path and the separation config still does not. * refactor(colocate): configure the allocator per role instead of at import AWEX_ACTOR_ALLOC_CONF existed because the colocate examples point rollout at ${actor.scheduling_spec}, so both roles shared one env_vars mapping while the actor wants expandable_segments and SGLang's memory saver cannot tolerate it. Applying it required rewriting PYTORCH_CUDA_ALLOC_CONF from the first lines of the top-level package, because the `from .infra` chain initializes CUDA and freezes the allocator config; that put argv sniffing and an environment mutation into every `import areal`, and the same block was duplicated in rpc_server. Give each colocated role its own scheduling_spec env_vars instead: the env then reaches the process through `srun --env`, before it starts, so no import-time hook is needed. Drop the mechanism and both copies. The mirror of it in the SGLang plugin never worked. It ran from the __main__ block, long after the module-level `from areal.utils import ...` had already frozen the allocator config, so it rewrote the variable while allocations stayed expandable. Replace it with an assertion that runs before any areal import and fails loudly, since a silently self-disabled memory saver surfaces much later as a colocate OOM or an invalid CUDA IPC target. * perf(mcore): build the gloo mirror group only when the engine can offload resolve_broadcast_target reads cpu_model_parallel_group only after an offloaded engine has handed the accelerator to rollout and device collectives are unusable. Creating it unconditionally cost one gloo new_group per data-parallel group at startup on every run, including separation runs that never offload. Gate it on the engine's offload config. Nothing changes on the consumer side: the attribute already defaults to None and resolve_broadcast_target falls back to the device group in that case. * fix(colocate): do not require a colocation strategy to select v1 AWEX The gate added earlier also demanded an actor-rollout colocation scheduling strategy. AWEX runs opt in through weight_update_mode and leave actor and rollout on the default separation strategy, so the gate went false for every such run: the controller never started the AWEX meta server nor passed its address to the rollout, each training worker then started a server of its own, and the rollout registered against a different one. The run waited on 'infer_conf' until the timeout with no error. Keep the controller-version check, which is what stops a v2 separation run from taking this path, and drop the strategy check. The truth table in the test was asserting the broken behaviour, so it is corrected alongside. * docs(examples): add the colocated AWEX GRPO config Mirrors the two-GPU setup the example test exercises: a Megatron actor and an SGLang rollout time-sharing both GPUs, with weights handed over through AWEX. The TMS switches live in the actor's own scheduling_spec env_vars rather than a global toggle: SGLang opens its own memory-saver regions for the colocated rollout, and an auto-opened region on the training side would nest inside them. * refactor: drop the unrelated rpc_server import reformat The parenthesised import left behind by removing the AWEX allocator environment knob carries no semantic change, so rpc_server.py no longer needs to appear in this change set at all. * refactor(recover): fold the colocate helpers into RecoverHandler Both helpers are only reachable from RecoverHandler.load(), so keeping them at module level widened the public surface of areal.utils.recover for no caller. They become private static methods next to the existing _ensure_recover_supported/_normalize_recover_engines pair.
…eal-project#1583) Workers are long-lived rpc_server processes, so any terminal state means they are gone. The status check only reacted to FAILED and CANCELLED, so a job that reached COMPLETED - for instance the batch script exiting 0 after a container FATAL - was still treated as healthy and the controller waited on workers that no longer existed. NODE_FAIL was not mapped at all. squeue also exits non-zero once a job leaves its window ("Invalid job id specified" right after completion), which is indistinguishable from a transient slurmctld error at the call site. Ask sacct in that case: a terminal state means the workers are gone, anything else is treated as transient and retried.
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
force-pushed
the
vmax/mask-incomplete-rollout-groups
branch
from
August 9, 2026 23:04
3e7fa87 to
9392e20
Compare
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>
…e-rollout-groups # Conflicts: # areal/trainer/ppo/actor.py
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>
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Fork-local focused review for upstream areal-project/AReaL#1563. Do not merge this PR; the upstream PR is the landing authority.
Stack
Summary
Noneand retain every usable slot exactly oncemin_usable_group_sizefrom group-relative reward or advantage normalization; batch-relative estimators retain a usable singletonThere is deliberately no
requires_peerprotocol noun. The estimator owns its minimum usable sample count; the rollout boundary carries only that normalized integer contract.Exact range
b0dbd4c423de3be706a3978fd9d1ced4678298e6..6b11d22af742baf18e3226dceb0ef6a2fd154a6cVerification
A real multi-GPU Megatron/Archon pipeline canary is not available on this host.