# PR #2216 完整报告

- 仓库：`THUDM/slime`
- 标题：feat: add backend-aware MUSA support
- 合并时间：2026-08-20 15:11
- 原文链接：http://prhub.com.cn/THUDM/slime/pull/2216

---

# 执行摘要

- 一句话：新增加速器抽象层，支持 MUSA 后端并保持 CUDA 兼容
- 推荐动作：值得精读，尤其 `slime/utils/accelerator/__init__.py` 的选择编排与 `musa_patch` 引导时序、`base.py` 的 `resolve_visible_device_id`（可见设备映射是厂商差异核心）、`reloadable_process_group.py` 的 backend 归一化。建议后续补充“后端接入指南”文档，并规划 MUSA/MCCL 端到端验证（至少一条 SMOKE 路径），以确认 `musa_patch` 引导、权重更新组 `cpu:gloo,musa:mccl` 与 teardown 顺序在真实硬件上的正确性。关注 `is_accelerator_backend` 与 `process_group_backend` 对复合 backend 字符串的处理。

# 功能与动机

PR body 指出多条共享路径直接假定 `torch.cuda`、`CUDA_VISIBLE_DEVICES` 与 NCCL，而 MUSA 通过 `torch.musa`、`MUSA_VISIBLE_DEVICES` 与 MCCL 暴露等价能力；没有后端边界时，每条共享运行时路径都不得不写入厂商条件分支并处理 import 顺序。作者在评论中补充：其基于 `main` 分支 commit `50f2d94466681258b2300753441e48540cc91e04` 适配摩尔线程 GPU，已能在 MUSA 环境跑通 Slime，并保持 CUDA 兼容，请求维护者评估合入。

# 实现拆解

1. **建立抽象契约**：新增 `slime/utils/accelerator/base.py`，定义 `Accelerator` 抽象基类，覆盖 `is_available`、`device`、`device_name`、`set_device`、`synchronize`、流 / 事件、内存统计、RNG、`resolve_visible_device_id` 等操作；厂商模块不会被 `base.py` 导入。
2. **实现委托与具体后端**：`torch_accelerator.py` 提供 `TorchAccelerator`，统一把 CUDA 风格 API 委托给 `torch.cuda` / `torch.musa` 命名空间；`cuda.py` 的 `CUDAAccelerator` 保持 `torch.cuda` + NCCL，`musa.py` 的 `MUSAAccelerator` 映射到 `torch.musa` + MCCL，并通过 `supports(capability)` 声明能力差异（如 MUSA 不支持 `sglang_fp8_utils`、`strict_fp32_logits`）。
3. **选择编排与 `musa_patch` 引导**：`slime/utils/accelerator/__init__.py` 维护 `_REGISTRY` 与 `_BackendRegistration`，支持 `register_accelerator()` 注入第三方后端。选择顺序为：显式 `SLIME_ACCELERATOR` > MUSA 环境变量请求（`MUSA_VISIBLE_DEVICES` / `MUSA_PATCH_PATH`）> 按 priority 自动选择。`musa_patch` 仅在 MUSA 被选定后、构造 backend 前引导一次，显式选择 CUDA 时不会导入。
4. **运行时路径接入**：将散落的 `torch.cuda` 直调替换为 `accelerator.*` 兼容 shim，涉及 `profile_utils.py`（`_create_torch_profiler` 按 `device_type` 构造 activity，`_TorchMemoryProfiler._memory_module` 抽象内存快照）、`sglang_engine.py`（删除 `_to_local_gpu_id`，改用 `accelerator.resolve_visible_device_id`）、`memory_utils.py`、`train_actor.py`、`hf_checkpoint_saver.py`、Megatron `update_weight` 模块、`logprob_utils.py`、`routing_replay.py`、`qwen3_5_vl.py` 以及 `convert_hf_to_torch_dist.py` / `convert_hf_to_int4_direct.py` 等。
5. **分布式进程组与配置配套**：`reloadable_process_group.py` 把 NCCL 专属生命周期泛化为加速器进程组（`_uses_accelerator_backend`、`_destroy_default_accelerator_process_group`），并在 `new_group` 包装中把逻辑 `nccl` 归一化为当前加速器通信后端，MUSA 映射为 `mccl`；MUSA 权重更新组使用 `cpu:gloo,musa:mccl` 混合后端。新增 `SLIME_ENABLE_EXPANDABLE_SEGMENTS`（0/1，默认 0）控制 allocator expandable segments。测试层面新增 `tests/test_accelerator.py`，扩展 `tests/test_reloadable_process_group_world.py`，并注册到 CI 的 CPU unit-test 矩阵。

