Implement text encoder offloading for Z-Image pipeline and resolve linter checks

Author: csgoogle

src/maxdiffusion/checkpointing/z_image_checkpointer.py

@@ -336,6 +336,7 @@ def _assemble(self, transformer, vae, vae_params, tokenizer, text_encoder, vae_s
         mesh=self.mesh,
         logical_axis_rules=self.config.logical_axis_rules,
         vae_shift_factor=vae_shift_factor,
+        offload_encoders=getattr(self.config, "offload_encoders", False),
     )
 
   def _place(self, params, shardings):

src/maxdiffusion/configs/base_zimage.yml

@@ -34,6 +34,8 @@ activations_dtype: 'bfloat16'
 
 # Maximum sequence length for the text encoder
 max_sequence_length: 512
+# offloads text encoder after text encoding to save memory.
+offload_encoders: True
 
 # Attention
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring

src/maxdiffusion/configs/base_zimage_turbo.yml

@@ -33,6 +33,8 @@ activations_dtype: 'bfloat16'
 
 # Maximum sequence length for the text encoder
 max_sequence_length: 512
+# offloads text encoder after text encoding to save memory.
+offload_encoders: True
 
 # Attention
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring

src/maxdiffusion/models/z_image/transformer_z_image.py

@@ -564,7 +564,6 @@ def __call__(
     if len(x) != len(cap_feats):
       raise ValueError("x and cap_feats must have one item per batch element.")
     adaln = self.t_embedder(t * self.t_scale).astype(x[0].dtype)
-    batch = len(x)
 
     # 1. Pad and patchify each element locally (on its device)
     caption_features, caption_positions, caption_masks = [], [], []

src/maxdiffusion/pipelines/z_image/z_image_pipeline.py

@@ -75,6 +75,7 @@ def __init__(
       mesh=None,
       logical_axis_rules=(),
       vae_shift_factor: Optional[float] = None,
+      offload_encoders: bool = False,
   ):
     self.transformer = transformer
     self.vae = vae
@@ -96,6 +97,12 @@ def __init__(
     self._text_encoder_split = nnx.split(self.text_encoder, nnx.Param, ...)
     self._transformer_split = nnx.split(self.transformer, nnx.Param, ...)
 
+    self.offload_encoders = offload_encoders
+    if self.offload_encoders:
+      # Save the original shardings of the text encoder parameters
+      graphdef, state, rest = self._text_encoder_split
+      self._text_encoder_shardings = jax.tree_util.tree_map(lambda x: x.sharding, state)
+
   def _decode_fn(self):
     """Jitted VAE decode, including the latent denormalization."""
     if self._jitted_decoder is None:
@@ -127,6 +134,15 @@ def encode_prompt(self, prompt: Union[str, list[str]], max_sequence_length: int
         chats, padding="max_length", max_length=max_sequence_length, truncation=True, return_tensors="np"
     )
     graphdef, state, rest = self._text_encoder_split
+
+    if self.offload_encoders:
+      # Restore parameters from CPU back to original device shardings
+      state = jax.tree_util.tree_map(
+          lambda x, sharding: jax.device_put(x, sharding),
+          state,
+          self._text_encoder_shardings
+      )
+
     with self.mesh if self.mesh is not None else nullcontext():
       embeddings = encode_step(
           graphdef,
@@ -136,6 +152,15 @@ def encode_prompt(self, prompt: Union[str, list[str]], max_sequence_length: int
           jnp.asarray(inputs["attention_mask"], jnp.int32),
       )
     lengths = np.asarray(inputs["attention_mask"]).sum(axis=-1)
+
+    if self.offload_encoders:
+      from maxdiffusion import max_logging
+      s_offload = time.perf_counter()
+      cpus = jax.devices("cpu")
+      offloaded_state = jax.tree_util.tree_map(lambda x: jax.device_put(x, device=cpus[0]), state)
+      self._text_encoder_split = (graphdef, offloaded_state, rest)
+      max_logging.log(f"Offloaded Qwen3 text encoder parameters to CPU in {(time.perf_counter() - s_offload):.4f}s.")
+
     return [jnp.asarray(embeddings[index, : lengths[index]], dtype=self.dtype) for index in range(len(lengths))]
 
   @staticmethod

src/maxdiffusion/pyconfig.py

@@ -204,6 +204,8 @@ def user_init(raw_keys):
       raw_keys["names_which_can_be_saved"] = []
     if "names_which_can_be_offloaded" not in raw_keys:
       raw_keys["names_which_can_be_offloaded"] = []
+    if "offload_encoders" not in raw_keys:
+      raw_keys["offload_encoders"] = False
 
     raw_keys["weights_dtype"] = jax.numpy.dtype(raw_keys["weights_dtype"])
     raw_keys["activations_dtype"] = jax.numpy.dtype(raw_keys["activations_dtype"])

src/maxdiffusion/tests/z_image/z_image_module_parity_test.py

@@ -127,19 +127,19 @@ def test_transformer_block_parity(self):
     )
 
   def test_full_transformer_parity(self):
-    config = dict(
-        all_patch_size=(2,),
-        all_f_patch_size=(1,),
-        in_channels=4,
-        dim=32,
-        n_layers=1,
-        n_refiner_layers=1,
-        n_heads=4,
-        n_kv_heads=4,
-        cap_feat_dim=8,
-        axes_dims=[2, 2, 4],
-        axes_lens=[64, 64, 64],
-    )
+    config = {
+        "all_patch_size": (2,),
+        "all_f_patch_size": (1,),
+        "in_channels": 4,
+        "dim": 32,
+        "n_layers": 1,
+        "n_refiner_layers": 1,
+        "n_heads": 4,
+        "n_kv_heads": 4,
+        "cap_feat_dim": 8,
+        "axes_dims": [2, 2, 4],
+        "axes_lens": [64, 64, 64],
+    }
     hf = HFZImageTransformer2DModel(**config).eval()
     local = copy_parameters(ZImageTransformer2DModel(rngs=self.rngs, attention_kernel="dot_product", **config), hf)
     images = [torch.randn(4, 1, 4, 4)]