diff --git a/README.md b/README.md index 5f8033f..2891463 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 | MUSA (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..fe9b002 --- /dev/null +++ b/docs/user_guide_musa/README.md @@ -0,0 +1,57 @@ +# 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 + +MUSA deployments may use two separate compatibility layers: + +- 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. diff --git a/docs/user_guide_musa/install_guidance.md b/docs/user_guide_musa/install_guidance.md new file mode 100644 index 0000000..3886a6e --- /dev/null +++ b/docs/user_guide_musa/install_guidance.md @@ -0,0 +1,81 @@ +# 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())' +``` + +A public image is: + +`registry.mthreads.com/mcctest/training-suite:v2.1.7.rc3-ut-verify` + +## 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..d34de62 --- /dev/null +++ b/docs/user_guide_musa/quick_start.md @@ -0,0 +1,108 @@ +# 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_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 \ + +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..e751e79 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,39 @@ 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, + 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 + 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..d674206 100755 --- a/verl_hardware_plugin/engines/__init__.py +++ b/verl_hardware_plugin/engines/__init__.py @@ -134,3 +134,18 @@ 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) + + 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) 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..b56c2c8 --- /dev/null +++ b/verl_hardware_plugin/engines/megatron_musa.py @@ -0,0 +1,45 @@ +# Copyright (c) 2026 BAAI. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""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__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +@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. + """ + + 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 5e16c8a..6c64182 100644 --- a/verl_hardware_plugin/platforms/__init__.py +++ b/verl_hardware_plugin/platforms/__init__.py @@ -78,3 +78,11 @@ 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_threads (musa)") + except Exception as e: + logger.debug("MUSA platform not registered: %s", e) diff --git a/verl_hardware_plugin/platforms/platform_musa.py b/verl_hardware_plugin/platforms/platform_musa.py new file mode 100644 index 0000000..fbf30e7 --- /dev/null +++ b/verl_hardware_plugin/platforms/platform_musa.py @@ -0,0 +1,130 @@ +# 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: + return _ensure_torch_musa() and 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: + 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() + + 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