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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user