From a5764b1ab1374be7055746113dfb2b88371bd2ef Mon Sep 17 00:00:00 2001 From: wkcc Date: Wed, 2 Sep 2026 11:09:28 +0800 Subject: [PATCH 1/8] feat(musa): add Moore Threads MUSA platform and engine support --- README.md | 2 + docs/user_guide_musa/README.md | 63 +++++++++ docs/user_guide_musa/install_guidance.md | 77 ++++++++++ docs/user_guide_musa/quick_start.md | 100 +++++++++++++ tests/test_plugin_registration.py | 50 +++++++ verl_hardware_plugin/engines/__init__.py | 19 +++ verl_hardware_plugin/engines/fsdp_musa.py | 56 ++++++++ verl_hardware_plugin/engines/megatron_musa.py | 28 ++++ verl_hardware_plugin/platforms/__init__.py | 10 ++ .../platforms/platform_musa.py | 131 ++++++++++++++++++ 10 files changed, 536 insertions(+) create mode 100644 docs/user_guide_musa/README.md create mode 100644 docs/user_guide_musa/install_guidance.md create mode 100644 docs/user_guide_musa/quick_start.md create mode 100644 verl_hardware_plugin/engines/fsdp_musa.py create mode 100644 verl_hardware_plugin/engines/megatron_musa.py create mode 100644 verl_hardware_plugin/platforms/platform_musa.py diff --git a/README.md b/README.md index 5f8033f..0c31a64 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,7 @@ The platforms and engines in this repository are **reference implementations** | Enflame GCU | GCU | ECCL / FlagCX | ✅ Example (requires vendor support) | [User Guide](docs/user_guide_enflame/README.md) | | Huawei NPU | Ascend 910B | HCCL | Built-in (verl core) | [Ascend Tutorial](https://github.com/verl-project/verl/tree/main/docs/ascend_tutorial) | | Iluvatar | BI-V150 (CUDA-compatible) | IXCCL | ✅ Supported | [User Guide](docs/user_guide_iluvatar/README.md) | +| Moore Threads | S5000 (CUDA-compatible) | MCCL | ✅ Supported | [User Guide](docs/user_guide_musa/README.md) | ## Installation @@ -89,6 +90,7 @@ Each hardware platform provides a standalone user guide (following the structure - **[MetaX GPU](docs/user_guide_metax/README.md)** — MetaX GPU user guide - **[FlagOS](docs/user_guide_flagos/README.md)** — FlagOS unified heterogeneous engine user guide ([NVIDIA](docs/user_guide_flagos/nvidia/README.md)) - **[Enflame GCU](docs/user_guide_enflame/README.md)** — Enflame GCU user guide +- **[Moore Threads GPU](docs/user_guide_musa/README.md)** — Moore Threads GPU user guide ### Developer Guides diff --git a/docs/user_guide_musa/README.md b/docs/user_guide_musa/README.md new file mode 100644 index 0000000..574564b --- /dev/null +++ b/docs/user_guide_musa/README.md @@ -0,0 +1,63 @@ +# VERL MUSA User Guide + +## Introduction + +This document describes how to use verl for reinforcement learning training on +Moore Threads MUSA accelerators. + +## Directory Structure + +```text +verl_hardware_plugin/ +├── engines +│ ├── fsdp_musa.py # FSDP engine support +│ └── megatron_musa.py # Megatron engine support +└── platforms + └── platform_musa.py # MUSA platform settings +``` + +```text +user_guide_musa/ +├── README.md # This file +├── install_guidance.md # Installation and environment setup +└── quick_start.md # GSM8K GRPO quick start +``` + +## Getting Started + +- [Installation Guide](./install_guidance.md) — prerequisites and environment setup +- [Quick Start](./quick_start.md) — run a GSM8K GRPO training job + +## Platform Summary + +| Item | Description | +|------|-------------| +| Device type | `musa` | +| Vendor identifier | `moore_threads` | +| Communication backend | `mccl` | +| Device visibility env var | `MUSA_VISIBLE_DEVICES` | +| Ray resource name | `GPU` | +| IPC support | Yes | + +## MUSA Migration Patches + +```bash +export VERL_PLATFORM=musa +export RAY_EXPERIMENTAL_NOSET_MUSA_VISIBLE_DEVICES=1 +``` + +MUSA deployments may use two separate compatibility layers: + +- The official Megatron code is adapted through the external + `megatron-lm-musa-patch`, selected with `MUSA_PATCH_PATH` (usually + `/home/megatron-lm-musa-patch` in the release image). +- VERL and SGLang runtime components may use an internal + `verl-musa-patch` overlay supplied by the deployment environment. Set + `VERL_MUSA_PATCH` to its path and add it to the workers' `PYTHONPATH` + (usually `/home/verl-musa-patch/` in the internal release image). This + overlay is not part of the public `verl-hardware-plugin` package. + +For Megatron/MCore training, add the required external patch directories to +`PYTHONPATH`. FSDP training does not require `MUSA_PATCH_PATH`; if the internal +runtime overlay is used, pass `VERL_MUSA_PATCH` to every Ray worker. See the +[Installation Guide](./install_guidance.md) for the complete setup. diff --git a/docs/user_guide_musa/install_guidance.md b/docs/user_guide_musa/install_guidance.md new file mode 100644 index 0000000..099484e --- /dev/null +++ b/docs/user_guide_musa/install_guidance.md @@ -0,0 +1,77 @@ +# MUSA Installation Guide + +## Prerequisites + +- A MUSA Docker image with the matching driver/runtime, `torch_musa`, MCCL, + SGLang, and other MUSA dependencies. +- Network access to download models and datasets. +- A VERL checkout and this plugin checkout. + +The standard MUSA images already include SGLang and the external Megatron-LM +MUSA patch (usually `/home/megatron-lm-musa-patch`). Other runtime dependencies +such as Ray are also normally pre-installed. Do not add CUDA versions of these +packages, as they may override the MUSA packages. + +## 1. Start the MUSA Docker Image + +Use the MUSA release image provided for your hardware. The exact image name and +device mounts depend on the driver release; the following is a generic example: + +```bash +docker_image="${MUSA_DOCKER_IMAGE:-}" +docker_name="${MUSA_DOCKER_NAME:-verl_musa}" + +docker container create \ + --name "${docker_name}" \ + --privileged \ + --net host \ + --pid=host \ + --shm-size 100g \ + --ulimit memlock=-1 \ + -v /home:/home \ + -it \ + "${docker_image}" \ + /bin/bash + +docker start -ai "${docker_name}" +``` + +Inside the container, verify that the pre-installed components are available: + +```bash +ls /home +python3 -c 'import torch; import sglang; print(torch.musa.is_available())' +``` + +## 2. Install verl and verl-hardware-plugin + +```bash +# Install verl +git clone https://github.com/verl-project/verl.git +cd verl +pip install -e . + +# Install verl-hardware-plugin +git clone https://github.com/verl-project/verl-hardware-plugin.git +cd verl-hardware-plugin +pip install -e . +``` + +## 3. Prepare Data and Models + +The baseline scripts use Qwen3-0.6B and GSM8K. Set `MODEL_DIR` and `DATA_DIR` +to the paths available in your environment, for example: + +```text +MODEL_DIR=/ipfs/models/Qwen/Qwen3-0.6B +DATA_DIR=/ipfs/models/gsm8k +``` + +## 4. Verify the Environment + +```bash +python3 -c 'import torch; print(torch.musa.is_available(), torch.musa.device_count())' +``` + +The output should show that MUSA is available and report the visible device +count. Then follow the [Quick Start](./quick_start.md) to run a VERL script. diff --git a/docs/user_guide_musa/quick_start.md b/docs/user_guide_musa/quick_start.md new file mode 100644 index 0000000..44f36cc --- /dev/null +++ b/docs/user_guide_musa/quick_start.md @@ -0,0 +1,100 @@ +# MUSA Quick Start + +This guide walks you through the GSM8K GRPO baseline on Moore Threads MUSA. +Complete the [Installation Guide](./install_guidance.md) first. + +**Baseline scenario:** Qwen3-0.6B + GSM8K + FSDP actor + SGLang rollout — see +[`scripts/baseline_grpo_gsm8k.sh`](../../scripts/baseline_grpo_gsm8k.sh). + +## 1. Prepare Data and Model + +The MUSA image normally provides the runtime dependencies. Set the model and +dataset directories to paths available in your environment: + +```bash +MODEL_DIR=/ipfs/models/Qwen/Qwen3-0.6B +DATA_DIR=/ipfs/models/gsm8k +``` + +## 2. Run the Baseline + +From the repository root: + +```bash + +export VERL_PLATFORM=musa +export VERL_USE_EXTERNAL_MODULES=verl_hardware_plugin +export VERL_MUSA_PATCH=/home/verl-musa-patch +export RAY_EXPERIMENTAL_NOSET_MUSA_VISIBLE_DEVICES=1 +export MUSA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 +export RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0 +export MCCL_LIB=/usr/local/musa/lib/libmccl.so +export LD_LIBRARY_PATH="/usr/local/musa/lib:${LD_LIBRARY_PATH:-}" +export VLLM_PATCH_MUSA_CUSTOM_OPS=1 +export SGLANG_MUSA_GRAPH_COMPAT=1 +export PYTHONPATH="${VERL_MUSA_PATCH}:${PYTHONPATH:-}" + +export INFER_BACKEND=sglang +export DATA_DIR=/ipfs/models/gsm8k +export MODEL_DIR=/ipfs/models/Qwen/Qwen3-0.6B + +exec bash "scripts/baseline_grpo_gsm8k.sh" \ + "+ray_kwargs.ray_init.runtime_env.env_vars.VERL_PLATFORM='musa'" \ + "+ray_kwargs.ray_init.runtime_env.env_vars.VERL_MUSA_PATCH='${VERL_MUSA_PATCH}'" \ + "+ray_kwargs.ray_init.runtime_env.env_vars.PYTHONPATH='${VERL_MUSA_PATCH}:${PYTHONPATH:-}'" \ + trainer.device=musa \ + +actor_rollout_ref.rollout.engine_kwargs.sglang.device=musa \ + +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=fa3 \ + +actor_rollout_ref.rollout.engine_kwargs.sglang.disable_piecewise_cuda_graph=True \ + "$@" + +``` + +The script passes the platform settings, `verl-musa-patch`, SGLang options, and +device environment to Ray workers through `runtime_env`. Shell exports alone +are not sufficient for Ray workers. + + + +## 3. Compare Results + +Compare `critic/rewards/mean` with the [NVIDIA reference run](https://swanlab.cn/@heavyrain/verl_grpo_gsm8k_math/runs/8h196r8o/chart). + +The baseline should: + +1. Complete all epochs without a crash or hang. +2. Show an upward reward trend within the first 20 steps. +3. Avoid a flat or collapsing reward curve during the first 100 steps. + +## 4. Quick Verification + +```bash +python3 -c 'import torch; print(torch.musa.is_available(), torch.musa.device_count())' +``` + +The output should show that MUSA is available and report the visible device +count. The logs should also contain `[VERL_MUSA_SITE]` bootstrap messages. + + +## Multi-Node Setup + +Start Ray on the head node and workers, then set `NNODES` and run the baseline: + +```bash +# Head node +ray start --head --port=6379 +export RAY_ADDRESS='auto' + +# Worker nodes +ray start --address=':6379' + +NNODES=2 bash scripts/baseline_grpo_gsm8k.sh +``` + +MUSA uses Ray's built-in `GPU` resource. Do not configure a custom `musa` +resource. + +## Next Steps + +- See [Installation Guide](./install_guidance.md) for image and dependency setup. +- See [development.md — Acceptance Baseline](../development.md#acceptance-baseline-for-new-hardware-adaptation) for the PR checklist. diff --git a/tests/test_plugin_registration.py b/tests/test_plugin_registration.py index ce3cbbd..97a10ff 100644 --- a/tests/test_plugin_registration.py +++ b/tests/test_plugin_registration.py @@ -58,6 +58,14 @@ def test_iluvatar_registered(self): cls = PlatformRegistry.get("iluvatar") assert cls is PlatformIluvatar + def test_musa_registered(self): + from verl.plugin.platform.platform_manager import PlatformRegistry + from verl_hardware_plugin.platforms.platform_musa import PlatformMUSA # noqa: F401 + + assert "musa" in PlatformRegistry.registered_names() + cls = PlatformRegistry.get("musa") + assert cls is PlatformMUSA + def test_xpu_detection_with_env(self): from verl.plugin.platform.platform_manager import _detect_platform_name from verl_hardware_plugin.platforms.platform_xpu import PlatformXPU # noqa: F401 @@ -140,6 +148,22 @@ def test_iluvatar_detection_with_env(self): with mock.patch.dict(os.environ, {"VERL_PLATFORM": "iluvatar"}): assert _detect_platform_name() == "iluvatar" + def test_musa_detection_with_env(self): + from verl.plugin.platform.platform_manager import _detect_platform_name + from verl_hardware_plugin.platforms.platform_musa import PlatformMUSA # noqa: F401 + + with _fresh_registries(): + with mock.patch.dict(os.environ, {"VERL_PLATFORM": "musa"}): + assert _detect_platform_name() == "musa" + + def test_musa_device_and_vendor_names(self): + from verl_hardware_plugin.platforms.platform_musa import PlatformMUSA + + platform = PlatformMUSA() + assert platform.device_name == "musa" + assert platform.vendor_name == "moore_threads" + assert platform.communication_backend_name() == "mccl" + class TestEngineRegistration: """Verify that engine classes register correctly.""" @@ -232,6 +256,32 @@ def test_megatron_iluvatar_engine_registered(self): is MegatronIluvatarEngineWithLMHead ) + def test_megatron_musa_engine_registered(self): + from verl.workers.engine.base import EngineRegistry + from verl_hardware_plugin.engines.megatron_musa import MegatronMUSAEngineWithLMHead + + assert ( + EngineRegistry._engines["language_model"]["megatron"][("musa", "moore_threads")] + is MegatronMUSAEngineWithLMHead + ) + + def test_fsdp_musa_engines_registered(self): + from verl.workers.engine.base import EngineRegistry + from verl_hardware_plugin.engines.fsdp_musa import ( + FSDPMUSAEngineWithLMHead, + FSDPMUSAEngineWithValueHead, + ) + + for backend in ("fsdp", "fsdp2"): + assert ( + EngineRegistry._engines["language_model"][backend][("musa", "moore_threads")] + is FSDPMUSAEngineWithLMHead + ) + assert ( + EngineRegistry._engines["value_model"][backend][("musa", "moore_threads")] + is FSDPMUSAEngineWithValueHead + ) + def test_fsdp_enflame_engines_registered(self): from verl.workers.engine.base import EngineRegistry from verl_hardware_plugin.engines.fsdp_enflame import ( diff --git a/verl_hardware_plugin/engines/__init__.py b/verl_hardware_plugin/engines/__init__.py index d8eafb2..39e997c 100755 --- a/verl_hardware_plugin/engines/__init__.py +++ b/verl_hardware_plugin/engines/__init__.py @@ -134,3 +134,22 @@ def register_all_engines(): logger.info("Registered engines: megatron_iluvatar") except Exception as e: logger.debug("Iluvatar Megatron engines not registered: %s", e) + + # Moore Threads MUSA engines (MCCL communication). + try: + from verl_hardware_plugin.engines import fsdp_musa # noqa: F401 + + logger.info("Registered engines: fsdp_musa") + except Exception as e: + logger.debug("MUSA FSDP engines not registered: %s", e) + if os.getenv("VERL_PLATFORM", "").strip().lower() == "musa": + raise + + try: + from verl_hardware_plugin.engines import megatron_musa # noqa: F401 + + logger.info("Registered engines: megatron_musa") + except Exception as e: + logger.debug("MUSA Megatron engine not registered: %s", e) + if os.getenv("VERL_PLATFORM", "").strip().lower() == "musa": + raise diff --git a/verl_hardware_plugin/engines/fsdp_musa.py b/verl_hardware_plugin/engines/fsdp_musa.py new file mode 100644 index 0000000..885d15f --- /dev/null +++ b/verl_hardware_plugin/engines/fsdp_musa.py @@ -0,0 +1,56 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""FSDP engines for Moore Threads MUSA devices.""" + +import logging +import os + +from verl.trainer.config import CheckpointConfig +from verl.workers.config import FSDPEngineConfig, FSDPOptimizerConfig, HFModelConfig +from verl.workers.engine.base import EngineRegistry +from verl.workers.engine.fsdp import FSDPEngineWithLMHead +from verl.workers.engine.fsdp.transformer_impl import FSDPEngineWithValueHead + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +@EngineRegistry.register( + model_type="language_model", + backend=["fsdp", "fsdp2"], + device="musa", + vendor="moore_threads", +) +class FSDPMUSAEngineWithLMHead(FSDPEngineWithLMHead): + """FSDP language-model engine for MUSA with MCCL communication.""" + + def __init__( + self, + model_config: HFModelConfig, + engine_config: FSDPEngineConfig, + optimizer_config: FSDPOptimizerConfig, + checkpoint_config: CheckpointConfig, + ): + super().__init__(model_config, engine_config, optimizer_config, checkpoint_config) + logger.info("FSDPMUSAEngineWithLMHead initialized") + + +@EngineRegistry.register( + model_type="value_model", + backend=["fsdp", "fsdp2"], + device="musa", + vendor="moore_threads", +) +class FSDPMUSAEngineWithValueHead(FSDPEngineWithValueHead): + """FSDP value-model engine for MUSA with MCCL communication.""" + + def __init__( + self, + model_config: HFModelConfig, + engine_config: FSDPEngineConfig, + optimizer_config: FSDPOptimizerConfig, + checkpoint_config: CheckpointConfig, + ): + super().__init__(model_config, engine_config, optimizer_config, checkpoint_config) + logger.info("FSDPMUSAEngineWithValueHead initialized") diff --git a/verl_hardware_plugin/engines/megatron_musa.py b/verl_hardware_plugin/engines/megatron_musa.py new file mode 100644 index 0000000..009fc22 --- /dev/null +++ b/verl_hardware_plugin/engines/megatron_musa.py @@ -0,0 +1,28 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Megatron engine registration for Moore Threads MUSA.""" + +import logging + +from verl.workers.engine.base import EngineRegistry + +logger = logging.getLogger(__name__) + +from verl.workers.engine.megatron.transformer_impl import ( # noqa: E402 + MegatronEngineWithLMHead, +) + + +@EngineRegistry.register( + model_type="language_model", + backend="megatron", + device="musa", + vendor="moore_threads", +) +class MegatronMUSAEngineWithLMHead(MegatronEngineWithLMHead): + """Megatron engine registration for Moore Threads MUSA. + + MUSA-specific lifecycle hooks are installed by ``verl_musa_patch`` before + this registry module is imported. + """ diff --git a/verl_hardware_plugin/platforms/__init__.py b/verl_hardware_plugin/platforms/__init__.py index 5e16c8a..baa9ae6 100644 --- a/verl_hardware_plugin/platforms/__init__.py +++ b/verl_hardware_plugin/platforms/__init__.py @@ -78,3 +78,13 @@ def register_all_platforms(): logger.info("Registered platform: iluvatar (cuda)") except Exception as e: logger.debug("Iluvatar platform not registered: %s", e) + + # Moore Threads MUSA — requires torch_musa and the MUSA runtime. + try: + from verl_hardware_plugin.platforms import platform_musa # noqa: F401 + + logger.info("Registered platform: moore (musa)") + except Exception as e: + logger.debug("MUSA platform not registered: %s", e) + if os.getenv("VERL_PLATFORM", "").strip().lower() == "musa": + raise diff --git a/verl_hardware_plugin/platforms/platform_musa.py b/verl_hardware_plugin/platforms/platform_musa.py new file mode 100644 index 0000000..e17c9e7 --- /dev/null +++ b/verl_hardware_plugin/platforms/platform_musa.py @@ -0,0 +1,131 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Moore Threads MUSA platform implementation.""" + +import logging +from contextlib import contextmanager +from types import ModuleType +from typing import Any, Optional + +import torch + +from verl.plugin.platform.platform_base import PlatformBase +from verl.plugin.platform.platform_manager import PlatformRegistry + +logger = logging.getLogger(__name__) + + +def _ensure_torch_musa() -> bool: + if hasattr(torch, "musa"): + return True + try: + import torch_musa # noqa: F401 + + return hasattr(torch, "musa") + except Exception as exc: + logger.debug("torch.musa is unavailable: %s", exc) + return False + + +@PlatformRegistry.register(platform="musa") +class PlatformMUSA(PlatformBase): + """Platform backend for Moore Threads accelerators.""" + + @property + def device_name(self) -> str: + return "musa" + + @property + def vendor_name(self) -> str: + return "moore_threads" + + @property + def device_module(self) -> ModuleType: + if not _ensure_torch_musa(): + raise RuntimeError("torch_musa is not installed or torch.musa is unavailable") + return torch.musa + + def is_available(self) -> bool: + return _ensure_torch_musa() and torch.musa.is_available() + + def is_platform_available(self, use_smi_check: bool = False) -> bool: + # The explicit VERL_PLATFORM=musa path is needed on Ray CPU actors, + # where torch.musa.is_available() can be false before device binding. + if not _ensure_torch_musa(): + return False + return True if use_smi_check else torch.musa.is_available() + + def current_device(self) -> int: + return torch.musa.current_device() + + def device_count(self) -> int: + return torch.musa.device_count() + + def set_device(self, device_index: int) -> None: + torch.musa.set_device(device_index) + + def synchronize(self, device_index: Optional[int] = None) -> None: + if device_index is None: + torch.musa.synchronize() + else: + torch.musa.synchronize(device_index) + + def manual_seed(self, seed: int) -> None: + torch.musa.manual_seed(seed) + + def manual_seed_all(self, seed: int) -> None: + torch.musa.manual_seed_all(seed) + + def set_allocator_settings(self, settings: str) -> None: + torch.musa.memory._set_allocator_settings(settings) + + def empty_cache(self) -> None: + torch.musa.empty_cache() + + def get_device_capability(self, device_index: int = 0) -> tuple[Optional[int], Optional[int]]: + if not self.is_available() or not hasattr(torch.musa, "get_device_capability"): + return None, None + result = torch.musa.get_device_capability(device_index) + return result if result is not None else (None, None) + + def communication_backend_name(self) -> str: + return "mccl" + + def visible_devices_envvar(self) -> str: + return "MUSA_VISIBLE_DEVICES" + + def ray_resource_name(self) -> str: + return "GPU" + + def ray_resource_options(self, num_gpus: float) -> dict[str, Any]: + return {"num_gpus": num_gpus} + + def ray_noset_envvars(self) -> list[str]: + return ["RAY_EXPERIMENTAL_NOSET_MUSA_VISIBLE_DEVICES"] + + def is_ipc_supported(self) -> bool: + return True + + @contextmanager + def nvtx_range(self, msg: str): + nvtx = getattr(torch.musa, "nvtx", None) + range_fn = getattr(nvtx, "range", None) + if range_fn is None: + yield + else: + with range_fn(msg): + yield + + def profiler_start(self) -> None: + start = getattr(getattr(torch.musa, "profiler", None), "start", None) + if start is not None: + start() + + def profiler_stop(self) -> None: + stop = getattr(getattr(torch.musa, "profiler", None), "stop", None) + if stop is not None: + stop() + + def cudart(self) -> Any: + return None From b91c18331878d9cc779d0240208d6789ce017ef6 Mon Sep 17 00:00:00 2001 From: wkcc Date: Wed, 2 Sep 2026 11:46:01 +0800 Subject: [PATCH 2/8] docs(musa): clarify platform and patch documentation --- README.md | 2 +- docs/user_guide_musa/README.md | 32 +++++++++++++------------------- 2 files changed, 14 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index 0c31a64..2891463 100644 --- a/README.md +++ b/README.md @@ -26,7 +26,7 @@ The platforms and engines in this repository are **reference implementations** | Enflame GCU | GCU | ECCL / FlagCX | ✅ Example (requires vendor support) | [User Guide](docs/user_guide_enflame/README.md) | | Huawei NPU | Ascend 910B | HCCL | Built-in (verl core) | [Ascend Tutorial](https://github.com/verl-project/verl/tree/main/docs/ascend_tutorial) | | Iluvatar | BI-V150 (CUDA-compatible) | IXCCL | ✅ Supported | [User Guide](docs/user_guide_iluvatar/README.md) | -| Moore Threads | S5000 (CUDA-compatible) | MCCL | ✅ Supported | [User Guide](docs/user_guide_musa/README.md) | +| Moore Threads | MUSA (CUDA-compatible) | MCCL | ✅ Supported | [User Guide](docs/user_guide_musa/README.md) | ## Installation diff --git a/docs/user_guide_musa/README.md b/docs/user_guide_musa/README.md index 574564b..fb8421b 100644 --- a/docs/user_guide_musa/README.md +++ b/docs/user_guide_musa/README.md @@ -33,7 +33,7 @@ user_guide_musa/ | Item | Description | |------|-------------| | Device type | `musa` | -| Vendor identifier | `moore_threads` | +| Vendor identifier | `Moore Threads` | | Communication backend | `mccl` | | Device visibility env var | `MUSA_VISIBLE_DEVICES` | | Ray resource name | `GPU` | @@ -41,23 +41,17 @@ user_guide_musa/ ## MUSA Migration Patches -```bash -export VERL_PLATFORM=musa -export RAY_EXPERIMENTAL_NOSET_MUSA_VISIBLE_DEVICES=1 -``` - MUSA deployments may use two separate compatibility layers: -- The official Megatron code is adapted through the external - `megatron-lm-musa-patch`, selected with `MUSA_PATCH_PATH` (usually - `/home/megatron-lm-musa-patch` in the release image). -- VERL and SGLang runtime components may use an internal - `verl-musa-patch` overlay supplied by the deployment environment. Set - `VERL_MUSA_PATCH` to its path and add it to the workers' `PYTHONPATH` - (usually `/home/verl-musa-patch/` in the internal release image). This - overlay is not part of the public `verl-hardware-plugin` package. - -For Megatron/MCore training, add the required external patch directories to -`PYTHONPATH`. FSDP training does not require `MUSA_PATCH_PATH`; if the internal -runtime overlay is used, pass `VERL_MUSA_PATCH` to every Ray worker. See the -[Installation Guide](./install_guidance.md) for the complete setup. +- MUSA support for the upstream Megatron/MCore implementation is provided by + the external `megatron-lm-musa-patch` compatibility layer. The patch is loaded + at runtime from the directory specified by `MUSA_PATCH_PATH` (usually + `/home/megatron-lm-musa-patch` in the release image); it adapts the + unmodified Megatron code for MUSA execution. + +- MUSA compatibility for VERL and SGLang runtime components is provided by the + deployment-specific `verl-musa-patch` compatibility layer. The patch is loaded + at runtime from the directory specified by `VERL_MUSA_PATCH` (usually + `/home/verl-musa-patch` in the release image) and made available to Ray workers + through `PYTHONPATH`; it adapts the VERL and SGLang runtime components for MUSA + execution. \ No newline at end of file From 31aac3e28ff9471704c1e74c92e40a6d5d98dbd6 Mon Sep 17 00:00:00 2001 From: wkcc Date: Wed, 2 Sep 2026 11:51:39 +0800 Subject: [PATCH 3/8] docs(musa): align vendor identifier with registry key --- docs/user_guide_musa/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/user_guide_musa/README.md b/docs/user_guide_musa/README.md index fb8421b..fe9b002 100644 --- a/docs/user_guide_musa/README.md +++ b/docs/user_guide_musa/README.md @@ -33,7 +33,7 @@ user_guide_musa/ | Item | Description | |------|-------------| | Device type | `musa` | -| Vendor identifier | `Moore Threads` | +| Vendor identifier | `moore_threads` | | Communication backend | `mccl` | | Device visibility env var | `MUSA_VISIBLE_DEVICES` | | Ray resource name | `GPU` | @@ -54,4 +54,4 @@ MUSA deployments may use two separate compatibility layers: at runtime from the directory specified by `VERL_MUSA_PATCH` (usually `/home/verl-musa-patch` in the release image) and made available to Ray workers through `PYTHONPATH`; it adapts the VERL and SGLang runtime components for MUSA - execution. \ No newline at end of file + execution. From 29441a9fbf66419cdf45696c291bd3882fa92b95 Mon Sep 17 00:00:00 2001 From: wukecong-mt Date: Thu, 3 Sep 2026 12:11:10 +0800 Subject: [PATCH 4/8] fix(musa): address review feedback for platform integration Address review feedback: - Align platform probing with other backends - Make allocator configuration non-fatal - Align the vendor identifier - Register the Megatron value-model engine --- tests/test_plugin_registration.py | 10 ++++++-- verl_hardware_plugin/engines/megatron_musa.py | 25 ++++++++++++++++--- verl_hardware_plugin/platforms/__init__.py | 2 +- .../platforms/platform_musa.py | 11 ++++---- 4 files changed, 35 insertions(+), 13 deletions(-) diff --git a/tests/test_plugin_registration.py b/tests/test_plugin_registration.py index 97a10ff..fda2306 100644 --- a/tests/test_plugin_registration.py +++ b/tests/test_plugin_registration.py @@ -164,7 +164,6 @@ def test_musa_device_and_vendor_names(self): assert platform.vendor_name == "moore_threads" assert platform.communication_backend_name() == "mccl" - class TestEngineRegistration: """Verify that engine classes register correctly.""" @@ -258,12 +257,19 @@ def test_megatron_iluvatar_engine_registered(self): def test_megatron_musa_engine_registered(self): from verl.workers.engine.base import EngineRegistry - from verl_hardware_plugin.engines.megatron_musa import MegatronMUSAEngineWithLMHead + from verl_hardware_plugin.engines.megatron_musa import ( + MegatronMUSAEngineWithLMHead, + MegatronMUSAEngineWithValueHead, + ) assert ( EngineRegistry._engines["language_model"]["megatron"][("musa", "moore_threads")] is MegatronMUSAEngineWithLMHead ) + assert ( + EngineRegistry._engines["value_model"]["megatron"][("musa", "moore_threads")] + is MegatronMUSAEngineWithValueHead + ) def test_fsdp_musa_engines_registered(self): from verl.workers.engine.base import EngineRegistry diff --git a/verl_hardware_plugin/engines/megatron_musa.py b/verl_hardware_plugin/engines/megatron_musa.py index 009fc22..b56c2c8 100644 --- a/verl_hardware_plugin/engines/megatron_musa.py +++ b/verl_hardware_plugin/engines/megatron_musa.py @@ -4,14 +4,13 @@ """Megatron engine registration for Moore Threads MUSA.""" import logging +import os from verl.workers.engine.base import EngineRegistry +from verl.workers.engine.megatron.transformer_impl import MegatronEngineWithLMHead, MegatronEngineWithValueHead logger = logging.getLogger(__name__) - -from verl.workers.engine.megatron.transformer_impl import ( # noqa: E402 - MegatronEngineWithLMHead, -) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) @EngineRegistry.register( @@ -26,3 +25,21 @@ class MegatronMUSAEngineWithLMHead(MegatronEngineWithLMHead): MUSA-specific lifecycle hooks are installed by ``verl_musa_patch`` before this registry module is imported. """ + + def initialize(self): + super().initialize() + logger.info("MegatronMUSAEngineWithLMHead initialized for MUSA") + + +@EngineRegistry.register( + model_type="value_model", + backend="megatron", + device="musa", + vendor="moore_threads", +) +class MegatronMUSAEngineWithValueHead(MegatronEngineWithValueHead): + """Megatron value-model engine registration for Moore Threads MUSA.""" + + def initialize(self): + super().initialize() + logger.info("MegatronMUSAEngineWithValueHead initialized for MUSA") diff --git a/verl_hardware_plugin/platforms/__init__.py b/verl_hardware_plugin/platforms/__init__.py index baa9ae6..8e8e69b 100644 --- a/verl_hardware_plugin/platforms/__init__.py +++ b/verl_hardware_plugin/platforms/__init__.py @@ -83,7 +83,7 @@ def register_all_platforms(): try: from verl_hardware_plugin.platforms import platform_musa # noqa: F401 - logger.info("Registered platform: moore (musa)") + logger.info("Registered platform: moore_threads (musa)") except Exception as e: logger.debug("MUSA platform not registered: %s", e) if os.getenv("VERL_PLATFORM", "").strip().lower() == "musa": diff --git a/verl_hardware_plugin/platforms/platform_musa.py b/verl_hardware_plugin/platforms/platform_musa.py index e17c9e7..fbf30e7 100644 --- a/verl_hardware_plugin/platforms/platform_musa.py +++ b/verl_hardware_plugin/platforms/platform_musa.py @@ -50,11 +50,7 @@ def is_available(self) -> bool: return _ensure_torch_musa() and torch.musa.is_available() def is_platform_available(self, use_smi_check: bool = False) -> bool: - # The explicit VERL_PLATFORM=musa path is needed on Ray CPU actors, - # where torch.musa.is_available() can be false before device binding. - if not _ensure_torch_musa(): - return False - return True if use_smi_check else torch.musa.is_available() + return _ensure_torch_musa() and torch.musa.is_available() def current_device(self) -> int: return torch.musa.current_device() @@ -78,7 +74,10 @@ def manual_seed_all(self, seed: int) -> None: torch.musa.manual_seed_all(seed) def set_allocator_settings(self, settings: str) -> None: - torch.musa.memory._set_allocator_settings(settings) + try: + torch.musa.memory._set_allocator_settings(settings) + except (AttributeError, RuntimeError): + logger.warning("torch.musa does not support _set_allocator_settings") def empty_cache(self) -> None: torch.musa.empty_cache() From 8d2a69594fd7fcbe73be84249c82ad8e04f22ef1 Mon Sep 17 00:00:00 2001 From: wukecong-mt Date: Thu, 3 Sep 2026 16:34:30 +0800 Subject: [PATCH 5/8] style: format plugin registration test --- tests/test_plugin_registration.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_plugin_registration.py b/tests/test_plugin_registration.py index fda2306..e751e79 100644 --- a/tests/test_plugin_registration.py +++ b/tests/test_plugin_registration.py @@ -164,6 +164,7 @@ def test_musa_device_and_vendor_names(self): assert platform.vendor_name == "moore_threads" assert platform.communication_backend_name() == "mccl" + class TestEngineRegistration: """Verify that engine classes register correctly.""" From b3ae2d2c3475eb060e41dc3c1123b7fb059368c6 Mon Sep 17 00:00:00 2001 From: wukecong-mt Date: Thu, 3 Sep 2026 16:40:54 +0800 Subject: [PATCH 6/8] docs(musa): pass runtime env vars in quick start --- docs/user_guide_musa/quick_start.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/user_guide_musa/quick_start.md b/docs/user_guide_musa/quick_start.md index 44f36cc..d34de62 100644 --- a/docs/user_guide_musa/quick_start.md +++ b/docs/user_guide_musa/quick_start.md @@ -40,8 +40,16 @@ export MODEL_DIR=/ipfs/models/Qwen/Qwen3-0.6B exec bash "scripts/baseline_grpo_gsm8k.sh" \ "+ray_kwargs.ray_init.runtime_env.env_vars.VERL_PLATFORM='musa'" \ + "+ray_kwargs.ray_init.runtime_env.env_vars.VERL_USE_EXTERNAL_MODULES='${VERL_USE_EXTERNAL_MODULES}'" \ "+ray_kwargs.ray_init.runtime_env.env_vars.VERL_MUSA_PATCH='${VERL_MUSA_PATCH}'" \ "+ray_kwargs.ray_init.runtime_env.env_vars.PYTHONPATH='${VERL_MUSA_PATCH}:${PYTHONPATH:-}'" \ + "+ray_kwargs.ray_init.runtime_env.env_vars.RAY_EXPERIMENTAL_NOSET_MUSA_VISIBLE_DEVICES='${RAY_EXPERIMENTAL_NOSET_MUSA_VISIBLE_DEVICES}'" \ + "+ray_kwargs.ray_init.runtime_env.env_vars.RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO='${RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO}'" \ + "+ray_kwargs.ray_init.runtime_env.env_vars.MUSA_VISIBLE_DEVICES='${MUSA_VISIBLE_DEVICES}'" \ + "+ray_kwargs.ray_init.runtime_env.env_vars.MCCL_LIB='${MCCL_LIB}'" \ + "+ray_kwargs.ray_init.runtime_env.env_vars.LD_LIBRARY_PATH='${LD_LIBRARY_PATH}'" \ + "+ray_kwargs.ray_init.runtime_env.env_vars.VLLM_PATCH_MUSA_CUSTOM_OPS='${VLLM_PATCH_MUSA_CUSTOM_OPS}'" \ + "+ray_kwargs.ray_init.runtime_env.env_vars.SGLANG_MUSA_GRAPH_COMPAT='${SGLANG_MUSA_GRAPH_COMPAT}'" \ trainer.device=musa \ +actor_rollout_ref.rollout.engine_kwargs.sglang.device=musa \ +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=fa3 \ From ba864767b78fecd26208da5e8bc58bed41bd604a Mon Sep 17 00:00:00 2001 From: wukecong-mt Date: Thu, 3 Sep 2026 18:19:29 +0800 Subject: [PATCH 7/8] fix(musa): make optional registration non-fatal --- verl_hardware_plugin/engines/__init__.py | 4 ---- verl_hardware_plugin/platforms/__init__.py | 2 -- 2 files changed, 6 deletions(-) diff --git a/verl_hardware_plugin/engines/__init__.py b/verl_hardware_plugin/engines/__init__.py index 39e997c..d674206 100755 --- a/verl_hardware_plugin/engines/__init__.py +++ b/verl_hardware_plugin/engines/__init__.py @@ -142,8 +142,6 @@ def register_all_engines(): logger.info("Registered engines: fsdp_musa") except Exception as e: logger.debug("MUSA FSDP engines not registered: %s", e) - if os.getenv("VERL_PLATFORM", "").strip().lower() == "musa": - raise try: from verl_hardware_plugin.engines import megatron_musa # noqa: F401 @@ -151,5 +149,3 @@ def register_all_engines(): logger.info("Registered engines: megatron_musa") except Exception as e: logger.debug("MUSA Megatron engine not registered: %s", e) - if os.getenv("VERL_PLATFORM", "").strip().lower() == "musa": - raise diff --git a/verl_hardware_plugin/platforms/__init__.py b/verl_hardware_plugin/platforms/__init__.py index 8e8e69b..6c64182 100644 --- a/verl_hardware_plugin/platforms/__init__.py +++ b/verl_hardware_plugin/platforms/__init__.py @@ -86,5 +86,3 @@ def register_all_platforms(): logger.info("Registered platform: moore_threads (musa)") except Exception as e: logger.debug("MUSA platform not registered: %s", e) - if os.getenv("VERL_PLATFORM", "").strip().lower() == "musa": - raise From 69829c0d37b2bd131de6d93b607e13ba1a525167 Mon Sep 17 00:00:00 2001 From: wukecong-mt Date: Fri, 4 Sep 2026 10:53:29 +0800 Subject: [PATCH 8/8] docs(musa): add public training image --- docs/user_guide_musa/install_guidance.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/user_guide_musa/install_guidance.md b/docs/user_guide_musa/install_guidance.md index 099484e..3886a6e 100644 --- a/docs/user_guide_musa/install_guidance.md +++ b/docs/user_guide_musa/install_guidance.md @@ -43,6 +43,10 @@ ls /home python3 -c 'import torch; import sglang; print(torch.musa.is_available())' ``` +A public image is: + +`registry.mthreads.com/mcctest/training-suite:v2.1.7.rc3-ut-verify` + ## 2. Install verl and verl-hardware-plugin ```bash