feat(musa): add Moore Threads MUSA platform and engine support - #13
Conversation
There was a problem hiding this comment.
🟡 Changes recommended
Platform auto-detection logic and the new MUSA documentation snippets contain issues that can cause incorrect platform selection and failing verification commands.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR extends verl-hardware-plugin with Moore Threads MUSA support by adding a new musa platform implementation, registering MUSA-specific engines (FSDP/FSDP2 and Megatron), and documenting installation/quick-start steps alongside registration tests.
Changes:
- Add
PlatformMUSAwith MCCL as the communication backend and Ray integration defaults for MUSA. - Register MUSA engines for FSDP/FSDP2 and Megatron under
(device="musa", vendor="moore_threads"). - Add MUSA registration tests and a dedicated user guide (install + quick start), plus README platform listing.
File summaries
| File | Description |
|---|---|
| verl_hardware_plugin/platforms/platform_musa.py | Introduces the MUSA platform adapter (device/vendor IDs, comm backend, Ray/env integration). |
| verl_hardware_plugin/platforms/init.py | Adds conditional import to register the MUSA platform. |
| verl_hardware_plugin/engines/fsdp_musa.py | Registers FSDP/FSDP2 engines for MUSA for LM/value models. |
| verl_hardware_plugin/engines/megatron_musa.py | Registers the Megatron LM engine for MUSA. |
| verl_hardware_plugin/engines/init.py | Adds conditional imports to register MUSA engines. |
| tests/test_plugin_registration.py | Adds tests validating MUSA platform and engine registrations. |
| README.md | Adds Moore Threads/MUSA to the supported platforms table and guide list. |
| docs/user_guide_musa/README.md | Adds a MUSA overview guide with key identifiers and patch layering notes. |
| docs/user_guide_musa/quick_start.md | Adds a GSM8K GRPO quick-start recipe for MUSA deployments. |
| docs/user_guide_musa/install_guidance.md | Adds MUSA environment and installation guidance. |
Review details
Suppressed comments (2)
docs/user_guide_musa/install_guidance.md:73
- Same issue as the earlier snippet:
torch.musamay not exist untiltorch_musais imported, so this verification command can fail unexpectedly.
python3 -c 'import torch; print(torch.musa.is_available(), torch.musa.device_count())'
docs/user_guide_musa/quick_start.md:72
- This quick verification snippet uses
torch.musawithout importingtorch_musafirst. Iftorch.musais registered only after importingtorch_musa(as implied by_ensure_torch_musa()in the platform implementation), this command can fail withAttributeError.
python3 -c 'import torch; print(torch.musa.is_available(), torch.musa.device_count())'
- Files reviewed: 10/10 changed files
- Comments generated: 6
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
|
||
| ## 3. Compare Results | ||
|
|
||
| Compare `critic/rewards/mean` with the [NVIDIA reference run](https://swanlab.cn/@heavyrain/verl_grpo_gsm8k_math/runs/8h196r8o/chart). |
There was a problem hiding this comment.
Is there a result on MUSA?
There was a problem hiding this comment.
I put this MUSA SwanLab link in this PR summary, at the very top, and the link is: https://swanlab.cn/@wukecong/verl_grpo_gsm8k_math/runs/z4ybajqy/chart
Address review feedback: - Align platform probing with other backends - Make allocator configuration non-fatal - Align the vendor identifier - Register the Megatron value-model engine
|
please check lint and copilot comments |
Summary
Validation
critic/rewards/meanshows the expected training trend.SwanLab result:
https://swanlab.cn/@wukecong/verl_grpo_gsm8k_math/runs/z4ybajqy/chart
NVIDIA reference:
https://swanlab.cn/@heavyrain/verl_grpo_gsm8k_math/runs/8h196r8o/chart
Tests: