[li-lizhe]

What

reduce_sum_across_dp_group() built the DP all_reduce tensor on CPU whenever CUDA was unavailable, which breaks every non-CUDA accelerator (Ascend NPU, XPU, ...) with RuntimeError: No backend type associated with device type cpu.

The function already receives the dp_mesh that owns the process group, and dp_mesh is guaranteed non-None at this point (the early return above handles None), so the device can be taken from the mesh instead of assuming CUDA. This keeps the CPU case working (torch.device("cpu") for a CPU mesh) and removes the CUDA-only branch.

-    if torch.cuda.is_available():
-        device = torch.device(f"cuda:{torch.cuda.current_device()}")
-    else:
-        device = torch.device("cpu")
+    device = torch.device(dp_mesh.device_type)
     tensor = torch.tensor([local_value], dtype=torch.int64, device=device)

Why it matters

Dataloader.get_state_dict() calls this helper, so on Ascend the dataloader state could not be saved/restored at all.

Verification

Ascend 910B4 (torch 2.15.0.dev + torch_npu), 2 processes via torchrun, real HCCL process group, calling the actual reduce_sum_across_dp_group() from this file (file copied byte-identically to the runner, md5 checked) with DeviceMesh("npu", [0, 1]):

before (unmodified file):

[rank0] current npu device=0 mesh.device_type=npu
[rank0] RAISED RuntimeError: No backend type associated with device type cpu
[rank1] current npu device=1 mesh.device_type=npu
[rank1] RAISED RuntimeError: No backend type associated with device type cpu

after (this PR):

[rank0] current npu device=0 mesh.device_type=npu
[rank1] current npu device=1 mesh.device_type=npu
[rank0] RESULT=6 expected=6
[rank1] RESULT=6 expected=6

Each rank keeps its own device (npu:0 / npu:1) and the collective result is correct (2 ranks x 3 = 6).

Not tested

• Only the npu mesh type was exercised; CUDA/XPU meshes were not run (the change is a pure device lookup).
• The full xtuner training loop was not run - only this helper under a real multi-process HCCL group. The only stubs were the sibling modules imported by dataloader.py (collator/packing/sampler/utils), which are unrelated to the function under test.

Fixes #2126

[li-lizhe]

Gentle ping on this one — it's been quiet for 8 days since the last activity, so I wanted to check whether there's anything else you'd like me to adjust or add before this can move forward. Happy to make changes if so. Thanks!

[li-lizhe]

@PengchengShi00 @jayhenry — would one of you be able to take a look at this when you get a chance?

DP all_reduce tensor is built on the wrong device

The dataloader creates its all_reduce tensor without a device argument, so it lands on the default device while the mesh may be on another one. It now takes the mesh device. Small and contained.

Happy to adjust if the approach doesn't fit.