Repository navigation
perf: accelerate Qwen3.5 OpenR1MM training - #454
Draft
yuanlehome wants to merge 5 commits into
Draft
yuanlehome wants to merge 5 commits into
yuanlehome wants to merge 5 commits into
Conversation
Reduce colocated GRPO overhead with a guarded inline-logprob path and BF16 expert gathering, asynchronous TP/GDN gathers, packed experts and TP-local direct loading. Enable DeepEP, CUTLASS grouped GEMM and Rust sgl-router for OpenR1MM; set global batch 512 and packing 8192. Avoid the intermediate ASCII copy when decoding R3 routing data. Preserve trainer offload/reload, process-group reconstruction and full weight synchronization. Add regression coverage for the optimized paths. Validation: - Full pre-commit checks; targeted CPU tests: 259 passed, 2 skipped - Five full batch512/packing8192 cycles without OOM or allocation retries - Packing8192 numerical differences accepted after frozen-input replay
Contributor
|
Thanks for contributing to Relax, @yuanlehome! 感谢你为 Relax 做出贡献！ 📚 Contribution guide / 贡献指南Describe the problem, your changes, and how you validated them. Keep each PR focused and run 请说明问题、改动和验证方式，保持 PR 聚焦，并在提交前运行 🛠️ CI commands / CI 指令
Put one command on the first line of a new PR comment. Rerun/cancel require PR authorship or repository write access. 在新 PR 评论的首行写一条指令。PR 作者或有仓库写权限的贡献者可以重跑、取消 CI。 🔎 Merge requirements / 合入条件GitHub：📝 请转为 Ready for review
|
There was a problem hiding this comment.
🔵 Needs a closer look
An unresolved training blocker needs correction, and distributed GPU weight-transfer changes require final human review.
1 open issue
What changed in this PR
Accelerates Qwen3.5 OpenR1MM training through guarded inline log-probabilities and optimized weight synchronization.
Changes:
- Adds asynchronous gathers, packed BF16 experts, and TP-local loading.
- Tunes the launcher for batch 512, packing 8192, and DeepEP.
- Adds regression coverage and reduces routing-data decoding overhead.
| File | Description |
|---|---|
| tests/backends/megatron/weight_update/test_qwen35_packed_experts.py | Tests packed expert ordering and lifetime. |
| tests/backends/megatron/weight_update/test_qwen35_direct_experts.py | Tests direct shards and receiver guards. |
| tests/backends/megatron/weight_update/test_expert_routing.py | Updates iterator fixtures. |
| tests/backends/megatron/weight_update/test_bf16_expert_gather.py | Tests gather layouts and group replacement. |
| tests/backends/megatron/weight_update/test_async_gdn_weight_gather.py | Tests asynchronous tensor reconstruction. |
| tests/backends/megatron/test_colocate_inline_log_probs.py | Tests eligibility and zero-KL inputs. |
| scripts/training/multimodal/run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh | Enables optimized training settings. |
| relax/engine/rollout/sglang_rollout.py | Removes intermediate ASCII encoding. |
| relax/backends/megatron/weight_update/update_weight_from_tensor.py | Integrates validated direct expert loading. |
| relax/backends/megatron/weight_update/qwen35_packed_experts.py | Adds packed expert export. |
| relax/backends/megatron/weight_update/qwen35_direct_experts.py | Adds TP-local shards and typed dispatch. |
| relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py | Integrates optimized gathers. |
| relax/backends/megatron/weight_update/common.py | Preserves GDN component ordering. |
| relax/backends/megatron/weight_update/bf16_expert_gather.py | Adds owner-balanced BF16 gathering. |
| relax/backends/megatron/loss.py | Supports zero-KL shapes without actor log-probabilities. |
| relax/backends/megatron/actor.py | Guards skipping the old-policy forward pass. |
Review effort: Balanced
| "max_staleness": 0, | ||
| "update_weights_interval": 1, | ||
| "partial_rollout": False, | ||
| "rl_partial_rollouts": False, |
Contributor
There was a problem hiding this comment.
在 7c216d3 中，这个条件移到了 inline_log_probs.py:32，问题仍未解决：真实 Namespace 不含 rl_partial_rollouts（SGLang 参数统一加 sglang_ 前缀），所以第 74 行总会将其判为不匹配。现在结果是静默返回 False，更新后的 Qwen 配方仍每轮执行本应跳过的 actor 前向。
请移除这项冗余检查，并增加字段缺省的回归用例；现有 partial_rollout=False 检查已覆盖 Relax 的开关。当前测试手工设置 rl_partial_rollouts=False，掩盖了真实入口的失败。
Select expert gathering, TP-local loading and inline log-prob reuse from runtime configuration and receiver capabilities. Consolidate expert layout and IPC handling, remove experiment-only switches, and retain bounded Bridge fallbacks for unsupported layouts. Preserve the accepted batch512, packing8192 and sgl-router configuration with regression coverage for loading, fallback and process-group reuse.
Use importorskip for weight-sync tests that require Megatron Core, Megatron Bridge or SGLang. Guard only the backend-dependent inline log-prob test so its pure CPU eligibility coverage still runs. Keep the full regression coverage active in training environments with the optional backends installed.
yuanlehome
marked this pull request as ready for review
October 19, 2026 12:38
Contributor
Nyanpasu 审查看板审查状态： 💬 已完成 · 有补充意见 审查版本： 0af6420 目标分支： main 已完成 0af6420 相对新基线的 25 文件范围及增量交互审查，未确认新缺陷。原 P2（不存在的参数使自动 logprob 复用失效）尚未修改，继续保留。
审查发现待处理
已解决或已取代
提交范围 · 接收 25 · 建议移出 0 · 待确认 0接收 25 个文件 · 建议移出 0 个文件 · 待确认 0 个文件。移出与待确认部分暂停深审，不代表审查通过。
Powered by Nyanpasu with gpt-6.1-sol xhigh, please check the suggestions carefully.
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||
Derive native BF16 expert shards from receiver MoE TP, EP and DP, and support Triton, DeepGEMM and FlashInfer CUTLASS independently of the A2A dispatcher. Preserve native gate/up order and fall back for unsupported weight layouts. Use scheduler-resolved defaults, disable shared-expert fusion by default, and clear visual embeddings during cache flush. Remove redundant recipe overrides and the opt-in CUTLASS grouped-GEMM setting. Validation: all-files pre-commit; 432 tests passed, 2 environment skips; 8 multi-GPU cases passed, including DeepEP normal and low-latency. The Qwen3.5 TP2 packing/copy microbenchmark showed no regression. FlashInfer A2A initialization remains blocked by the host Linux 5.4 kernel.
Split large BF16 expert projections into owner-balanced chunks within the sync buffer budget. Negotiate SGLang expert-slice loading automatically and copy into existing parameters while preserving the small-layer path. Merge main and retain both server lifecycle and expert-sync regression tests. Keep existing offload, process-group and IPC reader lifetimes. Validation: all-files pre-commit; 877 tests passed and 18 environment-dependent tests skipped after merging main. Prior multi-GPU validation covered Triton, DeepGEMM, CUTLASS, DeepEP, cross-process CUDA IPC and repeated updates with exact reference outputs.
yuanlehome
marked this pull request as draft
October 20, 2026 09:27
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
优化
run-qwen35-35B-A3B-8xgpu-openr1mm-r3的端到端训练性能，并将权重更新优化改为根据模型布局、训练拓扑和接收端能力自动选择。f3b4246补测了 12 组单层装载组件：收益随专家数量、MoE TP/EP/DP 和 runner 布局变化，结果及适用范围见下表。组件加速不等于完整同步或训练吞吐加速。Why
Qwen3.5-35B 的 40 层、每层 256 个 gated experts 会产生 30720 个 gate/up/down 专家条目。旧路径在转换后的细粒度元数据、flatten/repack、逐专家接收装载与切片上花费较多时间。仅调整 collective 或 batch，无法消除这些固定开销。
How
w13_weight/w2_weight，避开逐专家 HF loader 和中间 flatten。[up, gate]顺序。A2A backend 不再作为权重布局白名单；组合是否可运行仍由 SGLang 校验。条件不满足时回退兼容路径，不新增环境变量开关。flush_cache清理视觉 embedding cache，删除 recipe 的 VLM cache=0 覆盖。移除无效 reference 前向，符合条件时复用训练前向 logprob；保留 DeepEP 和 packing 调整。已移除本分支的 CUTLASS grouped-GEMM opt-in 修改。已有的两 chunk IPC overlap、raw iterator 的批量异步 gather 不计为本 PR 新增收益；本次也没有消除 Gloo/Ray/HTTP 控制通道或 EP 完整权重收集。
Testing
当前提交相关测试 432 passed、2 skipped；跳过项为 CPU 环境下的 CUDA-only MXFP4 检查及安装版本缺少
LoRAMerge。既有 GPU 正确性矩阵通过 8 个场景，覆盖 Triton / DeepGEMM / CUTLASS、EP、MoETP、MoEDP 和 DeepEP normal/low-latency：60 次 rank 级权重更新、234 次输出检查、54 次 CUDA Graph 检查通过。DeepEP normal 的另外 6 次因动态接收计数未执行 CUDA Graph 检查，仅验证 eager 路径。
本次新增 12 组装载性能测量全部完成，正确性检查通过，逐位比较接收权重并检查参数地址不变。
FlashInfer A2A 在当前 Linux 5.4 主机初始化时遇到
pidfd_getfd: ENOSYS，未形成可运行性能结论。DeepEP 的正确性通过不代表已测得其完整同步加速。通用化重构曾在 batch512 / packing8192 / sgl-router 下完成冻结输入和完整周期回归。当前最终提交未重跑完整训练、长期收敛或正式健康分验收；本次补测范围是下面明确列出的装载组件。
pre-commit run --all-filespassesTests pass (
pytest tests/) — 本轮执行针对性测试，未执行整个测试目录New tests added (if applicable)
Documentation updated (if applicable)
Type of Change