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(
):
with monkeypatch.context() as m:
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
m.delenv("HIP_VISIBLE_DEVICES", raising=False)
device = torch.device(f"cuda:{rank}")
torch.cuda.set_device(device)
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
......@@ -92,6 +93,7 @@ def eager_allreduce(
):
with monkeypatch.context() as m:
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
m.delenv("HIP_VISIBLE_DEVICES", raising=False)
device = torch.device(f"cuda:{rank}")
torch.cuda.set_device(device)
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