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

[li-lizhe]

Summary

funasr/bin/train.py, funasr/bin/train_ds.py and funasr/train_utils/trainer_ds.py bind every distributed process with CUDA-only calls:

torch.cuda.set_device(local_rank)      # train.py 85/164, train_ds.py 133
model = model.cuda(local_rank)         # train.py 131, trainer_ds.py 1030
device_id=torch.cuda.current_device()  # train.py 169
dist.init_process_group(backend=kwargs.get("backend", "nccl"), ...)

# end of every epoch
if device.type == "cuda":
    with torch.cuda.device(device):
        torch.cuda.empty_cache()

On a PyTorch build without CUDA — for example Ascend NPU — a distributed run cannot start:

• torch.cuda.set_device(0) -> AttributeError: module 'torch._C' has no attribute '_cuda_setDevice'
• init_process_group(backend="nccl") -> RuntimeError: Distributed package doesn't have NCCL built in
• the end-of-epoch cache release only fires for device.type == "cuda", so on NPU/XPU it is a silent no-op.

What changed

New helpers in funasr/train_utils/device_funcs.py (no new dependency):

• get_accelerator_device_type() — probes CUDA, then Ascend NPU, then XPU, "cpu" otherwise (same backend order AutoModel already uses for its device fallback).
• accelerator_device(local_rank) / set_accelerator_device(local_rank) — bind the process to local_rank on the active accelerator.
• get_default_dist_backend() — "hccl" on NPU, "xccl" on XPU (legacy "ccl" reserved for builds that report it), the historical "nccl" elsewhere. Capability-aware through torch.distributed.is_<backend>_available(), still overridable via backend=....
• clear_accelerator_cache(device) — dispatch empty_cache() to the accelerator that owns the device.

The three entry points now call these. CUDA behaviour is unchanged (nccl, cuda:<rank>, torch.cuda.empty_cache()).

Verification (Ascend 910B4, real machine)

Real 2-rank torchrun + HCCL, torch 2.15.0.dev20260917+cpu / torch_npu 2.15.0.dev20260917, physical cards 6-7:

step	before (upstream code)	after (this PR)
set_device	AttributeError: module 'torch._C' has no attribute '_cuda_setDevice' (both ranks)	ok, torch.npu.current_device() = 0 / 1
init_process_group	RuntimeError: Distributed package doesn't have NCCL built in	ok with backend=hccl
all_reduce, 2 ranks	not reached	sum 3.0 == expected 3.0
end-of-epoch cache release	skipped (device.type != "cuda")	ok

Also adds tests/test_distributed_device_setup.py (unittest, 7 cases, no accelerator needed): it stubs the accelerator modules and pins the resolution order, per-rank binding, backend choice (including the XPU default) and cache dispatch. python -m pytest tests/test_distributed_device_setup.py -> 7 passed on a CPU-only build.

Re-checked at the current head da0e18cd (2-rank torchrun on 2x Ascend 910B4, torch 2.14.0 + torch_npu 2.14.0, cards 6-7): device_type=npu, backend=hccl, rank 0/1 on npu:0/npu:1, all_reduce sum 3.0; python -m unittest tests.test_distributed_device_setup -> 7 passed. XPU hardware is still unavailable here, so the XCCL selection is unit-tested only.

Not tested

• CUDA hardware (not available here) — that path is unchanged.
• XPU hardware (not available here) — the new backend selection is covered by unit test only, not exercised on a device.
• Full FSDP / DeepSpeed end-to-end training runs.

[LauraGPT]

The CUDA compatibility path is sound in the scope I could independently exercise: at exact head ba0746849dde38ff8c115848165ebc700215a096, the six submitted device-helper tests pass, and a real H100 run with PyTorch 2.8.0+cu128 completed NCCL initialization, set_accelerator_device(0), an all-reduce, and DDP forward/backward on cuda:0.

One advertised accelerator path still needs a fix before merge. get_accelerator_device_type() can select "xpu", but get_default_dist_backend() returns "nccl" for every non-NPU device. Both training entry points therefore call init_process_group(backend="nccl") when XPU is active. In the same PyTorch 2.8 environment, torch.distributed.Backend.XCCL and is_xccl_available are present; a focused exact-head regression expecting the XPU default to be xccl fails with AssertionError: 'nccl' != 'xccl'.

Please select the repository-supported XPU backend (capability-aware if needed), retain the explicit backend= override, and add an XPU backend-selection case to tests/test_distributed_device_setup.py. The author's Ascend 2-rank evidence remains useful, but I have not independently reproduced NPU, XPU, full FSDP, or DeepSpeed hardware paths. Prepared with Codex assistance.

[li-lizhe]

Thanks for the review and for independently exercising the CUDA path at the exact head.

Fixed in da0e18cd: get_default_dist_backend() no longer returns "nccl" for everything that is not an NPU. It now resolves the backend per accelerator:

• Ascend NPU -> "hccl"
• XPU -> "xccl" (the former oneCCL name "ccl" is kept only for builds that positively confirm it through torch.distributed.is_ccl_available())
• CUDA / CPU -> "nccl" (unchanged)

Candidates are probed with torch.distributed.is_<backend>_available() where the build exposes it, and a non-CUDA accelerator never falls back to NCCL. The explicit backend=... override is untouched — both entry points still call kwargs.get("backend", get_default_dist_backend()).

tests/test_distributed_device_setup.py now covers the XPU default and the capability-aware fallback (7 cases; python -m unittest tests.test_distributed_device_setup -> OK on torch 2.14.0).

Re-verified on the Ascend hardware after the change (2-rank torchrun, torch 2.14.0 + torch_npu 2.14.0, physical cards 6-7): device_type=npu, backend=hccl, rank 0/1 bound to npu:0/npu:1, all_reduce sum 3.0 == expected, end-of-epoch cache release ok.

XPU hardware is still not available here, so the XCCL selection itself is unit-tested only, not exercised on a device.

[LauraGPT]

The backend-selection defect from the previous review is fixed at exact head da0e18cd9e3a4603562bb74bd66750d0d9022535: all seven submitted device-helper tests pass, including the XPU/XCCL selection and capability-aware fallback cases. I also reran the updated exact tree on a real NVIDIA H100 80GB HBM3 with PyTorch 2.8.0+cu128 / CUDA 12.8; NCCL initialization, set_accelerator_device(0), all-reduce, DDP forward/backward, and cache dispatch all completed on cuda:0.

The remaining merge gate is the newly changed XPU runtime path. The patch advertises XPU support and now chooses XCCL (or positively detected legacy CCL), but neither the author nor this review has exercised the exact head on Intel XPU hardware. Under the repository's current hardware-validation requirement, mock/unit backend selection and an unrelated CUDA GPU cannot substitute for the affected accelerator.

Please provide an exact-head Intel XPU run using the same entry-point order, preferably a 2-rank torchrun with XCCL, showing per-rank binding (xpu:0 / xpu:1), all-reduce, DDP forward/backward, and cache dispatch, plus exact PyTorch/oneAPI/device details. The reported exact-head 2-rank Ascend HCCL evidence is retained but was not independently reproduced here. This review does not cover full FSDP or DeepSpeed training. Prepared with Codex assistance.