关键文件：
- `slime/utils/accelerator/__init__.py`（模块 后端抽象；类别 source；类型 core-logic；符号 _BackendRegistration, register_accelerator, _append_musa_patch_path, _import_musa_patch）: 加速器抽象的核心入口，集中选择编排、注册表、环境变量判定与 musa_patch 引导时序，是整个 PR 的枢纽。
- `slime/utils/accelerator/base.py`（模块 后端契约；类别 source；类型 core-logic；符号 Accelerator, is_available, device, device_name）: 定义 Accelerator 契约与可见设备解析逻辑，是所有后端的实现基座，接口边界决定了后续新硬件接入成本。
- `slime/utils/accelerator/torch_accelerator.py`（模块 委托适配；类别 source；类型 core-logic；符号 TorchAccelerator, _module, accelerator_module, is_available）: 共享的 CUDA 风格命名空间委托适配器，CUDA 与 MUSA 都复用其设备名与内存 API，避免重复实现。
- `slime/utils/accelerator/musa.py`（模块 MUSA 后端；类别 source；类型 core-logic；符号 musa_module, is_musa_available, MUSAAccelerator, visible_devices_env）: MUSA 具体实现，负责通信后端映射、混合权重更新组、musa_patch 钩子与能力声明，是本 PR 的核心新增目标。
- `slime/utils/accelerator/cuda.py`（模块 CUDA 后端；类别 source；类型 core-logic；符号 CUDAAccelerator, _module, is_available, attach_oom_observer）: CUDA/ROCm 具体实现，保证显式选择 CUDA 时不加载 musa_patch 且持续使用 NCCL，是兼容性底线。
- `slime/utils/reloadable_process_group.py`（模块 进程组；类别 source；类型 core-logic；符号 _uses_accelerator_backend, _destroy_default_accelerator_process_group, _reload_default_process_group, new_group）: 将 WORLD 重建生命周期从 NCCL 特定逻辑泛化到加速器进程组，并新增 new_group 的 backend 归一化，属于分布式核心路径。
- `slime/utils/profile_utils.py`（模块 性能剖析；类别 source；类型 core-logic；符号 _create_torch_profiler, _TorchMemoryProfiler, _memory_module）: profiler activity 与内存快照抽象到加速器命名空间，避免 profiling 路径仅支持 torch.cuda。
- `slime/backends/sglang_utils/sglang_engine.py`（模块 SGLang 引擎；类别 source；类型 refactor；符号 _to_local_gpu_id, _compute_server_args）: 移除 CUDA_VISIBLE_DEVICES 专属的 _to_local_gpu_id，统一走 accelerator.resolve_visible_device_id，是 rollout 侧的兼容性关键点。
- `tests/test_accelerator.py`（模块 单测覆盖；类别 test；类型 test-coverage；符号 FakeAccelerator, test_cuda_selection_does_not_bootstrap_musa, test_selected_musa_bootstraps_patch_once, test_musa_backend_maps_devices_and_process_groups）: 对后端选择、musa_patch 引导时序、可见设备映射与 backend 映射进行 CPU 可运行覆盖，是本 PR 抽象层的主要验证。
- `tests/test_reloadable_process_group_world.py`（模块 单测覆盖；类别 test；类型 test-coverage；符号 _run_backend_normalization_worker, old_new_group, process_group_backend, test_accelerator_backend_detection）: 覆盖加速器 backend 检测与 new_group 归一化的进程级测试，是分布式改动的主要回归保障。

关键符号：register_accelerator, get_accelerator, initialize_accelerator, set_accelerator, reset_accelerator, Accelerator.resolve_visible_device_id, TorchAccelerator.device_name, MUSAAccelerator.weight_update_backend, MUSAAccelerator.post_import_torch, CUDAAccelerator.supports, _destroy_default_accelerator_process_group, _uses_accelerator_backend, new_group, _create_torch_profiler, _TorchMemoryProfiler._memory_module

## 关键源码片段

### `slime/utils/accelerator/__init__.py`

加速器抽象的核心入口，集中选择编排、注册表、环境变量判定与 musa_patch 引导时序，是整个 PR 的枢纽。

```python
# 加速器选择核心：先处理显式请求，再按可用性自动选择。
# `_REGISTRY` 中每个后端都带 `priority`，`priority` 越高越优先；
# 优先级相同时按名称字典序打破平局，保证多进程行为确定。
def get_accelerator() -> Accelerator:
    global _ACCELERATOR
    if _ACCELERATOR is not None:
        return _ACCELERATOR
    with _SELECTION_LOCK:
        if _ACCELERATOR is not None:
            return _ACCELERATOR
        _register_builtin_backends()
        requested = _requested_name()
        if requested is not None:
            # 用户显式指定 `SLIME_ACCELERATOR=cuda/musa` 时走 fail-fast 分支。
            _ACCELERATOR = _make_selected(requested, explicit=True)
            logger.info('Selected accelerator %s (explicit)', _ACCELERATOR.name)
            return _ACCELERATOR

        # 未显式请求时，遍历注册表，选出第一个可用后端。
        candidates = sorted(_REGISTRY.items(), key=lambda item: (-item[1].priority, item[0]))
        for name, registration in candidates:
            if registration.is_available():
                _ACCELERATOR = _make_selected(name, explicit=False)
                logger.info('Selected accelerator %s (auto)', _ACCELERATOR.name)
                return _ACCELERATOR
        registered = ', '.join(sorted(_REGISTRY))
        raise RuntimeError(
            'No usable accelerator was detected. '
            f'Registered backends: {registered}. '
            'Set SLIME_ACCELERATOR explicitly or install a supported accelerator runtime.'
        )


# 确认 MUSA 是否被请求：显式 `SLIME_ACCELERATOR=musa`，
# 或者存在 `MUSA_VISIBLE_DEVICES` / `MUSA_PATCH_PATH` 环境变量。
def _musa_requested() -> bool:
    configured = os.environ.get('SLIME_ACCELERATOR', '').lower()
    if configured and configured != 'auto':
        return configured == 'musa'
    return 'MUSA_VISIBLE_DEVICES' in os.environ or bool(os.environ.get('MUSA_PATCH_PATH'))


# 构造选定的后端。MUSA 分支会在校验可用性之前先尝试引入 `musa_patch`，
# 因为 `musa_patch` 可能负责把 `torch.musa` 挂到 `torch` 命名空间上。
def _make_selected(name: str, explicit: bool) -> Accelerator:
    _register_builtin_backends()
    entry = _REGISTRY.get(name)
    if entry is None:
        available = ', '.join(sorted(_REGISTRY))
        raise ValueError(f'Unknown accelerator {name!r}; registered backends: {available}')
    if name == 'musa':
        _bootstrap_musa_patch_if_needed()
    if explicit and not entry.is_available():
        if name == 'musa':
            detail = (
                'torch.musa is unavailable; install a MUSA-enabled PyTorch runtime and set MUSA_PATCH_PATH if required'
            )
        elif name == 'cuda':
            detail = 'torch.cuda.is_available() is false or no CUDA device is visible'
        else:
            detail = 'the backend availability check returned false'
        raise RuntimeError(f'Requested accelerator {name!r} is unavailable: {detail}')
    backend = entry.factory()
    if not backend.is_available():
        raise RuntimeError(f'Accelerator backend {name!r} was selected but is unavailable at runtime')
    return backend

```

### `slime/utils/accelerator/base.py`

定义 Accelerator 契约与可见设备解析逻辑，是所有后端的实现基座，接口边界决定了后续新硬件接入成本。

