From 21cc9b0f980957eb30a034d6dde4dca113ec5af6 Mon Sep 17 00:00:00 2001
From: Oliver Koenig <okoenig@nvidia.com>
Date: Tue, 3 Dec 2024 06:07:46 -0800
Subject: [PATCH] ADLR/megatron-lm!2416 - tests: Add barrier for destroy

---
 tests/unit_tests/conftest.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py
index f166a8179d..4833b30e33 100644
--- a/tests/unit_tests/conftest.py
+++ b/tests/unit_tests/conftest.py
@@ -19,6 +19,7 @@ def pytest_sessionfinish(session, exitstatus):
 def cleanup():
     yield
     if torch.distributed.is_initialized():
+        torch.distributed.barrier()
         torch.distributed.destroy_process_group()