diff --git a/tests/unit_tests/a2a_overlap/test_cuda_graphed_schedule_chunk_1f1b.py b/tests/unit_tests/a2a_overlap/test_cuda_graphed_schedule_chunk_1f1b.py
index 719bd5df18f..85586095bd7 100644
--- a/tests/unit_tests/a2a_overlap/test_cuda_graphed_schedule_chunk_1f1b.py
+++ b/tests/unit_tests/a2a_overlap/test_cuda_graphed_schedule_chunk_1f1b.py
@@ -333,6 +333,8 @@ def _run_test_helper(
 
         return loss_list
 
+    @pytest.mark.flaky
+    @pytest.mark.flaky_in_dev
     @pytest.mark.skipif(
         not (HAVE_TE and is_te_min_version("2.10.0")),
         reason="Partial CUDA graph support requires TransformerEngine version >= 2.10.0",
diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py
index 362d102200e..e251a3c1e7e 100644
--- a/tests/unit_tests/conftest.py
+++ b/tests/unit_tests/conftest.py
@@ -38,6 +38,14 @@ def pytest_sessionfinish(session, exitstatus):
         session.exitstatus = 0
 
 
+@pytest.fixture(scope="session", autouse=True)
+def cleanup():
+    yield
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+        torch.distributed.destroy_process_group()
+
+
 @pytest.fixture(scope="function", autouse=True)
 def set_env():
     if is_te_min_version("1.3"):
diff --git a/tests/unit_tests/transformer/test_cuda_graphs.py b/tests/unit_tests/transformer/test_cuda_graphs.py
index d4866b4839e..4696a3ed439 100644
--- a/tests/unit_tests/transformer/test_cuda_graphs.py
+++ b/tests/unit_tests/transformer/test_cuda_graphs.py
@@ -1013,6 +1013,8 @@ def _run_test_helper(
 
         return torch.tensor(loss_list)
 
+    @pytest.mark.flaky
+    @pytest.mark.flaky_in_dev
     @pytest.mark.skipif(
         not (HAVE_TE and is_te_min_version("2.10.0")),
         reason="Partial CUDA graph UT support requires TransformerEngine version >= 2.10.0",