Revert "Set all_reduce_token to None when exiting" (#6321)

pytorch · Jan 19, 2024 · ad669d5 · ad669d5
1 parent dce624c
commit ad669d5
Show file tree

Hide file tree

Showing 4 changed files with 1 addition and 29 deletions.
diff --git a/test/run_tests.sh b/test/run_tests.sh
@@ -234,7 +234,6 @@ function run_mp_op_tests {
   run_test "$CDIR/test_mp_save.py"
   run_test "$CDIR/test_mp_mesh_reduce.py"
   run_test "$CDIR/test_mp_sync_batch_norm.py"
-  run_test "$CDIR/test_mp_early_exit.py"
   run_pt_xla_debug "$CDIR/debug_tool/test_mp_pt_xla_debug.py"
   run_xla_backend_mp "$CDIR/test_torch_distributed_all_gather_xla_backend.py"
   run_xla_backend_mp "$CDIR/test_torch_distributed_all_reduce_xla_backend.py"

diff --git a/test/test_mp_early_exit.py b/test/test_mp_early_exit.py
diff --git a/test/test_zero1.py b/test/test_zero1.py
@@ -13,6 +13,7 @@
 class XlaZeRO1Test(TestCase):
 
   @unittest.skipIf(xr.device_type() == 'TPU', "Crash on TPU")
+  @unittest.skipIf(xr.device_type() == 'CUDA', "Crash on CUDA")
   def test_zero1(self):
     device = xm.xla_device()
 

diff --git a/torch_xla/__init__.py b/torch_xla/__init__.py
@@ -148,8 +148,6 @@ def _setup_tpu_vm_library_path() -> bool:
 
 
 def _prepare_to_exit():
-  device = _XLAC._xla_get_default_device()
-  _XLAC._set_all_reduce_token(device, None)
   _XLAC._prepare_to_exit()
   if int(os.environ.get('PT_XLA_DEBUG', '0')):
     _summarize_fn_tracker()