feat: add backend-aware MUSA support

.github/workflows/pr-test.yml

@@ -629,6 +629,14 @@ jobs:
               "num_gpus": 0,
               "test_file": "test_rollout_validation.py"
             },
+            {
+              "num_gpus": 0,
+              "test_file": "test_accelerator.py"
+            },
+            {
+              "num_gpus": 0,
+              "test_file": "test_update_weight_transport.py"
+            },
             {
               "num_gpus": 0,
               "test_file": "test_reloadable_process_group_world.py"

.github/workflows/pr-test.yml.j2

@@ -82,6 +82,8 @@
         {'test_file': 'test_rm_deepscaler.py', 'num_gpus': 0},
         {'test_file': 'test_sample.py', 'num_gpus': 0},
         {'test_file': 'test_rollout_validation.py', 'num_gpus': 0},
+        {'test_file': 'test_accelerator.py', 'num_gpus': 0},
+        {'test_file': 'test_update_weight_transport.py', 'num_gpus': 0},
         {'test_file': 'test_reloadable_process_group_world.py', 'num_gpus': 0},
         {'test_file': 'test_placement_group.py', 'num_gpus': 0},
         {'test_file': 'test_external_sglang_engines.py', 'num_gpus': 0},

slime/backends/megatron_utils/__init__.py

@@ -1,6 +1,6 @@
 import logging
 
-import torch
+from slime.utils import accelerator
 
 try:
     import deep_ep
@@ -12,32 +12,33 @@ def new_init(self, *args, **kwargs):
         if torch_memory_saver._impl is not None:
             torch_memory_saver._impl._binary_wrapper.cdll.tms_set_interesting_region(False)
         old_init(self, *args, **kwargs)
-        torch.cuda.synchronize()
+        accelerator.synchronize()
         if torch_memory_saver._impl is not None:
             torch_memory_saver._impl._binary_wrapper.cdll.tms_set_interesting_region(True)
 
     deep_ep.Buffer.__init__ = new_init
 except ImportError:
     logging.warning("deep_ep is not installed, some functionalities may be limited.")
 
-try:
-    from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import (
-        Qwen3VLMoETextRotaryEmbedding,
-        Qwen3VLTextRotaryEmbedding,
-    )
+if not accelerator.is_musa_environment():
+    try:
+        from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import (
+            Qwen3VLMoETextRotaryEmbedding,
+            Qwen3VLTextRotaryEmbedding,
+        )
 
-    def patch_rotary_embedding(cls):
-        _original_forward = cls.forward
+        def patch_rotary_embedding(cls):
+            _original_forward = cls.forward
 
-        def _patched_forward(self, *args, packed_seq_params=None, **kwargs):
-            return _original_forward(self, *args, **kwargs)
+            def _patched_forward(self, *args, packed_seq_params=None, **kwargs):
+                return _original_forward(self, *args, **kwargs)
 
-        cls.forward = _patched_forward
+            cls.forward = _patched_forward
 
-    patch_rotary_embedding(Qwen3VLTextRotaryEmbedding)
-    patch_rotary_embedding(Qwen3VLMoETextRotaryEmbedding)
-except ImportError:
-    pass
+        patch_rotary_embedding(Qwen3VLTextRotaryEmbedding)
+        patch_rotary_embedding(Qwen3VLMoETextRotaryEmbedding)
+    except ImportError:
+        pass
 
 logging.getLogger("megatron").setLevel(logging.WARNING)
 

slime/backends/megatron_utils/actor.py

@@ -13,7 +13,7 @@
 from transformers import AutoConfig, AutoTokenizer
 
 from slime.ray.train_actor import TrainRayActor
-from slime.utils import train_dump_utils
+from slime.utils import accelerator, train_dump_utils
 from slime.utils.data import process_rollout_data
 from slime.utils.distributed_utils import get_gloo_group
 from slime.utils.logging_utils import init_tracking
@@ -243,7 +243,7 @@ def _get_rollout_data(self, rollout_data_ref: Box) -> RolloutBatch:
         )
         # TODO: this is ugly, move to somewhere else?
         # move tokens to GPU in advance
