Run the Z-Image Qwen3 text encoder in NNX

Split Qwen3 checkpoint conversion out of the model definition into
models/qwen3_utils.py, mirroring transformer_z_image / z_image_utils.

Fix two bugs that made NNXFlaxQwen3Attention unusable: it applied no
causal mask, and it treated attention_mask as an additive mask rather
than the 0/1 padding mask its callers pass. Against HF Transformers the
port scored cosine 0.242; it now matches at 0.999996.

Declare param_dtype on the Qwen3 Linear/Embed layers. Flax stores
parameters in param_dtype -- float32 by default -- regardless of the
compute dtype, so a bfloat16 run held the 4B encoder in float32: 16.1 GB
per chip instead of 8.0 GB. Norm scales stay float32 deliberately, and
now say so in the module rather than being retyped after loading.

Z-Image builds its denoiser and text encoder through one NNX path in the
checkpointer; only the Diffusers VAE is still Linen. flux2klein keeps the
existing loader, so its numerics are unchanged.

Verified on v6e-8 (Z-Image-Turbo, 1024x1024, 9 steps, 8 images): 0.31s
per image, output within 0.024/255 of the previous Linen encoder.

Author: csgoogle

src/maxdiffusion/checkpointing/z_image_checkpointer.py

@@ -28,8 +28,8 @@
   3. fill it, either from safetensors or from the Orbax cache, placing each
      parameter directly into its target sharding.
 
-`_component` does 1 and 2 for the Flax modules; the denoiser needs its own
-version of the same because it is NNX rather than Linen.
+`_nnx_component` does 1 and 2 for the denoiser and the text encoder;
+`_linen_component` does the same for the Diffusers VAE, which is still Linen.
 """
 
 import inspect
@@ -50,7 +50,8 @@
 
 from .. import max_logging, max_utils
 from ..models import FlaxAutoencoderKL
-from ..models.qwen3_flax import FlaxQwen3Config, FlaxQwen3Model, load_qwen3_weights
+from ..models.qwen3_flax import FlaxQwen3Config, NNXFlaxQwen3Model
+from ..models.qwen3_utils import load_qwen3_weights
 from ..models.z_image.transformer_z_image import ZImageTransformer2DModel
 from ..models.z_image.z_image_utils import load_z_image_transformer
 from ..pipelines.z_image import ZImagePipeline
@@ -91,12 +92,25 @@ def _abstract_tree(shapes, shardings):
   )
 
 
-def _component(module, init_args, mesh, axis_rules):
-  """Abstract params of a Linen module, plus their target shardings.
+def _nnx_component(factory, mesh, axis_rules, seed=0):
+  """Build an NNX module abstractly: its graph, params and target shardings.
 
-  Shardings come from the module's own logical annotations. A module with no
-  annotations (`vae_flax.FlaxAutoencoderKL` has none) resolves to fully
-  replicated, which is exactly what it used to be given explicitly.
+  Shardings come from the module's own `nnx.with_partitioning` annotations,
+  resolved against the mesh.
+  """
+  graphdef, state, rest = nnx.split(nnx.eval_shape(factory, nnx.Rngs(jax.random.key(seed))), nnx.Param, ...)
+  shardings = None
+  if mesh is not None:
+    sharding_tree = nn.logical_to_mesh_sharding(nnx.get_partition_spec(state), mesh, axis_rules)
+    shardings = {path: variable.value for path, variable in nnx.to_flat_state(sharding_tree)}
+  return graphdef, state, rest, shardings
+
+
+def _linen_component(module, init_args, mesh, axis_rules):
+  """The Linen equivalent of `_nnx_component`, for the Diffusers VAE.
+
+  A module with no logical annotations (`vae_flax.FlaxAutoencoderKL` has none)
+  resolves to fully replicated.
   """
   abstract = jax.eval_shape(lambda: module.init(jax.random.key(0), *init_args))
   shapes = _plain(abstract["params"])
@@ -106,9 +120,15 @@ def _component(module, init_args, mesh, axis_rules):
   return shapes, dict(flatten_dict(_plain(mesh_shardings["params"])))
 
 
-def _transformer_component(transformer_config, config, mesh):
-  """The NNX equivalent of `_component`, plus the graph needed to rebuild it."""
+def _fill(graphdef, state, rest, params):
+  """Put loaded params into an abstract NNX state and rebuild the module."""
+  flat_state = dict(nnx.to_flat_state(state))
+  for path, value in flatten_dict(params).items():
+    flat_state[path].value = value
+  return nnx.merge(graphdef, nnx.from_flat_state(flat_state), rest)
+
 
+def _transformer_factory(transformer_config, config, mesh):
   def factory(rngs):
     return ZImageTransformer2DModel(
         rngs=rngs,
@@ -120,26 +140,15 @@ def factory(rngs):
         **transformer_config,
     )
 
-  graphdef, state, rest = nnx.split(nnx.eval_shape(factory, nnx.Rngs(jax.random.key(config.seed))), nnx.Param, ...)
-  shardings = None
-  if mesh is not None:
-    sharding_tree = nn.logical_to_mesh_sharding(nnx.get_partition_spec(state), mesh, config.logical_axis_rules)
-    shardings = {path: variable.value for path, variable in nnx.to_flat_state(sharding_tree)}
-  return graphdef, state, rest, shardings
-
-
-def _fill(graphdef, state, rest, params):
-  """Put loaded params into an abstract NNX state and rebuild the module."""
-  flat_state = dict(nnx.to_flat_state(state))
-  for path, value in flatten_dict(params).items():
-    flat_state[path].value = value
-  return nnx.merge(graphdef, nnx.from_flat_state(flat_state), rest)
+  return factory
 
 
 def create_z_image_transformer(model_id: str, config, mesh: Optional[jax.sharding.Mesh] = None):
   """Instantiate and stream a Diffusers Z-Image checkpoint into an NNX model."""
   transformer_config = ZImageTransformer2DModel.load_config(model_id, subfolder="transformer")
-  graphdef, state, rest, shardings = _transformer_component(transformer_config, config, mesh)
+  graphdef, state, rest, shardings = _nnx_component(
+      _transformer_factory(transformer_config, config, mesh), mesh, config.logical_axis_rules, config.seed
+  )
   params = load_z_image_transformer(model_id, state.to_pure_dict(), target_shardings=shardings)
   return _fill(graphdef, state, rest, params)
 
@@ -196,27 +205,31 @@ def build_vae(self, vae_config: dict):
     """VAE module, its abstract params and their target shardings."""
     vae = FlaxAutoencoderKL(**vae_config, dtype=self.config.activations_dtype, weights_dtype=self.config.weights_dtype)
     sample = jnp.ones((1, vae.config.in_channels, 64, 64), jnp.float32)
-    return vae, *_component(vae, (sample,), self.mesh, self._vae_axis_rules())
+    return vae, *_linen_component(vae, (sample,), self.mesh, self._vae_axis_rules())
 
   def build_text_encoder(self, text_encoder_config: dict):
-    """Qwen3 module, its abstract params and their target shardings."""
-    model = FlaxQwen3Model(
-        FlaxQwen3Config(
-            vocab_size=text_encoder_config["vocab_size"],
-            hidden_size=text_encoder_config["hidden_size"],
-            intermediate_size=text_encoder_config["intermediate_size"],
-            num_hidden_layers=text_encoder_config["num_hidden_layers"],
-            num_attention_heads=text_encoder_config["num_attention_heads"],
-            num_key_value_heads=text_encoder_config["num_key_value_heads"],
-            head_dim=text_encoder_config["head_dim"],
-            rms_norm_eps=text_encoder_config["rms_norm_eps"],
-            rope_theta=text_encoder_config["rope_theta"],
-            max_position_embeddings=text_encoder_config["max_position_embeddings"],
-            dtype=self.config.weights_dtype,
-        )
+    """Qwen3 graph, its abstract params and their target shardings."""
+    qwen3_config = FlaxQwen3Config(
+        vocab_size=text_encoder_config["vocab_size"],
+        hidden_size=text_encoder_config["hidden_size"],
+        intermediate_size=text_encoder_config["intermediate_size"],
+        num_hidden_layers=text_encoder_config["num_hidden_layers"],
+        num_attention_heads=text_encoder_config["num_attention_heads"],
+        num_key_value_heads=text_encoder_config["num_key_value_heads"],
+        head_dim=text_encoder_config["head_dim"],
+        rms_norm_eps=text_encoder_config["rms_norm_eps"],
+        rope_theta=text_encoder_config["rope_theta"],
+        max_position_embeddings=text_encoder_config["max_position_embeddings"],
+        dtype=self.config.weights_dtype,
+    )
+    # The norms declare float32 scales, so the abstract params already carry the
+    # right dtypes for both the loader and the Orbax restore target.
+    return _nnx_component(
+        lambda rngs: NNXFlaxQwen3Model(rngs=rngs, config=qwen3_config),
+        self.mesh,
+        self._text_encoder_axis_rules(),
+        self.config.seed,
     )
-    ids = jnp.ones((1, 8), jnp.int32)
-    return model, *_component(model, (ids, ids), self.mesh, self._text_encoder_axis_rules())
 
   def _vae_axis_rules(self):
     """`vae_logical_axis_rules` if the config declares them, as WAN does."""
@@ -248,18 +261,18 @@ def load_diffusers_pipeline(self) -> ZImagePipeline:
     )
     vae_params = self._place(vae_params, vae_shardings)
 
