From 156610a04e443735c26fc64c2257df5b9be16c9d Mon Sep 17 00:00:00 2001 From: tzhouam Date: Sun, 4 Oct 2026 04:43:03 +0800 Subject: [PATCH] Synchronize CPU ring workers before group teardown Signed-off-by: tzhouam (cherry picked from commit 4ad71d827040f255b936030badbae587428410d4) Signed-off-by: tzhouam --- tests/diffusion/distributed/test_comm.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/diffusion/distributed/test_comm.py b/tests/diffusion/distributed/test_comm.py index a8d8a032a..3c4f9ba03 100644 --- a/tests/diffusion/distributed/test_comm.py +++ b/tests/diffusion/distributed/test_comm.py @@ -230,6 +230,10 @@ def _run_ring_p2p( atol=atol, msg=f"[Rank {local_rank}] Ring P2P data mismatch", ) + if device_kind == "cpu": + # P2P completion only synchronizes neighbors. Keep fast CPU ranks + # alive until every rank has finished constructing and using groups. + torch.distributed.barrier() finally: destroy_distributed_env()