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 Ascend NPU, the historical "nccl" everywhere else including XPU (selecting XCCL for XPU is a separate follow-up, [dist] XPU distributed backend selection (XCCL / oneCCL) is unverified — needs XPU hardware #3777, because that path cannot be exercised without XPU hardware). 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, 8 cases, no accelerator needed): it stubs the accelerator modules and pins the resolution order, per-rank binding, backend choice (HCCL on NPU; XPU stays the historical "nccl") and cache dispatch. python -m unittest tests.test_distributed_device_setup -> 8 passed on a CPU-only build.

Re-checked at the current head 0a336ba5 (2-rank torchrun on an Ascend 910B2 container, torch 2.15.0.dev20260917+cpu / torch_npu 2.15.0.dev20260917+gitec69335, 4 devices): device_type=npu, backend=hccl, rank 0/1 on npu:0/npu:1, representative training step loss 0.835605 on both ranks, all_reduce sum 1.0, and the end-of-epoch cache release drops reserved memory 291504128 -> 2097152. python -m unittest tests.test_distributed_device_setup -> 8 passed.

Not tested

• CUDA hardware (not available here) — that path is unchanged.
• XPU hardware (not available here) — this PR no longer changes the XPU backend selection (it keeps the pre-PR "nccl" value), so there is nothing XPU-specific left to exercise; see [dist] XPU distributed backend selection (XCCL / oneCCL) is unverified — needs XPU hardware #3777.
• 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.

[li-lizhe]

Thanks for re-running this on the exact head, and for the H100 confirmation.

On the XPU gate: we do not have Intel XPU hardware available in our environment, so we cannot produce an exact-head XPU/XCCL 2-rank run ourselves - the Ascend 2-rank HCCL evidence is the only accelerator we can exercise end to end. Two ways to close this, your call:

1. Keep the XPU arm as it is now and treat it as best effort: it is strictly gated on a positive capability check (is_xccl_available(), with the legacy CCL name only when the build confirms it), and before this patch XPU + DDP could not start at all (the default was "nccl" on a build without NCCL). I can add an inline note saying the XPU branch is unverified on hardware, so nobody reads it as a tested path.

2. Narrow the PR to the Ascend NPU path only: revert the XPU entry to the pre-PR behaviour (leave get_default_dist_backend() at "nccl" for non-NPU) so that every runtime path this PR changes has exact-head real-device evidence, and I file the XPU/XCCL selection as a separate follow-up issue for whoever has XPU hardware.

Which do you prefer?

[li-lizhe]

Thanks for the review. We are taking option (2): this PR is now narrowed to the Ascend NPU path only, and the XPU/XCCL selection is filed separately as #3777.

Why: we do not have Intel XPU hardware available in our environment, so the XPU arm was the one runtime path this PR changed that we could not exercise on a device. The Ascend NPU path, by contrast, is verified end to end at the exact head. Narrowing means every runtime path this PR still changes now has real-device evidence, which matches the repository's hardware-validation requirement.

What changed at head 0a336ba5:

• _ACCELERATOR_DIST_BACKENDS no longer has an "xpu" entry, so an XPU run resolves to exactly the pre-PR value, "nccl" — identical to upstream behaviour for that device. Ascend NPU still resolves to "hccl", CUDA keeps "nccl", and an explicit backend=... still wins.
• The XPU/XCCL selection is filed as [dist] XPU distributed backend selection (XCCL / oneCCL) is unverified — needs XPU hardware #3777 for whoever has XPU hardware, including what evidence a fix would need.
• Tests updated: the previous XPU-default test is replaced by one that pins the narrowed behaviour (XPU stays "nccl" even when a build reports is_xccl_available()), plus one that pins HCCL as the NPU backend on a build that exposes no is_hccl_available() probe. 8 tests.

Re-verified at exact head 0a336ba5 on an Ascend 910B2 container (torch 2.15.0.dev20260917+cpu, torch_npu 2.15.0.dev20260917+gitec69335, 4 devices):

torchrun --nproc_per_node=2 ...
rank 0 local_rank 0 device_type npu backend hccl device npu:0 STEP_LOSS 0.835605 ALL_REDUCE_SUM 1.0 world_size 2
rank 1 local_rank 1 device_type npu backend hccl device npu:1 STEP_LOSS 0.835605 ALL_REDUCE_SUM 1.0 world_size 2
rank 0 CACHE reserved_before 20971520 peak 291504128 after_release 2097152 cpu_device_noop OK
rank 1 CACHE reserved_before 20971520 peak 291504128 after_release 2097152 cpu_device_noop OK
python -m unittest tests.test_distributed_device_setup  ->  8 tests OK

The remaining untested paths are unchanged from before: CUDA hardware (that path is untouched by this PR) and full FSDP / DeepSpeed end-to-end runs.

[LauraGPT]

The requested narrowing is complete at exact head 0a336ba5007e5b2e37ba904dd10b03913e0c9e23: XPU backend selection is restored to its historical behavior and moved to #3777, so this PR no longer advertises an unverified XCCL runtime path.

Fresh independent validation covered both the submitted head and synthetic merge 37ee1f6b321b894a2b05d5a051e7677e8a327ef9 over current main dddd5dc5dfdf759fc58c2a9c0c5a52ab60556922. The submitted eight device-helper tests passed on both trees, and affected-file compileall passed. Recursive tree reconciliation proved the merge changes exactly the five reviewed files; each changed blob matches the head and every other blob matches current main.

The synthetic merge also passed a real NVIDIA H100 80GB HBM3 run with PyTorch 2.8.0+cu128 / CUDA 12.8: NCCL initialization, cuda:0 binding, all-reduce, DDP forward/backward, and accelerator-cache dispatch all completed, with no pre/post GPU process remaining. The author's exact-head Ascend 910B2 two-rank HCCL evidence is retained; I did not independently reproduce NPU hardware. Intel XPU/XCCL remains explicitly unverified and held in #3777. This review does not claim full FSDP or DeepSpeed end-to-end coverage.

Prepared with Codex assistance.