-        device = torch.cuda.current_device()
+        device = accelerator.device()
         rollout_data["tokens"] = [
             t.to(device=device, dtype=torch.long, non_blocking=True) for t in rollout_data["tokens"]
         ]
@@ -349,7 +349,6 @@ def compute_log_prob(
         num_microbatches: list[int],
         store_prefix: str = "",
     ) -> dict[str, list[torch.Tensor]]:
-
         with timer(f"{store_prefix}log_probs"):
             return forward_only(
                 get_log_probs_and_entropy,

slime/backends/megatron_utils/data.py

@@ -9,7 +9,7 @@
 from megatron.core import mpu
 from megatron.core.packed_seq_params import PackedSeqParams
 
-from slime.utils import train_metric_utils
+from slime.utils import accelerator, train_metric_utils
 from slime.utils.flops_utils import calculate_fwd_flops
 from slime.utils.metric_utils import compute_pass_rate, compute_rollout_step
 from slime.utils.types import RolloutBatch
@@ -83,7 +83,7 @@ def get_batch(
             tokens = F.pad(tokens, (0, pad), value=pad_token_id)
             cu_seqlens_list.append(cu_seqlens_list[-1] + pad)
 
-        cu_seqlens = torch.tensor(cu_seqlens_list, dtype=torch.int, device=torch.cuda.current_device())
+        cu_seqlens = torch.tensor(cu_seqlens_list, dtype=torch.int, device=accelerator.device())
         tokens = tokens.chunk(cp_size, dim=0)[cp_rank]
     else:
         tokens = [slice_with_cp(t, pad_token_id) for t in tokens]
@@ -101,7 +101,7 @@ def get_batch(
             cu_seqlens.append(cu_seqlens[-1] + pad)
 
         # thd requires the cu_seqlens to be of the origin length
-        cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int).cuda() * cp_size
+        cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int, device=accelerator.device()) * cp_size
 
     max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
     packed_seq_params = PackedSeqParams(
@@ -537,5 +537,5 @@ def tensors_to_gpu(tensor_list, device=None):
     if tensor_list is None:
         return None
     if device is None:
-        device = torch.cuda.current_device()
+        device = accelerator.device()
     return [t.to(device=device, dtype=torch.float32) for t in tensor_list]

slime/backends/megatron_utils/hf_checkpoint_saver.py

@@ -8,6 +8,8 @@
 
 import torch
 
+from slime.utils import accelerator
+
 logger = logging.getLogger(__name__)
 
 _HF_WEIGHT_FILE_NAMES = {
@@ -256,9 +258,8 @@ def _write_pending_chunk(
     if pending_write is not None:
         shard_idx, named_tensors = pending_write
         writer.write(named_tensors, shard_idx=shard_idx)
-        if torch.cuda.is_available():
-            torch.cuda.ipc_collect()
-            torch.cuda.empty_cache()
+        accelerator.ipc_collect()
+        accelerator.empty_cache()
 
     return None
 

slime/backends/megatron_utils/loss.py

@@ -8,6 +8,7 @@
 from megatron.core import mpu
 from torch.utils.checkpoint import checkpoint
 
+from slime.utils import accelerator
 from slime.utils.distributed_utils import distributed_masked_whiten
 from slime.utils.misc import load_function
 from slime.utils.ppo_utils import (
@@ -82,7 +83,8 @@ def get_responses(
         `[R, V]` (policy) or `[R, 1]` (value) and `tokens_chunk` is shape `[R]`
         (1D int64), both aligned to response tokens for one sample.
     """
-    assert logits.dtype == torch.float32, f"{logits.dtype}"
+    if not accelerator.is_musa_environment():
+        assert logits.dtype == torch.float32, f"{logits.dtype}"
     assert len(logits.shape) == 3, f"{logits.shape}"
     assert logits.size(0) == 1, f"{logits.shape}"
     logits = logits.squeeze(0)
@@ -489,7 +491,8 @@ def get_log_probs_and_entropy(
     log-probabilities; entropy is always computed from the unmasked logits.
     """
     assert non_loss_data
-    assert logits.dtype == torch.float32, f"{logits.dtype}"
+    if not accelerator.is_musa_environment():
+        assert logits.dtype == torch.float32, f"{logits.dtype}"
     assert len(logits.shape) == 3, f"{logits.shape}"
     assert logits.size(0) == 1, f"{logits.shape}"
     logits = logits.squeeze(0)

slime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py

@@ -5,6 +5,8 @@
 import torch
 import torch.nn as nn
 
+from slime.utils import accelerator
+
 try:
     import fake_int4_quant_cuda
 except ImportError:
@@ -90,7 +92,7 @@ def from_linear(cls, linear, w_bit, group_size, init_only=False, scales=None, ze
             awq_linear.bias = linear.bias.clone().half()
 
         pack_num = 32 // awq_linear.w_bit
-        device = torch.device(f"cuda:{torch.cuda.current_device()}")
+        device = accelerator.device()
 
         repeat_scales = scales.to(device).t().repeat_interleave(group_size, 1)
         if isinstance(zeros, torch.Tensor):
@@ -283,7 +285,7 @@ def quantize_params_compressed_tensors(converted_named_params, quantization_conf
         qw, s, zp = pack_layer(param, group_size, is_symmetric)
         qweight_name = name.replace(".weight", ".weight_packed")
         scale_name = name.replace(".weight", ".weight_scale")
-        weight_shape = torch.tensor(param.shape, dtype=torch.int32, device="cuda")
+        weight_shape = torch.tensor(param.shape, dtype=torch.int32, device=accelerator.device())
         weight_shape_name = name.replace(".weight", ".weight_shape")
         if zp is not None:
             zp_name = name.replace(".weight", ".weight_zero_point")

slime/backends/megatron_utils/server/logprob_utils.py

@@ -12,6 +12,7 @@
 from slime.backends.megatron_utils.data import get_data_iterator
 from slime.backends.megatron_utils.loss import get_log_probs_and_entropy, get_responses
 from slime.backends.megatron_utils.model import forward_only
+from slime.utils import accelerator
 
 logging.getLogger().setLevel(logging.WARNING)
 
@@ -222,17 +223,17 @@ def get_label_token_log_probs_from_vocab_parallel_logits(
     return (local_selected_logits - global_max.to(reduction_dtype) - log_denom).to(logits_dtype)
 
 
-def _to_cuda_tensors(values, dtype: torch.dtype) -> list[torch.Tensor]:
-    return [torch.as_tensor(value, dtype=dtype, device=torch.cuda.current_device()) for value in values]
+def _to_accelerator_tensors(values, dtype: torch.dtype) -> list[torch.Tensor]:
+    return [torch.as_tensor(value, dtype=dtype, device=accelerator.device()) for value in values]
 
 
 def _prepare_rollout_data(rollout_data_ref):
     rollout_data = ray.get(rollout_data_ref[0].inner)
 
-    rollout_data["tokens"] = _to_cuda_tensors(rollout_data["tokens"], torch.long)
-    rollout_data["loss_masks"] = _to_cuda_tensors(rollout_data["loss_masks"], torch.int)
+    rollout_data["tokens"] = _to_accelerator_tensors(rollout_data["tokens"], torch.long)
+    rollout_data["loss_masks"] = _to_accelerator_tensors(rollout_data["loss_masks"], torch.int)
     if rollout_data.get("label_token_ids") is not None:
-        rollout_data["label_token_ids"] = _to_cuda_tensors(rollout_data["label_token_ids"], torch.long)
+        rollout_data["label_token_ids"] = _to_accelerator_tensors(rollout_data["label_token_ids"], torch.long)
         for idx, tensor in enumerate(rollout_data["label_token_ids"]):
             if tensor.dim() == 1 and tensor.numel() == 0:
                 rollout_data["label_token_ids"][idx] = tensor.reshape(0, 0)

slime/backends/megatron_utils/sglang.py

@@ -1,11 +1,15 @@
 # the file to manage all sglang deps in the megatron actor
-try:
-    from sglang.srt.layers.quantization.fp8_utils import quant_weight_ue8m0, transform_scale_ue8m0
-    from sglang.srt.model_loader.utils import should_deepgemm_weight_requant_ue8m0
-except ImportError:
-    quant_weight_ue8m0 = None
-    transform_scale_ue8m0 = None
-    should_deepgemm_weight_requant_ue8m0 = None
+from slime.utils import accelerator
+
+quant_weight_ue8m0 = None
+transform_scale_ue8m0 = None
+should_deepgemm_weight_requant_ue8m0 = None
+if not accelerator.is_musa_environment():
+    try:
+        from sglang.srt.layers.quantization.fp8_utils import quant_weight_ue8m0, transform_scale_ue8m0
+        from sglang.srt.model_loader.utils import should_deepgemm_weight_requant_ue8m0
+    except ImportError:
+        pass
 
 try:
     from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions

slime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py

@@ -1,7 +1,7 @@
 import dataclasses
 
 
-from slime.utils import megatron_bridge_utils
+from slime.utils import accelerator, megatron_bridge_utils
 from slime.utils.misc import chunk_named_params_by_size
 
 from ..megatron_to_hf import postprocess_hf_param
@@ -30,7 +30,7 @@ def _patched(self, task, converted_weights_dict):
         cpu_dict = {k: v.cpu() for k, v in converted_weights_dict.items()}
         result = _orig(self, task, cpu_dict)
         # Move merged result back to GPU for CUDA IPC serialization
-        return {k: v.cuda() for k, v in result.items()} if result else result
+        return {k: v.to(device=accelerator.device()) for k, v in result.items()} if result else result
 
     GPTOSSBridge.maybe_modify_converted_hf_weight = _patched
     GPTOSSBridge._cpu_cache_patched = True
@@ -100,7 +100,7 @@ def _handle_one(task):
         ), f"{weight_dict_key=} not in new_weight_dict ({task.vp_stage=}, {task.param_name=}, {list(new_weight_dict)=})"
 
         new_param_weight = new_weight_dict[weight_dict_key]
-        new_param_weight = new_param_weight.cuda()
+        new_param_weight = new_param_weight.to(device=accelerator.device())
         return dataclasses.replace(task, param_weight=new_param_weight)
 
     return _MapWithLen(_handle_one, vanilla_conversion_tasks)

slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py

@@ -7,6 +7,7 @@
 from megatron.core import mpu
 from tqdm import tqdm
 
+from slime.utils import accelerator
 from slime.utils.distributed_utils import get_gloo_group
 from slime.utils.types import ParamInfo
 
@@ -73,13 +74,13 @@ def _get_megatron_full_params(
         if dist.get_rank() == info.src_rank:
             params.append(
                 torch.nn.Parameter(
-                    megatron_local_weights[info.name].to(device=torch.cuda.current_device(), non_blocking=True),
+                    megatron_local_weights[info.name].to(device=accelerator.device(), non_blocking=True),
                     requires_grad=False,
                 )
             )
         else:
-            params.append(torch.empty(info.shape, dtype=info.dtype, device=torch.cuda.current_device()))
-    torch.cuda.synchronize()
+            params.append(torch.empty(info.shape, dtype=info.dtype, device=accelerator.device()))
+    accelerator.synchronize()
 
     # broadcast params across pp ranks
     if pp_size > 1:

slime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py

@@ -19,6 +19,7 @@
 from megatron.core import mpu
 from ray.actor import ActorHandle
 
+from slime.utils import accelerator
 from slime.utils.disk_delta import NUM_WORKERS, checksum, make_tensor_reader, overwrite_encode
 from slime.utils.distributed_utils import get_gloo_group
 
@@ -259,7 +260,7 @@ def collect(fut):
                 if use_pinned and nbytes <= max_bytes:
                     buf = free_q.get()  # blocks when all buffers are in flight -> backpressures the gather
                     buf[:nbytes].copy_(flat, non_blocking=True)
-                    torch.cuda.current_stream().synchronize()
+                    accelerator.current_stream().synchronize()
                     payload, pinned = buf, True
                 else:
                     payload, pinned = flat.cpu().numpy(), False
@@ -278,7 +279,7 @@ def _record_metrics(self) -> None:
         counts = torch.tensor(
             [self.changed_bytes, self.total_bytes, self.wire_bytes],
             dtype=torch.int64,
-            device=torch.cuda.current_device(),
+            device=accelerator.device(),
         )
         dist.all_reduce(counts)
         changed, total, wire = counts.tolist()