Args should default to true, removed tensor_list operator for cpus

axonn-ai · Feb 5, 2024 · 472d355 · 472d355
1 parent e3bfa64
commit 472d355
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/axonn/intra_layer/communication.py b/axonn/intra_layer/communication.py
@@ -48,7 +48,7 @@ def _gather(input_, dim, process_group=None, cache=False):
         tensor_list = [
             torch.empty_like(input_) for _ in range(dist.get_world_size(process_group))
         ]
-        tensor_list[rank] = input_
+        # tensor_list[rank] = input_
         dist.all_gather(tensor_list, input_, group=process_group)
 
         # Note: torch.cat already creates a contiguous tensor.

diff --git a/axonn/intra_layer/fully_connected.py b/axonn/intra_layer/fully_connected.py
@@ -253,8 +253,8 @@ def get_output_feature_size(self):
     def forward(
         self,
         x,
-        scatter_input=False,
-        gather_output=False,
+        scatter_input=True,
+        gather_output=True,
         cache_weights_in_all_gather=False,
     ):
         # gather weights from depth parallel group