Multi-GPU non-CUDA: unconditional torch.cuda.set_device() in worksplit-multigpu hot paths
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 6h
- Merged PRs (30d)
- 155
Description
Tracking issue for non-CUDA gaps in the multi-GPU code paths on the `worksplit-multigpu` branch. These do **not** affect single-GPU users (those paths are fixed in #14068) and only trigger when a workflow actually attaches `MultiGPU_WorkUnits` on a non-NVIDIA backend with ≥2 visible devices. Filing for visibility before the eventual merge into `master`; not urgent.
## Affected code
### 1. `comfy/multigpu.py::MultiGPUThreadPool._worker_loop` — [line 38](https://github.com/Comfy-Org/ComfyUI/blob/worksplit-multigpu/comfy/multigpu.py#L38)
`python
def _worker_loop(self, device: torch.device, work_q: queue.Queue, result_q: queue.Queue):
try:
torch.cuda.set_device(device) # <-- unconditional
except Exception as e:
...
`
The `try/except` only catches the error; the worker then enters a degraded loop that just re-raises on every submit. On XPU / NPU / MPS / DirectML, no usable work ever runs and every multi-GPU sampling step hits `Exception` immediately.
### 2. `comfy/samplers.py::_calc_cond_batch_multigpu::_handle_batch` — [line 476](https://github.com/Comfy-Org/ComfyUI/blob/worksplit-multigpu/comfy/samplers.py#L476)
`python
def _handle_batch(device, batch_tuple, results):
try:
# TODO: non-NVIDIA support -- guard with `if device.type == \"cuda\":` once
# we extend multigpu QA beyond CUDA. Unconditional call crashes on
# XPU/NPU/MPS/CPU/DirectML backends.
torch.cuda.set_device(device)
...
`
TODO is already in-source; just needs the gate.
### Related (cosmetic / async correctness)
3. `comfy/samplers.py` `_handle_batch` also has a TODO about `output_device`-side `torch.cuda.synchronize` being needed when extending beyond NVIDIA. Worth tracking together.
## Suggested fix
Wrap both `torch.cuda.set_device(device)` calls with:
`python
if device.type == \"cuda\":
torch.cuda.set_device(device)
elif device.type == \"xpu\":
torch.xpu.set_device(device)
elif device.type == \"npu\":
torch.npu.set_device(device)
# else: MPS / CPU / DirectML — no explicit per-thread device context required
`
And add the `torch.cuda.synchronize(output_device)` cross-stream sync once extending QA to non-NVIDIA.
## Out of scope (already fixed)
- AMD/ROCm + DirectML `unload_all_models()` silent no-op (single-GPU regression) — fixed in #14068.
- `--enable-dynamic-vram` `torch.cuda.device_count()` crash on non-NVIDIA (single-GPU) — fixed in #14068.
## References
- Code review thread: https://ampcode.com/threads/T-019e5164-c7b3-7429-988e-245a004436e0
- PR #14052 (threaded loader fix, merged)
- PR #14068 (single-GPU non-CUDA fixes)
Contributor guide
Research direction
Start with MultiGPUThreadPool._worker_loop in comfy/multigpu.py and _calc_cond_batch_multigpu::_handle_batch in comfy/samplers.py, focusing on the unconditional torch.cuda.set_device calls. Compare device handling for CUDA, XPU, and NPU, and review the output_device synchronization TODO. Done means multi-GPU work can run on supported non-CUDA backends without immediate exceptions, with the related synchronization addressed or explicitly scoped.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 55/100