fix(train): bind distributed training to the active accelerator

funasr/bin/train.py

@@ -32,6 +32,12 @@
 from funasr.download.download_model_from_hub import download_model
 from funasr.models.lora.utils import mark_only_lora_as_trainable
 from funasr.train_utils.set_all_random_seed import set_all_random_seed
+from funasr.train_utils.device_funcs import (
+    accelerator_device,
+    clear_accelerator_cache,
+    get_default_dist_backend,
+    set_accelerator_device,
+)
 from funasr.train_utils.load_pretrained_model import load_pretrained_model
 from funasr.utils.misc import prepare_model_dir
 from funasr.train_utils.model_summary import model_summary
@@ -81,8 +87,11 @@ def main(**kwargs):
     use_fsdp = kwargs.get("use_fsdp", False)
     # use_ddp = False if use_fsdp else use_fsdp
     if use_ddp or use_fsdp:
-        dist.init_process_group(backend=kwargs.get("backend", "nccl"), init_method="env://")
-        torch.cuda.set_device(local_rank)
+        dist.init_process_group(
+            backend=kwargs.get("backend", get_default_dist_backend()),
+            init_method="env://",
+        )
+        set_accelerator_device(local_rank)
 
     logging.info("Build model, frontend, tokenizer")
     device = kwargs.get("device", "cuda")
@@ -128,7 +137,7 @@ def main(**kwargs):
         logging.info(f"{model_summary(model)}")
 
     if use_ddp:
-        model = model.cuda(local_rank)
+        model = model.to(accelerator_device(local_rank))
         model = DDP(
             model,
             device_ids=[local_rank],
@@ -161,12 +170,12 @@ def custom_auto_wrap_policy(
 
         # Configure a custom `min_num_params`
         my_auto_wrap_policy = functools.partial(custom_auto_wrap_policy, min_num_params=int(1e5))
-        torch.cuda.set_device(local_rank)
+        set_accelerator_device(local_rank)
         model = FSDP(
             model,
             auto_wrap_policy=custom_auto_wrap_policy,
             mixed_precision=None,
-            device_id=torch.cuda.current_device(),
+            device_id=local_rank,
         )
     else:
         model = model.to(device=kwargs.get("device", "cuda"))
@@ -246,9 +255,7 @@ def custom_auto_wrap_policy(
             trainer.start_step = 0
 
             device = next(model.parameters()).device
-            if device.type == "cuda":
-                with torch.cuda.device(device):
-                    torch.cuda.empty_cache()
+            clear_accelerator_cache(device)
 
             time_escaped = (time.perf_counter() - time_slice_i) / 3600.0
             logging.info(

funasr/bin/train_ds.py

@@ -30,6 +30,11 @@
 from funasr.download.download_model_from_hub import download_model
 from funasr.models.lora.utils import mark_only_lora_as_trainable
 from funasr.train_utils.set_all_random_seed import set_all_random_seed
+from funasr.train_utils.device_funcs import (
+    clear_accelerator_cache,
+    get_default_dist_backend,
+    set_accelerator_device,
+)
 from funasr.train_utils.load_pretrained_model import load_pretrained_model
 from funasr.utils.misc import prepare_model_dir
 from funasr.train_utils.model_summary import model_summary
@@ -121,16 +126,17 @@ def main(**kwargs):
         deepspeed_config,
         trainer_conf,
     ) = _resolve_distributed_config(kwargs, world_size)
+    dist_backend = kwargs.get("backend", get_default_dist_backend())
     if use_deepspeed:
         logging.info(f"use_deepspeed: {use_deepspeed}")
-        deepspeed.init_distributed(dist_backend=kwargs.get("backend", "nccl"))
+        deepspeed.init_distributed(dist_backend=dist_backend)
     elif use_ddp or use_fsdp:
         logging.info(f"use_ddp: {use_ddp}, use_fsdp: {use_fsdp}")
         dist.init_process_group(
-            backend=kwargs.get("backend", "nccl"),
+            backend=dist_backend,
             init_method="env://",
         )
-        torch.cuda.set_device(local_rank)
+        set_accelerator_device(local_rank)
 
     # rank = dist.get_rank()
 
@@ -241,9 +247,7 @@ def main(**kwargs):
             trainer.start_step = 0
 
             device = next(model.parameters()).device
-            if device.type == "cuda":
-                with torch.cuda.device(device):
-                    torch.cuda.empty_cache()
+            clear_accelerator_cache(device)
 
             time_escaped = (time.perf_counter() - time_slice_i) / 3600.0
             logging.info(

funasr/train_utils/device_funcs.py

@@ -62,3 +62,124 @@ def force_gatherable(data, device):
     else:
         warnings.warn(f"{type(data)} may not be gatherable by DataParallel")
         return data
+
+
+_ACCELERATOR_DEVICE_TYPES = ("cuda", "npu", "xpu")
+
+#: Default ``torch.distributed`` backend for each accelerator, most preferred
+#: first. XCCL (oneCCL) is what current torch exposes for XPU; builds that
+#: predate the 2.6 rename know the same backend as ``"ccl"``.
+_ACCELERATOR_DIST_BACKENDS = {"npu": ("hccl",), "xpu": ("xccl", "ccl")}
+
+
+def _accelerator_module(device_type):
+    """Return the torch submodule owning ``device_type``, or ``None``.
+
+    ``torch_npu`` registers ``torch.npu`` lazily, so it is imported on demand;
+    builds without it simply have no ``torch.npu`` attribute. ``torch.cuda`` and
+    ``torch.xpu`` exist on every build, but their ``is_available()`` reports
+    ``False`` when that backend is missing.
+    """
+    if device_type == "npu":
+        try:
+            import torch_npu  # noqa: F401
+        except ImportError:
+            return None
+    return getattr(torch, device_type, None)
+
+
+def get_accelerator_device_type():
+    """Return the device type of the active accelerator, or ``"cpu"``.
+
+    Probes the same backends FunASR accepts for ``AutoModel`` device fallback
+    (CUDA, then Ascend NPU, then XPU). CPU-only torch builds report ``"cpu"`` so
+    callers can keep a device-agnostic, no-op codepath.
+    """
+    for device_type in _ACCELERATOR_DEVICE_TYPES:
+        module = _accelerator_module(device_type)
+        if module is None:
+            continue
+        try:
+            if module.is_available():
+                return device_type
+        except Exception:
+            continue
+    return "cpu"
+
+
+def accelerator_device(local_rank=0):
+    """Return the ``torch.device`` for ``local_rank`` on the active accelerator."""
+    device_type = get_accelerator_device_type()
+    if device_type == "cpu":
+        return torch.device("cpu")
+    return torch.device(device_type, int(local_rank))
+
+
+def set_accelerator_device(local_rank=0):
+    """Bind this process to ``local_rank`` on the active accelerator.
+
+    Replaces the CUDA-only ``torch.cuda.set_device(local_rank)``, which raises
+    (``module 'torch._C' has no attribute '_cuda_setDevice'``) on a build
+    without CUDA even when an Ascend NPU is present. Returns the resolved device
+    type; for ``"cpu"`` nothing is bound, matching the previous CPU behaviour.
+    """
+    device_type = get_accelerator_device_type()
+    if device_type not in _ACCELERATOR_DEVICE_TYPES:
+        return device_type
+    module = _accelerator_module(device_type)
+    set_device = getattr(module, "set_device", None)
+    if set_device is not None:
+        set_device(int(local_rank))
+    return device_type
+
+
+def _dist_backend_available(backend):
+    """Ask torch whether ``backend`` is usable: ``is_<backend>_available()``.
+
+    A build that does not expose the probe cannot confirm anything, so it
+    reports ``False`` and the caller keeps the accelerator's canonical backend
+    name instead of inventing availability.
+    """
+    distributed = getattr(torch, "distributed", None)
+    probe = getattr(distributed, "is_{}_available".format(backend), None)
+    if probe is None:
+        return False
+    try:
+        return bool(probe())
+    except Exception:
+        return False
+
+
+def get_default_dist_backend():
+    """Return the ``torch.distributed`` backend for the active accelerator.
+
+    CUDA keeps the historical ``"nccl"``, Ascend NPU uses HCCL and XPU uses
+    XCCL (``"ccl"`` on builds older than the oneCCL rename). Candidates are
+    probed with ``torch.distributed.is_<backend>_available()`` where available
+    so the selected backend matches the build, but a non-CUDA accelerator never
+    falls back to NCCL. This is only the default — an explicit ``backend=...``
+    still wins.
+    """
+    candidates = _ACCELERATOR_DIST_BACKENDS.get(get_accelerator_device_type())
+    if not candidates:
+        return "nccl"
+    for backend in candidates:
+        if _dist_backend_available(backend):
+            return backend
+    return candidates[0]
+
+
+def clear_accelerator_cache(device=None):
+    """Release cached allocator memory for ``device`` (no-op on CPU).
+
+    ``torch.cuda.empty_cache()`` is a silent no-op on a build without CUDA, so
+    the end-of-epoch release used by the training entry points did nothing on
+    Ascend NPU / XPU. Dispatch to the accelerator that owns ``device`` instead.
+    """
+    device_type = getattr(device, "type", None) or get_accelerator_device_type()
+    if device_type not in _ACCELERATOR_DEVICE_TYPES:
+        return
+    module = _accelerator_module(device_type)
+    empty_cache = getattr(module, "empty_cache", None)
+    if empty_cache is not None:
+        empty_cache()

funasr/train_utils/trainer_ds.py

@@ -11,7 +11,7 @@
 from pathlib import Path
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-from funasr.train_utils.device_funcs import to_device
+from funasr.train_utils.device_funcs import accelerator_device, to_device
 from funasr.train_utils.recursive_op import recursive_average
 from funasr.train_utils.average_nbest_models import average_checkpoints
 from funasr.train_utils.checkpoint_metrics import (
@@ -1027,7 +1027,7 @@ def warp_model(self, model, **kwargs):
 
         elif self.use_ddp:
             local_rank = int(os.environ.get("LOCAL_RANK", 0))
-            model = model.cuda(local_rank)
+            model = model.to(accelerator_device(local_rank))
             model = DDP(
                 model,
                 device_ids=[local_rank],

tests/test_distributed_device_setup.py

@@ -0,0 +1,113 @@
+import unittest
+from unittest import mock
+
+import torch
+
+from funasr.train_utils import device_funcs
+
+
+class _FakeAccelerator:
+    """Stand-in for ``torch.cuda`` / ``torch.npu`` / ``torch.xpu``."""
+
+    def __init__(self, available):
+        self._available = available
+        self.device_calls = []
+        self.empty_cache_calls = 0
+
+    def is_available(self):
+        return self._available
+
+    def set_device(self, index):
+        self.device_calls.append(index)
+
+    def empty_cache(self):
+        self.empty_cache_calls += 1
+
+
+def _fake_modules(**availability):
+    return {
+        device_type: _FakeAccelerator(availability.get(device_type, False))
+        for device_type in device_funcs._ACCELERATOR_DEVICE_TYPES
+    }
+
+
+class TestAcceleratorResolution(unittest.TestCase):
+    """The training entry points must bind devices without assuming CUDA."""
+
+    def _patch(self, **availability):
+        modules = _fake_modules(**availability)
+        patcher = mock.patch.object(
+            device_funcs,
+            "_accelerator_module",
+            side_effect=lambda device_type: modules.get(device_type),
+        )
+        patcher.start()
+        self.addCleanup(patcher.stop)
+        return modules
+
+    def test_device_type_prefers_cuda_then_npu_then_xpu(self):
+        self._patch(cuda=True, npu=True, xpu=True)
+        self.assertEqual(device_funcs.get_accelerator_device_type(), "cuda")
+        self._patch(npu=True, xpu=True)
+        self.assertEqual(device_funcs.get_accelerator_device_type(), "npu")
+        self._patch(xpu=True)
+        self.assertEqual(device_funcs.get_accelerator_device_type(), "xpu")
+        self._patch()
+        self.assertEqual(device_funcs.get_accelerator_device_type(), "cpu")
+
+    def test_accelerator_device_follows_active_backend(self):
+        self._patch()
+        self.assertEqual(device_funcs.accelerator_device(0), torch.device("cpu"))
+        self._patch(cuda=True)
+        self.assertEqual(device_funcs.accelerator_device(3), torch.device("cuda", 3))
+
+    def test_set_accelerator_device_binds_local_rank(self):
+        modules = self._patch(npu=True)
+        resolved = device_funcs.set_accelerator_device(3)
+        self.assertEqual(resolved, "npu")
+        self.assertEqual(modules["npu"].device_calls, [3])
+
+    def test_set_accelerator_device_is_noop_without_accelerator(self):
+        modules = self._patch()
+        self.assertEqual(device_funcs.set_accelerator_device(1), "cpu")
+        self.assertTrue(all(not module.device_calls for module in modules.values()))
+
+    def test_default_dist_backend_matches_accelerator(self):
+        self._patch(npu=True)
+        self.assertEqual(device_funcs.get_default_dist_backend(), "hccl")
+        self._patch(cuda=True)
+        self.assertEqual(device_funcs.get_default_dist_backend(), "nccl")
+        self._patch(xpu=True)
+        self.assertEqual(device_funcs.get_default_dist_backend(), "xccl")
+        self._patch()
+        self.assertEqual(device_funcs.get_default_dist_backend(), "nccl")
+
+    def test_xpu_backend_selection_is_capability_aware(self):
+        """XPU uses XCCL and never NCCL; only a build that confirms oneCCL's
+        legacy name keeps ``"ccl"``."""
+        self._patch(xpu=True)
+        with mock.patch.object(
+            device_funcs,
+            "_dist_backend_available",
+            side_effect=lambda backend: backend == "xccl",
+        ):
+            self.assertEqual(device_funcs.get_default_dist_backend(), "xccl")
+        with mock.patch.object(
+            device_funcs,
+            "_dist_backend_available",
+            side_effect=lambda backend: backend == "ccl",
+        ):
+            self.assertEqual(device_funcs.get_default_dist_backend(), "ccl")
+        with mock.patch.object(device_funcs, "_dist_backend_available", return_value=False):
+            self.assertEqual(device_funcs.get_default_dist_backend(), "xccl")
+
+    def test_clear_accelerator_cache_dispatches_to_active_backend(self):
+        modules = self._patch(npu=True)
+        device_funcs.clear_accelerator_cache(torch.device("cpu"))
+        self.assertEqual(modules["npu"].empty_cache_calls, 0)
+        device_funcs.clear_accelerator_cache(None)
+        self.assertEqual(modules["npu"].empty_cache_calls, 1)
+
+
+if __name__ == "__main__":
+    unittest.main()