Add put_params_on_devices helper to support multi-host SPMD parameter placement

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 maxdiffusion.max_utils import device_put_replicated
   from maxdiffusion.models.flux.util import (
       load_and_convert_flux_klein_weights,
       load_and_convert_vae_weights,
@@ -368,12 +369,20 @@ def unbox_fn(x):
       max_logging.log("\n" + "=" * 80)
       max_logging.log("🚀 Pinning all parameters to TPU HBM permanently...")
       max_logging.log("=" * 80 + "\n")
+      def put_params_on_devices(params_tree, shardings_tree):
+        def _put_leaf(param, sharding):
+          if hasattr(sharding, "is_fully_addressable") and sharding.is_fully_addressable:
+            return jax.device_put(param, sharding)
+          return device_put_replicated(param, sharding)
+
+        return jax.tree_util.tree_map(_put_leaf, params_tree, shardings_tree)
+
       max_logging.log("Putting params on TPU HBM...")
       with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
         try:
-          params = jax.device_put(params, transformer_shardings)
+          params = put_params_on_devices(params, transformer_shardings)
         except Exception as err:
-          max_logging.log("\n❌ jax.device_put(params, transformer_shardings) FAILED!")
+          max_logging.log("\n❌ put_params_on_devices(params, transformer_shardings) FAILED!")
           flat_p = flax.traverse_util.flatten_dict(params)
           flat_s = flax.traverse_util.flatten_dict(transformer_shardings)
           k_p = set(flat_p.keys())
@@ -383,9 +392,9 @@ def unbox_fn(x):
           sys.stdout.flush()
           raise err
         max_logging.log("Putting vae_params on TPU HBM...")
-        vae_params = jax.device_put(vae_params, vae_shardings)
+        vae_params = put_params_on_devices(vae_params, vae_shardings)
         max_logging.log("Putting qwen3_params on TPU HBM...")
-        qwen3_params = jax.device_put(qwen3_params, qwen3_shardings)
+        qwen3_params = put_params_on_devices(qwen3_params, qwen3_shardings)
       max_logging.log("All parameters placed on TPU HBM successfully!")
       gc.collect()
       jax.effects_barrier()