diff --git a/funasr/bin/train.py b/funasr/bin/train.py index d7c20f8f0..51e0c51e0 100644 --- a/funasr/bin/train.py +++ b/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( diff --git a/funasr/bin/train_ds.py b/funasr/bin/train_ds.py index 1b18fe1b6..f6d2dde23 100644 --- a/funasr/bin/train_ds.py +++ b/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( diff --git a/funasr/train_utils/device_funcs.py b/funasr/train_utils/device_funcs.py index dd7fc76a2..8cebc8572 100644 --- a/funasr/train_utils/device_funcs.py +++ b/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__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__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() diff --git a/funasr/train_utils/trainer_ds.py b/funasr/train_utils/trainer_ds.py index fc2f2b660..eb1ded1d2 100644 --- a/funasr/train_utils/trainer_ds.py +++ b/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], diff --git a/tests/test_distributed_device_setup.py b/tests/test_distributed_device_setup.py new file mode 100644 index 000000000..fce009498 --- /dev/null +++ b/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()