add variable for all-gather dtype

axonn-ai · Feb 26, 2024 · da66c0a · da66c0a
1 parent a9d38c2
commit da66c0a
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/axonn/intra_layer/__init__.py b/axonn/intra_layer/__init__.py
@@ -43,6 +43,7 @@ def gather(
 OVERLAP_REDUCE_SCATTER = False
 OVERLAP_ALL_REDUCE = False
 ALL_GATHER_ITERATOR = None
+ALL_GATHER_DTYPE = torch.bfloat16
 handles = []
 pending_grad_accumulations = []
 weights_cache = {}
@@ -99,10 +100,10 @@ def trigger_async_all_gathers(model):
                     assert weight.ndim == 1
                     output_shape = weight.shape[0] * world_size
                     all_gathered_weight = torch.empty(
-                        output_shape, dtype=weight.dtype, device=weight.device
+                        output_shape, dtype=ALL_GATHER_DTYPE, device=weight.device
                     )
                     handle = dist.all_gather_into_tensor(
-                        all_gathered_weight, weight, group=process_group, async_op=True
+                        all_gathered_weight, weight.to(ALL_GATHER_DTYPE), group=process_group, async_op=True
                     )
                 weights_cache[weight] = [all_gathered_weight, handle]
             yield