Repository: alpa-projects/alpa

Issue:
Rendezvous crashes while trying to access already killed NCCLUniqueIDStore
**Please describe the bug**

We often see the following error in our CI:
```
Epoch ... :   0%|          | 0/1 [00:00<?, ?it/s]
Training...:   0%|          | 0/100 [00:00<?, ?it/s]2023-03-27 10:09:39,484	ERROR worker.py:406 -- Unhandled error (suppress with 'RAY_IGNORE_UNHANDLED_ERRORS=1'): ray::MeshHostWorker.create_and_set_cross_mesh_communicators() (pid=1133, ip=172.31.220.237, repr=<alpa.device_mesh.MeshHostWorker object at 0x7f9f11e21c10>)
  File "/home/ray/jun-alpa/alpa/alpa/device_mesh.py", line 413, in create_and_set_cross_mesh_communicators
    g.create_and_set_xla_communicators(devices, key)
  File "/home/ray/jun-alpa/alpa/alpa/collective/collective_group/nccl_collective_group.py", line 454, in create_and_set_xla_communicators
    nccl_uid = self._rendezvous_nccl_uid(self.rank, comm_key,
  File "/home/ray/jun-alpa/alpa/alpa/collective/collective_group/nccl_collective_group.py", line 723, in _rendezvous_nccl_uid
    if rendezvous.get_access_counter() == max_counter:
  File "/home/ray/jun-alpa/alpa/alpa/collective/collective_group/base_collective_group.py", line 111, in get_access_counter
    return ray.get(self._store.get_access_counter.remote())
ray.exceptions.RayActorError: The actor died unexpectedly before finishing this task.
	class_name: NCCLUniqueIDStore
	actor_id: 1f4a015eb27d5a8f9a02b5df02000000
	pid: 4113
	name: 2ca6a4c123962c052614fee4d865b186
	namespace: alpa_default_space
	ip: 172.31.165.49
The actor is dead because it was killed by `ray.kill`.
2023-03-27 10:09:39,489	ERROR worker.py:406 -- Unhandled error (suppress with 'RAY_IGNORE_UNHANDLED_ERRORS=1'): ray::MeshHostWorker.create_and_set_cross_mesh_communicators() (pid=1147, ip=172.31.241.33, repr=<alpa.device_mesh.MeshHostWorker object at 0x7f4330afa790>)
  File "/home/ray/jun-alpa/alpa/alpa/device_mesh.py", line 413, in create_and_set_cross_mesh_communicators
    g.create_and_set_xla_communicators(devices, key)
  File "/home/ray/jun-alpa/alpa/alpa/collective/collective_group/nccl_collective_group.py", line 454, in create_and_set_xla_communicators
    nccl_uid = self._rendezvous_nccl_uid(self.rank, comm_key,
  File "/home/ray/jun-alpa/alpa/alpa/collective/collective_group/nccl_collective_group.py", line 723, in _rendezvous_nccl_uid
    if rendezvous.get_access_counter() == max_counter:
  File "/home/ray/jun-alpa/alpa/alpa/collective/collective_group/base_collective_group.py", line 111, in get_access_counter
    return ray.get(self._store.get_access_counter.remote())
ray.exceptions.RayActorError: The actor died unexpectedly before finishing this task.
	class_name: NCCLUniqueIDStore
	actor_id: 1f4a015eb27d5a8f9a02b5df02000000
	pid: 4113
	name: 2ca6a4c123962c052614fee4d865b186
	namespace: alpa_default_space
	ip: 172.31.165.49
The actor is dead because it was killed by `ray.kill`.
2023-03-27 10:09:39,496	ERROR worker.py:406 -- Unhandled error (suppress with 'RAY_IGNORE_UNHANDLED_ERRORS=1'): ray::MeshHostWorker.create_and_set_cross_mesh_communicators() (pid=14302, ip=172.31.220.187, repr=<alpa.device_mesh.MeshHostWorker object at 0x7f4cd1f0efd0>)
  File "/home/ray/jun-alpa/alpa/alpa/device_mesh.py", line 413, in create_and_set_cross_mesh_communicators
    g.create_and_set_xla_communicators(devices, key)
  File "/home/ray/jun-alpa/alpa/alpa/collective/collective_group/nccl_collective_group.py", line 454, in create_and_set_xla_communicators
    nccl_uid = self._rendezvous_nccl_uid(self.rank, comm_key,
  File "/home/ray/jun-alpa/alpa/alpa/collective/collective_group/nccl_collective_group.py", line 723, in _rendezvous_nccl_uid
    if rendezvous.get_access_counter() == max_counter:
  File "/home/ray/jun-alpa/alpa/alpa/collective/collective_group/base_collective_group.py", line 111, in get_access_counter
    return ray.get(self._store.get_access_counter.remote())
ray.exceptions.RayActorError: The actor died unexpectedly before finishing this task.
	class_name: NCCLUniqueIDStore
	actor_id: 1f4a015eb27d5a8f9a02b5df02000000
	pid: 4113
	name: 2ca6a4c123962c052614fee4d865b186
	namespace: alpa_default_space
	ip: 172.31.165.49
The actor is dead because it was killed by `ray.kill`.
```
The issue is that some Rendezvous deleted the underlying store while others are still trying to access it.

https://github.com/alpa-projects/alpa/blob/main/alpa/collective/collective_group/nccl_collective_group.py#L728

**Please describe the expected behavior**

have a feeling the NCCLUniqueIDStore doesn't have to be detached.
It can simply fate-share with the rank 0 workers, so nobody has to delete it.
we will also make sure the rank 0 workers is the last to get torn down maybe.

**To Reproduce**
Behavior is a little in-deterministic. Run a simple training job 10 times, and this will happen roughly 5 out of the 10 times.




Discussion:

Resolving patch:
diff --git a/alpa/collective/collective_group/base_collective_group.py b/alpa/collective/collective_group/base_collective_group.py
--- a/alpa/collective/collective_group/base_collective_group.py
+++ b/alpa/collective/collective_group/base_collective_group.py
@@ -112,7 +112,6 @@ def get_access_counter(self):
 
     def destroy_store(self):
         """Delete the named actor."""
-        ray.kill(self._store)
         self._store = None
 
 
diff --git a/alpa/collective/collective_group/nccl_collective_group.py b/alpa/collective/collective_group/nccl_collective_group.py
--- a/alpa/collective/collective_group/nccl_collective_group.py
+++ b/alpa/collective/collective_group/nccl_collective_group.py
@@ -569,8 +569,7 @@ def generate_nccl_uid():
         group_uid = nccl_util.get_nccl_unique_id()
         return group_uid
 
-    @staticmethod
-    def _generate_nccl_uid(key):
+    def _generate_nccl_uid(self, key):
         """Generate an NCCL unique ID for initializing communicators.
 
         The method will also create a KV store using Ray named actor and store
@@ -587,9 +586,9 @@ def _generate_nccl_uid(key):
         store_name = get_store_name(key)
         # Avoid a potential circular dependency in ray/actor.py
         from alpa.collective.util import NCCLUniqueIDStore  # pylint: disable=import-outside-toplevel
-        store = NCCLUniqueIDStore.options(
-            name=store_name, lifetime="detached").remote(store_name)
-        ray.get([store.set_id.remote(group_uid)])
+        self._store = NCCLUniqueIDStore.options(
+            name=store_name).remote(store_name)
+        ray.get([self._store.set_id.remote(group_uid)])
         return group_uid
 
     def _collective(self,
diff --git a/alpa/collective/collective_group/xla_nccl_collective_group.py b/alpa/collective/collective_group/xla_nccl_collective_group.py
--- a/alpa/collective/collective_group/xla_nccl_collective_group.py
+++ b/alpa/collective/collective_group/xla_nccl_collective_group.py
@@ -347,8 +347,7 @@ def generate_nccl_uid():
         group_uid = xla_nccl_util.get_nccl_unique_id()
         return group_uid
 
-    @staticmethod
-    def _generate_nccl_uid(key):
+    def _generate_nccl_uid(self, key):
         """Generate an NCCL unique ID for initializing communicators.
 
         The method will also create a KV store using Ray named actor and store
@@ -365,9 +364,9 @@ def _generate_nccl_uid(key):
         store_name = get_store_name(key)
         # Avoid a potential circular dependency in ray/actor.py
         from alpa.collective.util import NCCLUniqueIDStore  # pylint: disable=import-outside-toplevel
-        store = NCCLUniqueIDStore.options(
-            name=store_name, lifetime="detached").remote(store_name)
-        ray.get([store.set_id.remote(group_uid)])
+        self._store = NCCLUniqueIDStore.options(
+            name=store_name).remote(store_name)
+        ray.get([self._store.set_id.remote(group_uid)])
         return group_uid
 
     # unimplemented


Test patch: