Add multihost_utils.sync_global_devices to prevent premature process exit timeout

Author: amepas

src/maxdiffusion/generate_flux2klein.py

@@ -129,6 +129,7 @@ def main(argv):
   pyconfig.initialize(default_args)
 
   # Import modules after jax.distributed.initialize() has run via pyconfig.initialize()
+  from jax.experimental import multihost_utils
   from maxdiffusion.max_utils import device_put_replicated
   from maxdiffusion.models.flux.util import (
       load_and_convert_flux_klein_weights,
@@ -552,6 +553,7 @@ def _put_leaf(param, sharding):
     max_logging.log("\n=======================================================")
     max_logging.log(f"SUCCESS! Batched generation complete for {config.batch_size} images! 🎨🎉")
     max_logging.log("=======================================================\n")
+    multihost_utils.sync_global_devices("flux2klein_generation_complete")
 
 
 if __name__ == "__main__":