| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
feat(v2): support grouped reward normalization (#1619) * feat(v2): support grouped reward normalization Normalize per-prompt rollout groups at the data-proxy export boundary before trajectory tensors are merged and remotized. Preserve raw rewards for filtering and metrics while sharing the normalization implementation with v1. Keep reward metadata local for rollout filters and clear remote shards for rejected or failed trajectories so filtered v2 rollouts do not leak data-proxy storage. * fix(v2): bound rejected trajectory cleanup | 8 天前 | |
refactor(api): migrate allocation_mode to per-engine backend fields (#1044) * refactor(api): migrate allocation_mode to per-engine backend fields Replace the centralized `allocation_mode` string with explicit `backend` fields on `TrainEngineConfig` and `InferenceEngineConfig`. Each engine now owns its own backend+parallelism spec (e.g. `fsdp:d4`, `sglang:d4t2`), eliminating implicit auto-backend selection and the shared `AllocationMode` object. Key changes: - Add `backend` field to TrainEngineConfig and InferenceEngineConfig - Add `ModelAllocation.from_str()` for single-component parsing - Remove `AllocationMode` public export (replaced by `ModelAllocation`) - Rename internal `AllocationMode` to `_AllocationMode` for SPMD launcher backward compatibility with FutureWarning - Remove auto-backend selection — explicit backend prefix is now required - Controllers (`TrainController`, `RolloutController`) parse `backend` directly instead of receiving `alloc_mode` from trainers - `WeightUpdateMeta.alloc_mode` replaced by `gen_allocation` (single `ModelAllocation`) - Add `RWTrainer` and `ArchonRWEngine` for reward model training - Remove `get_model_update_meta()` helper (logic moved to trainers) - Update all YAML configs, examples, docs (EN+ZH), and tests BREAKING CHANGE: `AllocationMode` is removed from public API. Users must migrate to per-engine `backend` fields. SPMD launchers emit deprecation warnings. * chore(ci): fix backend specifier for vlm sft test * fix: fix bare dims for actor backends * chore(docs): fix reminder for bare allocation dims | 5 个月前 | |
chore(deps): upgrade runtime dependencies and CI workflow (#1206) * chore(deps): upgrade runtime dependencies and CI workflow Upgrade megatron-core, sglang, vllm (0.19.1), transformers, and related packages. Pin deepep/deepgemm commits and lock trackio version for reproducibility. Key changes: - Upgrade sglang, vllm, transformers, megatron-core versions - Pin vllm to 0.19.1, fix compatibility across backends - Optimize Dockerfile and slim Docker image - Add uv_sync.sh install script, fix uv install on Linux - Fix Archon Qwen3.5 precision and port range issues - Remove integration tests from PR CI (moved to nightly) - Sync GRPO integration test config Refs: upgrade-deps branch * chore: remove duplicated test | 3 个月前 | |
fix(infra): preserve LD_PRELOAD in local launchers (#1578) Avoid wrapping local targets with stdbuf when the effective target environment already defines LD_PRELOAD. GNU stdbuf appends libstdbuf to the value, which breaks TMS CDLL loading during offload. Preserve env=None versus explicit env={} semantics while retaining line buffering for targets without LD_PRELOAD. Refs: #1570 Signed-off-by: yaoyaoshiguonan <yaoyaoguonan@outlook.com> | 19 天前 | |
fix: align CP metrics and TP grad norm metadata (#1497) * fix: align CP metrics and TP grad norm metadata Two independent correctness fixes for the Megatron engine. 1. CP metrics alignment (stats_tracker): Per-key reduce_group overrides may be wider than the default DP export group (e.g. DP+CP for token-level metrics). Add a `key_sync_group` argument so override keys sync their key set / metadata and reduce over DP+CP, while default keys keep the DP alignment. Aggregation helpers now emit identity placeholder tensors (0 / +inf / -inf) for keys absent on a rank, so every rank still participates in the collective and avoids hangs / mismatched reductions when CP > 1. 2. TP grad norm metadata (megatron_engine): `_mark_duplicated_params` now also clears `param.tensor_model_parallel` on replicated (tp_size == 1, non-expert) params. Megatron's optimizer uses that attribute to decide which TP ranks contribute to grad norm/clipping, so leaving it True double-counts duplicated params when TP > 1. Clearing it is also consistent with `all_gather_param` and `hf_save`, which already treat non-TP params as replicated. Adds unit tests for both paths (tests/test_stats_tracker.py and a new case in tests/test_megatron_engine.py). * test: cover all_gather_param routing and duplicated-param grad norm Add two tests backing the TP grad-norm metadata fix: - tests/test_all_gather_param.py: assert `all_gather_param` returns the param as-is (no TP all-gather) when `tensor_model_parallel` is False or the name is in `duplicated_param_names`, and only all-gathers genuine TP-sharded params. Uses `pytest.importorskip` so it runs in the megatron CI env and skips gracefully elsewhere. - tests/test_grad_norm_duplicated.py: a CPU/gloo 2- and 4-rank test showing a replicated param is counted once (correct) when tensor_model_parallel is False, and inflated by sqrt(tp) when it is (incorrectly) left True. Mirrors Megatron's param_is_not_tensor_parallel_duplicate selection. * test: add 2-GPU grad-norm TP-invariance integration test End-to-end guard for the TP grad-norm metadata fix: run the real MegatronEngine at TP=1 and TP=2 on identical deterministic input via torchrun and assert the reported grad norm is TP-invariant. Double-counting replicated params (tensor_model_parallel left True) would inflate it as TP grows. Also asserts the fix demoted at least one real duplicated param at TP=2. Marked multi_gpu/slow and skipped without >= 2 GPUs. * refactor(stats): dedup identity tensors, avoid min/max GPU sync Address review feedback on the CP metrics fix: - Generalize the placeholder-tensor helper to `_placeholder_scalar(fill=...)` so the SUM/AVG (0.0) and MIN/MAX (+/-inf) empty-value branches share one device-aware constructor instead of four inline copies. - In `_min_of`/`_max_of`, reduce the per-shard list with `torch.stack(xs).min()` / `.max()` instead of Python `min()`/`max()`, which forced a CPU-GPU sync. - In the SCALAR branch, read values via `self.stats.get(key, [])` (a rank may learn a key only through metadata sync) and guard the `value / cnt` division against `cnt == 0` so a key absent everywhere yields 0.0 instead of NaN. - Add a regression test for the missing-on-this-rank SCALAR path. | 1 个月前 | |
feat(infra): add microservice-based training service (controller v2) (#1169) * feat(infra): add microservice-based training service (controller v2) Add GatewayTrainController that decomposes training into five HTTP microservices: guard (process manager), worker (engine container), data proxy (batch dispatcher), router (service registry), and gateway (API ingress). This enables training orchestration without requiring the scheduler's RPC infrastructure. Key changes: - Add GatewayTrainController with 7-step async initialization - Add guard /set_env endpoint for NCCL env propagation - Add worker, router, gateway, data proxy FastAPI/Flask services - Wire TrainEngineConfig.log_level to suppress HTTP access logs - Add create_process_group stubs to existing engines - Add comprehensive unit and integration tests * fix(infra): stabilize controller v2 training service Consolidate the controller v2 training-service follow-up work into one atomic infra commit. This keeps startup, routing, dispatch, recovery, and health-handling changes together as the post-9c70 stabilization series while restoring a clean branch history. * feat(infra): add microservice-based training service (controller v2) Add GatewayTrainController that decomposes training into five HTTP microservices: guard (process manager), worker (engine container), data proxy (batch dispatcher), router (service registry), and gateway (API ingress). This enables training orchestration without requiring the scheduler's RPC infrastructure. Key changes: - Add GatewayTrainController with 7-step async initialization - Add guard /set_env endpoint for NCCL env propagation - Add worker, router, gateway, data proxy FastAPI/Flask services - Wire TrainEngineConfig.log_level to suppress HTTP access logs - Add create_process_group stubs to existing engines - Add comprehensive unit and integration tests * fix(infra): stabilize controller v2 training service Consolidate the controller v2 training-service follow-up work into one atomic infra commit. This keeps startup, routing, dispatch, recovery, and health-handling changes together as the post-9c70 stabilization series while restoring a clean branch history. * chore: add SPDX headers to training service modules --------- Co-authored-by: Wentai Zhang <zhangwentai.zwt@antgroup.com> | 4 个月前 | |
feat(engine): enable model-owned THD for Qwen3-VL (#1608) Allow Qwen3-VL models to merge multimodal embeddings before packing while preserving existing padded-only and wrapper-owned paths. Key changes: - Route dense and MoE Qwen3-VL through model-owned THD - Reconstruct padded inputs and restore packed model outputs - Add routing, alignment, parity, and distributed forward coverage | 16 天前 | |
feat(v2): support grouped reward normalization (#1619) * feat(v2): support grouped reward normalization Normalize per-prompt rollout groups at the data-proxy export boundary before trajectory tensors are merged and remotized. Preserve raw rewards for filtering and metrics while sharing the normalization implementation with v1. Keep reward metadata local for rollout filters and clear remote shards for rejected or failed trajectories so filtered v2 rollouts do not leak data-proxy storage. * fix(v2): bound rejected trajectory cleanup | 8 天前 | |
refactor(tests): move tests from areal/tests to top-level tests directory (#944) * refactor(tests): move tests from areal/tests to top-level tests directory Move the test directory from areal/tests/ to a top-level tests/ directory to follow Python best practices and improve project structure. Key changes: - Move all test files from areal/tests/ to tests/ - Create __init__.py files for tests package and subdirectories - Create areal/utils/testing_utils.py with shared utilities: - get_model_path, get_dataset_path (model/dataset path resolution) - MODEL_PATHS, DENSE_MODEL_PATHS, MOE_MODEL_PATHS (test model configs) - load_archon_model (Archon model loading for tests) - Update areal/tools/profiling_utils to re-export from testing_utils - Update test utilities to import from areal/utils/testing_utils, keeping test-specific behavior (pytest.skip) as wrappers - Update all imports from 'areal.tests.' to 'tests.' - Update hardcoded torchrun script paths - Update workflow string references for dynamic class loading - Update CI workflow paths in .github/workflows/test-areal.yml - Update documentation references in CLAUDE.md, AGENTS.md, CONTRIBUTING.md, and all .claude/ and .opencode/ skill/agent files * minor fix * remove import side effect * fix(tests): add tests/utils.py module for test imports Re-export get_model_path and get_dataset_path from areal.utils.testing_utils to fix ModuleNotFoundError in tests that import from tests.utils. | 6 个月前 | |
feat: support model training in IPv6-only environment (#1072) * feat: support model training in IPv6-only environment --------- Co-authored-by: bingyechen <bingyechen@bytedance.com> Co-authored-by: root <root@dc05-p13-t0-n028.byted.org> | 5 个月前 | |
fix(hermes): preserve singleton online reward signal (#1474) * fix(ppo): warn when singleton groups erase rewards * fix(hermes): preserve singleton online reward signal * test(hermes): harden reward signal regression * style: format singleton reward warning test * fix(config): support mapping-shaped reward normalization | 1 个月前 | |
fix(tests): assert all_gather_param passthrough by storage identity (#1562) Tensor.data returns a fresh wrapper object on every access, so `out is param.data` compares two different wrappers and fails under current torch. Assert data_ptr and storage_offset instead, matching the same fix carried by the pending profiling PR so the two merge cleanly in either order. | 1 个月前 | |
feat(colocate): support AWEX colocated actor-rollout training (#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. | 26 天前 | |
refactor(api): migrate allocation_mode to per-engine backend fields (#1044) * refactor(api): migrate allocation_mode to per-engine backend fields Replace the centralized `allocation_mode` string with explicit `backend` fields on `TrainEngineConfig` and `InferenceEngineConfig`. Each engine now owns its own backend+parallelism spec (e.g. `fsdp:d4`, `sglang:d4t2`), eliminating implicit auto-backend selection and the shared `AllocationMode` object. Key changes: - Add `backend` field to TrainEngineConfig and InferenceEngineConfig - Add `ModelAllocation.from_str()` for single-component parsing - Remove `AllocationMode` public export (replaced by `ModelAllocation`) - Rename internal `AllocationMode` to `_AllocationMode` for SPMD launcher backward compatibility with FutureWarning - Remove auto-backend selection — explicit backend prefix is now required - Controllers (`TrainController`, `RolloutController`) parse `backend` directly instead of receiving `alloc_mode` from trainers - `WeightUpdateMeta.alloc_mode` replaced by `gen_allocation` (single `ModelAllocation`) - Add `RWTrainer` and `ArchonRWEngine` for reward model training - Remove `get_model_update_meta()` helper (logic moved to trainers) - Update all YAML configs, examples, docs (EN+ZH), and tests BREAKING CHANGE: `AllocationMode` is removed from public API. Users must migrate to per-engine `backend` fields. SPMD launchers emit deprecation warnings. * chore(ci): fix backend specifier for vlm sft test * fix: fix bare dims for actor backends * chore(docs): fix reminder for bare allocation dims | 5 个月前 | |
fix(api): return float zero after async reward timeout (#1541) AsyncRewardWrapper returned an integer zero when reward computation timed out, despite declaring a float result. The v1 OpenAI proxy only accepts dictionary or float rewards, so the fallback could raise ValueError instead of recording a zero reward. Return 0.0 from both wrapper-owned fallback paths and strengthen the existing timeout tests to verify the result type. Signed-off-by: Bo Yang <yb550079@antgroup.com> | 8 天前 | |
refactor(tests): move tests from areal/tests to top-level tests directory (#944) * refactor(tests): move tests from areal/tests to top-level tests directory Move the test directory from areal/tests/ to a top-level tests/ directory to follow Python best practices and improve project structure. Key changes: - Move all test files from areal/tests/ to tests/ - Create __init__.py files for tests package and subdirectories - Create areal/utils/testing_utils.py with shared utilities: - get_model_path, get_dataset_path (model/dataset path resolution) - MODEL_PATHS, DENSE_MODEL_PATHS, MOE_MODEL_PATHS (test model configs) - load_archon_model (Archon model loading for tests) - Update areal/tools/profiling_utils to re-export from testing_utils - Update test utilities to import from areal/utils/testing_utils, keeping test-specific behavior (pytest.skip) as wrappers - Update all imports from 'areal.tests.' to 'tests.' - Update hardcoded torchrun script paths - Update workflow string references for dynamic class loading - Update CI workflow paths in .github/workflows/test-areal.yml - Update documentation references in CLAUDE.md, AGENTS.md, CONTRIBUTING.md, and all .claude/ and .opencode/ skill/agent files * minor fix * remove import side effect * fix(tests): add tests/utils.py module for test imports Re-export get_model_path and get_dataset_path from areal.utils.testing_utils to fix ModuleNotFoundError in tests that import from tests.utils. | 6 个月前 | |
Add opt-in support for Hugging Face kernels (#1033) * feat: add opt-in huggingface kernels support * docs: explain how to enable kernels in training, link Hugging Face kernels docs * fix: move attn impl validation into fsdp utils --------- Co-authored-by: OpenAI Codex <codex@openai.com> | 5 个月前 | |
feat(colocate): support AWEX colocated actor-rollout training (#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. | 26 天前 | |
feat(awex): add separation AdamW delta weight transfer (#1623) * feat(config): add separation DTE topology gates * feat(awex): add separation AdamW delta weight transfer * docs(examples): add DTE separation GSM8K example * docs(examples): expand DTE GSM8K configuration * fix(awex): address separation DTE review feedback * fix(awex): harden separation delta transfer protocol * refactor(config): integrate delta transfer into train engine * refactor(config): simplify delta weight update switch * test(awex): isolate optional DTE dependency | 11 天前 | |
feat(engine): support Qwen3-VL with native AWEX colocate (#1605) * feat(engine): support Qwen3-VL with native AWEX colocate Preserve the complete multimodal Hugging Face config during colocate\nmetadata exchange so AWEX can resolve both the language model and vision\ntower sharding contracts.\n\nKey changes:\n- Publish composite Qwen3-VL inference configuration\n- Delegate nested config parsing to native AWEX\n- Add Dense and MoE colocate contract tests * fix(engine): preserve Qwen3-VL-MoE composite config SGLang stores only the text config on the Qwen3-VL-MoE runtime model.\nRead the original Hugging Face config from ModelRunner so AWEX also\nreceives the vision tower metadata required for weight sharding. | 14 天前 | |
feat(awex): add separation AdamW delta weight transfer (#1623) * feat(config): add separation DTE topology gates * feat(awex): add separation AdamW delta weight transfer * docs(examples): add DTE separation GSM8K example * docs(examples): expand DTE GSM8K configuration * fix(awex): address separation DTE review feedback * fix(awex): harden separation delta transfer protocol * refactor(config): integrate delta transfer into train engine * refactor(config): simplify delta weight update switch * test(awex): isolate optional DTE dependency | 11 天前 | |
feat(colocate): support AWEX colocated actor-rollout training (#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. | 26 天前 | |
feat(ppo): add CISPO loss surrogate (MiniMax-M1) (#1412) * feat(ppo): add CISPO advantage estimator (MiniMax-M1) PPO/GRPO clipping zeroes the gradient of any token whose importance-sampling ratio leaves the clip band: `min(r*A, clip(r)*A)` is constant in theta there. MiniMax-M1 (https://arxiv.org/abs/2506.13585, Eq. 4-5) observes those are disproportionately the low-probability "fork" tokens (`However`, `Wait`, ...) that steer reasoning -- exactly the tokens reasoning RL needs gradient on -- and instead clips the IS *weight* under stop-gradient while keeping the gradient on every token's `log pi_theta`. ScaleRL (arXiv:2510.13786 Eq. 4) adopts the same surrogate. AReaL had grpo/gspo/ppo/sapo but no CISPO. Per token, opt-in via `actor.use_cispo_loss=true`: ratio = exp(logprobs - proximal_logprobs) ratio_clipped = clip(ratio, 1 - eps_clip, 1 + eps_clip_higher) # stop-grad pg_loss = -sg(ratio_clipped) * advantages * logprobs Advantages are never clipped. The clip reuses the existing delta-from-1 `eps_clip` / `eps_clip_higher` plumbing (same convention as GSPO); CISPO is canonically single-sided, so the recommended setting is `eps_clip=1.0` (lower bound 0) with `eps_clip_higher=4.0` for the wide MiniMax-M1 range. Shape of the change -- additive, fits AReaL's existing estimator seam: - `functional.py`: `cispo_loss_fn` (+ export), reusing the token-mean reduction and the PPO-compatible stat schema so the existing clip-stat logging path applies unchanged; `clip_mask` reports band-exit on either side (under CISPO no clip zeroes the loss, so band-exit -- not loss-affecting clip -- is the meaningful metric). - `actor.py`: `grpo_loss_fn` dispatches to CISPO before SAPO/PPO. - `cli_args.py`: `PPOActorConfig.use_cispo_loss` with `__post_init__` validation (mutually exclusive with SAPO, requires `eps_clip_higher > 0`, token-level importance sampling only) + regenerated CLI docs (en/zh). Defaults unchanged (`use_cispo_loss=False`) -> byte-identical to before. `tests/test_cispo_loss.py` pins the two defining invariants, each across a PPO-like band (0.2/0.28) and the wide MiniMax band (1.0/4.0): the closed-form surrogate value + clip indices, and gradient routing (`logprobs.grad == -sg(clip(ratio)) * A / N`, zero gradient through the IS-ratio path -- mutation-verified to fail if the stop-gradient detach is dropped), plus config validation. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * feat(ppo): support decoupled loss for CISPO via rejection sampling The CISPO branch ignored use_decoupled_loss / rejection_sampling, silently dropping the off-policy correction in exactly the async/stale regime CISPO targets. CISPO already anchors its clipped IS ratio at pi_proximal; thread the behavior logp + rejection_sampling through cispo_loss_fn and rescale each token's surrogate by the detached pi_proximal/pi_behave weight, mirroring ppo_actor_loss_fn. Both factors are stop-gradient, so the estimator stays -sg(behave_imp_weight * ratio_clipped) * A * grad(log pi_theta). Addresses the maintainer review on the dispatch bypassing the decoupled check. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: EazyReal <8047065+EazyReal@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> | 2 个月前 | |
fix(reward): guard clevr_count_70k_reward_fn against scoring failures (#1430) * fix(reward): score non-string answers in clevr_count_70k_reward_fn clevr_count_70k_reward_fn did not str()-coerce its inputs or guard against errors, unlike the sibling reward fns (gsm8k, geometry3k). A non-string answer (e.g. an int) made ans.strip() raise AttributeError, which WorkflowExecutor catches and uses to reject the whole trajectory instead of scoring the sample — and even a matching completion was lost rather than scored 1.0. Coerce completions and answer to str and wrap the body in try/except, matching the sibling reward fns. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * fix(example): apply the same guard to the clevr GRPO example reward fn The example defines its own clevr_count_70k_reward_fn (referenced via workflow_kwargs) that still had the old logic. Mirror the built-in fix: str-coerce completions/answer and guard, so a non-string answer is scored rather than raising and dropping the trajectory. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> | 2 个月前 | |
refactor(tests): move tests from areal/tests to top-level tests directory (#944) * refactor(tests): move tests from areal/tests to top-level tests directory Move the test directory from areal/tests/ to a top-level tests/ directory to follow Python best practices and improve project structure. Key changes: - Move all test files from areal/tests/ to tests/ - Create __init__.py files for tests package and subdirectories - Create areal/utils/testing_utils.py with shared utilities: - get_model_path, get_dataset_path (model/dataset path resolution) - MODEL_PATHS, DENSE_MODEL_PATHS, MOE_MODEL_PATHS (test model configs) - load_archon_model (Archon model loading for tests) - Update areal/tools/profiling_utils to re-export from testing_utils - Update test utilities to import from areal/utils/testing_utils, keeping test-specific behavior (pytest.skip) as wrappers - Update all imports from 'areal.tests.' to 'tests.' - Update hardcoded torchrun script paths - Update workflow string references for dynamic class loading - Update CI workflow paths in .github/workflows/test-areal.yml - Update documentation references in CLAUDE.md, AGENTS.md, CONTRIBUTING.md, and all .claude/ and .opencode/ skill/agent files * minor fix * remove import side effect * fix(tests): add tests/utils.py module for test imports Re-export get_model_path and get_dataset_path from areal.utils.testing_utils to fix ModuleNotFoundError in tests that import from tests.utils. | 6 个月前 | |
refactor(infra): simplify RTensor to single-shard and adopt per-trajectory list pipeline (#1017) * refactor(infra): simplify RTensor to single-shard and adopt per-trajectory list pipeline * fix(infra): remove layout reliance of RTensor * fix(infra): refactor the train controller data dispatch logic * refactor(data): rename datapack -> seqpack and refactor functions in data utility * fix(infra): concurrent get rtensors * fix(engine): consolidate new list engine APIs --------- Co-authored-by: 博惟 <bowei.fw@antgroup.com> | 5 个月前 | |
fix(engine): engage Megatron deterministic mode before model build (#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. | 30 天前 | |
fix: make rollout sampling deterministic (#1625) | 12 天前 | |
feat(trainer): add dpo (#1190) * feat(trainer): add DPO trainer with FSDP backend Add Direct Preference Optimization (Rafailov et al. 2023) as a new trainer. The policy is directly optimized to prefer chosen over rejected responses via a contrastive loss on log-probability ratios against a frozen reference model, removing the need for a separately trained reward model. Reference logprobs are computed online each step by a colocated ref engine, following the PPO/GRPO pattern. FSDP is the supported backend; Megatron and Archon variants raise NotImplementedError as placeholders. Verified on Qwen2.5-7B-Base + Anthropic/hh-rlhf (1 epoch, no SFT): reward_accuracy rises from 0.50 to ~0.70 and reward_margin grows monotonically, matching the original DPO paper's HH-RLHF results. * fix(trainer): fix DPO config forwarding, require ref model, and correct IPO normalization Fixes several issues found during PR review of the DPO trainer: Key changes: - Create DPOEngineConfig(TrainEngineConfig) embedding beta and loss_type, fixing silent parameter drop in single-controller mode (as_controller never forwarded beta/loss_type to workers) - Make ref a required field in DPOConfig (ref_logprobs are required at runtime, so config should enforce this upfront) - Remove zero-ref fallback in compute_dpo_loss; use input_["ref_logprobs"] directly - Add IPO loss with per-token length normalization matching TRL author- confirmed convention (normalize per-sequence logratios by completion length before the squared loss) - Remove all ref-is-None guard branches from DPOTrainer - Update docs, YAML config, and tests for all changes Refs: #1190 --------- Co-authored-by: 博惟 <bowei.fw@antgroup.com> | 4 个月前 | |
feat(awex): add separation AdamW delta weight transfer (#1623) * feat(config): add separation DTE topology gates * feat(awex): add separation AdamW delta weight transfer * docs(examples): add DTE separation GSM8K example * docs(examples): expand DTE GSM8K configuration * fix(awex): address separation DTE review feedback * fix(awex): harden separation delta transfer protocol * refactor(config): integrate delta transfer into train engine * refactor(config): simplify delta weight update switch * test(awex): isolate optional DTE dependency | 11 天前 | |
refactor(tests): move tests from areal/tests to top-level tests directory (#944) * refactor(tests): move tests from areal/tests to top-level tests directory Move the test directory from areal/tests/ to a top-level tests/ directory to follow Python best practices and improve project structure. Key changes: - Move all test files from areal/tests/ to tests/ - Create __init__.py files for tests package and subdirectories - Create areal/utils/testing_utils.py with shared utilities: - get_model_path, get_dataset_path (model/dataset path resolution) - MODEL_PATHS, DENSE_MODEL_PATHS, MOE_MODEL_PATHS (test model configs) - load_archon_model (Archon model loading for tests) - Update areal/tools/profiling_utils to re-export from testing_utils - Update test utilities to import from areal/utils/testing_utils, keeping test-specific behavior (pytest.skip) as wrappers - Update all imports from 'areal.tests.' to 'tests.' - Update hardcoded torchrun script paths - Update workflow string references for dynamic class loading - Update CI workflow paths in .github/workflows/test-areal.yml - Update documentation references in CLAUDE.md, AGENTS.md, CONTRIBUTING.md, and all .claude/ and .opencode/ skill/agent files * minor fix * remove import side effect * fix(tests): add tests/utils.py module for test imports Re-export get_model_path and get_dataset_path from areal.utils.testing_utils to fix ModuleNotFoundError in tests that import from tests.utils. | 6 个月前 | |
feat(colocate): support AWEX colocated actor-rollout training (#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. | 26 天前 | |
feat(awex): add separation AdamW delta weight transfer (#1623) * feat(config): add separation DTE topology gates * feat(awex): add separation AdamW delta weight transfer * docs(examples): add DTE separation GSM8K example * docs(examples): expand DTE GSM8K configuration * fix(awex): address separation DTE review feedback * fix(awex): harden separation delta transfer protocol * refactor(config): integrate delta transfer into train engine * refactor(config): simplify delta weight update switch * test(awex): isolate optional DTE dependency | 11 天前 | |
refactor(tests): move tests from areal/tests to top-level tests directory (#944) * refactor(tests): move tests from areal/tests to top-level tests directory Move the test directory from areal/tests/ to a top-level tests/ directory to follow Python best practices and improve project structure. Key changes: - Move all test files from areal/tests/ to tests/ - Create __init__.py files for tests package and subdirectories - Create areal/utils/testing_utils.py with shared utilities: - get_model_path, get_dataset_path (model/dataset path resolution) - MODEL_PATHS, DENSE_MODEL_PATHS, MOE_MODEL_PATHS (test model configs) - load_archon_model (Archon model loading for tests) - Update areal/tools/profiling_utils to re-export from testing_utils - Update test utilities to import from areal/utils/testing_utils, keeping test-specific behavior (pytest.skip) as wrappers - Update all imports from 'areal.tests.' to 'tests.' - Update hardcoded torchrun script paths - Update workflow string references for dynamic class loading - Update CI workflow paths in .github/workflows/test-areal.yml - Update documentation references in CLAUDE.md, AGENTS.md, CONTRIBUTING.md, and all .claude/ and .opencode/ skill/agent files * minor fix * remove import side effect * fix(tests): add tests/utils.py module for test imports Re-export get_model_path and get_dataset_path from areal.utils.testing_utils to fix ModuleNotFoundError in tests that import from tests.utils. | 6 个月前 | |
feat(infra): allow colocation with offloading and disk weight updates (#1157) * feat(infra): allow colocation with offloading and disk weight updates Enable scheduler-level colocation to run actor/critic and engine processes on shared GPU allocations while preserving async update correctness. Key changes: - Add colocate/offload/disk-update config fields and scheduler plumbing - Harden RPC server and engine blueprint coordination for colocated roles - Extend trainer paths and tests for colocated evaluation dispatch behavior * fix(infra): restore default train-engine RPC broadcast Keep initialized TrainEngine RPC calls backward compatible across the guard and Ray servers so non-head model-parallel ranks continue to receive controller payloads without every call site opting in manually. * fix: enforce offload prerequisites for colocated training Fail fast when colocated or explicit train-engine offload would run without TMS support, and provision Ray workers with the same offload environment as the local and Slurm schedulers. --------- Co-authored-by: Wentai Zhang <zhangwentai.zwt@antgroup.com> | 4 个月前 | |
fix(trainer): run initial evaluation before first update (#1636) `eval_before_train` currently relies on the first scheduled evaluator check, which trainers perform only after one optimization step. The reported baseline can therefore contain updated weights. Run the one-shot evaluation separately so its metrics are logged before the first update without advancing periodic evaluation cadence. Skip the startup evaluation on recovery and clear legacy deferred triggers when loading evaluator state. Key changes: - Invoke the startup baseline from PPO, SFT, DPO, and reward trainers - Preserve periodic evaluation cadence and recovery behavior - Cover call ordering, logging steps, legacy state, and PPO offload Refs: #1232 Signed-off-by: Bo Yang <yb550079@antgroup.com> | 8 天前 | |
feat(colocate): support AWEX colocated actor-rollout training (#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. | 26 天前 | |
refactor(api): migrate allocation_mode to per-engine backend fields (#1044) * refactor(api): migrate allocation_mode to per-engine backend fields Replace the centralized `allocation_mode` string with explicit `backend` fields on `TrainEngineConfig` and `InferenceEngineConfig`. Each engine now owns its own backend+parallelism spec (e.g. `fsdp:d4`, `sglang:d4t2`), eliminating implicit auto-backend selection and the shared `AllocationMode` object. Key changes: - Add `backend` field to TrainEngineConfig and InferenceEngineConfig - Add `ModelAllocation.from_str()` for single-component parsing - Remove `AllocationMode` public export (replaced by `ModelAllocation`) - Rename internal `AllocationMode` to `_AllocationMode` for SPMD launcher backward compatibility with FutureWarning - Remove auto-backend selection — explicit backend prefix is now required - Controllers (`TrainController`, `RolloutController`) parse `backend` directly instead of receiving `alloc_mode` from trainers - `WeightUpdateMeta.alloc_mode` replaced by `gen_allocation` (single `ModelAllocation`) - Add `RWTrainer` and `ArchonRWEngine` for reward model training - Remove `get_model_update_meta()` helper (logic moved to trainers) - Update all YAML configs, examples, docs (EN+ZH), and tests BREAKING CHANGE: `AllocationMode` is removed from public API. Users must migrate to per-engine `backend` fields. SPMD launchers emit deprecation warnings. * chore(ci): fix backend specifier for vlm sft test * fix: fix bare dims for actor backends * chore(docs): fix reminder for bare allocation dims | 5 个月前 | |
feat: support model training in IPv6-only environment (#1072) * feat: support model training in IPv6-only environment --------- Co-authored-by: bingyechen <bingyechen@bytedance.com> Co-authored-by: root <root@dc05-p13-t0-n028.byted.org> | 5 个月前 | |
Fix grad norm hang when LoRA frozen ranks have no gradients (#1139) Ranks with no gradients (e.g. frozen non-LoRA params) previously returned 0.0 immediately, skipping the all_reduce. Ranks that do have gradients then hang waiting for the collective to complete. Move device init before the empty-grads check and make zero-grad ranks still participate in all_reduce with a zero-valued tensor. Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> | 4 个月前 | |
refactor(api): migrate allocation_mode to per-engine backend fields (#1044) * refactor(api): migrate allocation_mode to per-engine backend fields Replace the centralized `allocation_mode` string with explicit `backend` fields on `TrainEngineConfig` and `InferenceEngineConfig`. Each engine now owns its own backend+parallelism spec (e.g. `fsdp:d4`, `sglang:d4t2`), eliminating implicit auto-backend selection and the shared `AllocationMode` object. Key changes: - Add `backend` field to TrainEngineConfig and InferenceEngineConfig - Add `ModelAllocation.from_str()` for single-component parsing - Remove `AllocationMode` public export (replaced by `ModelAllocation`) - Rename internal `AllocationMode` to `_AllocationMode` for SPMD launcher backward compatibility with FutureWarning - Remove auto-backend selection — explicit backend prefix is now required - Controllers (`TrainController`, `RolloutController`) parse `backend` directly instead of receiving `alloc_mode` from trainers - `WeightUpdateMeta.alloc_mode` replaced by `gen_allocation` (single `ModelAllocation`) - Add `RWTrainer` and `ArchonRWEngine` for reward model training - Remove `get_model_update_meta()` helper (logic moved to trainers) - Update all YAML configs, examples, docs (EN+ZH), and tests BREAKING CHANGE: `AllocationMode` is removed from public API. Users must migrate to per-engine `backend` fields. SPMD launchers emit deprecation warnings. * chore(ci): fix backend specifier for vlm sft test * fix: fix bare dims for actor backends * chore(docs): fix reminder for bare allocation dims | 5 个月前 | |
fix(engine): avoid duplicating multimodal tensors (#1272) | 4 个月前 | |
fix(fsdp): maintain fp32 master weights for AdamW (#1292) (#1369) Fix issue #1292: torch.optim.AdamW silently inherits bf16 dtype from model.parameters() when actor.dtype=bfloat16, causing late-stage SFT loss to plateau ~3x higher than DS-Z3 / Megatron precision-aware optimizer paths. Decouple parameter storage dtype from forward/backward compute dtype so AdamW sees fp32 params natively (fp32 master weights), while FSDP2 MixedPrecisionPolicy(param_dtype=dtype) keeps forward/backward in bf16. Cast back to compute dtype on HF export and xccl weight sync to keep deployment artefacts and rollout broadcasts unchanged. Changes: - Add TrainEngineConfig.optimizer_dtype (default 'float32') with __post_init__ validation and fp32/bf16 alias canonicalization - FSDPEngine: load model in optimizer_dtype (was: dtype) so AdamW creates fp32 exp_avg/exp_avg_sq when default. FSDP2 MP policy unchanged - Add _cast_to_compute_dtype helper; use in _save_model_to_hf and _update_weights_from_distributed (xccl) to cast fp32 storage back to compute dtype before disk export / SGLang broadcast - Replace misleading adam_bf16 'less stable' warning with dtype-aware logic; explicitly warn when adam + optimizer_dtype=bfloat16 combination triggers #1292 - Regression test (1 / 2 GPU) covering 5 dtype invariants: storage / Adam state / forward / HF export / xccl cast - Document fp32 master tradeoffs and adam_bf16 OOM fallback in en + zh handling_oom.md; regenerate CLI reference Backward compatibility: setting optimizer_dtype='bfloat16' explicitly restores pre-fix behaviour bit-for-bit (storage bf16, Adam state bf16, all casts no-op). The fix is opt-out via config rather than breaking change. Verified on 2x RTX 4090 with Qwen3-0.6B, GSM8K, 466 steps: - First 2 steps bit-identical between BEFORE/AFTER (sanity) - Last-50 mean: BEFORE 0.3904 vs AFTER 0.3461 (gap +0.044, 11%) - All 3 dtype-invariant pytest cases PASSED Fixes #1292. | 3 个月前 | |
refactor(tests): move tests from areal/tests to top-level tests directory (#944) * refactor(tests): move tests from areal/tests to top-level tests directory Move the test directory from areal/tests/ to a top-level tests/ directory to follow Python best practices and improve project structure. Key changes: - Move all test files from areal/tests/ to tests/ - Create __init__.py files for tests package and subdirectories - Create areal/utils/testing_utils.py with shared utilities: - get_model_path, get_dataset_path (model/dataset path resolution) - MODEL_PATHS, DENSE_MODEL_PATHS, MOE_MODEL_PATHS (test model configs) - load_archon_model (Archon model loading for tests) - Update areal/tools/profiling_utils to re-export from testing_utils - Update test utilities to import from areal/utils/testing_utils, keeping test-specific behavior (pytest.skip) as wrappers - Update all imports from 'areal.tests.' to 'tests.' - Update hardcoded torchrun script paths - Update workflow string references for dynamic class loading - Update CI workflow paths in .github/workflows/test-areal.yml - Update documentation references in CLAUDE.md, AGENTS.md, CONTRIBUTING.md, and all .claude/ and .opencode/ skill/agent files * minor fix * remove import side effect * fix(tests): add tests/utils.py module for test imports Re-export get_model_path and get_dataset_path from areal.utils.testing_utils to fix ModuleNotFoundError in tests that import from tests.utils. | 6 个月前 | |
refactor(tests): move tests from areal/tests to top-level tests directory (#944) * refactor(tests): move tests from areal/tests to top-level tests directory Move the test directory from areal/tests/ to a top-level tests/ directory to follow Python best practices and improve project structure. Key changes: - Move all test files from areal/tests/ to tests/ - Create __init__.py files for tests package and subdirectories - Create areal/utils/testing_utils.py with shared utilities: - get_model_path, get_dataset_path (model/dataset path resolution) - MODEL_PATHS, DENSE_MODEL_PATHS, MOE_MODEL_PATHS (test model configs) - load_archon_model (Archon model loading for tests) - Update areal/tools/profiling_utils to re-export from testing_utils - Update test utilities to import from areal/utils/testing_utils, keeping test-specific behavior (pytest.skip) as wrappers - Update all imports from 'areal.tests.' to 'tests.' - Update hardcoded torchrun script paths - Update workflow string references for dynamic class loading - Update CI workflow paths in .github/workflows/test-areal.yml - Update documentation references in CLAUDE.md, AGENTS.md, CONTRIBUTING.md, and all .claude/ and .opencode/ skill/agent files * minor fix * remove import side effect * fix(tests): add tests/utils.py module for test imports Re-export get_model_path and get_dataset_path from areal.utils.testing_utils to fix ModuleNotFoundError in tests that import from tests.utils. | 6 个月前 | |
fix(utils): mask 2d sequence advantages (#1346) ## Summary AReaL's 2D padded sequence-level PPO/GSPO loss can include `loss_mask=False` padding positions when it averages token advantages into one sequence advantage. As a result, changing only masked padding values can silently change valid-token loss, gradients, and one-step updates. This patch applies `loss_mask` before the 2D per-sequence advantage reduction, then zeros masked positions after broadcasting the sequence advantage back to token shape. ## Concrete Minimal Example Boundary: ```python from areal.utils.functional import ppo_actor_loss_fn loss_mask = torch.tensor([[True, True, False, False]]) logprobs = torch.zeros(1, 4) proximal_logprobs = torch.zeros(1, 4) old_logprobs = torch.zeros(1, 4) clean_advantages = torch.tensor([[1.0, 1.0, 0.0, 0.0]]) contaminated_advantages = torch.tensor([[1.0, 1.0, 10.0, 10.0]]) ppo_actor_loss_fn( logprobs=logprobs, proximal_logprobs=proximal_logprobs, old_logprobs=old_logprobs, advantages=contaminated_advantages, eps_clip=0.2, loss_mask=loss_mask, importance_sampling_level="sequence", ) ``` Only positions where `loss_mask` is `False` change. Wrong value on current `main`: ```json { "clean_2d": { "loss": -1.0, "valid_logging_loss": [-1.0, -1.0], "grad_norm": 1.0, "update_delta_norm": 0.10000000149011612 }, "contaminated_2d": { "loss": -11.0, "valid_logging_loss": [-11.0, -11.0], "grad_norm": 11.0, "update_delta_norm": 1.100000023841858 } } ``` Fixed value on this branch: ```json { "clean_2d": { "loss": -1.0, "valid_logging_loss": [-1.0, -1.0], "grad_norm": 1.0, "update_delta_norm": 0.10000000149011612 }, "contaminated_2d": { "loss": -1.0, "valid_logging_loss": [-1.0, -1.0], "grad_norm": 1.0, "update_delta_norm": 0.10000000149011612 } } ``` ## Root Cause The packed 1D branch already does: ```python masked_advantages = torch.where(loss_mask, advantages, 0.0) ``` before reducing by sequence. The 2D padded branch used: ```python advantages.sum(dim=-1, keepdim=True) / loss_mask.sum(dim=-1, keepdim=True).clamp(min=1) ``` This divided by the valid-token count but still included padded or otherwise masked advantage values in the numerator. ## Fix The 2D branch now selects valid advantages before the sequence reduction and zeros the broadcast result on masked positions: ```python masked_advantages = torch.where(loss_mask, advantages, 0.0) advantages = ( masked_advantages.sum(dim=-1, keepdim=True) / seq_lengths ).expand_as(log_ratio) advantages = torch.where(loss_mask, advantages, 0.0) ``` ## Validation Recipe ```json { "bug_id": "AREAL-GSPO-2D-MASKED-ADV-LEAK", "validation_mode": "actual_areal_ppo_actor_loss_boundary_hook", "hooked_boundary": "areal.utils.functional.ppo_actor_loss_fn with importance_sampling_level='sequence'", "constructed_scenario": { "format": "2D padded batch", "loss_mask": [[true, true, false, false]], "clean_advantages": [[1.0, 1.0, 0.0, 0.0]], "contaminated_masked_advantages": [[1.0, 1.0, 10.0, 10.0]] }, "expected_invariant": "Changing advantage values where loss_mask is false must not change loss, valid-token logging loss, gradient, or one-step update.", "replaced_component": null } ``` Runner script: ```python import importlib import json import os import sys from pathlib import Path import torch target_repo = Path(os.environ["TARGET_REPO"]).resolve() sys.path.insert(0, str(target_repo)) functional_module = importlib.import_module("areal.utils.functional.functional") ppo_actor_loss_fn = functional_module.ppo_actor_loss_fn def run_case(advantages): loss_mask = torch.tensor([[True, True, False, False]]) theta = torch.nn.Parameter(torch.tensor(0.0)) optimizer = torch.optim.SGD([theta], lr=0.1) optimizer.zero_grad() logprobs = theta.expand_as(advantages) zeros = torch.zeros_like(advantages) loss, stat = ppo_actor_loss_fn( logprobs=logprobs, proximal_logprobs=zeros, old_logprobs=zeros, advantages=advantages, eps_clip=0.2, loss_mask=loss_mask, rejection_sampling=None, importance_sampling_level="sequence", cu_seqlens=None, ) loss.backward() before = float(theta.detach().item()) grad = float(theta.grad.detach().item()) optimizer.step() after = float(theta.detach().item()) return { "loss": float(loss.detach().item()), "valid_logging_loss": stat["loss"][loss_mask].detach().cpu().tolist(), "grad_norm": abs(grad), "update_delta_norm": abs(after - before), } clean_advantages = torch.tensor([[1.0, 1.0, 0.0, 0.0]]) contaminated_advantages = torch.tensor([[1.0, 1.0, 10.0, 10.0]]) print( json.dumps( { "clean_2d": run_case(clean_advantages), "contaminated_2d": run_case(contaminated_advantages), }, indent=2, sort_keys=True, ) ) ``` The hook imports `areal.utils.functional.functional` from the checkout under test and uses real PyTorch autograd plus `torch.optim.SGD`. It does not replace the target loss function. ## Checks Passed: ```bash git diff --check python3 -m ruff check areal/utils/functional/functional.py tests/test_functional.py python3 -m ruff format --check areal/utils/functional/functional.py tests/test_functional.py python3 -m pre_commit install --install-hooks python3 -m pre_commit run --files areal/utils/functional/functional.py tests/test_functional.py ``` ## Relevant Output Unfixed target output: ```json { "clean_2d": {"loss": -1.0, "grad_norm": 1.0, "update_delta_norm": 0.10000000149011612}, "contaminated_2d": {"loss": -11.0, "grad_norm": 11.0, "update_delta_norm": 1.100000023841858} } ``` Fixed-path hook output: ```json { "status": "fixed", "repair_commit": "9631a1d651cfd2870f18edc7213e4f932ce83869", "observed_fixed_behavior": { "clean_2d": {"loss": -1.0, "grad_norm": 1.0}, "contaminated_2d": {"loss": -1.0, "grad_norm": 1.0} } } ``` Local pytest note: ```text python3 -m pytest -q tests/test_functional.py::TestPPOActorLossFnSequenceLevel::test_sequence_level_2d_advantage_average_ignores_masked_values ``` is blocked in this local environment by a top-level AReaL import error: ```text ImportError: cannot import name 'DefaultStager' from 'torch.distributed.checkpoint.staging' ``` The targeted boundary hook avoids top-level `import areal` per the AReaL workflow and imports the concrete real module from the checkout under test. ## Related Work / Dedup - `inclusionAI/AReaL` redirects to `areal-project/AReaL`. - GitHub issue/PR search found related feature PRs, including `areal-project/AReaL#501` and `areal-project/AReaL#1088`, but no PR or issue describing this 2D padded masked-advantage bug. Current refreshed `main` still had the unmasked `advantages.sum(dim=-1, keepdim=True)` path. - Historical RL-Sentinel findings include same-family masked-signal bugs in other projects, but no exact AReaL `areal/utils/functional/functional.py` duplicate with this 2D padded advantage trigger. ## AI Assistance This draft and patch were prepared with AI assistance in an RL-Sentinel testing loop. | 3 个月前 | |
feat(trainer): add flexible GAE lambda strategies (#1600) * feat: add turn-level GAE support Treat each generated turn as a GAE timestep while preserving the token-level default and token-local KL regularization. Key changes: - propagate and validate token-aligned turn IDs - compute turn-level advantages without full CPU sequence scans - filter structural metadata at FSDP and Archon model boundaries - document the new selector and add focused regression tests * feat(trainer): add dynamic per-sample GAE lambda Allow GAE lambda to vary by trajectory using effective token or turn lengths while preserving static float behavior. Key changes: - Resolve custom lambda functions and keyword arguments from config - Add VAPO length-adaptive GAE with empty-trajectory handling - Validate per-sample lambda tensors and cover token and turn modes * perf(trainer): reduce GAE preprocessing overhead Hoist loop-invariant tensor work and bypass dynamic trajectory length construction when GAE lambda is static. * feat(trainer): add relative-position GAE lambda * fix(trainer): allow token lambda without turn metadata Keep custom token-level GAE lambda functions compatible with rollout workflows that do not emit turn IDs, while preserving the metadata requirement for turn-level GAE. * docs: regenerate CLI reference for GAE options Keep the generated configuration reference aligned with the current main branch after porting the AntCode GAE changes. * docs: document flexible GAE configuration Explain token- and turn-level recurrences, KL and critic semantics, dynamic lambda strategies, and custom workflow turn IDs in English and Chinese.\n\nFix CLI default rendering for dataclass factories and cover it with unit tests. * refactor(trainer): extract GAE helpers Keep PPOActor focused on training orchestration by moving GAE kernels, turn metadata validation, and lambda context construction into a dedicated module. --------- Co-authored-by: Wenhao Zhou <miumiu.zwh@antgroup.com> | 13 天前 | |
fix(inference): reject incomplete sampling evidence (#1554) Validate normalized token/logprob evidence before rollout accumulation, reject vLLM ambiguous sentinel values, and preserve the explicit SGLang abort-before-prefill result. Signed-off-by: EazyReal <8047065+EazyReal@users.noreply.github.com> Co-authored-by: EazyReal <8047065+EazyReal@users.noreply.github.com> | 1 个月前 | |
feat(colocate): support AWEX colocated actor-rollout training (#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. | 26 天前 | |
fix: align CP metrics and TP grad norm metadata (#1497) * fix: align CP metrics and TP grad norm metadata Two independent correctness fixes for the Megatron engine. 1. CP metrics alignment (stats_tracker): Per-key reduce_group overrides may be wider than the default DP export group (e.g. DP+CP for token-level metrics). Add a `key_sync_group` argument so override keys sync their key set / metadata and reduce over DP+CP, while default keys keep the DP alignment. Aggregation helpers now emit identity placeholder tensors (0 / +inf / -inf) for keys absent on a rank, so every rank still participates in the collective and avoids hangs / mismatched reductions when CP > 1. 2. TP grad norm metadata (megatron_engine): `_mark_duplicated_params` now also clears `param.tensor_model_parallel` on replicated (tp_size == 1, non-expert) params. Megatron's optimizer uses that attribute to decide which TP ranks contribute to grad norm/clipping, so leaving it True double-counts duplicated params when TP > 1. Clearing it is also consistent with `all_gather_param` and `hf_save`, which already treat non-TP params as replicated. Adds unit tests for both paths (tests/test_stats_tracker.py and a new case in tests/test_megatron_engine.py). * test: cover all_gather_param routing and duplicated-param grad norm Add two tests backing the TP grad-norm metadata fix: - tests/test_all_gather_param.py: assert `all_gather_param` returns the param as-is (no TP all-gather) when `tensor_model_parallel` is False or the name is in `duplicated_param_names`, and only all-gathers genuine TP-sharded params. Uses `pytest.importorskip` so it runs in the megatron CI env and skips gracefully elsewhere. - tests/test_grad_norm_duplicated.py: a CPU/gloo 2- and 4-rank test showing a replicated param is counted once (correct) when tensor_model_parallel is False, and inflated by sqrt(tp) when it is (incorrectly) left True. Mirrors Megatron's param_is_not_tensor_parallel_duplicate selection. * test: add 2-GPU grad-norm TP-invariance integration test End-to-end guard for the TP grad-norm metadata fix: run the real MegatronEngine at TP=1 and TP=2 on identical deterministic input via torchrun and assert the reported grad norm is TP-invariant. Double-counting replicated params (tensor_model_parallel left True) would inflate it as TP grows. Also asserts the fix demoted at least one real duplicated param at TP=2. Marked multi_gpu/slow and skipped without >= 2 GPUs. * refactor(stats): dedup identity tensors, avoid min/max GPU sync Address review feedback on the CP metrics fix: - Generalize the placeholder-tensor helper to `_placeholder_scalar(fill=...)` so the SUM/AVG (0.0) and MIN/MAX (+/-inf) empty-value branches share one device-aware constructor instead of four inline copies. - In `_min_of`/`_max_of`, reduce the per-shard list with `torch.stack(xs).min()` / `.max()` instead of Python `min()`/`max()`, which forced a CPU-GPU sync. - In the SCALAR branch, read values via `self.stats.get(key, [])` (a rank may learn a key only through metadata sync) and guard the `value / cnt` division against `cnt == 0` so a key absent everywhere yields 0.0 instead of NaN. - Add a regression test for the missing-on-this-rank SCALAR path. | 1 个月前 | |
feat(rollout): add grouped reward normalization controls (#1516) Add rollout-time reward normalization and incomplete group dropping for grouped InteractionWithTokenLogpReward workflows. Preserve original rewards for logging and dumps, and thread the controls through inference and training controllers. Add tests covering grouped reward normalization and parameter forwarding. Co-authored-by: chucai.dzq <chucai.dzq@alibaba-inc.com> | 1 个月前 | |
fix(hermes): preserve singleton online reward signal (#1474) * fix(ppo): warn when singleton groups erase rewards * fix(hermes): preserve singleton online reward signal * test(hermes): harden reward signal regression * style: format singleton reward warning test * fix(config): support mapping-shaped reward normalization | 1 个月前 | |
fix(inference): reject incomplete sampling evidence (#1554) Validate normalized token/logprob evidence before rollout accumulation, reject vLLM ambiguous sentinel values, and preserve the explicit SGLang abort-before-prefill result. Signed-off-by: EazyReal <8047065+EazyReal@users.noreply.github.com> Co-authored-by: EazyReal <8047065+EazyReal@users.noreply.github.com> | 1 个月前 | |
feat(engine): support direct engine construction via from_pretrained without config dataclass (#1140) * feat(engine): add from_pretrained factory method for FSDPEngine - Add from_pretrained method in class FSDPEngine - Integrate the engine created by from_pretrained method into test_train_engine.py - test_train_engine.py pass * style: auto-fix trailing whitespace and ruff issues * feat(engine): add from_pretrained factory method and test for RemoteSGLangEngine * fix: make create_process_group method parse backend config and add regression test. * fix: fix bug of backend parse and test_fsdp_engine_alloc_mode_construction * fix: add experiment_name and trial_name parameters for the factory classmethod of FSDPEngine. * fix: change the model parameter to tokenizer_path in RemoteSGLangEngine and fix data_parallel_size parse. * fix: fix test_fsdp_engine_alloc_mode_construction() * fix: delete test_fsdp_engine_alloc_mode_construction() because the lack of assertion. * style: Apply pre-commit formatting * fix: add experiment_name and trial_name in tests of FSDPEngine from_pretrained method * fix: add GPU check decorator for test_fsdp_engine_alloc_mode_construction() * feat: add from_pretrained method for RemotevllmEngine and test | 4 个月前 | |
feat(utils): add Karmarkar-Karp partitioning algorithm for sequence packing (#1151) Add KK (Largest Differencing Method) as an alternative to FFD for micro-batch allocation. KK produces more balanced partitions with lower max-min spread, beneficial for RL workloads with variable sequence lengths. Key changes: - Add _KKSet, _KKState, _kk_partition, kk_allocate in seqpack.py - Add packing_algorithm field to MicroBatchSpec (ffd/kk) - Wire KK allocation through dist_rollout and data utils - Add sequence_packing docs (en/zh) and CLI reference updates - Add comprehensive unit tests and torchrun benchmark Refs: #1151 | 4 个月前 | |
fix(docker): move venv out of /AReaL to avoid mount override (#1251) * fix(docker): move venv out of /AReaL to avoid mount override * fix: fix CI test bug * chore: pin ironclaw version * fix: fix ironcraw install bug * fix: fix zerocraw install bug * fix: fix multi gpu ci test bug * fix: fix uv pip install in dockerfile * chore: regenerate uv locks --------- Co-authored-by: 博惟 <bowei.fw@antgroup.com> | 4 个月前 | |
fix(infra): add two-phase teardown to prevent TCPStore race at shutdown (#1244) Problem: During teardown, rank-0 (TCPStore server owner) could exit before peer ranks finished their final NCCL abort, causing noisy "TCPStore.recvValue failed" / "Broken pipe" warnings on stderr. Root cause: All ranks were killed simultaneously without first coordinating a distributed barrier on the CPU (gloo) group to safely tear down NCCL communicators and the TCPStore. Solution: Implement a two-phase teardown protocol: Phase 1 - Engine destroy: call engine.destroy() on every worker concurrently. The engine-side destroy() now executes a CPU barrier (dist.barrier on a gloo process group) followed by dist.destroy_process_group(), ensuring all ranks leave the NCCL collective together. Phase 2 - Process kill: only after the barrier completes, kill the actual processes (Ray: remove placement groups; Slurm: scancel; Local: process tree cleanup). Changes: - engine (fsdp/megatron/archon): add _cpu_group + pre-destroy barrier - train_controller: two-phase destroy (engines first, then workers) - scheduler/ray: _cleanup_workers with ray.wait timeout + PG removal - scheduler/slurm: _destroy_engines_on_workers via HTTP before scancel - scheduler/local: graceful engine teardown before SIGKILL - scheduler_api: add reverse_order param to delete_workers interface - tests: updated test_train_controller, added test_local_scheduler Tested: DPO 4xH20 Ray scheduler - clean teardown, no TCPStore warnings. | 4 个月前 | |
refactor(tests): move tests from areal/tests to top-level tests directory (#944) * refactor(tests): move tests from areal/tests to top-level tests directory Move the test directory from areal/tests/ to a top-level tests/ directory to follow Python best practices and improve project structure. Key changes: - Move all test files from areal/tests/ to tests/ - Create __init__.py files for tests package and subdirectories - Create areal/utils/testing_utils.py with shared utilities: - get_model_path, get_dataset_path (model/dataset path resolution) - MODEL_PATHS, DENSE_MODEL_PATHS, MOE_MODEL_PATHS (test model configs) - load_archon_model (Archon model loading for tests) - Update areal/tools/profiling_utils to re-export from testing_utils - Update test utilities to import from areal/utils/testing_utils, keeping test-specific behavior (pytest.skip) as wrappers - Update all imports from 'areal.tests.' to 'tests.' - Update hardcoded torchrun script paths - Update workflow string references for dynamic class loading - Update CI workflow paths in .github/workflows/test-areal.yml - Update documentation references in CLAUDE.md, AGENTS.md, CONTRIBUTING.md, and all .claude/ and .opencode/ skill/agent files * minor fix * remove import side effect * fix(tests): add tests/utils.py module for test imports Re-export get_model_path and get_dataset_path from areal.utils.testing_utils to fix ModuleNotFoundError in tests that import from tests.utils. | 6 个月前 | |
feat(engine): support fixed warmup steps (#1597) Allow optimizer configs to select an absolute scheduler warmup while preserving the proportional fallback across FSDP, Megatron, and Archon. Validate shared scheduler boundaries and keep Megatron resume initialization independent of the active warmup config. Key changes: - Add fixed warmup resolution with cross-engine validation - Align Megatron decay and resume scheduler parameters - Add boundary and Megatron integration regression tests Co-authored-by: 峯回 <dh183333@antgroup.com> | 21 天前 | |
refactor: flatten sub-module imports to use parent package re-exports (#996) Add __init__.py with lazy re-exports (__getattr__ + __all__) for areal/api, areal/engine, areal/reward, areal/workflow, and __all__ for areal/dataset, then rewrite all external imports across the codebase to use the shorter parent-package form (e.g. `from areal.api import TrainEngine`). Key changes: - All re-exports are fully lazy via __getattr__ (no eager imports) - Flatten ~100 files across areal/, tests/, examples/ - Preserve cli_args deep imports (dozens of config classes, including SchedulingSpec which stays in cli_args) - Preserve intra-package relative imports to avoid circular deps - Reward submodules use sibling-relative imports (from . import ...) - Preserve non-exported symbols (DeviceRuntimeInfo, HttpRequest, etc.) Co-authored-by: Wentai Zhang <zhangwentai.zwt@antgroup.com> | 5 个月前 | |
feat(mcore): apply fp32 lm head forward when enabled (#1574) * 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. | 27 天前 | |
fix(megatron): release actor weights after async save (#1615) * fix(megatron): release actor weights after async save After MCore completes D2H staging and forks the writer, clear the parent-side tensor lists retained by the queued request. Preserve writer structure, byte payloads, results, and finalize callbacks so actor weights can be offloaded while checkpoint I/O continues. Signed-off-by: jiawei <jiaweibit@gmail.com> * docs(megatron): describe unsupported async layouts --------- Signed-off-by: jiawei <jiaweibit@gmail.com> | 12 天前 | |
fix(models): configure bridge provider for deterministic builds (#1603) Megatron-Bridge materializes a provider separate from the transformer config. Apply prebuild determinism to that provider before finalization so construction-time consumers inherit the requested settings. | 19 天前 | |
fix: align CP metrics and TP grad norm metadata (#1497) * fix: align CP metrics and TP grad norm metadata Two independent correctness fixes for the Megatron engine. 1. CP metrics alignment (stats_tracker): Per-key reduce_group overrides may be wider than the default DP export group (e.g. DP+CP for token-level metrics). Add a `key_sync_group` argument so override keys sync their key set / metadata and reduce over DP+CP, while default keys keep the DP alignment. Aggregation helpers now emit identity placeholder tensors (0 / +inf / -inf) for keys absent on a rank, so every rank still participates in the collective and avoids hangs / mismatched reductions when CP > 1. 2. TP grad norm metadata (megatron_engine): `_mark_duplicated_params` now also clears `param.tensor_model_parallel` on replicated (tp_size == 1, non-expert) params. Megatron's optimizer uses that attribute to decide which TP ranks contribute to grad norm/clipping, so leaving it True double-counts duplicated params when TP > 1. Clearing it is also consistent with `all_gather_param` and `hf_save`, which already treat non-TP params as replicated. Adds unit tests for both paths (tests/test_stats_tracker.py and a new case in tests/test_megatron_engine.py). * test: cover all_gather_param routing and duplicated-param grad norm Add two tests backing the TP grad-norm metadata fix: - tests/test_all_gather_param.py: assert `all_gather_param` returns the param as-is (no TP all-gather) when `tensor_model_parallel` is False or the name is in `duplicated_param_names`, and only all-gathers genuine TP-sharded params. Uses `pytest.importorskip` so it runs in the megatron CI env and skips gracefully elsewhere. - tests/test_grad_norm_duplicated.py: a CPU/gloo 2- and 4-rank test showing a replicated param is counted once (correct) when tensor_model_parallel is False, and inflated by sqrt(tp) when it is (incorrectly) left True. Mirrors Megatron's param_is_not_tensor_parallel_duplicate selection. * test: add 2-GPU grad-norm TP-invariance integration test End-to-end guard for the TP grad-norm metadata fix: run the real MegatronEngine at TP=1 and TP=2 on identical deterministic input via torchrun and assert the reported grad norm is TP-invariant. Double-counting replicated params (tensor_model_parallel left True) would inflate it as TP grows. Also asserts the fix demoted at least one real duplicated param at TP=2. Marked multi_gpu/slow and skipped without >= 2 GPUs. * refactor(stats): dedup identity tensors, avoid min/max GPU sync Address review feedback on the CP metrics fix: - Generalize the placeholder-tensor helper to `_placeholder_scalar(fill=...)` so the SUM/AVG (0.0) and MIN/MAX (+/-inf) empty-value branches share one device-aware constructor instead of four inline copies. - In `_min_of`/`_max_of`, reduce the per-shard list with `torch.stack(xs).min()` / `.max()` instead of Python `min()`/`max()`, which forced a CPU-GPU sync. - In the SCALAR branch, read values via `self.stats.get(key, [])` (a rank may learn a key only through metadata sync) and guard the `value / cnt` division against `cnt == 0` so a key absent everywhere yields 0.0 instead of NaN. - Add a regression test for the missing-on-this-rank SCALAR path. | 1 个月前 | |
feat(megatron): Qwen3.5 dense + MoE training/inference support via megatron-bridge (#1384) * feat(engine): delegate megatron live weight sync to bridge.export_hf_weights Add use_bridge_for_update_weights flag that routes the live weight update path through megatron-bridge.export_hf_weights instead of the hand-rolled convert_to_hf registry. Required for new model families (e.g. Qwen3.5) that don't have a registry entry. The bridge handles TP/EP/PP gather and HF layout transformation internally; AReaL keeps the bucketed broadcast loop unchanged. FP8 and LoRA paths fall back to the registry automatically. Also fix a latent device-context bug in _load_model_from_hf: megatron-bridge builds shard-index tensors via torch.arange() under the caller's `with self.device:` context, putting them on CUDA while HF weights are loaded on CPU. The resulting indexing error trips ChunkedMapping for any model with GDN/Mamba conv1d weights (e.g. Qwen3.5). Force CPU as the factory-op default just around the bridge.load_hf_weights call. Key changes: - New MegatronEngineConfig.use_bridge_for_update_weights flag - Refactor _update_weights_from_distributed into dispatch + _update_weights_via_registry helper - New _update_weights_via_bridge streams from bridge.export_hf_weights and reuses the bucket broadcast loop - Wrap bridge.load_hf_weights in `with torch.device("cpu"):` to prevent CUDA index / CPU tensor mismatch in ChunkedMapping * test(megatron): add Qwen3.5 distributed test scaffolding Add 1-GPU smoke + 5 multi-GPU tests (TP=2, PP=2, PP+VPP=2, DP=2 grad_norm invariance, DCP save/load) mirroring the Qwen3 dense set. All Qwen3.5 tests route through bridge_type=megatron-bridge because its GDN hybrid attention is only handled by megatron-bridge's model definitions (mbridge would substring-match qwen3 and emit wrong shapes). NOTE: these tests currently fail at engine.forward because megatron-core's GDN layer raises NotImplementedError on packed (THD) sequences. A follow-up will add a BSHD path mirroring verl's data_format switch; until then these tests document the expected coverage and act as a regression target. Key changes: - Add qwen3_5 to MODEL_PATHS in run_megatron_engine_distributed.py - Re-key MODEL_PATHS from areal.utils.testing_utils canonical paths so local-path overrides propagate from a single source - Wire bridge_type via _MODEL_BRIDGE_OVERRIDES (qwen3_5 → megatron-bridge) - Six test_qwen3_5_* tests in test_megatron_engine_distributed.py * feat(engine): add BSHD padded forward path and megatron-bridge patches for Qwen3.5 Qwen3.5's GDN (Gated Delta Net) layers reject packed (THD) sequences in megatron-core. Add a BSHD path that reconstructs [B, S] padded input from cu_seqlens inside packed_context_parallel_forward, mirroring the existing VLM 2D-reconstruction logic but for text-only models. Also add runtime monkey-patch for megatron-bridge PR #3143 (MTP shadow embedding missing word_embeddings attribute under sequence_parallel + tied embeddings). The patch lazily restores the attribute from the closure before _postprocess runs, avoiding the need to replace the full forward method. Key changes: - New MegatronEngineConfig.use_padded_seq flag (BSHD mode) - Generalize VLM 2D-reconstruction path in packed_context_parallel_forward to also fire on use_padded_seq=True for non-VLM models - CP>1 guard for use_padded_seq (same constraint VLM has) - New megatron_bridge_patches.py with PR #3143 workaround - New train_hf_save_load test_type in runner (replaces DCP for SSM models whose flattened_range tensors are unsupported by mcore DCP) - Qwen3.5 tests: 1-GPU, TP=2, PP=2, HF save/load all pass; VPP and grad_norm_mb_invariance skipped with documented reasons * fix(engine): register Qwen3.5 as vision model and fix non-contiguous broadcast Add qwen3_5 and qwen3_5_moe to VALID_VISION_MODELS so the engine loads the HF processor and passes pixel_values / image_grid_thw through the VLM forward path. Qwen3.5's base architecture (Qwen3_5ForConditionalGeneration) is inherently multimodal — there is no separate qwen3_5_vl model_type. Also fix a ValueError in _update_weights_via_bridge where bridge.export_hf_weights yields non-contiguous tensor views (from QKV split / gate-up chunk) that NCCL broadcast rejects. Call .contiguous() before bucketing. * feat(vllm): add gdn_prefill_backend to avoid FlashInfer GDN hang Qwen3.5 and other GDN hybrids default to vLLM's FlashInfer GDN prefill kernel, which hangs on SM90 — a runtime mbarrier deadlock (flashinfer #2623/#3329) and a JIT-compile deadlock (vLLM #41865/#39287), surfacing as shm_broadcast stall -> sample_tokens timeout -> EngineDeadError. Expose gdn_prefill_backend so configs can set "triton" (stable Triton/FLA kernel). None default emits no flag, so non-GDN models are unaffected. * test(megatron): add Qwen3.5-MoE expert-parallel + HF save/load tests Qwen3.5-35B-A3B megatron coverage via megatron-bridge, both running on 4 GPUs: - test_qwen3_5_moe_expert_parallel: PP2/TP2/EP2 forward + cross-rank logprob consistency. CP is unavailable for the GDN series (Megatron-LM #4043) and the full-attention layers cap TP<=2, so ranks are filled with PP at EP=2. - test_qwen3_5_moe_hf_save_load: save -> zero -> load -> compare round-trip validating MoE expert-weight conversion (TEGroupedLinear weight0..N + GLU linear_fc1 stride-2). The train step is skipped (_MODEL_SAVELOAD_SKIP_TRAIN) since a 35B-A3B optimizer state does not fit; the loaded HF weights are already non-trivial. The megatron-vs-FSDP logit comparison is skipped for this model (_MODEL_SKIP_FSDP_COMPARE): AReaL's FSDP engine materializes the full fp32 35B per rank on load, which cannot fit. * docs(example): add Qwen3.5-2B megatron geometry3k GRPO config * docs(engine): clarify megatron-bridge patch docstring + document gdn_prefill_backend choices * refactor(engine): auto-derive padded-seq layout from model type The padded (BSHD) vs packed (THD) forward layout is a hard architectural property of the model -- GDN/SSM kernels (the Qwen3.5 family) reject packed sequences -- not a user tunable. Exposing it as the `use_padded_seq` config field let it be mis-set and risked silent correctness or crash issues. Derive it from `model_type` instead so the layout can never disagree with the architecture. Also surface a startup warning when `use_bridge_for_update_weights=True` but a fallback condition (non-megatron-bridge, FP8/quantized, or LoRA) silently routes live weight sync through the registry path, so the effective behavior is visible in logs. Key changes: - Add requires_padded_seq(model_type) helper in engine/core/model.py - Derive self.use_padded_seq from model_type in MegatronEngine.initialize - Remove use_padded_seq from MegatronEngineConfig and regenerate CLI docs - Warn once when bridge weight-sync falls back to the registry path - Drop the test-runner override map and example yaml flag Refs: #1384 | 2 个月前 | |
perf(megatron): avoid synthetic BSHD padding rows (#1632) Assisted-by: OpenAI Codex <codex@openai.com> Signed-off-by: EazyReal <8047065+EazyReal@users.noreply.github.com> Co-authored-by: EazyReal <8047065+EazyReal@users.noreply.github.com> | 8 天前 | |
feat(engine): enable model-owned THD for Qwen3-VL (#1608) Allow Qwen3-VL models to merge multimodal embeddings before packing while preserving existing padded-only and wrapper-owned paths. Key changes: - Route dense and MoE Qwen3-VL through model-owned THD - Reconstruct padded inputs and restore packed model outputs - Add routing, alignment, parity, and distributed forward coverage | 16 天前 | |
fix(models): support FP32 operands with chunked LM head (#1594) * fix(models): support FP32 operands with chunked LM head Keep enable_fp32_lm_head orthogonal to chunked logits while avoiding repeated local vocab-weight casts across sequence chunks. Key changes: - Reuse one FP32 weight conversion per LM-head forward and backward - Preserve TP/SP gradient communication and FP32 main_grad accumulation - Cover flag combinations, gradients, and distributed execution * test(models): account for chunked bias reduction order Low-precision chunked backward sums bias gradients per chunk, while the full reference reduces all tokens at once. Compare the old BF16/FP16 path with an explicit relative tolerance and keep strict parity for FP32 operands. * refactor(models): scope FP32 projection helper Keep the full-sequence FP32 projection implementation private to the native linear class so unrelated call sites cannot invoke it accidentally. | 20 天前 | |
feat(awex): add separation AdamW delta weight transfer (#1623) * feat(config): add separation DTE topology gates * feat(awex): add separation AdamW delta weight transfer * docs(examples): add DTE separation GSM8K example * docs(examples): expand DTE GSM8K configuration * fix(awex): address separation DTE review feedback * fix(awex): harden separation delta transfer protocol * refactor(config): integrate delta transfer into train engine * refactor(config): simplify delta weight update switch * test(awex): isolate optional DTE dependency | 11 天前 | |
fix: only subtract in-range excluded ports in find_free_ports (#1436) available_range subtracted the full exclude_ports count regardless of whether those ports lay within [min_port, max_port]. Excluded ports outside the range deflated the availability count and could raise a spurious ValueError even when every in-range port was free. Intersect exclude_ports with the range before subtracting. Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> | 2 个月前 | |
refactor(api): migrate allocation_mode to per-engine backend fields (#1044) * refactor(api): migrate allocation_mode to per-engine backend fields Replace the centralized `allocation_mode` string with explicit `backend` fields on `TrainEngineConfig` and `InferenceEngineConfig`. Each engine now owns its own backend+parallelism spec (e.g. `fsdp:d4`, `sglang:d4t2`), eliminating implicit auto-backend selection and the shared `AllocationMode` object. Key changes: - Add `backend` field to TrainEngineConfig and InferenceEngineConfig - Add `ModelAllocation.from_str()` for single-component parsing - Remove `AllocationMode` public export (replaced by `ModelAllocation`) - Rename internal `AllocationMode` to `_AllocationMode` for SPMD launcher backward compatibility with FutureWarning - Remove auto-backend selection — explicit backend prefix is now required - Controllers (`TrainController`, `RolloutController`) parse `backend` directly instead of receiving `alloc_mode` from trainers - `WeightUpdateMeta.alloc_mode` replaced by `gen_allocation` (single `ModelAllocation`) - Add `RWTrainer` and `ArchonRWEngine` for reward model training - Remove `get_model_update_meta()` helper (logic moved to trainers) - Update all YAML configs, examples, docs (EN+ZH), and tests BREAKING CHANGE: `AllocationMode` is removed from public API. Users must migrate to per-engine `backend` fields. SPMD launchers emit deprecation warnings. * chore(ci): fix backend specifier for vlm sft test * fix: fix bare dims for actor backends * chore(docs): fix reminder for bare allocation dims | 5 个月前 | |
fix(CI): fix vlm_grpo CI OOM bug (#1438) Default fp32 master weights (PR #1369) push Qwen2.5-VL-3B beyond A100-40G capacity at the first AdamW step. Pin the test to bf16 storage with adam_bf16 (Kahan summation) so it runs on 40G runners. ## Description <!-- Provide a clear and concise description of what this PR does --> ## Related Issue <!-- Link to the issue this PR addresses. PRs should be related to a well-templated issue. --> Fixes #(issue) ## Type of Change <!-- Select ONE that best describes this PR --> - [x] 🐛 Bug fix - [ ] ✨ New feature - [ ] 💥 Breaking change - [ ] 📝 Documentation update - [ ] ♻️ Refactoring - [ ] ⚡ Performance improvement - [ ] ✅ Test coverage improvement ## Checklist <!-- Mark with 'x' what you've done --> - [ ] I have read the [Contributing Guide](https://github.com/areal-project/community/blob/main/CONTRIBUTING.md) - [ ] Pre-commit hooks pass (`pre-commit run --all-files`) - [ ] Relevant tests pass; new tests added for new functionality - [ ] Documentation updated (if applicable; built with `./docs/build_all.sh`) - [ ] Branch is up to date with `main` - [ ] Self-reviewed via `/review-pr` command - [ ] This PR was created by a coding agent via `/create-pr` - [ ] This PR is a breaking change **Breaking Change Details (if applicable):** <!-- Describe what breaks and how users should migrate --> ## Additional Context <!-- Add any other context, screenshots, logs, or explanations here --> ______________________________________________________________________ **Need help?** Check the [Contributing Guide](https://github.com/areal-project/community/blob/main/CONTRIBUTING.md) or ask in [GitHub Discussions](https://github.com/areal-project/AReaL/discussions)! | 2 个月前 | |
feat(fsdp): add per-layer optimizer step with streaming H2D/D2H pipeline (#983) Add PerLayerOptimWrapper that accelerates CPU-offloaded optimizer step by streaming optimizer states per-layer to device with async prefetch, instead of running Adam on CPU. Key design: - OptimKernel base class + AdamKernel for extensible optimizer support - ParamTransferState dataclass for typed per-layer transfer state - Pipelined H2D prefetch / compute / D2H offload via dedicated streams - record_stream() for CUDA allocator safety on cross-stream tensors - refresh_states() to re-pin/normalize after checkpoint load - Device-neutral naming (no hardcoded GPU references) Config: per_layer_optim_step, optim_step_prefetch_layers Requires: optimizer type 'adam' (AdamW); offload managed automatically Also: - Remove dead OptimizerConfig.offload field (was never consumed) - Replace DTensor._local_tensor with .to_local() (public API) - Fix docs to reference offload_params instead of optimizer.offload - Mark GPU-requiring tests with @pytest.mark.slow - Update CLI reference docs (en/zh) and handling_oom docs (en/zh) Co-authored-by: Wentai Zhang <zhangwentai.zwt@antgroup.com> | 5 个月前 | |
perf: reduce Megatron training memory peaks (#1555) * 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: #1555 | 26 天前 | |
perf: reduce Megatron training memory peaks (#1555) * 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: #1555 | 26 天前 | |
feat(trainer): add flexible GAE lambda strategies (#1600) * feat: add turn-level GAE support Treat each generated turn as a GAE timestep while preserving the token-level default and token-local KL regularization. Key changes: - propagate and validate token-aligned turn IDs - compute turn-level advantages without full CPU sequence scans - filter structural metadata at FSDP and Archon model boundaries - document the new selector and add focused regression tests * feat(trainer): add dynamic per-sample GAE lambda Allow GAE lambda to vary by trajectory using effective token or turn lengths while preserving static float behavior. Key changes: - Resolve custom lambda functions and keyword arguments from config - Add VAPO length-adaptive GAE with empty-trajectory handling - Validate per-sample lambda tensors and cover token and turn modes * perf(trainer): reduce GAE preprocessing overhead Hoist loop-invariant tensor work and bypass dynamic trajectory length construction when GAE lambda is static. * feat(trainer): add relative-position GAE lambda * fix(trainer): allow token lambda without turn metadata Keep custom token-level GAE lambda functions compatible with rollout workflows that do not emit turn IDs, while preserving the metadata requirement for turn-level GAE. * docs: regenerate CLI reference for GAE options Keep the generated configuration reference aligned with the current main branch after porting the AntCode GAE changes. * docs: document flexible GAE configuration Explain token- and turn-level recurrences, KL and critic semantics, dynamic lambda strategies, and custom workflow turn IDs in English and Chinese.\n\nFix CLI default rendering for dataclass factories and cover it with unit tests. * refactor(trainer): extract GAE helpers Keep PPOActor focused on training orchestration by moving GAE kernels, turn metadata validation, and lambda context construction into a dedicated module. --------- Co-authored-by: Wenhao Zhou <miumiu.zwh@antgroup.com> | 13 天前 | |
feat(ppo): report rejection-aware token and log-prob statistics (#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. | 27 天前 | |
perf: reduce Megatron training memory peaks (#1555) * 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: #1555 | 26 天前 | |
feat(api): add unified RejectionSamplingConfig for async training (#1088) Replace behave_imp_weight_cap/behave_imp_weight_mode with unified RejectionSamplingConfig supporting multiple metrics (ratio, kl_k1, kl_k2, kl_k3), levels (token/sequence), and actions (mask/clamp). Key changes: - Add RejectionSamplingConfig dataclass with comprehensive validation - Implement apply_rejection_sampling for 1D packed and 2D padded formats - Fix loss denominator scaling bug in mask mode (save count before filtering) - Use geometric mean for sequence-level ratio aggregation (matching GSPO) - Broadcast sequence-level geometric mean as uniform behave_imp_weight - Warn when use_decoupled_loss=True but rejection_sampling is None - Update ppo_actor_loss_fn and grpo_loss_fn to use new config - Migrate 40 example configs to new rejection_sampling field - Add 43 unit tests covering all modes, metrics, and edge cases Refs: #1052 | 4 个月前 | |
feat(scheduler): support grouped colocation in the Ray scheduler (#1575) * feat(scheduler): support grouped colocation in the Ray scheduler Colocation previously required the colocated role to match the target role's replica count, which rules out AWEX-style colocation where a few multi-GPU inference workers share the GPUs of many single-GPU trainer workers (e.g. 16 x 4-GPU SGLang servers over 64 x 1-GPU actors). When replica counts differ but the colocated role's total GPU demand exactly reuses the target role's GPUs, route worker creation to a grouped path: chunk each target node's physical GPUs into contiguous per-worker groups (never crossing nodes), pin one zero-GPU process launcher to each target node via node affinity, and start worker processes with explicit physical gpu_devices so CUDA_VISIBLE_DEVICES carries physical indices exactly like the Slurm launcher. Grouped roles own their workers and launchers, are not registered as colocated aliases, and create no placement groups, so readiness discovery and teardown follow the standard paths. * fix(scheduler): accept a device-free Ray driver node RayScheduler probed accelerators on the driver process only, so a CPU-only head node aborted with "does not support CPU-only clusters" even when every Ray worker exposed GPUs. Fall back to the advertised cluster resources when the driver itself owns no device. | 20 天前 | |
feat(scheduler): support grouped colocation in the Ray scheduler (#1575) * feat(scheduler): support grouped colocation in the Ray scheduler Colocation previously required the colocated role to match the target role's replica count, which rules out AWEX-style colocation where a few multi-GPU inference workers share the GPUs of many single-GPU trainer workers (e.g. 16 x 4-GPU SGLang servers over 64 x 1-GPU actors). When replica counts differ but the colocated role's total GPU demand exactly reuses the target role's GPUs, route worker creation to a grouped path: chunk each target node's physical GPUs into contiguous per-worker groups (never crossing nodes), pin one zero-GPU process launcher to each target node via node affinity, and start worker processes with explicit physical gpu_devices so CUDA_VISIBLE_DEVICES carries physical indices exactly like the Slurm launcher. Grouped roles own their workers and launchers, are not registered as colocated aliases, and create no placement groups, so readiness discovery and teardown follow the standard paths. * fix(scheduler): accept a device-free Ray driver node RayScheduler probed accelerators on the driver process only, so a CPU-only head node aborted with "does not support CPU-only clusters" even when every Ray worker exposed GPUs. Fall back to the advertised cluster resources when the driver itself owns no device. | 20 天前 | |
feat(infra): add HTTP-based Ray Scheduler (#1441) * fix(infra): preserve dataclass state over RPC * feat(engine): support headless vLLM server mode * fix(utils): use fixed Ray name resolve namespace * refactor(engine): expose backend server env builders * refactor(infra): remove Ray-native scheduler Drop the single-controller Ray scheduler path. Key changes: - Remove RayScheduler, Ray RPC actors, and Ray vLLM remote launcher - Make single-controller trainers dispatch only local or slurm schedulers - Simplify RTensor and RPC serialization to the HTTP backend * feat(infra): add HTTP-based Ray scheduler Add a Ray-backed scheduler that allocates accelerator placement groups while keeping worker and inference engine traffic on the existing HTTP RPC path, with batched launcher operations and multi-node rollout support. Key changes: - Add RayScheduler and RayWorkerProcessLauncher for Ray-managed HTTP workers - Wire scheduler.type=ray into infra exports, trainer initialization, logging, docs, and examples - Batch worker startup and status checks by Ray launcher to reduce per-worker actor calls - Split multi-node rollout backend launch and cleanup into a dedicated coordinator - Tighten Ray launcher lifecycle handling for worker shutdown, placement groups, and backend process cleanup * test(infra): add Ray scheduler tests * fix(trainer): allow proxy workers with Ray scheduler * fix(infra): reuse existing Ray cluster on init * fix(api): restore deprecated Ray placement config --------- Co-authored-by: Ge Shi <utashih@gmail.com> | 1 个月前 | |
fix(sft): report CP-invariant token-count stats (#1242) (#1249) * fix(sft): report CP-invariant token-count stats (#1242) After PR #1223 introduced CP-local loss, compute_packed_sft_loss started recording `n_tokens`, `n_valid_tokens` and `prompt_tokens` using the CP-split `loss_mask` / `logprobs`. These denominators are summed only across the DP group at export time, so the reported values under-count by the CP factor (e.g. ~4x smaller with CP=4). The ratios reported via `stat(..., denominator=...)` (loss, ppl, vocab_*) remain correct because numerator and denominator scale together, so the issue is easy to miss. Preserve the pre-CP-split loss_mask as `_global_loss_mask` when the CP-local path constructs its inputs, and use it as the denominator for the token-count metrics so they are invariant to the CP topology. Keep separate `n_tokens_local` / `n_valid_tokens_local` denominators with CP-local shapes for the CP-local tensors (`logprobs`, `vocab_*`), since `stats_tracker.stat` requires matching shapes. Verified on 64-GPU (CP=2) and 128-GPU (CP=4) SPMD runs with the same Qwen3-30B + swe_distilled_1000 setup: step-1 `n_tokens` matches exactly across both topologies (6,320,900), and equals `CP * n_tokens_local` on each, confirming the fix is topology-invariant. Fixes #1242 Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * style: fix 'Loggin stats' typo Noticed by Copilot reviewer on PR #1249. * fix(sft): make CP-local SFT loss/ppl/vocab_* stats CP-invariant After PR #1223 introduced CP-local loss, ratio metrics in compute_packed_sft_loss (loss/entropy/ppl/vocab_*) started reporting the average over rank 0's CP slice rather than the global average, because both numerator and denominator were CP-local tensors and export only reduced across DP. Counts were already fixed in PR #1242, but the ratios were not -- empirically loss/ppl/max can drift up to 25% (e.g. ppl/max 2.57 vs 3.20) on long-context SFT where prompt / completion tokens land in different CP slices. Fix the reporting without reintroducing the expensive logits all-gather that #1223 removed: at export time, the per-key reduce sees only scalar numerator/denominator (already .sum()-ed), so all-reducing across DP + CP costs a few bytes -- not 37GB of logits. Key changes: - stats_tracker: add per-key reduce_group override (kw-only) on denominator/scalar/stat; _avg/_min/_max/_sum/SCALAR honor it via _effective_reduce_group; reset clears it; fix latent reduce_types pop bug in single-key export path. - megatron_engine: in CP-local forward_step, expose _cp_reduce_group (CP) and _cp_dp_reduce_group (DP+CP) on cp_inputs. - lm_engine: use _cp_reduce_group to all-reduce per-sequence seqlogp/valid-count so ppl is CP-invariant; use _cp_dp_reduce_group as reduce_group for loss/entropy/vocab_* stats so the global mean is reported. Verified on a 64-GPU CP=2 SFT replay (Qwen3-30B-A3B + scale-swe data, seed=1, BS=128): with the fix, loss/avg, ppl/avg/max, vocab_*/avg match the pre-CP-local reference run to within 0.06%; grad_norm and n_tokens/n_valid_tokens are unchanged (the latter remains as fixed in PR #1242). Refs: #1242 * fix: reassemble CP packed sequences * chore: revert unnecessary changes --------- Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com> Co-authored-by: 博惟 <bowei.fw@antgroup.com> | 3 个月前 | |
docs(api): correct obsolete no_save_optim guidance (#1609) The no_save_optim help text told users the flag was "required when using use_distributed_optimizer with Megatron (flattened_range incompatibility)". That workaround was obsoleted by #1468, which switched the Megatron checkpointer to dp_reshardable optimizer sharding, so saving optimizer state works again on the pinned megatron-core 0.17.0. Following the stale advice now costs users silently: recovery resumes with a freshly initialized optimizer, continuing at full learning rate with reset Adam moments. Reword both optimizer-skip flags to describe the actual trade-off, and add tests pinning the behavior the help text now promises. Key changes: - Reword no_save_optim/no_load_optim help in RecoverConfig - Assert both flags default to off, guarding against reintroducing the workaround as a default - Assert the help text no longer documents the flags as required - Cover RecoverHandler threading both flags into SaveLoadMeta.with_optim - Regenerate docs/{en,zh}/cli_reference.md Refs: #1341, #1468 | 16 天前 | |
feat(api): add unified RejectionSamplingConfig for async training (#1088) Replace behave_imp_weight_cap/behave_imp_weight_mode with unified RejectionSamplingConfig supporting multiple metrics (ratio, kl_k1, kl_k2, kl_k3), levels (token/sequence), and actions (mask/clamp). Key changes: - Add RejectionSamplingConfig dataclass with comprehensive validation - Implement apply_rejection_sampling for 1D packed and 2D padded formats - Fix loss denominator scaling bug in mask mode (save count before filtering) - Use geometric mean for sequence-level ratio aggregation (matching GSPO) - Broadcast sequence-level geometric mean as uniform behave_imp_weight - Warn when use_decoupled_loss=True but rejection_sampling is None - Update ppo_actor_loss_fn and grpo_loss_fn to use new config - Migrate 40 example configs to new rejection_sampling field - Add 43 unit tests covering all modes, metrics, and edge cases Refs: #1052 | 4 个月前 | |
fix(ppo): coerce ppo_n_minibatches to 1 for reuse_train_logp (#1457) reuse_train_logp requires ppo_n_minibatches=1 so the training forward pass still reflects the policy that produced the rollout. Previously an invalid combination raised ValueError; instead warn and force ppo_n_minibatches=1, making clear this reduces the PPO update to a single optimizer step. Remove obsolete tests that asserted behavior no longer relevant to this follow-up. | 2 个月前 | |
fix(ppo): handle variable-size trajectory groups in reward normalization (#1454) Use actual trajectory group sizes when applying group-level reward and advantage normalization so failed or filtered rollout samples do not cause fixed-size slices to cross prompt groups. Pass trajectory metadata through batched_call instead of injecting a sentinel batch key, rename the Normalization argument to group_sizes, and zero singleton leave-one-out groups because they have no peer baseline. Add tests for variable-size groups, singleton leave-one-out behavior, validation, 2D advantage normalization, and batched_call metadata forwarding. | 2 个月前 | |
fix(infra): add two-phase teardown to prevent TCPStore race at shutdown (#1244) Problem: During teardown, rank-0 (TCPStore server owner) could exit before peer ranks finished their final NCCL abort, causing noisy "TCPStore.recvValue failed" / "Broken pipe" warnings on stderr. Root cause: All ranks were killed simultaneously without first coordinating a distributed barrier on the CPU (gloo) group to safely tear down NCCL communicators and the TCPStore. Solution: Implement a two-phase teardown protocol: Phase 1 - Engine destroy: call engine.destroy() on every worker concurrently. The engine-side destroy() now executes a CPU barrier (dist.barrier on a gloo process group) followed by dist.destroy_process_group(), ensuring all ranks leave the NCCL collective together. Phase 2 - Process kill: only after the barrier completes, kill the actual processes (Ray: remove placement groups; Slurm: scancel; Local: process tree cleanup). Changes: - engine (fsdp/megatron/archon): add _cpu_group + pre-destroy barrier - train_controller: two-phase destroy (engines first, then workers) - scheduler/ray: _cleanup_workers with ray.wait timeout + PG removal - scheduler/slurm: _destroy_engines_on_workers via HTTP before scancel - scheduler/local: graceful engine teardown before SIGKILL - scheduler_api: add reverse_order param to delete_workers interface - tests: updated test_train_controller, added test_local_scheduler Tested: DPO 4xH20 Ray scheduler - clean teardown, no TCPStore warnings. | 4 个月前 | |
fix(infra): pause proxy workers during weight updates (#1618) Proxy engines keep issuing stale requests during rollout weight updates. Pause them before normal workers and resume them before dispatch restarts. | 13 天前 | |
fix(rollout): attribute output tokens to the serving version (#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. | 1 个月前 | |
feat(infra): add HTTP-based Ray Scheduler (#1441) * fix(infra): preserve dataclass state over RPC * feat(engine): support headless vLLM server mode * fix(utils): use fixed Ray name resolve namespace * refactor(engine): expose backend server env builders * refactor(infra): remove Ray-native scheduler Drop the single-controller Ray scheduler path. Key changes: - Remove RayScheduler, Ray RPC actors, and Ray vLLM remote launcher - Make single-controller trainers dispatch only local or slurm schedulers - Simplify RTensor and RPC serialization to the HTTP backend * feat(infra): add HTTP-based Ray scheduler Add a Ray-backed scheduler that allocates accelerator placement groups while keeping worker and inference engine traffic on the existing HTTP RPC path, with batched launcher operations and multi-node rollout support. Key changes: - Add RayScheduler and RayWorkerProcessLauncher for Ray-managed HTTP workers - Wire scheduler.type=ray into infra exports, trainer initialization, logging, docs, and examples - Batch worker startup and status checks by Ray launcher to reduce per-worker actor calls - Split multi-node rollout backend launch and cleanup into a dedicated coordinator - Tighten Ray launcher lifecycle handling for worker shutdown, placement groups, and backend process cleanup * test(infra): add Ray scheduler tests * fix(trainer): allow proxy workers with Ray scheduler * fix(infra): reuse existing Ray cluster on init * fix(api): restore deprecated Ray placement config --------- Co-authored-by: Ge Shi <utashih@gmail.com> | 1 个月前 | |
refactor(api): migrate allocation_mode to per-engine backend fields (#1044) * refactor(api): migrate allocation_mode to per-engine backend fields Replace the centralized `allocation_mode` string with explicit `backend` fields on `TrainEngineConfig` and `InferenceEngineConfig`. Each engine now owns its own backend+parallelism spec (e.g. `fsdp:d4`, `sglang:d4t2`), eliminating implicit auto-backend selection and the shared `AllocationMode` object. Key changes: - Add `backend` field to TrainEngineConfig and InferenceEngineConfig - Add `ModelAllocation.from_str()` for single-component parsing - Remove `AllocationMode` public export (replaced by `ModelAllocation`) - Rename internal `AllocationMode` to `_AllocationMode` for SPMD launcher backward compatibility with FutureWarning - Remove auto-backend selection — explicit backend prefix is now required - Controllers (`TrainController`, `RolloutController`) parse `backend` directly instead of receiving `alloc_mode` from trainers - `WeightUpdateMeta.alloc_mode` replaced by `gen_allocation` (single `ModelAllocation`) - Add `RWTrainer` and `ArchonRWEngine` for reward model training - Remove `get_model_update_meta()` helper (logic moved to trainers) - Update all YAML configs, examples, docs (EN+ZH), and tests BREAKING CHANGE: `AllocationMode` is removed from public API. Users must migrate to per-engine `backend` fields. SPMD launchers emit deprecation warnings. * chore(ci): fix backend specifier for vlm sft test * fix: fix bare dims for actor backends * chore(docs): fix reminder for bare allocation dims | 5 个月前 | |
feat(engine): enable model-owned THD for Qwen3-VL (#1608) Allow Qwen3-VL models to merge multimodal embeddings before packing while preserving existing padded-only and wrapper-owned paths. Key changes: - Route dense and MoE Qwen3-VL through model-owned THD - Reconstruct padded inputs and restore packed model outputs - Add routing, alignment, parity, and distributed forward coverage | 16 天前 | |
feat(infra): add HTTP-based Ray Scheduler (#1441) * fix(infra): preserve dataclass state over RPC * feat(engine): support headless vLLM server mode * fix(utils): use fixed Ray name resolve namespace * refactor(engine): expose backend server env builders * refactor(infra): remove Ray-native scheduler Drop the single-controller Ray scheduler path. Key changes: - Remove RayScheduler, Ray RPC actors, and Ray vLLM remote launcher - Make single-controller trainers dispatch only local or slurm schedulers - Simplify RTensor and RPC serialization to the HTTP backend * feat(infra): add HTTP-based Ray scheduler Add a Ray-backed scheduler that allocates accelerator placement groups while keeping worker and inference engine traffic on the existing HTTP RPC path, with batched launcher operations and multi-node rollout support. Key changes: - Add RayScheduler and RayWorkerProcessLauncher for Ray-managed HTTP workers - Wire scheduler.type=ray into infra exports, trainer initialization, logging, docs, and examples - Batch worker startup and status checks by Ray launcher to reduce per-worker actor calls - Split multi-node rollout backend launch and cleanup into a dedicated coordinator - Tighten Ray launcher lifecycle handling for worker shutdown, placement groups, and backend process cleanup * test(infra): add Ray scheduler tests * fix(trainer): allow proxy workers with Ray scheduler * fix(infra): reuse existing Ray cluster on init * fix(api): restore deprecated Ray placement config --------- Co-authored-by: Ge Shi <utashih@gmail.com> | 1 个月前 | |
fix: emit PEFT-standard disk LoRA adapter keys so vLLM can load them (incl. #1577) (#1579) * fix: emit PEFT-standard disk LoRA adapter keys so vLLM can load them, and align sglang best-effort unload test assertion (#1577) * test(inference-service): raise VLM controller init timeout to 600s to avoid vLLM startup timeout | 28 天前 | |
feat(colocate): support AWEX colocated actor-rollout training (#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. | 26 天前 | |
feat: support pp for Sglang (#1162) * Introduces pipeline parallelism (PP) support for the SGLang inference backend, enabling AReaL to train with `pp_size > 1` on the inference side across all three training engines (Megatron, FSDP, Archon). Create a separate NCCL weight update group per PP stage (per-PP-rank group). Each group contains only the TP workers at the corresponding PP rank plus one training rank, allowing the rendezvous to complete within a single PP stage without cross -stage blocking. | 3 个月前 | |
refactor: move 5 experimental modules into areal/v2 for 2.0 release (#1448) * refactor: move 5 experimental modules into areal/v2 for 2.0 release Move agent_service, inference_service, training_service, weight_update, and cli from areal/experimental/ to areal/v2/, and rewrite every reference (Python imports, `python -m` invocations, console-script entry points, CODEOWNERS, docs, review-pr signals) to the new path. - 5 directories migrated via `git mv` (history preserved) - 70 files modified across areal/, tests/, examples/, pyproject{,vllm}.toml, .github/CODEOWNERS, ROADMAP.md, and tooling docs - `areal/v2/__init__.py` added so v2 is an importable package - Build config (tool.uv.build-backend with module-root="") auto-discovers the new package — no pyproject changes beyond the entry point Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * style: apply pre-commit auto-fixes after areal/v2 move CI pre-commit job reformatted 54 files automatically: - .github/CODEOWNERS: realign owner columns (32-col) after shorter /areal/v2/ paths broke the previous 40-col alignment - areal/**, tests/**, examples/**, docs/**: ruff isort reorders `areal.v2.*` imports into their new alphabetical slot (between `areal.engine` and `areal.infra`) - markdown/yaml whitespace normalized by mdformat / ruff-format Pure formatting; no logic change. `pre-commit run --all-files` is now green locally. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * fix(tests): point py<3.12 conftest stubs at areal/v2/ Two test conftest.py stubs (Python 3.10/3.11 compat) still had the `areal/experimental/{inference_service,weight_update}` path as the namespace package's __path__, broken by the v2 move. The sed pass missed them because the paths were comma-separated `os.path.join` args (`"areal", "experimental", "X"`), not slash-form path strings. Additionally insert an `areal.v2` stub between `areal` and the leaf package so the parent→child attribute wiring loop (which uses `name.rsplit(".", 1)`) can find a parent module in sys.modules. Without it `setattr(parent, child, ...)` silently no-ops and `unittest.mock.patch` traversal breaks on the new path. Spotted by gemini-code-assist on PR #1448. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor(tests): mirror areal/v2 move under tests/v2 Move agent_service, inference_service, training_service, and weight_update test directories from tests/experimental/ to tests/v2/ to mirror the source-tree layout. tests/experimental/ retains archon/ and openai/ (still-experimental modules). - 50 files migrated via `git mv` (history preserved) - tests/v2/__init__.py added - 9 files rewritten for the new dotted/slashed test paths: `tests.experimental.{4 modules}` → `tests.v2.{4 modules}` `tests/experimental/{4 modules}` → `tests/v2/{4 modules}` (covers integration_utils imports, fake_train_engine engine_class, pytest invocation strings in docstrings, and the tests/v2/weight_update/torchrun/run_nccl_weight_transfer.py path) conftest stubs (Python <3.12 namespace shims) keep working because _REPO_ROOT is computed via "..", "..", ".." from the conftest file — same depth under tests/v2/X/ as under tests/experimental/X/. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com> | 2 个月前 | |
feat(scheduler): honor reservation, exclusive and user env overrides on Slurm (#1584) SchedulingSpec had no way to reach sbatch's --reservation or --exclusive, so runs that need a reserved partition had to be launched by hand outside the scheduler. Explicit env_vars were also lost: the scheduler applied AReaL's forwarding and the thread-count defaults on top of them, so a spec asking for a specific OMP_NUM_THREADS or allocator config silently got the framework value. Snapshot the user's mapping and re-apply it last, which also lets colocated roles carry per-role settings that must differ between actor and rollout. | 21 天前 | |
refactor: flatten sub-module imports to use parent package re-exports (#996) Add __init__.py with lazy re-exports (__getattr__ + __all__) for areal/api, areal/engine, areal/reward, areal/workflow, and __all__ for areal/dataset, then rewrite all external imports across the codebase to use the shorter parent-package form (e.g. `from areal.api import TrainEngine`). Key changes: - All re-exports are fully lazy via __getattr__ (no eager imports) - Flatten ~100 files across areal/, tests/, examples/ - Preserve cli_args deep imports (dozens of config classes, including SchedulingSpec which stays in cli_args) - Preserve intra-package relative imports to avoid circular deps - Reward submodules use sibling-relative imports (from . import ...) - Preserve non-exported symbols (DeviceRuntimeInfo, HttpRequest, etc.) Co-authored-by: Wentai Zhang <zhangwentai.zwt@antgroup.com> | 5 个月前 | |
fix(scheduler): treat every terminal Slurm state as a dead worker (#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. | 26 天前 | |
fix: per-sample version tracking with loss_mask filter and multi-turn… (#1408) ## Summary - head_version/tail_version now per-sample, filtered by loss_mask==1 - fixes head_version always being -1 due to input token version placeholders - adds version_rle field (run-length encoded per-token version sequence) - adds _split_trajectory_for_dump helper for correct multi-turn prompt_end - adds segments field for multi-turn agent trajectory analysis ## Context The previous `_dump_trajectory` had three issues: 1. `head_version = min(versions)` always returned -1 because input tokens use -1 as placeholder 2. `prompt_end = seqlen - sum(mask)` is incorrect for multi-turn agent rollouts where loss_mask is interleaved 3. No per-token version granularity was persisted for cross-version analysis --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> | 2 个月前 | |
fix(infra): correct staleness capacity inflation after recovery (#1345) * fix(infra): correct staleness capacity inflation after checkpoint recovery StalenessManager's accepted counter started at 0 while the version was restored to a high value by the recovery path. This caused the capacity formula to yield (max_staleness + recovered_version + 1) * batch_size instead of the intended (max_staleness + 1) * batch_size, allowing a burst of rollout submissions and unbounded staleness growth. Add on_version_recovered() to StalenessManager and call it from rl_trainer after recover completes. The trainer accesses the staleness manager directly via the known concrete type (RolloutController in single-controller mode, workflow_executor in SPMD mode). * fix(infra): clarify staleness recovery semantics and use public APIs Address review feedback on the staleness manager recovery path: - Document that on_version_recovered is expected to be called with running == 0 and explain the bound when it is not. - Reach the manager through the public staleness_manager properties on RolloutController and WorkflowExecutor instead of the private _staleness_manager attribute, avoiding coupling to internal layout. - Extend tests with the version=0 no-op case and a parametrized case with in-flight rollouts to verify accepted is set correctly. --------- Co-authored-by: fenghui <dh183333@antgroup.com> | 3 个月前 | |
fix: align CP metrics and TP grad norm metadata (#1497) * fix: align CP metrics and TP grad norm metadata Two independent correctness fixes for the Megatron engine. 1. CP metrics alignment (stats_tracker): Per-key reduce_group overrides may be wider than the default DP export group (e.g. DP+CP for token-level metrics). Add a `key_sync_group` argument so override keys sync their key set / metadata and reduce over DP+CP, while default keys keep the DP alignment. Aggregation helpers now emit identity placeholder tensors (0 / +inf / -inf) for keys absent on a rank, so every rank still participates in the collective and avoids hangs / mismatched reductions when CP > 1. 2. TP grad norm metadata (megatron_engine): `_mark_duplicated_params` now also clears `param.tensor_model_parallel` on replicated (tp_size == 1, non-expert) params. Megatron's optimizer uses that attribute to decide which TP ranks contribute to grad norm/clipping, so leaving it True double-counts duplicated params when TP > 1. Clearing it is also consistent with `all_gather_param` and `hf_save`, which already treat non-TP params as replicated. Adds unit tests for both paths (tests/test_stats_tracker.py and a new case in tests/test_megatron_engine.py). * test: cover all_gather_param routing and duplicated-param grad norm Add two tests backing the TP grad-norm metadata fix: - tests/test_all_gather_param.py: assert `all_gather_param` returns the param as-is (no TP all-gather) when `tensor_model_parallel` is False or the name is in `duplicated_param_names`, and only all-gathers genuine TP-sharded params. Uses `pytest.importorskip` so it runs in the megatron CI env and skips gracefully elsewhere. - tests/test_grad_norm_duplicated.py: a CPU/gloo 2- and 4-rank test showing a replicated param is counted once (correct) when tensor_model_parallel is False, and inflated by sqrt(tp) when it is (incorrectly) left True. Mirrors Megatron's param_is_not_tensor_parallel_duplicate selection. * test: add 2-GPU grad-norm TP-invariance integration test End-to-end guard for the TP grad-norm metadata fix: run the real MegatronEngine at TP=1 and TP=2 on identical deterministic input via torchrun and assert the reported grad norm is TP-invariant. Double-counting replicated params (tensor_model_parallel left True) would inflate it as TP grows. Also asserts the fix demoted at least one real duplicated param at TP=2. Marked multi_gpu/slow and skipped without >= 2 GPUs. * refactor(stats): dedup identity tensors, avoid min/max GPU sync Address review feedback on the CP metrics fix: - Generalize the placeholder-tensor helper to `_placeholder_scalar(fill=...)` so the SUM/AVG (0.0) and MIN/MAX (+/-inf) empty-value branches share one device-aware constructor instead of four inline copies. - In `_min_of`/`_max_of`, reduce the per-shard list with `torch.stack(xs).min()` / `.max()` instead of Python `min()`/`max()`, which forced a CPU-GPU sync. - In the SCALAR branch, read values via `self.stats.get(key, [])` (a rank may learn a key only through metadata sync) and guard the `value / cnt` division against `cnt == 0` so a key absent everywhere yields 0.0 instead of NaN. - Add a regression test for the missing-on-this-rank SCALAR path. | 1 个月前 | |
feat(megatron): add CP-safe vocab stats and MoE config support (#1460) Add Megatron context-parallel output gathering for forward-only paths and plumb vocabulary logits statistics through PPO/DPO/SFT losses. Support per-key reduce groups in StatsTracker so CP-local loss and vocab statistics can reduce across the appropriate DP/CP group without changing the default reduction group for unrelated stats. Expose Megatron/MoE configuration knobs for router fusion, auxiliary-loss-free balancing, router z-loss, FP32 lm_head output, and fused cross entropy. Update Bailing MoE defaults and regenerate CLI docs. | 2 个月前 | |
feat(utils): add Trackio experiment tracking backend (#1113) Integrate Trackio (Hugging Face) as a new experiment tracking option alongside existing WandB, SwanLab, and TensorBoard backends. Key changes: - Add TrackioConfig dataclass with mode/project/name/space_id fields - Integrate trackio init/log/finish lifecycle in StatsLogger - Add trackio.log() fallback in logging.py helper function - Add trackio to pyproject.toml dependencies - Update CLI docs generator to include TrackioConfig - Add unit tests for config and StatsLogger integration | 4 个月前 | |
feat(rollout): add grouped reward normalization controls (#1516) Add rollout-time reward normalization and incomplete group dropping for grouped InteractionWithTokenLogpReward workflows. Preserve original rewards for logging and dumps, and thread the controls through inference and training controllers. Add tests covering grouped reward normalization and parameter forwarding. Co-authored-by: chucai.dzq <chucai.dzq@alibaba-inc.com> | 1 个月前 | |
fix(CI): fix vlm_grpo CI OOM bug (#1438) Default fp32 master weights (PR #1369) push Qwen2.5-VL-3B beyond A100-40G capacity at the first AdamW step. Pin the test to bf16 storage with adam_bf16 (Kahan summation) so it runs on 40G runners. ## Description <!-- Provide a clear and concise description of what this PR does --> ## Related Issue <!-- Link to the issue this PR addresses. PRs should be related to a well-templated issue. --> Fixes #(issue) ## Type of Change <!-- Select ONE that best describes this PR --> - [x] 🐛 Bug fix - [ ] ✨ New feature - [ ] 💥 Breaking change - [ ] 📝 Documentation update - [ ] ♻️ Refactoring - [ ] ⚡ Performance improvement - [ ] ✅ Test coverage improvement ## Checklist <!-- Mark with 'x' what you've done --> - [ ] I have read the [Contributing Guide](https://github.com/areal-project/community/blob/main/CONTRIBUTING.md) - [ ] Pre-commit hooks pass (`pre-commit run --all-files`) - [ ] Relevant tests pass; new tests added for new functionality - [ ] Documentation updated (if applicable; built with `./docs/build_all.sh`) - [ ] Branch is up to date with `main` - [ ] Self-reviewed via `/review-pr` command - [ ] This PR was created by a coding agent via `/create-pr` - [ ] This PR is a breaking change **Breaking Change Details (if applicable):** <!-- Describe what breaks and how users should migrate --> ## Additional Context <!-- Add any other context, screenshots, logs, or explanations here --> ______________________________________________________________________ **Need help?** Check the [Contributing Guide](https://github.com/areal-project/community/blob/main/CONTRIBUTING.md) or ask in [GitHub Discussions](https://github.com/areal-project/AReaL/discussions)! | 2 个月前 | |
perf: reduce Megatron training memory peaks (#1555) * 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: #1555 | 26 天前 | |
fix(trainer): run initial evaluation before first update (#1636) `eval_before_train` currently relies on the first scheduled evaluator check, which trainers perform only after one optimization step. The reported baseline can therefore contain updated weights. Run the one-shot evaluation separately so its metrics are logged before the first update without advancing periodic evaluation cadence. Skip the startup evaluation on recovery and clear legacy deferred triggers when loading evaluator state. Key changes: - Invoke the startup baseline from PPO, SFT, DPO, and reward trainers - Preserve periodic evaluation cadence and recovery behavior - Cover call ordering, logging steps, legacy state, and PPO offload Refs: #1232 Signed-off-by: Bo Yang <yb550079@antgroup.com> | 8 天前 | |
refactor(api): migrate allocation_mode to per-engine backend fields (#1044) * refactor(api): migrate allocation_mode to per-engine backend fields Replace the centralized `allocation_mode` string with explicit `backend` fields on `TrainEngineConfig` and `InferenceEngineConfig`. Each engine now owns its own backend+parallelism spec (e.g. `fsdp:d4`, `sglang:d4t2`), eliminating implicit auto-backend selection and the shared `AllocationMode` object. Key changes: - Add `backend` field to TrainEngineConfig and InferenceEngineConfig - Add `ModelAllocation.from_str()` for single-component parsing - Remove `AllocationMode` public export (replaced by `ModelAllocation`) - Rename internal `AllocationMode` to `_AllocationMode` for SPMD launcher backward compatibility with FutureWarning - Remove auto-backend selection — explicit backend prefix is now required - Controllers (`TrainController`, `RolloutController`) parse `backend` directly instead of receiving `alloc_mode` from trainers - `WeightUpdateMeta.alloc_mode` replaced by `gen_allocation` (single `ModelAllocation`) - Add `RWTrainer` and `ArchonRWEngine` for reward model training - Remove `get_model_update_meta()` helper (logic moved to trainers) - Update all YAML configs, examples, docs (EN+ZH), and tests BREAKING CHANGE: `AllocationMode` is removed from public API. Users must migrate to per-engine `backend` fields. SPMD launchers emit deprecation warnings. * chore(ci): fix backend specifier for vlm sft test * fix: fix bare dims for actor backends * chore(docs): fix reminder for bare allocation dims | 5 个月前 | |
feat(trainer): add flexible GAE lambda strategies (#1600) * feat: add turn-level GAE support Treat each generated turn as a GAE timestep while preserving the token-level default and token-local KL regularization. Key changes: - propagate and validate token-aligned turn IDs - compute turn-level advantages without full CPU sequence scans - filter structural metadata at FSDP and Archon model boundaries - document the new selector and add focused regression tests * feat(trainer): add dynamic per-sample GAE lambda Allow GAE lambda to vary by trajectory using effective token or turn lengths while preserving static float behavior. Key changes: - Resolve custom lambda functions and keyword arguments from config - Add VAPO length-adaptive GAE with empty-trajectory handling - Validate per-sample lambda tensors and cover token and turn modes * perf(trainer): reduce GAE preprocessing overhead Hoist loop-invariant tensor work and bypass dynamic trajectory length construction when GAE lambda is static. * feat(trainer): add relative-position GAE lambda * fix(trainer): allow token lambda without turn metadata Keep custom token-level GAE lambda functions compatible with rollout workflows that do not emit turn IDs, while preserving the metadata requirement for turn-level GAE. * docs: regenerate CLI reference for GAE options Keep the generated configuration reference aligned with the current main branch after porting the AntCode GAE changes. * docs: document flexible GAE configuration Explain token- and turn-level recurrences, KL and critic semantics, dynamic lambda strategies, and custom workflow turn IDs in English and Chinese.\n\nFix CLI default rendering for dataclass factories and cover it with unit tests. * refactor(trainer): extract GAE helpers Keep PPOActor focused on training orchestration by moving GAE kernels, turn metadata validation, and lambda context construction into a dedicated module. --------- Co-authored-by: Wenhao Zhou <miumiu.zwh@antgroup.com> | 13 天前 | |
refactor(tests): move tests from areal/tests to top-level tests directory (#944) * refactor(tests): move tests from areal/tests to top-level tests directory Move the test directory from areal/tests/ to a top-level tests/ directory to follow Python best practices and improve project structure. Key changes: - Move all test files from areal/tests/ to tests/ - Create __init__.py files for tests package and subdirectories - Create areal/utils/testing_utils.py with shared utilities: - get_model_path, get_dataset_path (model/dataset path resolution) - MODEL_PATHS, DENSE_MODEL_PATHS, MOE_MODEL_PATHS (test model configs) - load_archon_model (Archon model loading for tests) - Update areal/tools/profiling_utils to re-export from testing_utils - Update test utilities to import from areal/utils/testing_utils, keeping test-specific behavior (pytest.skip) as wrappers - Update all imports from 'areal.tests.' to 'tests.' - Update hardcoded torchrun script paths - Update workflow string references for dynamic class loading - Update CI workflow paths in .github/workflows/test-areal.yml - Update documentation references in CLAUDE.md, AGENTS.md, CONTRIBUTING.md, and all .claude/ and .opencode/ skill/agent files * minor fix * remove import side effect * fix(tests): add tests/utils.py module for test imports Re-export get_model_path and get_dataset_path from areal.utils.testing_utils to fix ModuleNotFoundError in tests that import from tests.utils. | 6 个月前 | |
refactor(tests): move tests from areal/tests to top-level tests directory (#944) * refactor(tests): move tests from areal/tests to top-level tests directory Move the test directory from areal/tests/ to a top-level tests/ directory to follow Python best practices and improve project structure. Key changes: - Move all test files from areal/tests/ to tests/ - Create __init__.py files for tests package and subdirectories - Create areal/utils/testing_utils.py with shared utilities: - get_model_path, get_dataset_path (model/dataset path resolution) - MODEL_PATHS, DENSE_MODEL_PATHS, MOE_MODEL_PATHS (test model configs) - load_archon_model (Archon model loading for tests) - Update areal/tools/profiling_utils to re-export from testing_utils - Update test utilities to import from areal/utils/testing_utils, keeping test-specific behavior (pytest.skip) as wrappers - Update all imports from 'areal.tests.' to 'tests.' - Update hardcoded torchrun script paths - Update workflow string references for dynamic class loading - Update CI workflow paths in .github/workflows/test-areal.yml - Update documentation references in CLAUDE.md, AGENTS.md, CONTRIBUTING.md, and all .claude/ and .opencode/ skill/agent files * minor fix * remove import side effect * fix(tests): add tests/utils.py module for test imports Re-export get_model_path and get_dataset_path from areal.utils.testing_utils to fix ModuleNotFoundError in tests that import from tests.utils. | 6 个月前 | |
perf(megatron): avoid synthetic BSHD padding rows (#1632) Assisted-by: OpenAI Codex <codex@openai.com> Signed-off-by: EazyReal <8047065+EazyReal@users.noreply.github.com> Co-authored-by: EazyReal <8047065+EazyReal@users.noreply.github.com> | 8 天前 | |
feat(colocate): support AWEX colocated actor-rollout training (#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. | 26 天前 | |
feat(models): shard vision encoder across Ulysses SP ranks (#929) When using Ulysses Sequence Parallelism with VLMs, every SP rank redundantly runs the full vision encoder. This adds a `shard_vision_across_sp` option that distributes whole images across SP ranks, runs ViT locally, and all-gathers the embeddings - eliminating redundant computation while preserving gradient correctness via all_reduce(SUM) in backward. Key changes: - Add vision_sp_shard.py with greedy contiguous image assignment, padded all-gather, and custom autograd backward - Support Qwen2-VL, Qwen2.5-VL, and Qwen3-VL (incl. deepstack) - Add shard_vision_across_sp flag to FSDPEngineConfig (effective only when context_parallel_size > 1) - Monkey-patch VisionTransformer.forward via table-driven registry - Add 31 CPU-only unit tests covering all helper functions, deepstack unpacking, patching idempotency, and integration --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Wentai Zhang <zhangwentai.zwt@antgroup.com> | 5 个月前 | |
fix: make rollout sampling deterministic (#1625) | 12 天前 | |
fix(engine): gate per-PP weight sync on SGLang backend (#1564) The train_pp_size == gen_pp_size restriction is required only by SGLang's per-PP-stage weight synchronization, whose rollout-side request forms one NCCL group per PP stage. The guard was keyed on gen_pp_size > 1 alone, so it also fired for vLLM, which joins a single flat group spanning all inference workers and does not need a 1:1 mapping between training and inference PP stages. This blocked valid configurations such as Megatron PP=4 training with vLLM PP=2 rollout. Gate the per-PP path on the SGLang backend so vLLM falls back to the single-group path (already used for gen_pp == 1), where each training PP head broadcasts the parameters it owns to all vLLM workers. Key changes: - megatron_engine: gate per-PP branch on gen_backend == "sglang" - archon_weight_sync: apply the same backend gating in the dispatcher - add tests/test_vllm_pp_mismatch.py covering vLLM PP-mismatch success and preserved SGLang mismatch failure Refs: #1560 | 1 个月前 | |
refactor(tests): move tests from areal/tests to top-level tests directory (#944) * refactor(tests): move tests from areal/tests to top-level tests directory Move the test directory from areal/tests/ to a top-level tests/ directory to follow Python best practices and improve project structure. Key changes: - Move all test files from areal/tests/ to tests/ - Create __init__.py files for tests package and subdirectories - Create areal/utils/testing_utils.py with shared utilities: - get_model_path, get_dataset_path (model/dataset path resolution) - MODEL_PATHS, DENSE_MODEL_PATHS, MOE_MODEL_PATHS (test model configs) - load_archon_model (Archon model loading for tests) - Update areal/tools/profiling_utils to re-export from testing_utils - Update test utilities to import from areal/utils/testing_utils, keeping test-specific behavior (pytest.skip) as wrappers - Update all imports from 'areal.tests.' to 'tests.' - Update hardcoded torchrun script paths - Update workflow string references for dynamic class loading - Update CI workflow paths in .github/workflows/test-areal.yml - Update documentation references in CLAUDE.md, AGENTS.md, CONTRIBUTING.md, and all .claude/ and .opencode/ skill/agent files * minor fix * remove import side effect * fix(tests): add tests/utils.py module for test imports Re-export get_model_path and get_dataset_path from areal.utils.testing_utils to fix ModuleNotFoundError in tests that import from tests.utils. | 6 个月前 | |
perf: reduce Megatron training memory peaks (#1555) * 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: #1555 | 26 天前 | |
fix(engine): eagerly init HCCL subgroups to fix ref compute_logp on NPU (#1254) * feat(engine): add warmup_process_groups helper NCCL/HCCL communicators initialize lazily on first use. On Ascend NPU, deferring that until a collective runs during training can fail with HCCP init errors when multiple colocated engines race on overlapping subgroups. Introduce a small helper that forces eager communicator creation via a dummy all-reduce while all ranks are aligned. Refs: #1099 * fix(engine): eagerly init FSDPEngine HCCL subgroups Warm up dp/sp/mp groups right after they are minted so the first collective from ref.compute_logp doesn't race with the actor's HCCL init on Ascend NPU. Refs: #1099 * fix(engine): eagerly init MegatronEngine HCCL subgroups Apply the same warmup to the context+model-parallel and data-parallel groups created in MegatronEngine to avoid the lazy-init HCCL race on Ascend NPU when the ref engine is colocated with the actor. Refs: #1099 * fix(archon): eagerly init ArchonEngine HCCL subgroups Warm up dp/pp_cp_tp/tp/cp groups in ArchonEngine.create_process_group so colocated engines don't trigger the HCCL lazy-init race on NPU. Refs: #1099 * fix(engine): tolerate missing LOCAL_RANK in warmup helper Custom launchers (non-torchrun) may not export LOCAL_RANK. Fall back to the device the caller has already configured via set_device instead of raising KeyError. Also simplify the dedup pass to dict.fromkeys. Addresses review feedback on #1254. * test(engine): cover LOCAL_RANK fallback in warmup helper Assert set_device is called with LOCAL_RANK when present, and that the helper falls back to current_device() when it is unset. | 4 个月前 | |
refactor: flatten sub-module imports to use parent package re-exports (#996) Add __init__.py with lazy re-exports (__getattr__ + __all__) for areal/api, areal/engine, areal/reward, areal/workflow, and __all__ for areal/dataset, then rewrite all external imports across the codebase to use the shorter parent-package form (e.g. `from areal.api import TrainEngine`). Key changes: - All re-exports are fully lazy via __getattr__ (no eager imports) - Flatten ~100 files across areal/, tests/, examples/ - Preserve cli_args deep imports (dozens of config classes, including SchedulingSpec which stays in cli_args) - Preserve intra-package relative imports to avoid circular deps - Reward submodules use sibling-relative imports (from . import ...) - Preserve non-exported symbols (DeviceRuntimeInfo, HttpRequest, etc.) Co-authored-by: Wentai Zhang <zhangwentai.zwt@antgroup.com> | 5 个月前 | |
feat(v2): support grouped reward normalization (#1619) * feat(v2): support grouped reward normalization Normalize per-prompt rollout groups at the data-proxy export boundary before trajectory tensors are merged and remotized. Preserve raw rewards for filtering and metrics while sharing the normalization implementation with v1. Keep reward metadata local for rollout filters and clear remote shards for rejected or failed trajectories so filtered v2 rollouts do not leak data-proxy storage. * fix(v2): bound rejected trajectory cleanup | 8 天前 | |
refactor(tests): move tests from areal/tests to top-level tests directory (#944) * refactor(tests): move tests from areal/tests to top-level tests directory Move the test directory from areal/tests/ to a top-level tests/ directory to follow Python best practices and improve project structure. Key changes: - Move all test files from areal/tests/ to tests/ - Create __init__.py files for tests package and subdirectories - Create areal/utils/testing_utils.py with shared utilities: - get_model_path, get_dataset_path (model/dataset path resolution) - MODEL_PATHS, DENSE_MODEL_PATHS, MOE_MODEL_PATHS (test model configs) - load_archon_model (Archon model loading for tests) - Update areal/tools/profiling_utils to re-export from testing_utils - Update test utilities to import from areal/utils/testing_utils, keeping test-specific behavior (pytest.skip) as wrappers - Update all imports from 'areal.tests.' to 'tests.' - Update hardcoded torchrun script paths - Update workflow string references for dynamic class loading - Update CI workflow paths in .github/workflows/test-areal.yml - Update documentation references in CLAUDE.md, AGENTS.md, CONTRIBUTING.md, and all .claude/ and .opencode/ skill/agent files * minor fix * remove import side effect * fix(tests): add tests/utils.py module for test imports Re-export get_model_path and get_dataset_path from areal.utils.testing_utils to fix ModuleNotFoundError in tests that import from tests.utils. | 6 个月前 | |
refactor(tests): move tests from areal/tests to top-level tests directory (#944) * refactor(tests): move tests from areal/tests to top-level tests directory Move the test directory from areal/tests/ to a top-level tests/ directory to follow Python best practices and improve project structure. Key changes: - Move all test files from areal/tests/ to tests/ - Create __init__.py files for tests package and subdirectories - Create areal/utils/testing_utils.py with shared utilities: - get_model_path, get_dataset_path (model/dataset path resolution) - MODEL_PATHS, DENSE_MODEL_PATHS, MOE_MODEL_PATHS (test model configs) - load_archon_model (Archon model loading for tests) - Update areal/tools/profiling_utils to re-export from testing_utils - Update test utilities to import from areal/utils/testing_utils, keeping test-specific behavior (pytest.skip) as wrappers - Update all imports from 'areal.tests.' to 'tests.' - Update hardcoded torchrun script paths - Update workflow string references for dynamic class loading - Update CI workflow paths in .github/workflows/test-areal.yml - Update documentation references in CLAUDE.md, AGENTS.md, CONTRIBUTING.md, and all .claude/ and .opencode/ skill/agent files * minor fix * remove import side effect * fix(tests): add tests/utils.py module for test imports Re-export get_model_path and get_dataset_path from areal.utils.testing_utils to fix ModuleNotFoundError in tests that import from tests.utils. | 6 个月前 |
| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
| 8 天前 | ||
| 5 个月前 | ||
| 3 个月前 | ||
| 19 天前 | ||
| 1 个月前 | ||
| 4 个月前 | ||
| 16 天前 | ||
| 8 天前 | ||
| 6 个月前 | ||
| 5 个月前 | ||
| 1 个月前 | ||
| 1 个月前 | ||
| 26 天前 | ||
| 5 个月前 | ||
| 8 天前 | ||
| 6 个月前 | ||
| 5 个月前 | ||
| 26 天前 | ||
| 11 天前 | ||
| 14 天前 | ||
| 11 天前 | ||
| 26 天前 | ||
| 2 个月前 | ||
| 2 个月前 | ||
| 6 个月前 | ||
| 5 个月前 | ||
| 30 天前 | ||
| 12 天前 | ||
| 4 个月前 | ||
| 11 天前 | ||
| 6 个月前 | ||
| 26 天前 | ||
| 11 天前 | ||
| 6 个月前 | ||
| 4 个月前 | ||
| 8 天前 | ||
| 26 天前 | ||
| 5 个月前 | ||
| 5 个月前 | ||
| 4 个月前 | ||
| 5 个月前 | ||
| 4 个月前 | ||
| 3 个月前 | ||
| 6 个月前 | ||
| 6 个月前 | ||
| 3 个月前 | ||
| 13 天前 | ||
| 1 个月前 | ||
| 26 天前 | ||
| 1 个月前 | ||
| 1 个月前 | ||
| 1 个月前 | ||
| 1 个月前 | ||
| 4 个月前 | ||
| 4 个月前 | ||
| 4 个月前 | ||
| 4 个月前 | ||
| 6 个月前 | ||
| 21 天前 | ||
| 5 个月前 | ||
| 27 天前 | ||
| 12 天前 | ||
| 19 天前 | ||
| 1 个月前 | ||
| 2 个月前 | ||
| 8 天前 | ||
| 16 天前 | ||
| 20 天前 | ||
| 11 天前 | ||
| 2 个月前 | ||
| 5 个月前 | ||
| 2 个月前 | ||
| 5 个月前 | ||
| 26 天前 | ||
| 26 天前 | ||
| 13 天前 | ||
| 27 天前 | ||
| 26 天前 | ||
| 4 个月前 | ||
| 20 天前 | ||
| 20 天前 | ||
| 1 个月前 | ||
| 3 个月前 | ||
| 16 天前 | ||
| 4 个月前 | ||
| 2 个月前 | ||
| 2 个月前 | ||
| 4 个月前 | ||
| 13 天前 | ||
| 1 个月前 | ||
| 1 个月前 | ||
| 5 个月前 | ||
| 16 天前 | ||
| 1 个月前 | ||
| 28 天前 | ||
| 26 天前 | ||
| 3 个月前 | ||
| 2 个月前 | ||
| 21 天前 | ||
| 5 个月前 | ||
| 26 天前 | ||
| 2 个月前 | ||
| 3 个月前 | ||
| 1 个月前 | ||
| 2 个月前 | ||
| 4 个月前 | ||
| 1 个月前 | ||
| 2 个月前 | ||
| 26 天前 | ||
| 8 天前 | ||
| 5 个月前 | ||
| 13 天前 | ||
| 6 个月前 | ||
| 6 个月前 | ||
| 8 天前 | ||
| 26 天前 | ||
| 5 个月前 | ||
| 12 天前 | ||
| 1 个月前 | ||
| 6 个月前 | ||
| 26 天前 | ||
| 4 个月前 | ||
| 5 个月前 | ||
| 8 天前 | ||
| 6 个月前 | ||
| 6 个月前 |