```python
class Accelerator(abc.ABC):
    # 后端中性契约：只包含 Slime 运行时会用到的设备与内存操作。
    name: str
    device_type: str
    communication_backend_name: str

    @abc.abstractmethod
    def is_available(self) -> bool:
        # 返回该后端在本机是否可执行。
        ...

    @abc.abstractmethod
    def device(self, index=None) -> torch.device:
        # 返回本地设备索引对应的 torch.device。
        ...

    @abc.abstractmethod
    def synchronize(self, device=None) -> None:
        # 同步单个设备或当前设备上的全部工作。
        ...

    # 可见设备映射：把物理设备号换算成被 `CUDA_VISIBLE_DEVICES` /
    # `MUSA_VISIBLE_DEVICES` 截断后的本地序号，是厂商差异的集中点。
    def resolve_visible_device_id(self, physical_device_id) -> int:
        raw_value = str(physical_device_id).strip()
        visible = os.environ.get(self.visible_devices_env)
        if not visible:
            # 没有可见设备过滤时，物理号就是本地号。
            return int(float(raw_value))

        ids = [item.strip() for item in visible.split(',') if item.strip()]
        # 先尝试字符串精确匹配（支持 CUDA UUID 与 MUSA 物理号）。
        if raw_value in ids:
            return ids.index(raw_value)

        try:
            value = int(float(raw_value))
        except ValueError:
            value = None
        # 再尝试数值匹配：既允许物理号，也允许已经是本地号的输入。
        if value is not None and str(value) in ids:
            return ids.index(str(value))
        if value is not None and 0 <= value < len(ids):
            return value
        raise RuntimeError(
            f'Device id {raw_value} is not valid under {self.visible_devices_env}={visible}. '
            f'Expected one of {ids} (physical) or 0..{len(ids) - 1} (local).'
        )

```

### `slime/utils/accelerator/musa.py`

MUSA 具体实现，负责通信后端映射、混合权重更新组、musa_patch 钩子与能力声明，是本 PR 的核心新增目标。

```python
class MUSAAccelerator(TorchAccelerator):
    name = 'musa'
    device_type = 'musa'
    communication_backend_name = 'mccl'

    @property
    def visible_devices_env(self) -> str:
        # MUSA 使用 `MUSA_VISIBLE_DEVICES`，而不是 `CUDA_VISIBLE_DEVICES`。
        return 'MUSA_VISIBLE_DEVICES'

    def _module(self) -> Any:
        module = musa_module()
        if module is None:
            raise RuntimeError('MUSA backend requires a runtime that exposes torch.musa')
        return module

    # 权重更新组使用混合后端：CPU 元数据走 `gloo`，MUSA tensor 走 `mccl`。
    def weight_update_backend(self, default: str = 'nccl') -> str:
        return 'cpu:gloo,musa:mccl' if default == 'nccl' else default

    # MCCL 的 `init_process_group` 不接受 device_id 参数，因此这里返回 `None`。
    def distributed_device_id(self, index=None) -> None:
        return None

    # 提供 `post_import_torch` 钩子，供 `musa_patch` 在 torch 导入完成后打补丁。
    def post_import_torch(self) -> None:
        try:
            module = importlib.import_module('musa_patch')
        except ModuleNotFoundError as exc:
            if exc.name == 'musa_patch':
                return
            raise RuntimeError(f'musa_patch failed because dependency {exc.name!r} is missing') from exc
        callback = getattr(module, 'patch_after_import_torch', None)
        if callback is not None:
            callback()

    def supports(self, capability: str) -> bool:
        # 明确声明哪些 CUDA 专属能力在 MUSA 上不可用，调用方据此降级。
        if capability in {'nvml_affinity', 'sglang_fp8_utils', 'strict_fp32_logits'}:
            return False
        if capability == 'requires_cpu_initialization':
            # MUSA 运行时常要求先以 CPU 初始化分布式后端。
            return True
        if capability == 'amp':
            amp = getattr(self._module(), 'amp', None)
            return callable(getattr(amp, 'autocast', None))
        if capability == 'bf16':
            checker = getattr(self._module(), 'is_bf16_supported', None)
            return bool(checker and checker())
        return super().supports(capability)

```

# 评论区精华

> 作者 ForAxel：I adapted Slime for **Moore Threads GPUs**based on commit `50f2d94466681258b2300753441e48540cc91e04` from the `main` branch. The changes allow Slime to run successfully in the MUSA environment while preserving compatibility with CUDA. Could you please review whether this implementation is suitable for merging into `main`? If any additional tests or experimental results would help with the review, please feel free to contact me.

