Add global exception hook and post_param_placement_sync barrier

Author: amepas

src/maxdiffusion/generate_flux2klein.py

@@ -138,6 +138,16 @@ def main(argv):
   )
   from maxdiffusion.pipelines.flux.flux2klein_pipeline import FlaxFlux2KleinPipeline
 
+  def handle_exception(exc_type, exc_value, exc_traceback):
+    rank = jax.process_index()
+    cnt = jax.process_count()
+    print(f"\n❌ [HOST {rank}/{cnt}] UNCAUGHT FATAL EXCEPTION: {exc_value}", flush=True)
+    import traceback
+    traceback.print_exception(exc_type, exc_value, exc_traceback)
+    sys.stdout.flush()
+
+  sys.excepthook = handle_exception
+
   config = pyconfig.config
   os.makedirs(config.output_dir, exist_ok=True)
 
@@ -399,6 +409,9 @@ def _put_leaf(param, sharding):
       max_logging.log("All parameters placed on TPU HBM successfully!")
       gc.collect()
       jax.effects_barrier()
+      from jax.experimental import multihost_utils
+      multihost_utils.sync_global_devices("post_param_placement_sync")
+      print(f"[HOST {jax.process_index()}/{jax.process_count()}] Passed post_param_placement_sync barrier successfully! ✅", flush=True)
 
   load_time = time.time() - t_load_start
   max_logging.log(f" -> [TIMING] Total Model Loading & Device Placement: {load_time:.2f} seconds ⏱️\n")