-    text_encoder, shapes, text_encoder_shardings = self.build_text_encoder(
+    graphdef, state, rest, shardings = self.build_text_encoder(
         AutoConfig.from_pretrained(model_id, subfolder="text_encoder").to_dict()
     )
-    text_encoder_params = load_qwen3_weights(_text_encoder_dir(model_id), shapes, target_shardings=text_encoder_shardings)
+    params = load_qwen3_weights(_text_encoder_dir(model_id), state.to_pure_dict(), target_shardings=shardings)
+    text_encoder = _fill(graphdef, state, rest, params)
 
     return self._assemble(
         transformer,
         vae,
         vae_params,
         AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", use_fast=True),
         text_encoder,
-        text_encoder_params,
         raw_vae_config.get("shift_factor", _DEFAULT_VAE_SHIFT_FACTOR),
     )
 
@@ -278,9 +291,14 @@ def _restore_pipeline(self, orbax_dir: str) -> Optional[ZImagePipeline]:
         return None
 
       max_logging.log(f"Loading Z-Image pipeline from orbax checkpoint step {step} in {orbax_dir}")
-      graphdef, state, rest, shardings = _transformer_component(metadata["transformer_config"], self.config, self.mesh)
+      graphdef, state, rest, shardings = _nnx_component(
+          _transformer_factory(metadata["transformer_config"], self.config, self.mesh),
+          self.mesh,
+          self.config.logical_axis_rules,
+          self.config.seed,
+      )
       vae, vae_shapes, vae_shardings = self.build_vae(metadata["vae_config"])
-      text_encoder, text_shapes, text_shardings = self.build_text_encoder(metadata["text_encoder_config"])
+      text_graphdef, text_state, text_rest, text_shardings = self.build_text_encoder(metadata["text_encoder_config"])
 
       # Shapes come from the models themselves, so every item lands directly in
       # the sharding its component was built for -- no host staging.
@@ -290,7 +308,7 @@ def _restore_pipeline(self, orbax_dir: str) -> Optional[ZImagePipeline]:
               **{
                   "transformer_state": ocp.args.StandardRestore(_abstract_tree(state.to_pure_dict(), shardings)),
                   "vae_state": ocp.args.StandardRestore(_abstract_tree(vae_shapes, vae_shardings)),
-                  "text_encoder_state": ocp.args.StandardRestore(_abstract_tree(text_shapes, text_shardings)),
+                  "text_encoder_state": ocp.args.StandardRestore(_abstract_tree(text_state.to_pure_dict(), text_shardings)),
               }
           ),
       )
@@ -300,22 +318,20 @@ def _restore_pipeline(self, orbax_dir: str) -> Optional[ZImagePipeline]:
           vae,
           restored["vae_state"],
           self._load_tokenizer(orbax_dir),
-          text_encoder,
-          restored["text_encoder_state"],
+          _fill(text_graphdef, text_state, text_rest, restored["text_encoder_state"]),
           metadata["vae_shift_factor"],
       )
     except Exception as e:  # pylint: disable=broad-except
       max_logging.log(f"Failed to load orbax checkpoint from {orbax_dir}, falling back to Diffusers: {e}")
       return None
 
-  def _assemble(self, transformer, vae, vae_params, tokenizer, text_encoder, text_encoder_params, vae_shift_factor):
+  def _assemble(self, transformer, vae, vae_params, tokenizer, text_encoder, vae_shift_factor):
     return ZImagePipeline(
         transformer,
         vae,
         vae_params,
         tokenizer,
         text_encoder,
-        text_encoder_params,
         dtype=self.config.activations_dtype,
         mesh=self.mesh,
         logical_axis_rules=self.config.logical_axis_rules,
@@ -370,7 +386,9 @@ def _save_pipeline(self, orbax_dir: str, pipeline: ZImagePipeline):
               **{
                   "transformer_state": ocp.args.StandardSave(transformer_state.to_pure_dict()),
                   "vae_state": ocp.args.StandardSave(_plain(pipeline.vae_params)),
-                  "text_encoder_state": ocp.args.StandardSave(_plain(pipeline.text_encoder_params)),
+                  "text_encoder_state": ocp.args.StandardSave(
+                      nnx.split(pipeline.text_encoder, nnx.Param, ...)[1].to_pure_dict()
+                  ),
                   CONFIG_ITEM: ocp.args.JsonSave(json.loads(json.dumps(metadata, default=str))),
               }
           ),

src/maxdiffusion/generate_flux2klein.py

@@ -37,7 +37,8 @@
 
 from maxdiffusion.models.flux.transformers.transformer_flux_flax import Flux2KleinTransformer2DModel
 from maxdiffusion.models.vae_flax import FlaxAutoencoderKL
-from maxdiffusion.models.qwen3_flax import FlaxQwen3Config, FlaxQwen3Model, load_and_convert_qwen3_weights
+from maxdiffusion.models.qwen3_flax import FlaxQwen3Config, FlaxQwen3Model
+from maxdiffusion.models.qwen3_utils import load_and_convert_qwen3_weights
 from maxdiffusion.schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler