Onboard Flux.2-klein-4B multi-host TPU pod support (TP=4, DP=2, B=16)

Author: amepas

src/maxdiffusion/configs/base_flux2klein.yml

@@ -270,9 +270,6 @@ width: 1024
 batch_size: 4
 interactive: False
 
-# 4B Architecture Dimensions
-depth: 20 # num_single_layers
-num_double_layers: 5
-hidden_size: 3072
-num_attention_heads: 24
+# Note: Architecture dimensions (depth, num_double_layers, num_attention_heads) are
+# automatically inferred from pretrained_model_name_or_path (transformer/config.json).
 

src/maxdiffusion/generate_flux2klein.py

@@ -31,6 +31,7 @@
 
 from maxdiffusion import pyconfig
 from maxdiffusion import max_logging
+from maxdiffusion import max_utils
 from maxdiffusion.max_utils import create_device_mesh
 from maxdiffusion.train_utils import transformer_engine_context
 
@@ -176,10 +177,18 @@ def main(argv):
     repo_id = "black-forest-labs/FLUX.2-klein-9B" if depth_val == 24 else "black-forest-labs/FLUX.2-klein-4B"
   max_logging.log(f"Target model detected: {repo_id}")
 
-  from huggingface_hub import snapshot_download
+  import glob
+
+  cache_pattern = os.path.expanduser(f"~/.cache/huggingface/hub/models--{repo_id.replace('/', '--')}/snapshots/*")
+  snapshots = sorted(glob.glob(cache_pattern))
+  if snapshots:
+    snapshot_dir = snapshots[-1]
+  else:
+    from huggingface_hub import snapshot_download
+
+    snapshot_dir = snapshot_download(repo_id=repo_id)
 
-  max_logging.log(f"Resolving snapshot directory for model '{repo_id}' from HF Hub...")
-  snapshot_dir = snapshot_download(repo_id=repo_id)
+  max_logging.log(f"Host {jax.process_index()} using HF snapshot directory: {snapshot_dir}")
   safetensors_path = os.path.join(snapshot_dir, "transformer")
   vae_safetensors_path = os.path.join(snapshot_dir, "vae", "diffusion_pytorch_model.safetensors")
   text_encoder_path = os.path.join(snapshot_dir, "text_encoder")
@@ -326,7 +335,7 @@ def qwen3_init_fn():
   # 8. Load weights on Host CPU
   max_logging.log("Loading parameters on Host CPU...")
   t_load_start = time.time()
-  cpu_device = jax.devices("cpu")[0]
+  cpu_device = jax.local_devices(backend="cpu")[0]
   with jax.default_device(cpu_device):
     with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
       import flax.linen.spmd as flax_spmd
@@ -355,9 +364,9 @@ def unbox_fn(x):
 
       if config.weights_dtype == "bfloat16":
         max_logging.log("Casting JAX parameters to bfloat16 in-place...")
-        cast_dict_to_bfloat16_inplace(params, device=cpu_device, exclude_keywords=("norm",))
-        cast_dict_to_bfloat16_inplace(vae_params, device=cpu_device, exclude_keywords=("norm",))
-        cast_dict_to_bfloat16_inplace(qwen3_params, device=cpu_device, exclude_keywords=("norm",))
+        cast_dict_to_bfloat16_inplace(params, exclude_keywords=("norm",))
+        cast_dict_to_bfloat16_inplace(vae_params, exclude_keywords=("norm",))
+        cast_dict_to_bfloat16_inplace(qwen3_params, exclude_keywords=("norm",))
         vae_bn_mean = vae_bn_mean.astype(jnp.bfloat16)
         vae_bn_std = vae_bn_std.astype(jnp.bfloat16)
 
@@ -371,7 +380,7 @@ def unbox_fn(x):
       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 = jax.tree_util.tree_map(max_utils.device_put_replicated, params, transformer_shardings)
         except Exception as err:
           max_logging.log("\n❌ jax.device_put(params, transformer_shardings) FAILED!")
           flat_p = flax.traverse_util.flatten_dict(params)
@@ -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 = jax.tree_util.tree_map(max_utils.device_put_replicated, vae_params, vae_shardings)
         max_logging.log("Putting qwen3_params on TPU HBM...")
-        qwen3_params = jax.device_put(qwen3_params, qwen3_shardings)
+        qwen3_params = jax.tree_util.tree_map(max_utils.device_put_replicated, qwen3_params, qwen3_shardings)
       max_logging.log("All parameters placed on TPU HBM successfully!")
       gc.collect()
       jax.effects_barrier()

src/maxdiffusion/models/flux/util.py

@@ -389,7 +389,6 @@ def cast_dict_to_bfloat16_inplace(d, device=None, exclude_keywords=None, parent_
   """
   import gc
   import jax.numpy as jnp
-  import jax
 
   for k, v in list(d.items()):
     current_key = f"{parent_key}.{k}" if parent_key else str(k)
@@ -399,10 +398,7 @@ def cast_dict_to_bfloat16_inplace(d, device=None, exclude_keywords=None, parent_
       is_excluded = exclude_keywords and any(kw.lower() in current_key.lower() for kw in exclude_keywords)
       target_dtype = jnp.float32 if is_excluded else jnp.bfloat16
 
-      if device is not None:
-        d[k] = jax.device_put(jnp.array(v, dtype=target_dtype), device=device)
-      else:
-        d[k] = jnp.array(v, dtype=target_dtype)
+      d[k] = v.astype(target_dtype)
       if hasattr(d[k], "block_until_ready"):
         d[k].block_until_ready()
       del v