fix(rocm): stop overwriting ROCR_VISIBLE_DEVICES in apply_gpu_ids (#6123)

* fix(rocm): stop overwriting ROCR_VISIBLE_DEVICES in apply_gpu_ids

ROCR_VISIBLE_DEVICES uses HSA agent-level indexing, not physical GPU
indices. Setting it to a bare integer breaks multi-GPU ROCm systems
where the parent already set ROCR_VISIBLE_DEVICES=0,1: narrowing to
1 causes torch.cuda.is_available() to return False in the training
worker, producing a misleading 'no HIP accelerator' error even on a
correctly configured ROCm host.

HIP_VISIBLE_DEVICES is sufficient for GPU selection on ROCm.
Leave ROCR_VISIBLE_DEVICES inherited from the parent environment.

* test(rocm): update apply_gpu_ids test to assert ROCR_VISIBLE_DEVICES is not overwritten
This commit is contained in:
Leo Borcherding 2026-06-11 10:39:36 -05:00 committed by GitHub
parent 181288e118
commit 3964f44f02
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 14 additions and 4 deletions

View file

@ -1818,7 +1818,13 @@ def apply_gpu_ids(gpu_ids) -> None:
)
if _is_rocm:
os.environ["HIP_VISIBLE_DEVICES"] = value
os.environ["ROCR_VISIBLE_DEVICES"] = value
# ROCR_VISIBLE_DEVICES operates at the HSA agent level and uses
# different indexing semantics to HIP_VISIBLE_DEVICES. Setting it
# to a physical GPU index breaks multi-GPU ROCm systems where the
# parent already set ROCR_VISIBLE_DEVICES (e.g. "0,1"): narrowing
# to "1" causes torch.cuda.is_available() to return False in the
# worker subprocess. HIP_VISIBLE_DEVICES is sufficient for GPU
# selection on ROCm -- leave ROCR_VISIBLE_DEVICES inherited.
_visible_gpu_count = None
if _is_rocm:
logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s' (rocm)", value)

View file

@ -1614,14 +1614,18 @@ class TestApplyGpuIdsRocmFallback:
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert 'getattr(_torch.version, "hip", None)' in func_body
def test_apply_gpu_ids_sets_hip_and_rocr_visible_devices(self):
"""apply_gpu_ids should set both HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES on ROCm."""
def test_apply_gpu_ids_sets_hip_but_not_rocr_visible_devices(self):
"""apply_gpu_ids should set HIP_VISIBLE_DEVICES but leave ROCR_VISIBLE_DEVICES inherited.
ROCR_VISIBLE_DEVICES uses HSA agent-level indexing, not physical GPU indices.
Overwriting it breaks multi-GPU ROCm systems (see issue #6118).
"""
hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
source = hw_path.read_text(encoding = "utf-8")
func_start = source.find("def apply_gpu_ids")
func_body = source[func_start : source.find("\ndef ", func_start + 1)]
assert 'os.environ["HIP_VISIBLE_DEVICES"] = value' in func_body
assert 'os.environ["ROCR_VISIBLE_DEVICES"] = value' in func_body
assert 'os.environ["ROCR_VISIBLE_DEVICES"] = value' not in func_body
def test_apply_gpu_ids_rocm_fallback_is_guarded_by_try_except(self):
"""torch import in apply_gpu_ids must be wrapped in try/except so a missing torch never crashes."""