PR body 中明确了验证边界：不提供 MUSA kernel/compiler/vendor runtime，不承诺已有 CUDA 专属 kernel 可移植；`musa_patch` 仍是可选外部依赖；聚焦测试全部 CPU 可运行，真实 CUDA/NCCL 与 MUSA/MCCL 端到端验证留作独立集成。本次未出现 reviewer 对抗性讨论，属于“提出方案 + 等待维护者确认”的状态，最终 PR 由 zhuzilin 合入。

- 摩尔线程 MUSA 适配合并请求 (question): 未在本次 review 评论中看到维护者正式回复；PR 最终由 zhuzilin 合并。

# 风险与影响

- 风险：
 - **进程组全局 monkeypatch**：`reloadable_process_group.py` 的 `new_group` 会全局改写 `dist.new_group(..., backend='nccl')` 调用；若显式选择 CUDA 而环境存在 MUSA 变量，选择编排依赖 `SLIME_ACCELERATOR=cuda` 的显式优先级，该全局副作用缺少真实 NCCL 下的回归验证。
 - **`musa_patch` 失败路径**：`_import_musa_patch` 对“musa_patch 自身缺失”返回 `False`，但对依赖缺失或初始化异常会抛 `RuntimeError`；只设置了 `MUSA_PATCH_PATH` 的普通环境可能因此启动失败，需要更友好的降级。
 - **profiler 回归**：`_TorchMemoryProfiler` 依赖 `accelerator.memory_module()`，后端不支持时静默跳过；CUDA 路径的内存快照与 OOM dump 行为必须与旧版一致，否则影响排障能力。
 - **可见设备映射语义**：`base.py` 的 `resolve_visible_device_id` 取代 `sglang_engine._to_local_gpu_id`，两者都需兼容数字物理号、本地号与 UUID 串；`int(float(raw_value))` 对 `GPU-...` 会走 UUID 匹配分支，边界行为需测试覆盖。
 - **跨模块改动范围**：39 个文件、6 个模块，CUDA 默认路径虽保持不变，但缺少 MUSA/MCCL 端到端硬件验证，allocator 与通信 teardown 顺序只能在真实硬件上确认。
 - **配置项严格校验**：`SLIME_ENABLE_EXPANDABLE_SEGMENTS` 只接受 `0` 或 `1`，非法值会抛 `ValueError`，部署脚本需避免误传。
 - 影响：对 MUSA（摩尔线程）用户，这是首次获得官方支持，可通过 `SLIME_ACCELERATOR=musa` 或 MUSA 环境变量运行训练、rollout 与权重更新全链路，前提是安装 MUSA-enabled PyTorch 与可选 `musa_patch`。对 CUDA/ROCm 用户，默认行为保持兼容，但若环境导出 `MUSA_VISIBLE_DEVICES` 等变量，应显式设置 `SLIME_ACCELERATOR=cuda` 以恢复优先。对系统与团队，新增抽象层成为跨训练、rollout、profiling、模型、转换的公共基座，未来接入新加速器只需实现 `Accelerator` 并注册；但本次改动同时触及 39 个文件，回归风险需通过集中测试与预发布验证管控。对 CI，新增 CPU 单元测试覆盖选择与映射逻辑，但真实硬件集成验证缺位。
 - 风险标记：跨模块改动（39 文件）, 进程组全局 monkeypatch, musa_patch 失败路径未完全降级 , 缺少真实 MUSA/MCCL 端到端测试 , CUDA 默认行为需保持兼容

# 关联脉络

- PR #2267 Fix model convert when use latest megatron: 本 PR 与 #2267 都修改 tools/convert_hf_to_torch_dist.py，共同演进 Megatron 模型转换兼容性路径。
- PR #2271 fix transform_ue8m0 in fp8 convert: 两 PR 都修改 hf_checkpoint_saver.py，同属 Megatron-HF 检查点保存链路的适配工作。
- PR #2262 feat(glm5): align Megatron DeepEP training with SGLang rollout: 本 PR 修改 routing_replay.py（tests 中有 test_routing_replay_uses_selected_backend_current_device），与 #2262 的确定性路由 / 对齐路径相关。
- PR #2286 fix: improve compatibility with older SGLang versions: 两 PR 都涉及 SGLang 后端适配；本 PR 同时处理 sglang_engine.py 的 GPU 可见设备映射与 sglang-router 旧版本兼容。