| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
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> | 4 天前 | |
feat: add Bailing V3 SWE SFT support (#1598) * feat: add Bailing V3 SWE SFT support Port the reviewed internal implementation to the public main branch while preserving newer upstream engine behavior. Key changes: - Add Bailing V3 KDA, gated MLA, and MoE model support - Add SWE SFT dataset loading and cache handling - Add focused model, loader, and dataset tests Refs: inclusionAI/AReaL#2188 Signed-off-by: chucai.dzq <chucai.dzq@alibaba-inc.com> * refactor(dataset): split SWE SFT loader into modules Separate message processing, tokenization, pipeline orchestration, and CLI code so each concern can evolve without growing a single dataset module. Signed-off-by: chucai.dzq <chucai.dzq@alibaba-inc.com> * fix(dataset): align SWE SFT with adaptive templates Port the follow-up SWE SFT fixes from swe-dev so Bailing V3 adaptive chat templates use structural assistant masks and consistent thinking modes. Signed-off-by: chucai.dzq <chucai.dzq@alibaba-inc.com> * test(dataset): restore modules after SWE loader tests Prevent collection-time stubs from replacing the real areal.dataset package for subsequent data-service tests. * refactor(engine): remove precision dump hooks Keep Bailing V3 support focused on production training behavior by removing out-of-scope routing and log-probability dump paths. * test(models): run zigzag coverage in CI Place the CP zigzag unit tests under the root test pattern used by the GCP unit-test workflow. * fix: preserve Bailing V3 HF export metadata Keep runtime model metadata valid across fast, native, direct, and in-place mbridge exports while retaining the production local-source fallback. Key changes: - snapshot and validate HF config metadata before exporters overwrite it - preserve source assets and support native mbridge finalization - forward SWE preprocessing kwargs through the remote dataset path - add config round-trip, Saver, and controller regression coverage Refs: #1598 Constraint: Preserve swe-dev Bailing V3 architecture-based bridge dispatch Confidence: high Scope-risk: moderate Not-tested: Real multi-rank Bailing V3 HF save/load canary * test(engine): initialize native save fixture correctly Use the MegatronEngine backing process-group fields so the native mbridge finalization test can exercise the real cpu_group property. Refs: #1598 Constraint: Keep the production save path unchanged Confidence: high Scope-risk: narrow Not-tested: Full suite rerun pending on GCP * refactor(dataset): remove unused SWE augmentation options Keep the public SWE SFT loader focused on the pair and trajectory behavior exercised by production recipes, without shipping disabled sampling and truncation branches. Key changes: - remove random thinking variants and ratio balancing - remove task-notification truncation and the dead trajectory-copy CLI - preserve canonical pair, pre-split, and trajectory outputs with tests Refs: #1598 Constraint: Preserve tracked swe-dev production defaults and trajectory mode Rejected: Broad SWE substring dispatch | reintroduces c84db0bbf false matches Confidence: high Scope-risk: moderate Not-tested: Real production JSONL end-to-end run * fix(dataset): harden SWE SFT preprocessing Use structural assistant masks so literal template delimiters cannot silently drop supervised tokens. Reject unsupervised rows before training and pass data-worker cache topology explicitly. Key changes: - Make Qwen and Bailing template patches idempotent and fail closed - Filter malformed and all-zero masks, including pre-tokenized data - Coordinate shared caches with explicit data-worker rank metadata - Restore the split SWE preprocessing CLI and add regression coverage Refs: #1598 * fix: harden Bailing V3 and SWE preprocessing Keep the pull request focused on Bailing V3 SWE SFT correctness while retaining cache invalidation and export metadata fixes requested in review. Generic Hugging Face checkpoint publication hardening belongs in a separate change. --------- Signed-off-by: chucai.dzq <chucai.dzq@alibaba-inc.com> Co-authored-by: 楚财 <chucai.dzq@alibaba-inc.com> | 8 天前 | |
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> | 4 天前 | |
perf: reduce duplicated VLM rollout processing and transfers (#1671) * perf(infra): add group-scoped processor cache * perf(workflow): reuse processor results across VLM rollout groups * perf(experimental): cache agent multimodal prompt processing * perf(experimental): share processor cache across proxy sessions * perf(experimental): add grouped multimodal tensor references * perf(experimental): share tensor references across proxy sessions * perf(infra): preserve shared RTensor aliases * perf(infra): preserve aliases across staged engine RPCs * perf(trainer): avoid multimodal advantage RPC transfers * test(experimental): complete multimodal export fixtures * fix(experimental): prevent shared tensor refs from surviving store reuse * fix(infra): cancel rollout siblings before group finalization * fix(infra): retain source tensors during alias-preserving remotize | 4 天前 | |
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> | 4 天前 | |
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> | 4 天前 | |
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 个月前 | |
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 | 1 个月前 | |
perf: reduce duplicated VLM rollout processing and transfers (#1671) * perf(infra): add group-scoped processor cache * perf(workflow): reuse processor results across VLM rollout groups * perf(experimental): cache agent multimodal prompt processing * perf(experimental): share processor cache across proxy sessions * perf(experimental): add grouped multimodal tensor references * perf(experimental): share tensor references across proxy sessions * perf(infra): preserve shared RTensor aliases * perf(infra): preserve aliases across staged engine RPCs * perf(trainer): avoid multimodal advantage RPC transfers * test(experimental): complete multimodal export fixtures * fix(experimental): prevent shared tensor refs from surviving store reuse * fix(infra): cancel rollout siblings before group finalization * fix(infra): retain source tensors during alias-preserving remotize | 4 天前 | |
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> | 4 天前 | |
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> | 4 天前 | |
perf: reduce duplicated VLM rollout processing and transfers (#1671) * perf(infra): add group-scoped processor cache * perf(workflow): reuse processor results across VLM rollout groups * perf(experimental): cache agent multimodal prompt processing * perf(experimental): share processor cache across proxy sessions * perf(experimental): add grouped multimodal tensor references * perf(experimental): share tensor references across proxy sessions * perf(infra): preserve shared RTensor aliases * perf(infra): preserve aliases across staged engine RPCs * perf(trainer): avoid multimodal advantage RPC transfers * test(experimental): complete multimodal export fixtures * fix(experimental): prevent shared tensor refs from surviving store reuse * fix(infra): cancel rollout siblings before group finalization * fix(infra): retain source tensors during alias-preserving remotize | 4 天前 | |
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 个月前 | |
chore: enforce license (#1171) | 4 个月前 |