Skip to content

GitLab

Unverified Commit 0632ed87 authored by Ryan Rock's avatar Ryan Rock Committed by GitHub
Browse files

[AMD][CI] Fix test_custom_allreduce for A100 testgroup (#34735) 


Signed-off-by: default avatarRyan Rock <ryan.rock@amd.com>
parent aaefc58e
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
tests/distributed/test_custom_all_reduce.py
View file @ 0632ed87
...@@ -33,6 +33,7 @@ def graph_allreduce( ...@@ -33,6 +33,7 @@ def graph_allreduce(
): ):
with monkeypatch.context() as m: with monkeypatch.context() as m:
m.delenv("CUDA_VISIBLE_DEVICES", raising=False) m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
m.delenv("HIP_VISIBLE_DEVICES", raising=False)
device = torch.device(f"cuda:{rank}") device = torch.device(f"cuda:{rank}")
torch.cuda.set_device(device) torch.cuda.set_device(device)
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
...@@ -92,6 +93,7 @@ def eager_allreduce( ...@@ -92,6 +93,7 @@ def eager_allreduce(
): ):
with monkeypatch.context() as m: with monkeypatch.context() as m:
m.delenv("CUDA_VISIBLE_DEVICES", raising=False) m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
m.delenv("HIP_VISIBLE_DEVICES", raising=False)
device = torch.device(f"cuda:{rank}") device = torch.device(f"cuda:{rank}")
torch.cuda.set_device(device) torch.cuda.set_device(device)
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment