| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
chore: enforce license (#1171) | 4 个月前 | |
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. | 1 个月前 | |
feat: add MTP training with speculative decoding rollout (#1659) * feat: add MTP training with NEXTN speculative decoding rollout Train the built-in MTP head jointly with GRPO/SFT on Megatron and use it during rollout through SGLang NEXTN speculative decoding. Online weight synchronization updates both the target model and the built-in MTP draft runner, ensuring that speculative decoding uses the latest RL-trained MTP weights instead of a stale initialization. MTP is trained through Megatron-Core's auxiliary cross-entropy path with a configurable loss scaling factor of `0.1` by default. AReaL supplies independent MTP label and loss-mask channels while keeping the main forward path logits-based. Shared output weights, backbone hidden states, and embedding inputs are detached from the MTP loss graph, so the backbone receives only the policy/SFT gradient while the MTP-specific parameters learn from future-token supervision. For packed THD training with context parallelism, AReaL applies the same per-sequence zigzag CP split and rank-local repacking to MTP labels and loss masks as it does to input IDs. Megatron-Core's CP-aware rolling then aligns future-token targets across CP ranks without crossing packed sequence boundaries. For online rollout, a focused compatibility bridge for `sglang==0.5.10.post1` receives each distributed weight bucket once and applies it to both the built-in MTP draft runner and the target runner. It supports both SGLang speculative worker layouts, handles SGLang's internal `NEXTN`-to-`EAGLE` normalization, and leaves external EAGLE draft models untouched. Draft-weight CPU backup can be enabled so the server remains available while updated weights arrive online. End-to-end validation was performed with Qwen3.5-2B on Geometry3K GRPO. The training-side MTP weights changed during optimization, and the same updated tensors were loaded into all SGLang draft runners. Key changes: - Add `enable_mtp_training` and `mtp_loss_scaling_factor` to the Megatron engine configuration. MTP training implies retaining the model's MTP layers and is incompatible with `lm_head_loss_chunk_size`. - Feed independent MTP labels and loss masks through Megatron forward passes while preserving the main logits-based loss path. - Patch Megatron-Core `GPTModel` and Megatron-Bridge `Qwen3VLGPTModel` forwarding so Qwen3.5 text and multimodal batches can train MTP. - Align MTP labels and masks with padded or packed execution layouts and prevent targets from crossing sequence or padding boundaries. - Support MTP training with `CP > 1` for wrapper-owned packed THD by applying the same per-sequence zigzag split and rank-local repacking to input IDs, MTP labels, and MTP loss masks. - Reuse Megatron-Core's packed, CP-aware rolling semantics to align future-token supervision across CP ranks. - Isolate MTP gradients from shared output weights, embeddings, and backbone hidden states. - Report the auxiliary `mtp_loss` in training statistics. - Add SGLang speculative-decoding configuration passthrough for `NEXTN`, speculative steps, EAGLE top-k, draft-token count, external draft-model path, and draft-weight CPU backup. - Add an SGLang distributed weight-update bridge that updates both target and built-in MTP draft runners from the same received tensors. - Support both SGLang Spec v1 and Spec v2 draft-runner layouts, while failing fast for unsupported SGLang versions, missing draft runners, unsupported load formats, and inference pipeline parallelism. - Record `rollout/spec_accept_rate` and `rollout/spec_accept_length` from SGLang response metadata. - Add a Qwen3.5-2B Geometry3K GRPO example with MTP training and NEXTN rollout enabled. - Add unit coverage for padded and packed MTP label/mask layouts, CP zigzag alignment, multimodal forwarding, NEXTN/EAGLE routing, Spec v1/v2 compatibility, and draft/target online weight updates. - Regenerate the English and Chinese CLI reference documentation. Current limitations: - MTP training with `CP > 1` is supported only for wrapper-owned packed THD. Padded BSHD, VLM, and model-owned THD execution still require `CP=1`. - The current gradient-isolation implementation supports a single MTP prediction layer. Multi-layer MTP gradient propagation is not supported. - The SGLang distributed MTP update bridge requires inference pipeline parallel size `1`. - The compatibility bridge is intentionally pinned to `sglang==0.5.10.post1` because it relies on version-specific internal weight-update and draft-runner APIs. * test: use HttpGenerationResult in rollout version race test * fix(engine): align MTP masks and detach untied output weights Keep next-token-aligned masks unchanged before MCore's per-layer roll, and detach internal output-layer weights for untied models. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(engine): unify Megatron main and auxiliary loss scaling Let MCore apply the optimizer loss scale to both the main backward path and separately seeded MTP/MoE auxiliary gradients. This prevents FP16 optimizer unscaling from suppressing auxiliary updates. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(models): reject unsupported multilayer MTP training Fail before model construction when MTP training requests more than one prediction layer, whose gradients are not fully supported yet. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: huaqingyuan <huaqingyuan@tencent.com> Co-authored-by: Cursor <cursoragent@cursor.com> | 6 天前 | |
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 个月前 | |
feat: add MTP training with speculative decoding rollout (#1659) * feat: add MTP training with NEXTN speculative decoding rollout Train the built-in MTP head jointly with GRPO/SFT on Megatron and use it during rollout through SGLang NEXTN speculative decoding. Online weight synchronization updates both the target model and the built-in MTP draft runner, ensuring that speculative decoding uses the latest RL-trained MTP weights instead of a stale initialization. MTP is trained through Megatron-Core's auxiliary cross-entropy path with a configurable loss scaling factor of `0.1` by default. AReaL supplies independent MTP label and loss-mask channels while keeping the main forward path logits-based. Shared output weights, backbone hidden states, and embedding inputs are detached from the MTP loss graph, so the backbone receives only the policy/SFT gradient while the MTP-specific parameters learn from future-token supervision. For packed THD training with context parallelism, AReaL applies the same per-sequence zigzag CP split and rank-local repacking to MTP labels and loss masks as it does to input IDs. Megatron-Core's CP-aware rolling then aligns future-token targets across CP ranks without crossing packed sequence boundaries. For online rollout, a focused compatibility bridge for `sglang==0.5.10.post1` receives each distributed weight bucket once and applies it to both the built-in MTP draft runner and the target runner. It supports both SGLang speculative worker layouts, handles SGLang's internal `NEXTN`-to-`EAGLE` normalization, and leaves external EAGLE draft models untouched. Draft-weight CPU backup can be enabled so the server remains available while updated weights arrive online. End-to-end validation was performed with Qwen3.5-2B on Geometry3K GRPO. The training-side MTP weights changed during optimization, and the same updated tensors were loaded into all SGLang draft runners. Key changes: - Add `enable_mtp_training` and `mtp_loss_scaling_factor` to the Megatron engine configuration. MTP training implies retaining the model's MTP layers and is incompatible with `lm_head_loss_chunk_size`. - Feed independent MTP labels and loss masks through Megatron forward passes while preserving the main logits-based loss path. - Patch Megatron-Core `GPTModel` and Megatron-Bridge `Qwen3VLGPTModel` forwarding so Qwen3.5 text and multimodal batches can train MTP. - Align MTP labels and masks with padded or packed execution layouts and prevent targets from crossing sequence or padding boundaries. - Support MTP training with `CP > 1` for wrapper-owned packed THD by applying the same per-sequence zigzag split and rank-local repacking to input IDs, MTP labels, and MTP loss masks. - Reuse Megatron-Core's packed, CP-aware rolling semantics to align future-token supervision across CP ranks. - Isolate MTP gradients from shared output weights, embeddings, and backbone hidden states. - Report the auxiliary `mtp_loss` in training statistics. - Add SGLang speculative-decoding configuration passthrough for `NEXTN`, speculative steps, EAGLE top-k, draft-token count, external draft-model path, and draft-weight CPU backup. - Add an SGLang distributed weight-update bridge that updates both target and built-in MTP draft runners from the same received tensors. - Support both SGLang Spec v1 and Spec v2 draft-runner layouts, while failing fast for unsupported SGLang versions, missing draft runners, unsupported load formats, and inference pipeline parallelism. - Record `rollout/spec_accept_rate` and `rollout/spec_accept_length` from SGLang response metadata. - Add a Qwen3.5-2B Geometry3K GRPO example with MTP training and NEXTN rollout enabled. - Add unit coverage for padded and packed MTP label/mask layouts, CP zigzag alignment, multimodal forwarding, NEXTN/EAGLE routing, Spec v1/v2 compatibility, and draft/target online weight updates. - Regenerate the English and Chinese CLI reference documentation. Current limitations: - MTP training with `CP > 1` is supported only for wrapper-owned packed THD. Padded BSHD, VLM, and model-owned THD execution still require `CP=1`. - The current gradient-isolation implementation supports a single MTP prediction layer. Multi-layer MTP gradient propagation is not supported. - The SGLang distributed MTP update bridge requires inference pipeline parallel size `1`. - The compatibility bridge is intentionally pinned to `sglang==0.5.10.post1` because it relies on version-specific internal weight-update and draft-runner APIs. * test: use HttpGenerationResult in rollout version race test * fix(engine): align MTP masks and detach untied output weights Keep next-token-aligned masks unchanged before MCore's per-layer roll, and detach internal output-layer weights for untied models. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(engine): unify Megatron main and auxiliary loss scaling Let MCore apply the optimizer loss scale to both the main backward path and separately seeded MTP/MoE auxiliary gradients. This prevents FP16 optimizer unscaling from suppressing auxiliary updates. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(models): reject unsupported multilayer MTP training Fail before model construction when MTP training requests more than one prediction layer, whose gradients are not fully supported yet. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: huaqingyuan <huaqingyuan@tencent.com> Co-authored-by: Cursor <cursoragent@cursor.com> | 6 天前 | |
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> | 18 天前 | |
feat(infra): add n_gpus_per_node abstract property to Scheduler API (#1275) Move n_gpus_per_node from GatewayControllerConfig into the Scheduler base class so that any consumer can query GPU topology without depending on a specific config dataclass. Key changes: - Add abstract property to Scheduler base class - Implement in LocalScheduler, RayScheduler, SlurmScheduler - Update GatewayInferenceController to read from scheduler | 4 个月前 | |
fix(ppo): track trajectory truncation explicitly (#1637) Dynamic padding width identifies the longest trajectory in a batch, not whether generation exhausted its token budget. Propagate the inference stop reason and use it consistently for reward masking, GAE bootstrapping, and truncation metrics. Select bootstrap values from each trajectory's final valid token for both token- and turn-level GAE, filter the metadata at model boundaries, and retain the legacy heuristic for custom workflows that do not provide it. Signed-off-by: Bo Yang <yb550079@antgroup.com> | 15 天前 |
| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
| 4 个月前 | ||
| 1 个月前 | ||
| 6 天前 | ||
| 1 个月前 | ||
| 6 天前 | ||
| 18 天前 | ||
| 4 个月前 | ||
| 15 天前 |