Gradient bucketing using a pre-defined bucket size cap

pytorch · Mar 1, 2024 · fdb0f9e · fdb0f9e
1 parent 3bee0a7
commit fdb0f9e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/torch_xla/core/xla_model.py b/torch_xla/core/xla_model.py
@@ -41,6 +41,8 @@
 _ALLREDUCE_BUCKET_CAP_MB = 50
 
 XLA_LIB = Library("xla", "DEF")
+# Default bucket size for all-reduce
+_ALLREDUCE_BUCKET_CAP_MB = 50
 
 
 def _init_world_size_ordinal():
@@ -1052,7 +1054,6 @@ def reduce_gradients(optimizer, groups=None, pin_layout=True):
     gradients = reversed(_fetch_gradients(optimizer))
     bucketed_allreduce(gradients)
 
-
 def optimizer_step(optimizer,
                    barrier=False,
                    optimizer_args={},