fix(nnx): optimize scan memory by excluding read-only parameters from scan outputs

Author: hsuan-lun-chiang

src/maxtext/layers/nnx_decoders.py

@@ -356,13 +356,17 @@ def layer_fn(carry, scanned_vars):
           **kwargs,
       )
       new_carry = layer_out[0] if isinstance(layer_out, tuple) else layer_out
-      return new_carry, nnx.state(layer)
+      # Avoid returning and stacking read-only parameters inside the scan body.
+      # This prevents huge unnecessary memory allocation.
+      _, _, updated_state = nnx.split(layer, nnx.Param, ...)
+      return new_carry, updated_state
 
     final_carry, scanned_state = jax.lax.scan(layer_fn, inputs, (params, state))
 
     if scan_axis != 0:
       scanned_params, scanned_other = scanned_state.split(nnx.Param, ...)
-      scanned_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), scanned_params)
+      if scanned_params:
+        scanned_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), scanned_params)
       scanned_state = nnx.State.merge(scanned_params, scanned_other)
 
     nnx.update(self.scanned_layers, scanned_state)
@@ -1031,7 +1035,10 @@ def layer_fn(carry, scanned_vars):
         returned_params = updated_params
         new_current_state = nnx.State.merge(returned_params, updated_state)
       else:
-        new_current_state = nnx.state(layer)
+        # Avoid returning and stacking read-only parameters inside the scan body.
+        # This prevents huge unnecessary memory allocation.
+        _, _, updated_state = nnx.split(layer, nnx.Param, ...)
+        new_current_state = updated_state
 
       if use_kv:
         return new_carry, (new_current_state, updated_kv)
@@ -1077,7 +1084,8 @@ def layer_fn(carry, scanned_vars):
 
       if scan_axis != 0:
         new_params, new_rest = scanned_state.split(nnx.Param, ...)
-        new_params = maxtext_utils_nnx.nnx_sync_moveaxis(new_params, 0, scan_axis)
+        if new_params:
+          new_params = maxtext_utils_nnx.nnx_sync_moveaxis(new_params, 0, scan_axis)
         scanned_state = nnx.merge_state(new_params, new_rest)
 
       returned_kv_stacked = None