feat: add backend-aware MUSA support

[ForAxel]

Add backend-aware MUSA support

Summary

This PR adds backend-aware MUSA (Moore Threads GPU) support to Slime while preserving the existing CUDA execution path.

The implementation is based on the MUSA adaptations required by the current Megatron-LM and SGLang integration. It centralizes device and distributed-runtime differences instead of scattering MUSA-specific conditionals throughout the training code.

Motivation

Slime currently assumes CUDA/NCCL in several training, checkpoint, profiling, Ray, SGLang, and weight-update paths. On MUSA, these assumptions cause one or more of the following problems:

• CUDA device creation and synchronization are used directly.
• NCCL is selected where MCCL is required.
• Ray-visible device IDs are resolved using CUDA_VISIBLE_DEVICES only.
• CUDA-only memory and profiler APIs are unavailable.
• Distributed weight updates use a collective protocol that does not match SGLang's MUSA receive path.
• The MUSA compatibility patch must be imported before Megatron or SGLang initialization.

What Changed

1. Backend abstraction and bootstrap

• Add slime/utils/accelerator.py with common helpers for:
  • device type and device construction;
  • current device and device selection;
  • synchronization and cache management;
  • IPC collection and memory statistics;
  • visible-device ID resolution;
  • process-group backend selection.
• Import the optional musa_patch compatibility package before Megatron/SGLang initialization in the main entry points and independently imported plugin/tool entry points.
• Only attempt to import musa_patch when a MUSA environment is detected. CUDA processes do not import or execute the MUSA patch.
• Map the default NCCL selection to MCCL for MUSA. Weight-update groups use cpu:gloo,musa:mccl so CPU metadata communication and MUSA tensor communication use the appropriate backends.

2. Ray and device visibility

• Resolve local device IDs from either CUDA_VISIBLE_DEVICES or MUSA_VISIBLE_DEVICES.
• Preserve Ray's existing CUDA/ROCm behavior and add support for preventing Ray from overriding MUSA-visible devices.
• Skip NVIDIA NVML NUMA affinity setup on MUSA.

3. Megatron and model paths

• Replace CUDA-bound temporary tensors, synchronization, IPC collection, and device transfers in Megatron actor, data, checkpoint, log-probability, quantization, and model-plugin paths.
• Use MUSA-compatible CPU initialization and process-group setup when required.
• Skip CUDA-only Qwen-VL rotary patching and SGLang FP8 helper imports in MUSA environments.
• Keep the existing CUDA dtype validation and CUDA-specific model behavior unchanged.

4. Weight synchronization

• Make distributed weight-update process groups use the selected backend rather than hard-coded NCCL.
• Preserve the existing asynchronous broadcast, synchronous broadcast, and Gloo transport modes.
• Implement the p2p-broadcast sender protocol required by the SGLang MUSA receiver: the training rank sends each tensor to engine rank 1, after which SGLang performs the intra-engine broadcast.
• Route all weight-update temporary tensors and stream synchronization through the accelerator abstraction.

5. Process-group and profiling support

• Extend reloadable process-group handling to recognize both NCCL and MCCL-backed accelerator groups.
• Make profiler activity selection and memory snapshot hooks backend-aware.
• Gracefully skip unsupported accelerator-specific memory snapshot hooks instead of failing unrelated training paths.

CUDA Compatibility

The CUDA path remains the default when no MUSA environment is requested:

• accelerator.device_type() resolves to cuda when CUDA is available and MUSA is not selected.
• Default process-group and weight-update backends remain NCCL.
• CUDA_VISIBLE_DEVICES resolution remains supported.
• CUDA profiler, memory, synchronization, and model behavior are preserved.
• musa_patch is not imported in a normal CUDA environment.

If MUSA_VISIBLE_DEVICES or MUSA_PATCH_PATH is explicitly configured, the process is treated as a MUSA process by design.

Testing

The following tests pass locally:

PYTHONPATH=. pytest -q tests/test_accelerator.py tests/test_reloadable_process_group_world.py
14 passed

PYTHONPATH=/home/Megatron-LM:. pytest -q tests/test_update_weight_transport.py
2 passed

PYTHONPATH=. python tests/test_accelerator.py
8 passed

PYTHONPATH=/home/Megatron-LM:. python tests/test_update_weight_transport.py
2 passed

Additional validation passes:

• Black formatting check.
• Isort import-order check.
• Python bytecode compilation with compileall.
• git diff --cached --check.

The new CPU-only tests are registered in both the workflow template and generated GitHub Actions workflow.

End-to-end MUSA validation

The integration was also validated with a full post-training run on Moore Threads S5000 GPUs using the Qwen3-30B-A3B-Instruct-2507 model. The training job completed the rollout, reward computation, distributed weight synchronization, and Megatron optimization stages successfully. Evaluation pass rate increased clearly as the training step advanced, confirming that the MUSA execution path not only runs end to end but also preserves the expected learning behavior.

Scope and Non-Goals

• This PR does not modify CUDA kernels, compiler toolchains, or kernel-launch syntax.
• CUDA-only INT4 and FlashQLA features remain explicitly CUDA-only and are not silently redirected to MUSA.
• DeepSeek-specific model changes and reward-function changes are not part of this MUSA runtime adaptation.
• The external musa_patch package remains an environment prerequisite for the current MUSA Megatron/SGLang stack and is configured through MUSA_PATCH_PATH.