Fix Device Index for ROCm Ray Workers in MoE Benchmark (#38108)

Signed-off-by: Liwen <53441624+li-liwen@users.noreply.github.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
This commit is contained in:
Liwen
2026-03-28 08:27:11 +00:00
committed by GitHub
parent 58a249bc61
commit 171775f306

View File

@@ -627,9 +627,8 @@ class BenchmarkWorker:
need_device_guard = True
with (
torch.accelerator.device_index(self.device_id)
if need_device_guard
else nullcontext()
# Ray restricts each worker to one GPU; use local index 0
torch.accelerator.device_index(0) if need_device_guard else nullcontext()
):
for idx, config in enumerate(tqdm(search_space)):
try: