Repository: alpa-projects/alpa

Issue:
[PERF] Move cached property out of `DistributedArray`
Background
---
To manage a tensor distributed on a device mesh, we use `DistributedArray` to represent this concept on the driver side. The class implements apis to load, save, and retrieve data from devices to the driver. Because the tensor is sharded on the device mesh, the `DistributedArray` should know how to get all parts of the data, see [this property](https://github.com/alpa-projects/alpa/blob/main/alpa/device_mesh.py#L1638-L1643). The `DistributedArray` is read-only, so we cache the property.
However, we found `DistributedArray` exceptionally light. In each iteration(a.k.a. generating a token), the driver needs to get a tiny tensor from devices. Although the tensor's sharding spec is the same for each iteration, the `DistributedArray` object differs, and the property needs to be repeatedly computed for each iter.
In this issue, we want to move computing `one_replica_buffer_ids` from `DistributedArray` to `PhysicalDeviceMesh` or even just a util function, because once the mesh shape and sharding spec are kept, we can know how to get and assemble shards to cover the whole tensor.

TODO
---
- [ ] Move `compute_one_replica_ids` to `PhysicalDeviceMesh` or `util`. It only uses sharding spec and mesh shape as inputs, regardless of which array it computes;
- [ ] Support `ReplicatedDistributedArray`.


Discussion:
Hey @ZYHowell, I'd like to take a stab at fixing this issue unless someone else is already working on it. 

Looking at a couple of tests (`runtime/test_device_mesh.py` and `pipeline_parallel/test_mlp.py` mainly) and setting breakpoints, I do see a lot of calls to the `compute_one_replica_ids` but I wasn't sure if that's what you had in mind for this issue. Could you point me to the bit in code that causes the repeated calls you mentioned above? 

>  In each iteration(a.k.a. generating a token), the driver needs to get a tiny tensor from devices.

And to make sure I understand this correctly - we are able to move this property out of `DistributedArray` because given a fixed `physical_mesh` and a fixed `sharding_spec`, any calls to shard an array using [`shard_args_to_arrays`](https://github.com/alpa-projects/alpa/blob/main/alpa/device_mesh.py#L1317) will shard the arrays similarly and hence we do not need to recompute the `one_replica_buffer_ids` for this given (`physical_mesh`, `sharding_spec`) pair  every time [`_value`](https://github.com/alpa-projects/alpa/blob/main/alpa/device_mesh.py#L1649) is called?


> Hey @ZYHowell, I'd like to take a stab at fixing this issue unless someone else is already working on it.
> 
> Looking at a couple of tests (`runtime/test_device_mesh.py` and `pipeline_parallel/test_mlp.py` mainly) and setting breakpoints, I do see a lot of calls to the `compute_one_replica_ids` but I wasn't sure if that's what you had in mind for this issue. Could you point me to the bit in code that causes the repeated calls you mentioned above?
> 
> > In each iteration(a.k.a. generating a token), the driver needs to get a tiny tensor from devices.
> 
> And to make sure I understand this correctly - we are able to move this property out of `DistributedArray` because given a fixed `physical_mesh` and a fixed `sharding_spec`, any calls to shard an array using [`shard_args_to_arrays`](https://github.com/alpa-projects/alpa/blob/main/alpa/device_mesh.py#L1317) will shard the arrays similarly and hence we do not need to recompute the `one_replica_buffer_ids` for this given (`physical_mesh`, `sharding_spec`) pair every time [`_value`](https://github.com/alpa-projects/alpa/blob/main/alpa/device_mesh.py#L1649) is called?

Hello @Vatshank,

Thank you for asking if someone is working on the issue. I have been working on it for a couple of weeks, so would it be okay if I continue? I am very sorry for not commenting here earlier for an official assignment.

Regarding your questions, I will try my best to answer them in a few days.


Resolving patch:
diff --git a/alpa/device_mesh.py b/alpa/device_mesh.py
--- a/alpa/device_mesh.py
+++ b/alpa/device_mesh.py
@@ -638,6 +638,7 @@ class PhysicalDeviceMesh(ABC):
     num_devices_per_host: int
     mesh_id: int
     operation_executables: dict
+    one_replica_ids: dict
 
     def get_signature(self) -> str:
         """Return a signature string that contains the mesh shape and GPU
@@ -648,6 +649,27 @@ def get_signature(self) -> str:
         ret = ret.replace(" ", "-")
         return ret
 
+    def _compute_one_replica_ids(self, indices, aval_shape, sharding_spec):
+        # Tuple (aval_shape, sharding_spec) is 1-1 mapped to indices
+        # used to compute one_replica_ids
+        if (aval_shape, sharding_spec) in self.one_replica_ids:
+            return self.one_replica_ids[(aval_shape, sharding_spec)]
+
+        one_replica_indices = []
+        one_replica_host_local_ids = []
+        seen_index_hashes = set()
+        for i, index in enumerate(indices):
+            hashed_index = _hashable_index(index)
+            if hashed_index not in seen_index_hashes:
+                one_replica_indices.append(i)
+                one_replica_host_local_ids.append(
+                    divmod(i, self.num_devices_per_host))
+                seen_index_hashes.add(hashed_index)
+        self.one_replica_ids[(
+            aval_shape,
+            sharding_spec)] = one_replica_indices, one_replica_host_local_ids
+        return one_replica_indices, one_replica_host_local_ids
+
     @property
     def shape(self):
         return self.num_hosts, self.num_devices_per_host
@@ -845,6 +867,7 @@ def __init__(self, devices: Sequence["Device"] = None):
         self.mesh_id = -1
         self.device_strs = []
         self.operation_executables = {}
+        self.one_replica_ids = {}
 
         self.backend = xb.get_backend(global_config.backend)
 
@@ -974,6 +997,7 @@ def __init__(self,
         self.workers = None
         self.service_server = None
         self.operation_executables = {}
+        self.one_replica_ids = {}
         self.namespace = namespace
 
         if devices is not None:
@@ -1508,8 +1532,6 @@ def __init__(self,
         self.shape = self.aval.shape
         self.dtype = self.aval.dtype
         self._npy_value = None
-        self._one_replica_host_local_ids = None
-        self._one_replica_buffer_ids = None
         self._fetched_np_buffers = None
         self._fetched_np_buffers_ref = None
         self.skip_shard_args_check = False
@@ -1616,34 +1638,16 @@ def load(cls, path: str, aval: ShapedArray, device_mesh: PhysicalDeviceMesh,
         return DistributedArray(device_mesh, aval, sharding_spec, ary_ref,
                                 indices)
 
-    def _compute_one_replica_ids(self):
-        one_replica_indices = []
-        one_replica_host_local_ids = []
-        seen_index_hashes = set()
-        for i, index in enumerate(self.indices):
-            hashed_index = _hashable_index(index)
-            if hashed_index not in seen_index_hashes:
-                one_replica_indices.append(i)
-                one_replica_host_local_ids.append(
-                    divmod(i, self.device_mesh.num_devices_per_host))
-                seen_index_hashes.add(hashed_index)
-        self._one_replica_buffer_ids = one_replica_indices
-        self._one_replica_host_local_ids = one_replica_host_local_ids
-
-    # TODO(yonghao): to make ._value faster(in reorder buffer), cache different
-    # buffers with the same mesh shape and sharding spec.
     @property
     def one_replica_buffer_ids(self):
         """Indices of buffers containing one complete copy of the array data."""
-        if self._one_replica_buffer_ids is None:
-            self._compute_one_replica_ids()
-        return self._one_replica_buffer_ids
+        return self.device_mesh._compute_one_replica_ids(
+            self.indices, self.aval.shape, self.sharding_spec)[0]
 
     @property
     def one_replica_host_local_ids(self):
-        if self._one_replica_host_local_ids is None:
-            self._compute_one_replica_ids()
-        return self._one_replica_host_local_ids
+        return self.device_mesh._compute_one_replica_ids(
+            self.indices, self.aval.shape, self.sharding_spec)[1]
 
     @property
     def _value(self):


Test patch: