Comfy-Org / Comfy-Org/ComfyUI

Multi-GPU non-CUDA: unconditional torch.cuda.set_device() in worksplit-multigpu hot paths

Open
#14,069 1 comment 1 reaction 0 assignees View on GitHub
multigpu non-cuda
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.