ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: fix ROCm visibility precedence + narrow ROCm child env

Browse source 
Two reviewer-flagged correctness bugs in the AMD GPU probe path.

1) ROCm visibility precedence was reversed. torch.cuda enumerates GPUs
   relative to HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES on ROCm builds,
   but the probe's env-var lookup checked CUDA_VISIBLE_DEVICES first. With
   CUDA_VISIBLE_DEVICES=0,1 and HIP_VISIBLE_DEVICES=6,7 the probe returned
   [(0, ...), (1, ...)] when torch's view was actually [(6, ...), (7, ...)].
   The wrong physical IDs flowed downstream into CUDA_VISIBLE_DEVICES for
   the llama-server subprocess, pinning it to GPUs 0,1 instead of 6,7.

   Fix: branch on torch.version.hip. On ROCm, prefer HIP > ROCR > CUDA
   (matches torch's own ordering). On NVIDIA, use CUDA only -- ignoring
   any HIP/ROCR vars the parent happens to have set.

2) Child env narrowing only set CUDA_VISIBLE_DEVICES. On ROCm, llama-server
   honors HIP/ROCR; if the parent shell exported HIP_VISIBLE_DEVICES=4,5
   and the selector picked just GPU 4, the child still saw both because
   we never narrowed HIP/ROCR. Now we set all three on ROCm so the AMD
   subprocess actually sees the planned subset.

Both branches verified via temp/pr_simulation/sim_5172_rocm_precedence.py
(7/7 cases pass), including the reviewer's verbatim R5 case
(CVD=0,1 + HIP/ROCR=6,7).
This commit is contained in:
Daniel Han 2026-04-24 15:50:25 +00:00
commit eff55fb8db
1 changed files with 28 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

41
studio/backend/core/inference/llama_cpp.py
View file

@ -609,19 +609,23 @@ class LlamaCppBackend:
return []
if not hasattr(torch.cuda, "mem_get_info"):
return []
# torch.cuda enumerates GPUs RELATIVE to CUDA_VISIBLE_DEVICES
# (or HIP_VISIBLE_DEVICES on ROCm). Downstream we feed these
# IDs back into CUDA_VISIBLE_DEVICES for the llama-server
# subprocess, so we must translate visible ordinals back to
# physical indices first; otherwise launching with
# ``CUDA_VISIBLE_DEVICES=2,3`` would get rewritten to
# torch.cuda enumerates GPUs RELATIVE to the visibility mask.
# On NVIDIA builds the mask is CUDA_VISIBLE_DEVICES; on AMD
# ROCm builds it is HIP_VISIBLE_DEVICES (or ROCR_VISIBLE_DEVICES
# if HIP is unset). Downstream we feed these IDs back into the
# llama-server subprocess as CVD, so we must translate visible
# ordinals back to physical indices first; otherwise launching
# with ``CUDA_VISIBLE_DEVICES=2,3`` would get rewritten to
# ``CUDA_VISIBLE_DEVICES=0,1`` and target the wrong GPUs.
physical_ids: Optional[list[int]] = None
cvd = (
os.environ.get("CUDA_VISIBLE_DEVICES")
or os.environ.get("HIP_VISIBLE_DEVICES")
or os.environ.get("ROCR_VISIBLE_DEVICES")
)
if getattr(torch.version, "hip", None) is not None:
cvd = (
os.environ.get("HIP_VISIBLE_DEVICES")
or os.environ.get("ROCR_VISIBLE_DEVICES")
or os.environ.get("CUDA_VISIBLE_DEVICES")
)
else:
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
if cvd and cvd.strip():
try:
physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()]
@ -1792,9 +1796,20 @@ class LlamaCppBackend:
f"{new_ld}:{existing_ld}" if existing_ld else new_ld
)
# Pin to selected GPU(s) via CUDA_VISIBLE_DEVICES
# Pin to selected GPU(s). On ROCm, llama-server (and any torch
# in the subprocess) honors HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES;
# narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child seeing
# the full HIP/ROCR set the parent inherited.
if gpu_indices is not None:
env["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in gpu_indices)
pinned = ",".join(str(i) for i in gpu_indices)
env["CUDA_VISIBLE_DEVICES"] = pinned
try:
import torch as _torch
if getattr(_torch.version, "hip", None) is not None:
env["HIP_VISIBLE_DEVICES"] = pinned
env["ROCR_VISIBLE_DEVICES"] = pinned
except Exception:
pass
self._stdout_lines = []
self._process = subprocess.Popen(