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

fix(studio/rocm): dedup rocminfo gfx tokens and honor disabled visibility

Browse source 
Follow-up to 8793aef0. rocminfo / hipinfo emit each gfx target multiple
times per GPU (Name, ISA triple, marketing name), so the prior re.findall
indexing returned the wrong device when HIP_VISIBLE_DEVICES picked GPU 1
on a mixed-arch host -- the helper picked the second occurrence of GPU 0
instead. Collapse to unique tokens (insertion-ordered) before indexing.

Also handle HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES values of '' and
'-1' as "no AMD visible" (matches the rest of Studio's visibility code),
returning None so the planner does not pick a Lemonade asset for a
hidden GPU.
This commit is contained in:
Daniel Han 2026-05-19 09:01:11 +00:00
commit 25d4ab63a4
1 changed files with 21 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

29
studio/install_llama_prebuilt.py
View file

@ -2633,16 +2633,29 @@ def _pick_rocm_gfx_target(out: str) -> str | None:
(e.g. Strix Halo gfx1151 + discrete RX 7900 gfx1100). Respect (e.g. Strix Halo gfx1151 + discrete RX 7900 gfx1100). Respect
HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES so the asset matches what HIP HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES so the asset matches what HIP
actually runs on. Falls back to the first GPU when no env var is set. actually runs on. Falls back to the first GPU when no env var is set.
rocminfo / hipinfo print the same gfx token multiple times per GPU
(Name, ISA, marketing-name), so we de-duplicate consecutive-or-repeated
matches before indexing -- otherwise HIP_VISIBLE_DEVICES=1 picks the
second occurrence of GPU 0 instead of GPU 1. Empty / "-1" env values
mean no AMD GPU is visible to HIP and yield None.
""" """
_tokens = re.findall(r"gfx[1-9][0-9a-z]{2,3}", out.lower()) raw = re.findall(r"gfx[1-9][0-9a-z]{2,3}", out.lower())
if not _tokens: if not raw:
return None return None
_vis = ( # Dict preserves insertion order on Python 3.7+; collapses duplicates.
os.environ.get("HIP_VISIBLE_DEVICES") _tokens = list(dict.fromkeys(raw))
or os.environ.get("ROCR_VISIBLE_DEVICES") _vis_raw = None
or "" for _env in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"):
) _val = os.environ.get(_env)
if _vis: if _val is not None:
_vis_raw = _val
break
if _vis_raw is not None:
_vis = _vis_raw.strip()
# Empty or "-1" means "no AMD GPU visible" (matches the rest of Studio).
if _vis == "" or _vis == "-1":
return None
_first = _vis.split(",")[0].strip() _first = _vis.split(",")[0].strip()
try: try:
_idx = int(_first) _idx = int(_first)