Improve log message and print only on hosts that broadcast. Also bloc…

…k subtree instead of tree. PiperOrigin-RevId: 658166778
google · Aug 1, 2024 · b37e71a · b37e71a
1 parent 6a7fdf6
commit b37e71a
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/checkpoint/orbax/checkpoint/multihost/multislice_utils.py b/checkpoint/orbax/checkpoint/multihost/multislice_utils.py
@@ -232,9 +232,11 @@ def pre_jit(x, per_replica_sharding):
         out_shardings=out_sharding,
     )(in_tree_sharded)
     out_tree.extend(out_subtree)
-    jax.block_until_ready(out_tree)
+    jax.block_until_ready(out_subtree)
     start = end
-  logging.info('Number of broadcasts: %d', num_broadcasts)
+
+  if is_source:
+    logging.info('Total number of broadcasts: %d', num_broadcasts)
   return tuple(out_tree), num_broadcasts