From 807cbd1b478c73ac1d1c6d3d4e719446592043ba Mon Sep 17 00:00:00 2001 From: foraxel Date: Fri, 17 Jul 2026 19:28:27 +0800 Subject: [PATCH 1/2] feat: add backend-aware MUSA support --- .github/workflows/pr-test.yml | 8 + .github/workflows/pr-test.yml.j2 | 2 + slime/backends/megatron_utils/__init__.py | 33 +-- slime/backends/megatron_utils/actor.py | 5 +- slime/backends/megatron_utils/data.py | 8 +- .../megatron_utils/hf_checkpoint_saver.py | 7 +- slime/backends/megatron_utils/loss.py | 7 +- .../quantizer_compressed_tensors.py | 6 +- .../megatron_utils/server/logprob_utils.py | 11 +- slime/backends/megatron_utils/sglang.py | 18 +- .../hf_weight_iterator_bridge.py | 6 +- .../hf_weight_iterator_direct.py | 7 +- .../update_weight_from_disk_delta.py | 5 +- .../update_weight_from_distributed.py | 51 +++-- .../update_weight_from_tensor.py | 21 +- slime/backends/sglang_utils/sglang_engine.py | 24 +-- slime/ray/train_actor.py | 15 +- slime/ray/utils.py | 10 +- slime/utils/accelerator.py | 203 ++++++++++++++++++ slime/utils/memory_utils.py | 14 +- slime/utils/profile_utils.py | 52 ++++- slime/utils/reloadable_process_group.py | 42 ++-- slime/utils/routing_replay.py | 6 +- slime/utils/tensor_backper.py | 10 +- slime_plugins/megatron_bridge/glm4v_moe.py | 10 +- .../models/flash_dot_product_attention.py | 8 +- slime_plugins/models/qwen3_5.py | 6 +- slime_plugins/models/qwen3_next.py | 6 +- tests/test_accelerator.py | 129 +++++++++++ tests/test_reloadable_process_group_world.py | 12 +- tests/test_update_weight_transport.py | 88 ++++++++ tools/convert_hf_to_torch_dist.py | 19 +- tools/convert_to_hf.py | 6 +- train.py | 4 + train_async.py | 4 + 35 files changed, 706 insertions(+), 157 deletions(-) create mode 100644 slime/utils/accelerator.py create mode 100644 tests/test_accelerator.py create mode 100644 tests/test_update_weight_transport.py diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 45606dd532..74ce349726 100644 --- a/.github/workflows/pr-test.yml +++ b/.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" diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index f58b064830..72a850dc97 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.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}, diff --git a/slime/backends/megatron_utils/__init__.py b/slime/backends/megatron_utils/__init__.py index b1936ae692..db20738914 100644 --- a/slime/backends/megatron_utils/__init__.py +++ b/slime/backends/megatron_utils/__init__.py @@ -2,6 +2,8 @@ import torch +from slime.utils import accelerator + try: import deep_ep from torch_memory_saver import torch_memory_saver @@ -12,7 +14,7 @@ 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) @@ -20,24 +22,25 @@ def new_init(self, *args, **kwargs): 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) diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index fd50276454..8cf97239a2 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/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, diff --git a/slime/backends/megatron_utils/data.py b/slime/backends/megatron_utils/data.py index 51b008d111..6267b61af0 100644 --- a/slime/backends/megatron_utils/data.py +++ b/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] diff --git a/slime/backends/megatron_utils/hf_checkpoint_saver.py b/slime/backends/megatron_utils/hf_checkpoint_saver.py index a1c6979d12..e5093357bd 100644 --- a/slime/backends/megatron_utils/hf_checkpoint_saver.py +++ b/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 diff --git a/slime/backends/megatron_utils/loss.py b/slime/backends/megatron_utils/loss.py index adfa5d67e6..07aaa5e979 100644 --- a/slime/backends/megatron_utils/loss.py +++ b/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) diff --git a/slime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py b/slime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py index ca69df8e65..82d50c4e37 100644 --- a/slime/backends/megatron_utils/megatron_to_hf/processors/quantizer_compressed_tensors.py +++ b/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") diff --git a/slime/backends/megatron_utils/server/logprob_utils.py b/slime/backends/megatron_utils/server/logprob_utils.py index 61c3fca5f2..110aa163d0 100644 --- a/slime/backends/megatron_utils/server/logprob_utils.py +++ b/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) diff --git a/slime/backends/megatron_utils/sglang.py b/slime/backends/megatron_utils/sglang.py index 97c82a31cd..7555206d0f 100644 --- a/slime/backends/megatron_utils/sglang.py +++ b/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 diff --git a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py b/slime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py index 60a85fdf79..4007bb3c36 100644 --- a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py +++ b/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) diff --git a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py b/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py index e257e42200..88ce484e97 100644 --- a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py +++ b/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: diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py b/slime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py index 2632835096..ef9ef343e6 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py +++ b/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() diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py index 1ba987f06b..877e5c1c26 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os import socket import time from argparse import Namespace @@ -13,6 +14,7 @@ from ray.actor import ActorHandle from tqdm import tqdm +from slime.utils import accelerator from slime.utils.distributed_utils import get_gloo_group, init_process_group from slime.utils.http_utils import _wrap_ipv6 @@ -22,8 +24,8 @@ class UpdateWeightFromDistributed: """ - Update distributed engines via NCCL. Each PP rank: group "slime-pp_{pp_rank}", - only DP=TP=0 broadcasts. Non-expert (TP) and expert (EP) params separate. + Update distributed engines through a device process group. Each PP rank: group "slime-pp_{pp_rank}", + only DP=TP=0 transfers. Non-expert (TP) and expert (EP) params separate. Subclasses override ``_send_weights`` / ``_on_chunk`` to inject per-mode behaviour. """ @@ -63,7 +65,7 @@ def connect_rollout_engines( engine_parallel_configs: Sequence[Mapping[str, object]] | None = None, ) -> None: """ - Create NCCL "slime-pp_{pp_rank}" if PP source (DP=TP=0). Lock prevents concurrent broadcasts. + Create "slime-pp_{pp_rank}" if PP source (DP=TP=0). Lock prevents concurrent transfers. """ self.rollout_engines = rollout_engines self.rollout_engine_lock = rollout_engine_lock @@ -217,7 +219,7 @@ def _ep_gather_and_convert(self, named_tensors: list[tuple[str, torch.Tensor]]) handles = [] for i, (_name, param) in enumerate(named_tensors): params = [ - torch.empty_like(param.data, device=torch.cuda.current_device()) + torch.empty_like(param.data, device=accelerator.device()) for _ in range(mpu.get_expert_model_parallel_world_size()) ] handle = dist.all_gather(params, param.data, group=mpu.get_expert_model_parallel_group(), async_op=True) @@ -244,9 +246,9 @@ def _update_bucket_weights_from_distributed( load_format: str | None = None, ) -> None: """ - Lock → broadcast → clear → unlock → pbar++. Lock prevents NCCL deadlock. + Lock → transfer → clear → unlock → pbar++. Lock prevents communication deadlock. """ - # lock the rollout engines to prevent dead lock on broadcast. + # Lock the rollout engines to prevent communication deadlock. while not ray.get(self.rollout_engine_lock.acquire.remote()): time.sleep(0.1) @@ -272,11 +274,11 @@ def connect_rollout_engines_from_distributed( engine_gpu_counts: Sequence[int] | None = None, ) -> dist.ProcessGroup: """ - Create NCCL group: training rank 0 + all engine GPUs. Blocks until joined. + Create a device process group: training rank 0 + all engine GPUs. Blocks until joined. ``engine_gpu_counts`` gives the number of GPUs per engine. When engines have heterogeneous TP sizes (e.g. prefill TP=2, decode TP=4), each engine - occupies a different number of ranks in the NCCL group. + occupies a different number of ranks in the process group. """ if engine_gpu_counts is None: engine_gpu_counts = [args.rollout_num_gpus_per_engine] * len(rollout_engines) @@ -292,6 +294,7 @@ def connect_rollout_engines_from_distributed( for c in engine_gpu_counts: cumulative.append(cumulative[-1] + c) + backend = accelerator.weight_update_backend() refs = [ engine.init_weights_update_group.remote( master_address=master_address, @@ -299,12 +302,12 @@ def connect_rollout_engines_from_distributed( rank_offset=cumulative[i] + 1, world_size=world_size, group_name=group_name, - backend="nccl", + backend=backend, ) for i, engine in enumerate(rollout_engines) ] model_update_groups = init_process_group( - backend="nccl", + backend=backend, init_method=f"tcp://{_wrap_ipv6(master_address)}:{master_port}", world_size=world_size, rank=0, @@ -316,7 +319,7 @@ def connect_rollout_engines_from_distributed( def disconnect_rollout_engines_from_distributed(args, group_name, model_update_groups, rollout_engines): """ - Destroy NCCL on training and engines. + Destroy the weight-update process group on training and engines. """ refs = [engine.destroy_weights_update_group.remote(group_name) for engine in rollout_engines] dist.destroy_process_group(model_update_groups) @@ -332,8 +335,15 @@ def update_weights_from_distributed( load_format: str | None = None, ) -> list[ObjectRef]: """ - Send metadata (Ray), broadcast tensors (NCCL rank 0 → engines). + Send metadata through Ray and tensors through the configured transport. """ + comm_mode = os.environ.get("UPDATE_MODE", "broadcast") + if comm_mode == "p2p-broadcast": + group_rank = group.rank() + global_rank = dist.get_global_rank(group, group_rank) + else: + global_rank = None + refs = [ engine.update_weights_from_distributed.remote( names=[name for name, _ in converted_named_tensors], @@ -347,8 +357,21 @@ def update_weights_from_distributed( ] handles = [] - for _, param in converted_named_tensors: - handles.append(dist.broadcast(param.data, 0, group=group, async_op=True)) + if comm_mode == "broadcast": + for _, param in converted_named_tensors: + handles.append(dist.broadcast(param.data, 0, group=group, async_op=True)) + elif comm_mode == "gloo-broadcast": + for _, param in converted_named_tensors: + dist.broadcast(param.data.cpu(), 0, group=group, async_op=False) + elif comm_mode == "sync-broadcast": + for _, param in converted_named_tensors: + dist.broadcast(param.data, 0, group=group, async_op=False) + elif comm_mode == "p2p-broadcast": + if global_rank == 0: + for tag, (_, param) in enumerate(converted_named_tensors): + dist.send(param.data, 1, group=group, tag=tag) + else: + raise NotImplementedError(f"Unknown UPDATE_MODE={comm_mode!r}") for handle in handles: handle.wait() diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py index 877febb7ca..784342d7f5 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -11,6 +11,7 @@ from ray.actor import ActorHandle from tqdm import tqdm +from slime.utils import accelerator from slime.utils.distributed_utils import get_gloo_group from slime.utils.types import ParamInfo @@ -31,7 +32,7 @@ def _build_flattened_tensor_data( ) -> dict[str, Any]: if not named_tensors: return { - "flattened_tensor": torch.empty(0, dtype=torch.uint8, device=torch.cuda.current_device()), + "flattened_tensor": torch.empty(0, dtype=torch.uint8, device=accelerator.device()), "metadata": [], } @@ -266,12 +267,12 @@ def _update_expert_weights( refs, long_lived_tensors = self._send_hf_params(hf_named_tensors) ray.get(refs) dist.barrier(group=get_gloo_group()) - torch.cuda.synchronize() + accelerator.synchronize() del refs, long_lived_tensors, hf_named_tensors - torch.cuda.ipc_collect() - torch.cuda.empty_cache() + accelerator.ipc_collect() + accelerator.empty_cache() del staging_buffers - torch.cuda.empty_cache() + accelerator.empty_cache() @torch.no_grad() def update_weights(self) -> None: @@ -306,8 +307,8 @@ def update_weights(self) -> None: # then release CUDA IPC cache entries whose consumers (sglang engines) # have already closed their IPC handles. del refs, long_lived_tensors, hf_named_tensors - torch.cuda.ipc_collect() - torch.cuda.empty_cache() + accelerator.ipc_collect() + accelerator.empty_cache() if self._expert_transfer_plan: self._update_expert_weights(megatron_local_weights) @@ -316,8 +317,8 @@ def update_weights(self) -> None: dist.barrier(group=get_gloo_group()) # After the barrier all engines have returned, so every rank's last-chunk # IPC handles are now released by the consumers. Clean them up. - torch.cuda.ipc_collect() - torch.cuda.empty_cache() + accelerator.ipc_collect() + accelerator.empty_cache() # int4/fp4 post_process if self.rank == 0: @@ -426,6 +427,6 @@ def _send_to_colocated_engine( def _empty_flattened_tensor_data(): return { - "flattened_tensor": torch.empty(0, dtype=torch.uint8, device=torch.cuda.current_device()), + "flattened_tensor": torch.empty(0, dtype=torch.uint8, device=accelerator.device()), "metadata": [], } diff --git a/slime/backends/sglang_utils/sglang_engine.py b/slime/backends/sglang_utils/sglang_engine.py index c89af61e77..8a50990b98 100644 --- a/slime/backends/sglang_utils/sglang_engine.py +++ b/slime/backends/sglang_utils/sglang_engine.py @@ -6,6 +6,10 @@ import time from urllib.parse import quote +from slime.utils import accelerator + +# isort: split + import requests import sglang_router from packaging.version import parse @@ -30,24 +34,6 @@ def get_base_gpu_id(args, rank): return start_index -def _to_local_gpu_id(physical_gpu_id: int) -> int: - cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if not cvd: - return physical_gpu_id # no remapping - # CUDA_VISIBLE_DEVICES can be like "4,5,6,7" - visible = [int(x) for x in cvd.split(",") if x.strip() != ""] - # In a remapped process, valid torch device indices are 0..len(visible)-1 - if physical_gpu_id in visible: - return visible.index(physical_gpu_id) - # If we're already getting local IDs, allow them - if 0 <= physical_gpu_id < len(visible): - return physical_gpu_id - raise RuntimeError( - f"GPU id {physical_gpu_id} is not valid under CUDA_VISIBLE_DEVICES={cvd}. " - f"Expected one of {visible} (physical) or 0..{len(visible)-1} (local)." - ) - - def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process: if getattr(server_args, "encoder_only", False): from sglang.srt.disaggregation.encode_server import launch_server_process as sglang_launch_server_process @@ -552,7 +538,7 @@ def _compute_server_args( nnodes = max(1, _gpus_per_engine // args.num_gpus_per_node) node_rank = rank % nnodes base = base_gpu_id if base_gpu_id is not None else get_base_gpu_id(args, rank) - base = _to_local_gpu_id(base) + base = accelerator.resolve_visible_device_id(base) kwargs = { "model_path": args.hf_checkpoint, "trust_remote_code": True, diff --git a/slime/ray/train_actor.py b/slime/ray/train_actor.py index a8ba6ddc64..5c1ae5983b 100644 --- a/slime/ray/train_actor.py +++ b/slime/ray/train_actor.py @@ -10,6 +10,7 @@ import slime.utils.eval_config from slime.ray.ray_actor import RayActor +from slime.utils import accelerator from slime.utils.distributed_utils import init_gloo_group from slime.utils.logging_utils import configure_logger from slime.utils.memory_utils import clear_memory, print_memory @@ -18,11 +19,7 @@ def get_local_gpu_id(): - cvd = os.environ.get("CUDA_VISIBLE_DEVICES", None) - if cvd is None: - return ray.get_gpu_ids()[0] - else: - return cvd.split(",").index(str(ray.get_gpu_ids()[0])) + return accelerator.resolve_visible_device_id(ray.get_gpu_ids()[0]) class TrainRayActor(RayActor): @@ -56,9 +53,9 @@ def init(self, args, role, with_ref=False, with_opd_teacher=False): torch.serialization.add_safe_globals([slime.utils.eval_config.EvalDatasetConfig]) local_rank = int(os.environ.get("LOCAL_RANK", 0)) - torch.cuda.set_device(f"cuda:{local_rank}") + accelerator.set_device(local_rank) - backend = args.distributed_backend + backend = accelerator.process_group_backend(args.distributed_backend) dist.init_process_group( backend=backend, @@ -70,7 +67,9 @@ def init(self, args, role, with_ref=False, with_opd_teacher=False): args.world_size = dist.get_world_size() try: - if torch.version.hip is not None: + if accelerator.is_musa_available(): + logger.info("Detected MUSA environment, skipping NVML NUMA affinity setup") + elif torch.version.hip is not None: logger.info("Detected ROCm/HIP environment, skipping NUMA affinity setup") # will find the coresponding API to implement ROCm version as below else: diff --git a/slime/ray/utils.py b/slime/ray/utils.py index 7b3a3d5262..50bd35ee1b 100644 --- a/slime/ray/utils.py +++ b/slime/ray/utils.py @@ -2,8 +2,9 @@ import os import ray -import torch + from slime.ray.ray_actor import RayActor +from slime.utils import accelerator # Refer to # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/nvidia_gpu.py#L95-L96 @@ -15,6 +16,7 @@ # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/intel_gpu.py#L97-L98 NOSET_VISIBLE_DEVICES_ENV_VARS_LIST = [ "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_MUSA_VISIBLE_DEVICES", "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES", "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES", "RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES", @@ -38,9 +40,9 @@ def ray_noset_visible_devices(env_vars=os.environ): def get_physical_gpu_id(): - device = torch.cuda.current_device() - props = torch.cuda.get_device_properties(device) - return str(props.uuid) + device = accelerator.current_device() + props = accelerator.get_device_properties(device) + return str(getattr(props, "uuid", device)) @ray.remote diff --git a/slime/utils/accelerator.py b/slime/utils/accelerator.py new file mode 100644 index 0000000000..fa2d72fdea --- /dev/null +++ b/slime/utils/accelerator.py @@ -0,0 +1,203 @@ +import importlib +import logging +import os +import sys +from types import ModuleType +from typing import Any + +logger = logging.getLogger(__name__) + +_MUSA_PATCH_IMPORTED = False + + +def _append_musa_patch_path() -> None: + patch_path = os.environ.get("MUSA_PATCH_PATH") + if patch_path and patch_path not in sys.path: + sys.path.append(patch_path) + + +def _import_musa_patch() -> bool: + _append_musa_patch_path() + try: + importlib.import_module("musa_patch") + except ImportError: + return False + except Exception as exc: + logger.warning("Failed to import musa_patch: %s", exc) + return False + return True + + +import torch + + +def _musa() -> ModuleType | None: + return getattr(torch, "musa", None) + + +def is_musa_available() -> bool: + musa = _musa() + return bool(musa is not None and getattr(musa, "is_available", lambda: False)()) + + +def is_musa_environment() -> bool: + return is_musa_available() or "MUSA_VISIBLE_DEVICES" in os.environ or bool(os.environ.get("MUSA_PATCH_PATH")) + + +def _require_musa_available() -> None: + if is_musa_environment() and not is_musa_available(): + raise RuntimeError( + "MUSA environment was requested via MUSA_VISIBLE_DEVICES or MUSA_PATCH_PATH, " + "but torch.musa is unavailable. Check the MUSA runtime and musa_patch import order." + ) + + +def device_type() -> str: + if is_musa_available(): + return "musa" + if torch.cuda.is_available(): + return "cuda" + return "cpu" + + +def device_name(index: int | None = None) -> str: + current_type = device_type() + if current_type == "cpu": + return "cpu" + if index is None: + index = current_device() + return f"{current_type}:{index}" + + +def device(index: int | None = None) -> torch.device: + return torch.device(device_name(index)) + + +def accelerator_module() -> Any: + if is_musa_available(): + return _musa() + return torch.cuda + + +def set_device(index: int | str | torch.device) -> None: + if device_type() != "cpu": + accelerator_module().set_device(index) + + +def current_device() -> int | str: + if device_type() == "cpu": + return "cpu" + return accelerator_module().current_device() + + +def synchronize(device_arg: int | str | torch.device | None = None) -> None: + if device_type() == "cpu": + return + module = accelerator_module() + if device_arg is None: + module.synchronize() + else: + module.synchronize(device_arg) + + +def empty_cache() -> None: + if device_type() != "cpu": + accelerator_module().empty_cache() + + +def ipc_collect() -> None: + if device_type() == "cpu": + return + ipc_collect_fn = getattr(accelerator_module(), "ipc_collect", None) + if ipc_collect_fn is not None: + ipc_collect_fn() + + +def mem_get_info(device_arg: int | str | torch.device | None = None) -> tuple[int, int]: + if device_type() == "cpu": + raise RuntimeError("Accelerator memory info is unavailable on CPU") + if device_arg is None: + device_arg = current_device() + return accelerator_module().mem_get_info(device_arg) + + +def memory_allocated(device_arg: int | str | torch.device | None = None) -> int: + if device_type() == "cpu": + return 0 + return accelerator_module().memory_allocated(device_arg) + + +def memory_reserved(device_arg: int | str | torch.device | None = None) -> int: + if device_type() == "cpu": + return 0 + return accelerator_module().memory_reserved(device_arg) + + +def get_device_properties(device_arg: int | str | torch.device | None = None) -> Any: + if device_type() == "cpu": + return None + if device_arg is None: + device_arg = current_device() + return accelerator_module().get_device_properties(device_arg) + + +def visible_devices_env_key() -> str: + if "MUSA_VISIBLE_DEVICES" in os.environ or is_musa_available(): + return "MUSA_VISIBLE_DEVICES" + return "CUDA_VISIBLE_DEVICES" + + +def resolve_visible_device_id(physical_device_id: int | float | str) -> int: + env_key = visible_devices_env_key() + visible_devices = os.environ.get(env_key) + physical_device_id = int(float(physical_device_id)) + if not visible_devices: + return physical_device_id + + visible = [int(device_id) for device_id in visible_devices.split(",") if device_id.strip()] + if physical_device_id in visible: + return visible.index(physical_device_id) + if 0 <= physical_device_id < len(visible): + return physical_device_id + raise RuntimeError( + f"Device id {physical_device_id} is not valid under {env_key}={visible_devices}. " + f"Expected one of {visible} (physical) or 0..{len(visible) - 1} (local)." + ) + + +def process_group_backend(default: str = "nccl") -> str: + if default == "nccl": + if is_musa_available(): + return "mccl" + _require_musa_available() + return default + + +def weight_update_backend(default: str = "nccl") -> str: + if default == "nccl": + if is_musa_available(): + return "cpu:gloo,musa:mccl" + _require_musa_available() + return default + + +def current_stream(): + if device_type() == "cpu": + return None + return accelerator_module().current_stream() + + +def _try_import_musa_patch() -> bool: + global _MUSA_PATCH_IMPORTED + if _MUSA_PATCH_IMPORTED: + return True + if not is_musa_environment(): + return False + + _MUSA_PATCH_IMPORTED = _import_musa_patch() + if not _MUSA_PATCH_IMPORTED and is_musa_environment(): + logger.warning("musa_patch is not importable; continuing without it") + return _MUSA_PATCH_IMPORTED + + +_try_import_musa_patch() diff --git a/slime/utils/memory_utils.py b/slime/utils/memory_utils.py index d4b2d89321..4bad36b51f 100644 --- a/slime/utils/memory_utils.py +++ b/slime/utils/memory_utils.py @@ -5,28 +5,30 @@ import torch import torch.distributed as dist +from slime.utils import accelerator + logger = logging.getLogger(__name__) def clear_memory(clear_host_memory: bool = False): - torch.cuda.synchronize() + accelerator.synchronize() gc.collect() - torch.cuda.empty_cache() + accelerator.empty_cache() if clear_host_memory: torch._C._host_emptyCache() def available_memory(): - device = torch.cuda.current_device() - free, total = torch.cuda.mem_get_info(device) + device = accelerator.current_device() + free, total = accelerator.mem_get_info(device) vm = psutil.virtual_memory() return { "gpu": str(device), "total_GB": _byte_to_gb(total), "free_GB": _byte_to_gb(free), "used_GB": _byte_to_gb(total - free), - "allocated_GB": _byte_to_gb(torch.cuda.memory_allocated(device)), - "reserved_GB": _byte_to_gb(torch.cuda.memory_reserved(device)), + "allocated_GB": _byte_to_gb(accelerator.memory_allocated(device)), + "reserved_GB": _byte_to_gb(accelerator.memory_reserved(device)), "host_total_GB": _byte_to_gb(vm.total), "host_available_GB": _byte_to_gb(vm.available), "host_used_GB": _byte_to_gb(vm.used), diff --git a/slime/utils/profile_utils.py b/slime/utils/profile_utils.py index 504d1ce868..7f1fbceb59 100644 --- a/slime/utils/profile_utils.py +++ b/slime/utils/profile_utils.py @@ -5,6 +5,7 @@ import torch +from slime.utils import accelerator from slime.utils.memory_utils import print_memory logger = logging.getLogger(__name__) @@ -58,7 +59,13 @@ def _profile_simple_loop(iterator, args, name): def _create_torch_profiler(args, name): + activities = [torch.profiler.ProfilerActivity.CPU] + activity_name = accelerator.device_type().upper() + if hasattr(torch.profiler.ProfilerActivity, activity_name): + activities.append(getattr(torch.profiler.ProfilerActivity, activity_name)) + return torch.profiler.profile( + activities=activities, schedule=torch.profiler.schedule( # TODO the train_actor and train_log_probs ones may need to have different args to control step wait=max(args.profile_step_start - 1, 0), @@ -101,30 +108,59 @@ def stop(self): class _TorchMemoryProfiler(_BaseMemoryProfiler): + def __init__(self, args): + super().__init__(args) + self._recording = False + + @staticmethod + def _memory_module(): + return getattr(accelerator.accelerator_module(), "memory", None) + def start(self): logger.info("Attach OOM dump memory history.") - - torch.cuda.memory._record_memory_history( + memory_module = self._memory_module() + if memory_module is None or not hasattr(memory_module, "_record_memory_history"): + logger.warning("Accelerator memory history is unavailable; skip torch memory profiler.") + return + if not hasattr(memory_module, "_dump_snapshot"): + logger.warning("Accelerator memory snapshot is unavailable; skip torch memory profiler.") + return + + memory_module._record_memory_history( max_entries=1000000, - # record stack information for the trace events - # trace_alloc_record_context=True, stacks="all", ) + self._recording = True def oom_observer(device, alloc, device_alloc, device_free): logger.info( f"Observe OOM, will dump snapshot to {self._path_dump}. ({device=} {alloc=} {device_alloc=} {device_free=}; stacktrace is as follows)" ) traceback.print_stack() - torch.cuda.memory._dump_snapshot(self._path_dump) + memory_module._dump_snapshot(str(self._path_dump)) print_memory("when oom") - torch._C._cuda_attach_out_of_memory_observer(oom_observer) + if accelerator.is_musa_available(): + musa_c = getattr(torch.musa, "_MUSAC", None) + attach_oom_observer = getattr(musa_c, "_musa_attach_out_of_memory_observer", None) + else: + attach_oom_observer = getattr(torch._C, "_cuda_attach_out_of_memory_observer", None) + if attach_oom_observer is not None: + attach_oom_observer(oom_observer) + else: + logger.warning("Accelerator OOM observer is unavailable; memory snapshot on OOM is disabled.") def stop(self): + if not self._recording: + return logger.info(f"Dump memory snapshot to: {self._path_dump}") - torch.cuda.memory._dump_snapshot(self._path_dump) - torch.cuda.memory._record_memory_history(enabled=None) + memory_module = self._memory_module() + if memory_module is None or not hasattr(memory_module, "_dump_snapshot"): + logger.warning("Accelerator memory snapshot is unavailable; skip dump.") + return + memory_module._dump_snapshot(str(self._path_dump)) + memory_module._record_memory_history(enabled=None) + self._recording = False class _MemrayMemoryProfiler(_BaseMemoryProfiler): diff --git a/slime/utils/reloadable_process_group.py b/slime/utils/reloadable_process_group.py index 5bc574dace..a48d6e22ef 100644 --- a/slime/utils/reloadable_process_group.py +++ b/slime/utils/reloadable_process_group.py @@ -9,6 +9,7 @@ import torch.distributed as dist from torch.distributed.distributed_c10d import PrefixStore, _get_default_group, _get_default_store +from slime.utils import accelerator from slime.utils.distributed_utils import get_gloo_group, init_gloo_group, set_gloo_group from slime.utils.memory_utils import available_memory, clear_memory, print_memory @@ -26,11 +27,11 @@ class _DefaultProcessGroupState: rank: int world_size: int generation: int = 0 - nccl_world_destroyed: bool = False + accelerator_world_destroyed: bool = False def register_default_process_group(timeout: timedelta) -> None: - """Register the NCCL WORLD group so it can be destroyed and rebuilt. + """Register the accelerator WORLD group so it can be destroyed and rebuilt. Keeping a reference to the rendezvous store is intentional. It keeps the rank-0 TCPStore alive after ``destroy_process_group()`` and lets every @@ -58,8 +59,9 @@ def register_default_process_group(timeout: timedelta) -> None: ) -def _uses_nccl(backend: str) -> bool: - return "nccl" in backend.lower() +def _uses_accelerator_backend(backend: str) -> bool: + backend = backend.lower() + return "nccl" in backend or "mccl" in backend def _new_default_process_group(state: _DefaultProcessGroupState, backend: str) -> None: @@ -74,9 +76,9 @@ def _new_default_process_group(state: _DefaultProcessGroupState, backend: str) - ) -def _destroy_default_nccl_process_group() -> None: +def _destroy_default_accelerator_process_group() -> None: state = default_process_group_states.get(os.getpid()) - if state is None or state.nccl_world_destroyed or not _uses_nccl(state.backend): + if state is None or state.accelerator_world_destroyed or not _uses_accelerator_backend(state.backend): return # Use the CPU group as an out-of-band synchronization point. PyTorch @@ -87,7 +89,7 @@ def _destroy_default_nccl_process_group() -> None: _new_default_process_group(state, backend="gloo") set_gloo_group(_get_default_group()) - state.nccl_world_destroyed = True + state.accelerator_world_destroyed = True logger.info( "Destroyed default %s WORLD process group and initialized a temporary Gloo WORLD (generation %s)", state.backend, @@ -97,18 +99,18 @@ def _destroy_default_nccl_process_group() -> None: def _reload_default_process_group() -> None: state = default_process_group_states.get(os.getpid()) - if state is None or not state.nccl_world_destroyed: + if state is None or not state.accelerator_world_destroyed: return - # WORLD uses Gloo while the NCCL WORLD is destroyed, so this barrier does - # not allocate CUDA or recreate an NCCL communicator before all ranks are ready. + # WORLD uses Gloo while the accelerator WORLD is destroyed, so this barrier + # does not recreate an accelerator communicator before all ranks are ready. dist.barrier() dist.destroy_process_group() set_gloo_group(None) _new_default_process_group(state, backend=state.backend) init_gloo_group() - state.nccl_world_destroyed = False + state.accelerator_world_destroyed = False logger.info( "Reloaded default WORLD process group with backend %s (generation %s)", state.backend, @@ -145,9 +147,17 @@ def monkey_patch_torch_dist(): dist.old_new_group = old_new_group def new_group(*args, **kwargs): - group = old_new_group(*args, **kwargs) explicit_backend = args[2] if len(args) >= 3 else kwargs.get("backend") backend = str(explicit_backend) if explicit_backend is not None else str(dist.get_backend()) + normalized_backend = accelerator.process_group_backend(backend) + if normalized_backend != backend: + if len(args) >= 3: + args = (*args[:2], normalized_backend, *args[3:]) + else: + kwargs = {**kwargs, "backend": normalized_backend} + backend = normalized_backend + + group = old_new_group(*args, **kwargs) # Before WORLD is registered, preserve the historical behavior of # leaving CPU groups and singleton groups untouched. Afterwards every @@ -441,16 +451,16 @@ def bound_device_id(self, dev): def destroy_process_groups(): - """Destroy registered subgroups and replace NCCL WORLD with a temporary Gloo WORLD.""" + """Destroy registered subgroups and replace accelerator WORLD with a temporary Gloo WORLD.""" state = default_process_group_states.get(os.getpid()) - if state is not None and not state.nccl_world_destroyed and _uses_nccl(state.backend): + if state is not None and not state.accelerator_world_destroyed and _uses_accelerator_backend(state.backend): dist.barrier(group=get_gloo_group()) ReloadableProcessGroup.destroy_process_groups() - _destroy_default_nccl_process_group() + _destroy_default_accelerator_process_group() def reload_process_groups(): - """Restore NCCL WORLD and recreate all registered subgroups.""" + """Restore accelerator WORLD and recreate all registered subgroups.""" _reload_default_process_group() ReloadableProcessGroup.reload_process_groups() diff --git a/slime/utils/routing_replay.py b/slime/utils/routing_replay.py index 8647281664..0782542358 100644 --- a/slime/utils/routing_replay.py +++ b/slime/utils/routing_replay.py @@ -1,6 +1,8 @@ import os import torch +from slime.utils import accelerator + ROUTING_REPLAY = None @@ -28,12 +30,12 @@ def record(self, top_indices): def pop_forward(self): top_indices = self.top_indices_list[self.forward_index] self.forward_index += 1 - return top_indices.to(torch.cuda.current_device()) + return top_indices.to(accelerator.device()) def pop_backward(self): top_indices = self.top_indices_list[self.backward_index] self.backward_index += 1 - return top_indices.to(torch.cuda.current_device()) + return top_indices.to(accelerator.device()) def clear(self): self.forward_index = 0 diff --git a/slime/utils/tensor_backper.py b/slime/utils/tensor_backper.py index 2fc2a6359d..7d876a4be5 100644 --- a/slime/utils/tensor_backper.py +++ b/slime/utils/tensor_backper.py @@ -4,6 +4,8 @@ import torch +from slime.utils import accelerator + _SourceGetter = Callable[[], Iterable[tuple[str, torch.Tensor]]] @@ -58,7 +60,7 @@ def backup(self, tag: str) -> None: if name not in backup_dict: backup_dict[name] = torch.empty_like(param, device=torch.device("cpu"), pin_memory=True) backup_dict[name].copy_(param.detach(), non_blocking=True) - torch.cuda.synchronize() + accelerator.synchronize() @torch.no_grad() def copy(self, *, src_tag: str, dst_tag: str): @@ -71,7 +73,7 @@ def restore(self, tag: str) -> None: for name, param in self._source_getter(): assert name in backup_dict param.copy_(backup_dict[name], non_blocking=True) - torch.cuda.synchronize() + accelerator.synchronize() class _TensorBackuperNoop(TensorBackuper): @@ -94,12 +96,12 @@ def get(self, tag: str): def backup(self, tag: str) -> None: assert tag == self._single_tag self._backup_hash_dict = _compute_hash_dict(dict(self._source_getter())) - torch.cuda.synchronize() + accelerator.synchronize() def restore(self, tag: str) -> None: assert tag == self._single_tag assert _compute_hash_dict(dict(self._source_getter())) == self._backup_hash_dict - torch.cuda.synchronize() + accelerator.synchronize() def _compute_hash_dict(tensors: dict[str, torch.Tensor]): diff --git a/slime_plugins/megatron_bridge/glm4v_moe.py b/slime_plugins/megatron_bridge/glm4v_moe.py index 2f78285782..70d60b2d1d 100644 --- a/slime_plugins/megatron_bridge/glm4v_moe.py +++ b/slime_plugins/megatron_bridge/glm4v_moe.py @@ -17,6 +17,10 @@ from copy import deepcopy from dataclasses import dataclass, field +from slime.utils import accelerator + +# isort: split + import torch from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge @@ -206,8 +210,8 @@ def __init__( self.vision_model.requires_grad_(False) self.vision_model.eval() hook_hf_module_setattr_for_tp_grad_sync(self.vision_model) - if torch.cuda.is_available(): - self.vision_model = self.vision_model.to("cuda") + if accelerator.device_type() != "cpu": + self.vision_model = self.vision_model.to(accelerator.device()) # Language model -- standard Megatron GPT with M-RoPE self.language_model = MCoreGPTModel( @@ -440,7 +444,7 @@ def forward( # Non-first PP stage: allocate buffer with correct shape if cu_seqlens is not None: T = cu_seqlens[-1].item() - position_ids = torch.zeros(3, 1, T, dtype=torch.long, device=torch.cuda.current_device()) + position_ids = torch.zeros(3, 1, T, dtype=torch.long, device=accelerator.device()) else: raise NotImplementedError( "Non-THD position_ids broadcast not yet supported for non-first PP stages" diff --git a/slime_plugins/models/flash_dot_product_attention.py b/slime_plugins/models/flash_dot_product_attention.py index 06aef6dc3a..40cf6b1abb 100644 --- a/slime_plugins/models/flash_dot_product_attention.py +++ b/slime_plugins/models/flash_dot_product_attention.py @@ -7,6 +7,10 @@ import math +from slime.utils import accelerator + +# isort: split + import torch from flash_attn.flash_attn_interface import flash_attn_varlen_func from megatron.core.packed_seq_params import PackedSeqParams @@ -75,7 +79,7 @@ def __init__( elif config.softmax_type == "off-by-one": self.softmax_offset = torch.zeros( num_heads_per_partition, - device=torch.cuda.current_device(), + device=accelerator.device(), dtype=config.params_dtype, ) elif config.softmax_type == "learnable": @@ -84,7 +88,7 @@ def __init__( torch.nn.Parameter( torch.empty( num_heads_per_partition, - device=torch.cuda.current_device(), + device=accelerator.device(), dtype=config.params_dtype, ) ), diff --git a/slime_plugins/models/qwen3_5.py b/slime_plugins/models/qwen3_5.py index 294c9d97a9..c466268af0 100644 --- a/slime_plugins/models/qwen3_5.py +++ b/slime_plugins/models/qwen3_5.py @@ -1,5 +1,9 @@ import copy +from slime.utils import accelerator + +# isort: split + import torch import torch.nn as nn import torch.nn.functional as F @@ -77,7 +81,7 @@ def __init__(self, config, layer_idx: int, args=None): self.head_v_dim, eps=self.layer_norm_epsilon, activation=self.activation, - device=torch.cuda.current_device(), + device=accelerator.device(), dtype=config.dtype if config.dtype is not None else torch.get_default_dtype(), ) diff --git a/slime_plugins/models/qwen3_next.py b/slime_plugins/models/qwen3_next.py index 683cb28063..8d513c9c7a 100644 --- a/slime_plugins/models/qwen3_next.py +++ b/slime_plugins/models/qwen3_next.py @@ -1,5 +1,9 @@ import copy +from slime.utils import accelerator + +# isort: split + import torch import torch.nn as nn import torch.nn.functional as F @@ -70,7 +74,7 @@ def __init__(self, config, layer_idx: int, args=None): self.head_v_dim, eps=self.layer_norm_epsilon, activation=self.activation, - device=torch.cuda.current_device(), + device=accelerator.device(), dtype=config.dtype if config.dtype is not None else torch.get_default_dtype(), ) diff --git a/tests/test_accelerator.py b/tests/test_accelerator.py new file mode 100644 index 0000000000..5db8ab6525 --- /dev/null +++ b/tests/test_accelerator.py @@ -0,0 +1,129 @@ +import importlib +from types import SimpleNamespace + +import pytest + +from slime.utils import accelerator + +NUM_GPUS = 0 + + +@pytest.mark.unit +def test_resolve_cuda_visible_device_id(monkeypatch): + monkeypatch.delenv("MUSA_VISIBLE_DEVICES", raising=False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,6") + monkeypatch.setattr(accelerator, "is_musa_available", lambda: False) + + assert accelerator.resolve_visible_device_id(4) == 0 + assert accelerator.resolve_visible_device_id(6.0) == 1 + assert accelerator.resolve_visible_device_id(1) == 1 + + +@pytest.mark.unit +def test_resolve_musa_visible_device_id(monkeypatch): + monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "2,5") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") + + assert accelerator.visible_devices_env_key() == "MUSA_VISIBLE_DEVICES" + assert accelerator.resolve_visible_device_id("5") == 1 + + +@pytest.mark.unit +def test_invalid_visible_device_id_reports_active_environment(monkeypatch): + monkeypatch.delenv("MUSA_VISIBLE_DEVICES", raising=False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,6") + monkeypatch.setattr(accelerator, "is_musa_available", lambda: False) + + with pytest.raises(RuntimeError, match="CUDA_VISIBLE_DEVICES=4,6"): + accelerator.resolve_visible_device_id(7) + + +@pytest.mark.unit +def test_musa_process_group_backends(monkeypatch): + monkeypatch.setattr(accelerator, "is_musa_available", lambda: True) + + assert accelerator.process_group_backend() == "mccl" + assert accelerator.weight_update_backend() == "cpu:gloo,musa:mccl" + assert accelerator.process_group_backend("gloo") == "gloo" + + +@pytest.mark.unit +def test_requested_musa_runtime_does_not_fall_back_to_nccl(monkeypatch): + monkeypatch.setenv("MUSA_VISIBLE_DEVICES", "0") + monkeypatch.setattr(accelerator, "is_musa_available", lambda: False) + + with pytest.raises(RuntimeError, match="torch.musa is unavailable"): + accelerator.process_group_backend() + + +@pytest.mark.unit +def test_musa_patch_import_is_idempotent(monkeypatch): + import_count = 0 + + def import_patch_once(): + nonlocal import_count + import_count += 1 + return True + + monkeypatch.setattr(accelerator, "_MUSA_PATCH_IMPORTED", False) + monkeypatch.setattr(accelerator, "_import_musa_patch", import_patch_once) + monkeypatch.setattr(accelerator, "is_musa_environment", lambda: True) + + assert accelerator._try_import_musa_patch() + assert accelerator._try_import_musa_patch() + assert import_count == 1 + + +@pytest.mark.unit +def test_import_does_not_reapply_patch_after_import(monkeypatch): + reapplied = False + + def patch_after_import_torch(): + nonlocal reapplied + reapplied = True + + monkeypatch.setattr( + accelerator.importlib, + "import_module", + lambda name: SimpleNamespace(patch_after_import_torch=patch_after_import_torch), + ) + + assert accelerator._import_musa_patch() + assert not reapplied + + +@pytest.mark.unit +def test_cuda_environment_does_not_import_musa_patch(monkeypatch): + monkeypatch.delenv("MUSA_VISIBLE_DEVICES", raising=False) + monkeypatch.delenv("MUSA_PATCH_PATH", raising=False) + monkeypatch.setattr(accelerator, "_MUSA_PATCH_IMPORTED", False) + monkeypatch.setattr(accelerator, "is_musa_available", lambda: False) + monkeypatch.setattr( + accelerator, + "_import_musa_patch", + lambda: (_ for _ in ()).throw(AssertionError("musa_patch must not be imported in a CUDA environment")), + ) + + assert not accelerator._try_import_musa_patch() + + +@pytest.mark.unit +def test_accelerator_import_bootstraps_musa_patch(monkeypatch): + imported_modules = [] + + with monkeypatch.context() as patch: + patch.setenv("MUSA_VISIBLE_DEVICES", "0") + patch.setattr( + importlib, + "import_module", + lambda name: imported_modules.append(name) or SimpleNamespace(), + ) + + importlib.reload(accelerator) + + assert imported_modules == ["musa_patch"] + importlib.reload(accelerator) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_reloadable_process_group_world.py b/tests/test_reloadable_process_group_world.py index cb1e842f02..a6ded451c8 100644 --- a/tests/test_reloadable_process_group_world.py +++ b/tests/test_reloadable_process_group_world.py @@ -9,6 +9,12 @@ NUM_GPUS = 0 +@pytest.mark.unit +@pytest.mark.parametrize("backend", ["nccl", "mccl", "cpu:gloo,musa:mccl"]) +def test_accelerator_backend_detection(backend): + assert rpg._uses_accelerator_backend(backend) + + @pytest.mark.unit def test_register_default_process_group_captures_rendezvous_state(monkeypatch): timeout = timedelta(minutes=7) @@ -27,7 +33,7 @@ def test_register_default_process_group_captures_rendezvous_state(monkeypatch): assert state.store == "rendezvous-store" assert state.rank == 3 assert state.world_size == 8 - assert not state.nccl_world_destroyed + assert not state.accelerator_world_destroyed @pytest.mark.unit @@ -71,7 +77,7 @@ def init_process_group(**kwargs): rpg.destroy_process_groups() - assert state.nccl_world_destroyed + assert state.accelerator_world_destroyed assert state.generation == 1 assert events == [ ("barrier", "canonical-gloo"), @@ -95,7 +101,7 @@ def init_process_group(**kwargs): events.clear() rpg.reload_process_groups() - assert not state.nccl_world_destroyed + assert not state.accelerator_world_destroyed assert state.generation == 2 assert events == [ ("barrier", "WORLD"), diff --git a/tests/test_update_weight_transport.py b/tests/test_update_weight_transport.py new file mode 100644 index 0000000000..354715821c --- /dev/null +++ b/tests/test_update_weight_transport.py @@ -0,0 +1,88 @@ +import pytest +import torch + +pytest.importorskip("megatron") + +from slime.backends.megatron_utils.update_weight import update_weight_from_distributed + + +NUM_GPUS = 0 + + +class _RemoteMethod: + def __init__(self): + self.calls = [] + + def remote(self, **kwargs): + self.calls.append(kwargs) + return kwargs + + +class _Engine: + def __init__(self): + self.update_weights_from_distributed = _RemoteMethod() + + +class _Group: + def rank(self): + return 0 + + +def test_p2p_weight_update_matches_engine_receive_protocol(monkeypatch): + monkeypatch.setenv("UPDATE_MODE", "p2p-broadcast") + monkeypatch.setattr(update_weight_from_distributed.dist, "get_global_rank", lambda group, rank: rank) + + sends = [] + monkeypatch.setattr( + update_weight_from_distributed.dist, + "send", + lambda tensor, dst, group, tag: sends.append((tensor, dst, group, tag)), + ) + monkeypatch.setattr( + update_weight_from_distributed.dist, + "broadcast", + lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("unexpected collective broadcast")), + ) + + group = _Group() + engine = _Engine() + tensors = [("a", torch.ones(2)), ("b", torch.ones(3))] + + refs = update_weight_from_distributed.update_weights_from_distributed("slime-pp_0", group, 1, [engine], tensors) + + assert refs == engine.update_weights_from_distributed.calls + assert [(dst, send_group, tag) for _, dst, send_group, tag in sends] == [ + (1, group, 0), + (1, group, 1), + ] + + +def test_default_weight_update_keeps_async_broadcast(monkeypatch): + monkeypatch.delenv("UPDATE_MODE", raising=False) + + waited = [] + + class _Handle: + def wait(self): + waited.append(True) + + broadcasts = [] + monkeypatch.setattr( + update_weight_from_distributed.dist, + "broadcast", + lambda tensor, src, group, async_op: broadcasts.append((tensor, src, group, async_op)) or _Handle(), + ) + + group = _Group() + tensors = [("a", torch.ones(2)), ("b", torch.ones(3))] + update_weight_from_distributed.update_weights_from_distributed("slime-pp_0", group, 1, [_Engine()], tensors) + + assert [(src, broadcast_group, async_op) for _, src, broadcast_group, async_op in broadcasts] == [ + (0, group, True), + (0, group, True), + ] + assert waited == [True, True] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py index 13c1939130..b4e04fd39c 100644 --- a/tools/convert_hf_to_torch_dist.py +++ b/tools/convert_hf_to_torch_dist.py @@ -2,6 +2,10 @@ import os import shutil +from slime.utils import accelerator + +# isort: split + import torch import torch.distributed as dist from megatron.core.enums import ModelType @@ -87,6 +91,7 @@ def ceildiv(a, b): def main(): if torch.version.hip: import megatron.core.dist_checkpointing.strategies.filesystem_async as filesystem_async_module + from slime.utils.rocm_checkpoint_writer import ROCmFileSystemWriterAsync filesystem_async_module.FileSystemWriterAsync = ROCmFileSystemWriterAsync @@ -99,23 +104,25 @@ def main(): local_rank = int(os.getenv("LOCAL_RANK") or os.getenv("SLURM_LOCALID") or 0) global_rank = int(os.getenv("RANK") or os.getenv("SLURM_PROCID") or 0) - torch.cuda.set_device(local_rank) + accelerator.set_device(local_rank) os.environ.setdefault("WORLD_SIZE", str(world_size)) os.environ.setdefault("RANK", str(global_rank)) os.environ.setdefault("LOCAL_RANK", str(local_rank)) os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "12355") dist.init_process_group( - backend="nccl", + backend=accelerator.process_group_backend(), world_size=world_size, rank=global_rank, - device_id=torch.device(f"cuda:{local_rank}"), + device_id=None if accelerator.is_musa_available() else torch.device(accelerator.device_name(local_rank)), ) args = get_args() init(args) # if using AMD gpus, we have to do the conversion in cpu - if hasattr(torch.version, "hip") and torch.version.hip is not None: + if accelerator.is_musa_available(): + assert args.use_cpu_initialization, "MUSA requires --use_cpu_initialization=True" + elif hasattr(torch.version, "hip") and torch.version.hip is not None: assert args.use_cpu_initialization, "AMD GPU requires --use_cpu_initialization=True" model = get_model(get_model_provider_func(args), ModelType.encoder_or_decoder, wrap_with_ddp=False) @@ -130,9 +137,9 @@ def main(): model[0] = model[0].cpu() print_memory("after loading model") - torch.cuda.synchronize() + accelerator.synchronize() gc.collect() - torch.cuda.empty_cache() + accelerator.empty_cache() save_checkpoint(1, model, None, None, 0) diff --git a/tools/convert_to_hf.py b/tools/convert_to_hf.py index 9f16850f0c..19a2744ea4 100644 --- a/tools/convert_to_hf.py +++ b/tools/convert_to_hf.py @@ -1,3 +1,7 @@ +from slime.utils import accelerator + +# isort: split + import torch import torch.distributed as dist from megatron.core import mpu @@ -57,7 +61,7 @@ def main(args): param = param_ break else: - param = torch.empty(info.shape, dtype=info.dtype, device=torch.cuda.current_device()) + param = torch.empty(info.shape, dtype=info.dtype, device=accelerator.device()) if pp_size > 1: if info.src_rank in dist.get_process_group_ranks(mpu.get_pipeline_model_parallel_group()): diff --git a/train.py b/train.py index 9cb2968866..3cd5a0acdf 100644 --- a/train.py +++ b/train.py @@ -1,3 +1,7 @@ +from slime.utils import accelerator # noqa: F401 + +# isort: split + import ray from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models diff --git a/train_async.py b/train_async.py index 9141b612fe..7c66ba9196 100644 --- a/train_async.py +++ b/train_async.py @@ -1,3 +1,7 @@ +from slime.utils import accelerator # noqa: F401 + +# isort: split + import ray from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models From 9bd1a736400021231fb4ffceb07cbbb8ac020665 Mon Sep 17 00:00:00 2001 From: foraxel Date: Mon, 20 Jul 2026 11:27:05 +0800 Subject: [PATCH 2/2] fix(ci): make MUSA weight-update tests self-contained --- slime/backends/megatron_utils/__init__.py | 2 - slime/backends/sglang_utils/sglang_engine.py | 1 - tests/test_empty_colocated_weight_bucket.py | 4 + tests/test_update_weight_transport.py | 79 ++++++++++++++++++-- 4 files changed, 78 insertions(+), 8 deletions(-) diff --git a/slime/backends/megatron_utils/__init__.py b/slime/backends/megatron_utils/__init__.py index db20738914..4a56ca5ab1 100644 --- a/slime/backends/megatron_utils/__init__.py +++ b/slime/backends/megatron_utils/__init__.py @@ -1,7 +1,5 @@ import logging -import torch - from slime.utils import accelerator try: diff --git a/slime/backends/sglang_utils/sglang_engine.py b/slime/backends/sglang_utils/sglang_engine.py index 8a50990b98..1b8c885167 100644 --- a/slime/backends/sglang_utils/sglang_engine.py +++ b/slime/backends/sglang_utils/sglang_engine.py @@ -2,7 +2,6 @@ import ipaddress import logging import multiprocessing -import os import time from urllib.parse import quote diff --git a/tests/test_empty_colocated_weight_bucket.py b/tests/test_empty_colocated_weight_bucket.py index 7e5db7e0db..ea06f8194a 100644 --- a/tests/test_empty_colocated_weight_bucket.py +++ b/tests/test_empty_colocated_weight_bucket.py @@ -67,6 +67,9 @@ def _install_fake_deps(monkeypatch): update_weight_pkg.__path__ = [str(REPO_ROOT / "slime" / "backends" / "megatron_utils" / "update_weight")] slime_utils_pkg = types.ModuleType("slime.utils") slime_utils_pkg.__path__ = [str(REPO_ROOT / "slime" / "utils")] + accelerator_mod = types.ModuleType("slime.utils.accelerator") + accelerator_mod.device = lambda: "cuda:0" + accelerator_mod.ipc_collect = lambda: None dist_mod = types.ModuleType("torch.distributed") @@ -133,6 +136,7 @@ def gather_object(obj, object_gather_list, dst, group): monkeypatch.setitem(sys.modules, "slime.backends.megatron_utils", megatron_utils_pkg) monkeypatch.setitem(sys.modules, "slime.backends.megatron_utils.update_weight", update_weight_pkg) monkeypatch.setitem(sys.modules, "slime.utils", slime_utils_pkg) + monkeypatch.setitem(sys.modules, "slime.utils.accelerator", accelerator_mod) monkeypatch.setitem(sys.modules, "torch", torch_mod) monkeypatch.setitem(sys.modules, "torch.distributed", dist_mod) monkeypatch.setitem(sys.modules, "ray", ray_mod) diff --git a/tests/test_update_weight_transport.py b/tests/test_update_weight_transport.py index 354715821c..f9cdea8172 100644 --- a/tests/test_update_weight_transport.py +++ b/tests/test_update_weight_transport.py @@ -1,13 +1,15 @@ +import importlib.util +import sys +import types +from pathlib import Path + import pytest import torch -pytest.importorskip("megatron") - -from slime.backends.megatron_utils.update_weight import update_weight_from_distributed - - NUM_GPUS = 0 +REPO_ROOT = Path(__file__).resolve().parents[1] + class _RemoteMethod: def __init__(self): @@ -28,7 +30,73 @@ def rank(self): return 0 +def _load_update_weight_module(monkeypatch): + slime_pkg = types.ModuleType("slime") + slime_pkg.__path__ = [str(REPO_ROOT / "slime")] + backends_pkg = types.ModuleType("slime.backends") + backends_pkg.__path__ = [str(REPO_ROOT / "slime" / "backends")] + megatron_utils_pkg = types.ModuleType("slime.backends.megatron_utils") + megatron_utils_pkg.__path__ = [str(REPO_ROOT / "slime" / "backends" / "megatron_utils")] + update_weight_pkg = types.ModuleType("slime.backends.megatron_utils.update_weight") + update_weight_pkg.__path__ = [str(REPO_ROOT / "slime" / "backends" / "megatron_utils" / "update_weight")] + utils_pkg = types.ModuleType("slime.utils") + utils_pkg.__path__ = [str(REPO_ROOT / "slime" / "utils")] + + megatron_mod = types.ModuleType("megatron") + megatron_core_mod = types.ModuleType("megatron.core") + megatron_core_mod.mpu = types.ModuleType("megatron.core.mpu") + + ray_mod = types.ModuleType("ray") + ray_mod.ObjectRef = object + ray_actor_mod = types.ModuleType("ray.actor") + ray_actor_mod.ActorHandle = object + + accelerator_mod = types.ModuleType("slime.utils.accelerator") + distributed_utils_mod = types.ModuleType("slime.utils.distributed_utils") + distributed_utils_mod.get_gloo_group = lambda: None + distributed_utils_mod.init_process_group = lambda **kwargs: None + http_utils_mod = types.ModuleType("slime.utils.http_utils") + http_utils_mod._wrap_ipv6 = lambda address: address + + megatron_to_hf_mod = types.ModuleType("slime.backends.megatron_utils.megatron_to_hf") + megatron_to_hf_mod.convert_to_hf = lambda *args, **kwargs: [] + common_mod = types.ModuleType("slime.backends.megatron_utils.update_weight.common") + common_mod.all_gather_param = lambda *args, **kwargs: None + common_mod.named_params_and_buffers = lambda *args, **kwargs: [] + + fake_modules = { + "slime": slime_pkg, + "slime.backends": backends_pkg, + "slime.backends.megatron_utils": megatron_utils_pkg, + "slime.backends.megatron_utils.megatron_to_hf": megatron_to_hf_mod, + "slime.backends.megatron_utils.update_weight": update_weight_pkg, + "slime.backends.megatron_utils.update_weight.common": common_mod, + "slime.utils": utils_pkg, + "slime.utils.accelerator": accelerator_mod, + "slime.utils.distributed_utils": distributed_utils_mod, + "slime.utils.http_utils": http_utils_mod, + "megatron": megatron_mod, + "megatron.core": megatron_core_mod, + "ray": ray_mod, + "ray.actor": ray_actor_mod, + } + for name, module in fake_modules.items(): + monkeypatch.setitem(sys.modules, name, module) + + module_name = "slime.backends.megatron_utils.update_weight.update_weight_from_distributed" + module_path = ( + REPO_ROOT / "slime" / "backends" / "megatron_utils" / "update_weight" / "update_weight_from_distributed.py" + ) + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, module_name, module) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + def test_p2p_weight_update_matches_engine_receive_protocol(monkeypatch): + update_weight_from_distributed = _load_update_weight_module(monkeypatch) monkeypatch.setenv("UPDATE_MODE", "p2p-broadcast") monkeypatch.setattr(update_weight_from_distributed.dist, "get_global_rank", lambda group, rank: rank) @@ -58,6 +126,7 @@ def test_p2p_weight_update_matches_engine_receive_protocol(monkeypatch): def test_default_weight_update_keeps_async_broadcast(monkeypatch): + update_weight_from_distributed = _load_update_weight_module(monkeypatch) monkeypatch.delenv("UPDATE_MODE", raising=False) waited = []