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,16 @@ 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 +333,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 +362,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 +378,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 +390,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

@@ -118,18 +118,14 @@ def validate_flax_state_dict(expected_pytree: dict, new_pytree: dict):
   expected_pytree: dict - a pytree that comes from initializing the model.
   new_pytree: dict - a pytree that has been created from pytorch weights.
   """
-  expected_flat = (
-      flatten_dict(expected_pytree) if not isinstance(next(iter(expected_pytree.keys()), None), tuple) else expected_pytree
-  )
+  expected_flat = flatten_dict(expected_pytree) if not isinstance(next(iter(expected_pytree.keys()), None), tuple) else expected_pytree
   new_flat = flatten_dict(new_pytree) if not isinstance(next(iter(new_pytree.keys()), None), tuple) else new_pytree
 
   if len(expected_flat.keys()) != len(new_flat.keys()):
     set1 = set(expected_flat.keys())
     set2 = set(new_flat.keys())
     missing_keys = set1 ^ set2
-    max_logging.log(
-        f"Missing or extra parameter keys count mismatch ({len(expected_flat)} expected vs {len(new_flat)} converted): {missing_keys}"
-    )
+    max_logging.log(f"Missing or extra parameter keys count mismatch ({len(expected_flat)} expected vs {len(new_flat)} converted): {missing_keys}")
 
   for key in expected_flat.keys():
     if key in new_flat.keys():
@@ -273,7 +269,6 @@ def load_flow_model(name: str, eval_shapes: dict, device: str, hf_download: bool
       jax.clear_caches()
   return flax_state_dict
 
-
 # -----------------------------------------------------------------------------
 # Latent Packing & Unpacking Helpers
 # -----------------------------------------------------------------------------
@@ -399,10 +394,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
@@ -450,9 +442,7 @@ def convert_and_transpose_tensor(tensor, transpose=False):
     return jnp.array(tensor, dtype=target_dtype)
 
   # Global layers
-  params["context_embedder"]["kernel"] = convert_and_transpose_tensor(
-      pt_state_dict.pop("context_embedder.weight"), transpose=True
-  )
+  params["context_embedder"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop("context_embedder.weight"), transpose=True)
   params["x_embedder"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop("x_embedder.weight"), transpose=True)
   params["double_stream_modulation_img"]["kernel"] = convert_and_transpose_tensor(
       pt_state_dict.pop("double_stream_modulation_img.linear.weight"), transpose=True
@@ -466,9 +456,7 @@ def convert_and_transpose_tensor(tensor, transpose=False):
   params["proj_out"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop("proj_out.weight"), transpose=True)
 
   # norm_out
-  params["norm_out"]["linear"]["kernel"] = convert_and_transpose_tensor(
-      pt_state_dict.pop("norm_out.linear.weight"), transpose=True
-  )
+  params["norm_out"]["linear"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop("norm_out.linear.weight"), transpose=True)
 
   # time_text_embed (Timestep Embedding)
   if "time_guidance_embed.timestep_embedder.linear_1.weight" in pt_state_dict:
@@ -506,30 +494,18 @@ def convert_and_transpose_tensor(tensor, transpose=False):
     jax_db["attn"]["e_qkv"]["kernel"] = jnp.array(np.concatenate([add_q, add_k, add_v], axis=1), dtype=target_dtype)
 
     # Projections out
-    jax_db["attn"]["i_proj"]["kernel"] = convert_and_transpose_tensor(
-        pt_state_dict.pop(prefix + "attn.to_out.0.weight"), transpose=True
-    )
-    jax_db["attn"]["e_proj"]["kernel"] = convert_and_transpose_tensor(
-        pt_state_dict.pop(prefix + "attn.to_add_out.weight"), transpose=True
-    )
+    jax_db["attn"]["i_proj"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "attn.to_out.0.weight"), transpose=True)
+    jax_db["attn"]["e_proj"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "attn.to_add_out.weight"), transpose=True)
 
     # Norm scales
     jax_db["attn"]["query_norm"]["scale"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "attn.norm_q.weight"))
     jax_db["attn"]["key_norm"]["scale"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "attn.norm_k.weight"))
-    jax_db["attn"]["encoder_query_norm"]["scale"] = convert_and_transpose_tensor(
-        pt_state_dict.pop(prefix + "attn.norm_added_q.weight")
-    )
-    jax_db["attn"]["encoder_key_norm"]["scale"] = convert_and_transpose_tensor(
-        pt_state_dict.pop(prefix + "attn.norm_added_k.weight")
-    )
+    jax_db["attn"]["encoder_query_norm"]["scale"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "attn.norm_added_q.weight"))
+    jax_db["attn"]["encoder_key_norm"]["scale"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "attn.norm_added_k.weight"))
 
     # SwiGLU MLPs
-    jax_db["ff"]["linear_in"]["kernel"] = convert_and_transpose_tensor(
-        pt_state_dict.pop(prefix + "ff.linear_in.weight"), transpose=True
-    )
-    jax_db["ff"]["linear_out"]["kernel"] = convert_and_transpose_tensor(
-        pt_state_dict.pop(prefix + "ff.linear_out.weight"), transpose=True
-    )
+    jax_db["ff"]["linear_in"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "ff.linear_in.weight"), transpose=True)
+    jax_db["ff"]["linear_out"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "ff.linear_out.weight"), transpose=True)
     jax_db["ff_context"]["linear_in"]["kernel"] = convert_and_transpose_tensor(
         pt_state_dict.pop(prefix + "ff_context.linear_in.weight"), transpose=True
     )
@@ -544,12 +520,8 @@ def convert_and_transpose_tensor(tensor, transpose=False):
     s_prefix = f"single_transformer_blocks.{block_idx}."
 
     # Joint projections
-    jax_sb["linear1"]["kernel"] = convert_and_transpose_tensor(
-        pt_state_dict.pop(s_prefix + "attn.to_qkv_mlp_proj.weight"), transpose=True
-    )
-    jax_sb["linear2"]["kernel"] = convert_and_transpose_tensor(
-        pt_state_dict.pop(s_prefix + "attn.to_out.weight"), transpose=True
-    )
+    jax_sb["linear1"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop(s_prefix + "attn.to_qkv_mlp_proj.weight"), transpose=True)
+    jax_sb["linear2"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop(s_prefix + "attn.to_out.weight"), transpose=True)
 
     # Norm scales
     jax_sb["attn"]["query_norm"]["scale"] = convert_and_transpose_tensor(pt_state_dict.pop(s_prefix + "attn.norm_q.weight"))
@@ -585,15 +557,11 @@ def get_pytorch_weight_tensor(key):
   max_logging.log("Mapping VAE decoder weights to JAX parameters...")
 
   # post_quant_conv
-  jax_params["post_quant_conv"]["kernel"] = jnp.array(
-      get_pytorch_weight_tensor("post_quant_conv.weight").transpose(2, 3, 1, 0)
-  )
+  jax_params["post_quant_conv"]["kernel"] = jnp.array(get_pytorch_weight_tensor("post_quant_conv.weight").transpose(2, 3, 1, 0))
   jax_params["post_quant_conv"]["bias"] = jnp.array(get_pytorch_weight_tensor("post_quant_conv.bias"))
 
   # decoder.conv_in
-  jax_params["decoder"]["conv_in"]["kernel"] = jnp.array(
-      get_pytorch_weight_tensor("decoder.conv_in.weight").transpose(2, 3, 1, 0)
-  )
+  jax_params["decoder"]["conv_in"]["kernel"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_in.weight").transpose(2, 3, 1, 0))
   jax_params["decoder"]["conv_in"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_in.bias"))
 
   # decoder.mid_block
@@ -657,17 +625,13 @@ def get_pytorch_weight_tensor(key):
       upsampler_jax = up_block_jax["upsamplers_0"]
       upsampler_pt = f"{up_block_pt}.upsamplers.0"
 
-      upsampler_jax["conv"]["kernel"] = jnp.array(
-          get_pytorch_weight_tensor(f"{upsampler_pt}.conv.weight").transpose(2, 3, 1, 0)
-      )
+      upsampler_jax["conv"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{upsampler_pt}.conv.weight").transpose(2, 3, 1, 0))
       upsampler_jax["conv"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{upsampler_pt}.conv.bias"))
 
   # decoder.conv_norm_out & conv_out
   jax_params["decoder"]["conv_norm_out"]["scale"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_norm_out.weight"))
   jax_params["decoder"]["conv_norm_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_norm_out.bias"))
-  jax_params["decoder"]["conv_out"]["kernel"] = jnp.array(
-      get_pytorch_weight_tensor("decoder.conv_out.weight").transpose(2, 3, 1, 0)
-  )
+  jax_params["decoder"]["conv_out"]["kernel"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_out.weight").transpose(2, 3, 1, 0))
   jax_params["decoder"]["conv_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_out.bias"))
 
   jax_params = jax.tree_util.tree_map(