Internal change.

PiperOrigin-RevId: 686564643
google · Oct 16, 2024 · 71344e4 · 71344e4
1 parent 8ef917e
commit 71344e4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/checkpoint/orbax/checkpoint/experimental/emergency/checkpoint_manager.py b/checkpoint/orbax/checkpoint/experimental/emergency/checkpoint_manager.py
@@ -1097,6 +1097,8 @@ def create_zeros(shape_dtype_tup):
       zeros_pytree = create_zeros(tuple(shape_dtypes))
       in_tree = tuple(zeros_pytree)
 
+    multihost.sync_global_processes('local_restore_pre_broadcast')
+
     start_broadcast = time.time()
     shared_states, _ = multislice.broadcast_one_replica_to_all(
         in_tree,