Update

Author: hsuan-lun-chiang

src/maxtext/layers/nnx_decoders.py

@@ -463,19 +463,51 @@ def _update_leaf(leaf):
 
   def _apply_layer_with_remat(self, layer: nnx.Module, y: jax.Array, policy: Any, prevent_cse: bool, **kwargs):
     """Helper to cleanly apply jax.checkpoint to a single unscanned layer or block."""
+    """Runs the layer stack using nnx.scan."""
+    policy = self.get_remat_policy()
+    prevent_cse = maxtext_utils.should_prevent_cse_in_remat(self.config)
+    graphdef, params, state = nnx.split(layers, nnx.Param, ...)
 
-    graphdef, state = nnx.split(layer)
+    scan_axis = self.config.param_scan_axis
+    if scan_axis != 0:
+      params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0), params)
+
+    layer_cls = layers.__class__
+    sig = inspect.signature(layer_cls.__call__)
+    valid_kwargs = {k: v for k, v in kwargs.items() if k in sig.parameters or "kwargs" in sig.parameters}
 
-    def pure_layer_fn(state_in, y_in):
-      merged_layer = nnx.merge(graphdef, state_in)
-      out = merged_layer(y_in, **kwargs)
-      return out, nnx.state(merged_layer)
+    def _extract_matching_state(template, full):
+      if isinstance(template, nnx.State):
+        return nnx.State({k: _extract_matching_state(v, full[k]) for k, v in template.items()})
+      elif isinstance(template, dict):
+        return {k: _extract_matching_state(v, full[k]) for k, v in template.items()}
+      return full
 
-    checkpointed_fn = jax.checkpoint(pure_layer_fn, policy=policy, prevent_cse=prevent_cse)
-    out, new_state = checkpointed_fn(state, y)
-    nnx.update(layer, new_state)
+    def layer_fn(carry, scanned_vars):
+      current_params, current_state = scanned_vars
 
-    return out
+      if self.config.parameter_memory_host_offload:
+        current_params = jax.tree.map(lambda x: jax.device_put(x, max_utils.device_space()), current_params)
+
+      layer = nnx.merge(graphdef, current_params, current_state)
+      layer_out = layer(carry, *args, **valid_kwargs)
+      new_carry = layer_out[0] if isinstance(layer_out, tuple) else layer_out
+      
+      new_full_state = nnx.state(layer)
+      new_current_state = _extract_matching_state(current_state, new_full_state)
+      
+      # ONLY return non-param state to prevent memory duplication of weights
+      return new_carry, new_current_state
+
+    layer_fn = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
+
+    final_carry, scanned_other = jax.lax.scan(layer_fn, x_in, (params, state))
+
+    if scan_axis != 0:
+      params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), params)
+      
+    scanned_state = nnx.State.merge(params, scanned_other)
+    return final_carry, nnx.merge(graphdef, scanned_state)
 
   def _apply_layers_sequentially(self, layers, x_in, *args, length: int, **kwargs):
     """Runs the layer stack using nnx.scan."""
@@ -885,8 +917,6 @@ def _extract_slice(x, idx, axis):
 
     # Run the single layer
     out = single_layer(
-        y, *args, decoder_input_tokens=kwargs.get("decoder_input_tokens"), **kwargs.get("layer_kwargs", {})
-    )
     y = out[0] if isinstance(out, tuple) else out
 
     # Re-merge the updated state back into the specific slice of the stack
@@ -944,6 +974,7 @@ def _apply_scanned_chunk(self, y, current_idx, next_boundary, layer_stack, *args
 
     return y
 
+
   def _apply_interleaved_scanned_layers(self, y, layer_stack, start_idx, end_idx, engram_indices, *args, **kwargs):
     """Applies a mix of scanned standard layers and unscanned Engram layers."""
     current_idx = start_idx