Skip to content

feat(musa): add Moore Threads MUSA platform and engine support - #13

Merged
physics31415926 merged 8 commits into
verl-project:mainfrom
wukecong-mt:musa-support-pr
Sep 4, 2026
Merged

feat(musa): add Moore Threads MUSA platform and engine support#13
physics31415926 merged 8 commits into
verl-project:mainfrom
wukecong-mt:musa-support-pr

Conversation

@wukecong-mt

Copy link
Copy Markdown
Contributor

Summary

  • Add Moore Threads MUSA platform registration with MCCL communication support.
  • Add FSDP/FSDP2 and Megatron engine registrations for MUSA.
  • Add MUSA plugin registration tests.
  • Add MUSA installation and quick-start documentation.
  • Keep vendor-specific runtime patches outside the public hardware plugin.

Validation

  • MUSA 8-device GSM8K GRPO baseline completed.
  • Reward and accuracy meet the acceptance criteria.
  • critic/rewards/mean shows 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:

pytest -q tests/test_plugin_registration.py -k musa
ruff check verl_hardware_plugin tests/test_plugin_registration.py

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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 PlatformMUSA with 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.musa may not exist until torch_musa is 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.musa without importing torch_musa first. If torch.musa is registered only after importing torch_musa (as implied by _ensure_torch_musa() in the platform implementation), this command can fail with AttributeError.
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.

Comment thread verl_hardware_plugin/platforms/platform_musa.py
Comment thread verl_hardware_plugin/platforms/platform_musa.py Outdated
Comment thread docs/user_guide_musa/install_guidance.md
Comment thread docs/user_guide_musa/quick_start.md
Comment thread verl_hardware_plugin/engines/megatron_musa.py Outdated
Comment thread verl_hardware_plugin/platforms/__init__.py Outdated

## 3. Compare Results

Compare `critic/rewards/mean` with the [NVIDIA reference run](https://swanlab.cn/@heavyrain/verl_grpo_gsm8k_math/runs/8h196r8o/chart).

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a result on MUSA?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread verl_hardware_plugin/engines/__init__.py Outdated
Comment thread verl_hardware_plugin/engines/megatron_musa.py
Comment thread docs/user_guide_musa/install_guidance.md
Address review feedback:

- Align platform probing with other backends
- Make allocator configuration non-fatal
- Align the vendor identifier
- Register the Megatron value-model engine
@physics31415926

Copy link
Copy Markdown
Collaborator

please check lint and copilot comments

@physics31415926
physics31415926 merged commit 60802f1 into verl-project:main Sep 4, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants