Slight refactor of _split_local_dist_buffers()

Summary: Iterators enable this so no more bookkeeping on buffer indices.

Reviewed By: gajjanag

Differential Revision: D74159431

fbshipit-source-id: ab448d231d6d6d9bee3bbb7b801f8bd2cdffd859

Author: tsunghsienlee

distributed_shampoo/utils/shampoo_ddp_distributor.py

@@ -299,18 +299,10 @@ def _split_local_dist_buffers(
             )
             split_tensors_list.append(split_tensors)
 
-        # Obtain ordered buffer ranks containing (view of local buffer, rank).
-        splitted_local_dist_buffers = []
-        buffer_indices = [0] * len(
-            local_dist_buffers
-        )  # index counter for each rank for obtaining right buffer
-        for _, rank in buffer_size_ranks:
-            splitted_local_dist_buffers.append(
-                split_tensors_list[rank][buffer_indices[rank]]
-            )
-            buffer_indices[rank] += 1
-
-        return tuple(splitted_local_dist_buffers)
+        split_tensors_iterators = list(map(iter, split_tensors_list))
+        return tuple(
+            next(split_tensors_iterators[rank]) for _, rank in buffer_size_ranks
+        )
 
     def _construct_distributed_buffers(
         self,

distributed_shampoo/utils/shampoo_hsdp_distributor.py

@@ -492,18 +492,10 @@ def _split_local_dist_buffers(
             )
             split_tensors_list.append(split_tensors)
 
-        # Obtain ordered buffer ranks containing (view of local buffer, rank).
-        splitted_local_dist_buffers = []
-        buffer_indices = [0] * len(
-            local_dist_buffers
-        )  # index counter for each rank for obtaining right buffer
-        for _, rank in buffer_size_ranks:
-            splitted_local_dist_buffers.append(
-                split_tensors_list[rank][buffer_indices[rank]]
-            )
-            buffer_indices[rank] += 1
-
-        return tuple(splitted_local_dist_buffers)
+        split_tensors_iterators = list(map(iter, split_tensors_list))
+        return tuple(
+            next(split_tensors_iterators[rank]) for _, rank in buffer_size_ranks
+        )
 
     def _construct_distributed_buffers(
         self,

distributed_shampoo/utils/shampoo_hybrid_shard_distributor.py

@@ -428,18 +428,10 @@ def _split_local_dist_buffers(
             )
             split_tensors_list.append(split_tensors)
 
-        # Obtain ordered buffer ranks containing (view of local buffer, rank).
-        splitted_local_dist_buffers = []
-        buffer_indices = [0] * len(
-            local_dist_buffers
-        )  # index counter for each rank for obtaining right buffer
-        for _, rank in buffer_size_ranks:
-            splitted_local_dist_buffers.append(
-                split_tensors_list[rank][buffer_indices[rank]]
-            )
-            buffer_indices[rank] += 1
-
-        return tuple(splitted_local_dist_buffers)
+        split_tensors_iterators = list(map(iter, split_tensors_list))
+        return tuple(
+            next(split_tensors_iterators[rank]) for _, rank in buffer_size_ranks
+        )
 
     def _construct_distributed_buffers(
         self,