Skip to content

[Do not submit] test4 #4433

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
aireenmei wants to merge 6 commits into
base: main
Choose a base branch
from
Draft

[Do not submit] test4 #4433

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
9a3b4fc
Rewrite scannable block for Gemma4
aireenmei Jul 11, 2026
c9645be
add global scan
aireenmei Jul 13, 2026
6a873e3
fix full remat case where remat_policy=None
aireenmei Jul 13, 2026
391679a
update test
aireenmei Jul 14, 2026
62649ce
avoid offloading for global layer
aireenmei Jul 14, 2026
65d2d1a
fix linen path
aireenmei Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 18 additions & 31 deletions src/maxtext/layers/decoders.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -411,36 +411,15 @@ def get_remat_policy(self):
"kv_proj",
"qkv_proj",
)
elif cfg.remat_policy == "qkv_proj_offloaded":
elif cfg.remat_policy in ("qkv_proj_offloaded", "minimal_offloaded", "custom"):
# minimal_offloaded offloads all except context. All three share a single
# source of truth for their save/offload name lists (see
# maxtext_utils.get_offload_remat_names) so that offloading configured via
# `custom` resolves identically to the named presets.
save_names, offload_names = maxtext_utils.get_offload_remat_names(cfg)
policy = jax.checkpoint_policies.save_and_offload_only_these_names(
names_which_can_be_saved=[],
names_which_can_be_offloaded=["query_proj", "value_proj", "key_proj", "kv_proj"],
offload_src="device",
offload_dst="pinned_host",
)
elif cfg.remat_policy == "minimal_offloaded":
# offload all except context
policy = jax.checkpoint_policies.save_and_offload_only_these_names(
names_which_can_be_saved=[],
names_which_can_be_offloaded=[
"query_proj",
"value_proj",
"key_proj",
"kv_proj",
"qkv_proj",
"out_proj",
"mlpwi_0",
"mlpwi_1",
"mlpwi",
"mlpwo",
],
offload_src="device",
offload_dst="pinned_host",
)
elif cfg.remat_policy == "custom":
policy = jax.checkpoint_policies.save_and_offload_only_these_names(
names_which_can_be_saved=cfg.tensors_on_device,
names_which_can_be_offloaded=cfg.tensors_to_offload,
names_which_can_be_saved=save_names,
names_which_can_be_offloaded=offload_names,
offload_src="device",
offload_dst="pinned_host",
)
Expand Down Expand Up @@ -1433,7 +1412,12 @@ def _apply_gemma4_scanned_blocks(
if num_full_blocks > 0:
ScannableBlockToLinen = gemma4.Gemma4ScannableBlockToLinen
policy = self.get_remat_policy()
RemattedGemma4Block = self.set_remat_policy([ScannableBlockToLinen], policy)[0]
# Gemma4ScannableBlock rematerializes its own local (scanned) and global
# layers when apply_internal_remat=True, so we do NOT wrap it in
# block-level remat here (that would double-rematerialize and make XLA
# treat the whole block as one unit). Unrolling the block scan lets XLA
# free each block's activations across iterations instead of keeping the
# block live as a unit.

kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan(
kv_caches, num_full_blocks, block_pattern_len, stack=True
Expand All @@ -1456,7 +1440,7 @@ def _apply_gemma4_scanned_blocks(

# For a fully scanned block, apply it inside a nn.scan over the calculated number of full blocks
y, returned_kv_cache = nn.scan(
RemattedGemma4Block,
ScannableBlockToLinen,
variable_axes={
"params": cfg.param_scan_axis,
"cache": 0,
Expand All @@ -1467,6 +1451,7 @@ def _apply_gemma4_scanned_blocks(
split_rngs={"params": True, "dropout": cfg.enable_dropout},
in_axes=in_axes_tuple,
length=num_full_blocks,
unroll=num_full_blocks,
metadata_params={
nn.PARTITION_NAME: "layers",
"abstract_init": False,
Expand All @@ -1477,6 +1462,8 @@ def _apply_gemma4_scanned_blocks(
quant=self.quant,
model_mode=model_mode,
num_of_layers=block_pattern_len,
remat_policy_fn=policy,
apply_internal_remat=True,
name="scanned_blocks",
)(
y, *broadcast_args
Expand Down
161 changes: 64 additions & 97 deletions src/maxtext/layers/nnx_decoders.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@
ShardMode,
)
from maxtext.layers import initializers, linears, mhc, normalizations, quantizations
from maxtext.layers import nnx_wrappers
from maxtext.layers import nnx_scan, nnx_wrappers
from maxtext.layers.attentions import Attention
from maxtext.layers.embeddings import Embed, PositionalEmbedding, attend_on_embedding
from maxtext.layers.normalizations import RMSNorm
Expand Down Expand Up @@ -646,9 +646,19 @@ def _init_scanned_gemma4(self, decoder_block_classes, rngs, mesh):
attention_pattern_length = len(gemma4.GEMMA4_ATTENTION_PATTERN)
scan_length = config.num_decoder_layers // attention_pattern_length
num_remaining_layers = config.num_decoder_layers % attention_pattern_length
layer_kwargs = {"num_of_layers": attention_pattern_length}

rem_layer_kwargs = {"num_of_layers": num_remaining_layers}
policy = self.get_remat_policy()
# The pure-NNX decoder skips block-level remat (skip_block_remat=True below),
# so the block rematerializes its own local/global layers instead.
layer_kwargs = {
"num_of_layers": attention_pattern_length,
"remat_policy_fn": policy,
"apply_internal_remat": True,
}
rem_layer_kwargs = {
"num_of_layers": num_remaining_layers,
"remat_policy_fn": policy,
"apply_internal_remat": True,
}

RemattedGemma4Block = gemma4.Gemma4ScannableBlock

Expand Down Expand Up @@ -838,95 +848,20 @@ def _create_scanned_layers(
rngs: nnx.Rngs,
**layer_kwargs,
):
"""Creates a scanned stack of layers using jax.lax.scan for memory-efficient initialization."""
if length == 0:
return None
scan_axis = self.config.param_scan_axis

# Fork rngs to get per-layer RNG states for scanning
try:
forked_rngs = rngs.fork(split=length)
except: # pylint: disable=bare-except
pass

rngs_graphdef, rngs_state = nnx.split(forked_rngs)

first_rng_state = jax.tree.map(lambda x: x[0], rngs_state)
ref_rngs = nnx.merge(rngs_graphdef, first_rng_state)
ref_layer = decoder_layer_class(
config=self.config,
mesh=self.mesh,
quant=self.quant,
model_mode=self.model_mode,
rngs=ref_rngs,
**layer_kwargs,
return nnx_scan.create_scanned_layers(
lambda layer_rngs: decoder_layer_class(
config=self.config,
mesh=self.mesh,
model_mode=self.model_mode,
quant=self.quant,
rngs=layer_rngs,
**layer_kwargs,
),
length=length,
param_scan_axis=self.config.param_scan_axis,
metadata_axis_name=metadata_axis_name,
rngs=rngs,
)
layer_graphdef, _, _ = nnx.split(ref_layer, nnx.Param, ...)
del ref_layer

def scan_body(carry, rng_state_slice):
layer_rngs = nnx.merge(rngs_graphdef, rng_state_slice)
layer = decoder_layer_class(
config=self.config,
mesh=self.mesh,
quant=self.quant,
model_mode=self.model_mode,
rngs=layer_rngs,
**layer_kwargs,
)
_, params, rest = nnx.split(layer, nnx.Param, ...)
return carry, (params, rest)

_, (stacked_params, stacked_rest) = jax.lax.scan(scan_body, None, rngs_state)

if scan_axis != 0:
stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0), stacked_params)

def _add_scan_metadata(state, axis):
def _update_leaf(leaf):
if hasattr(leaf, "replace") and hasattr(leaf, "value"):
replace_kwargs = {}
if hasattr(leaf, "get_metadata"):
replace_kwargs.update(leaf.get_metadata())

replace_kwargs[nnx.PARTITION_NAME] = metadata_axis_name
replace_kwargs["param_scan_axis"] = axis

for key in [
"sharding",
"out_sharding",
"kernel_axes",
"sharding_names",
]:
val = getattr(leaf, key, None)
if val is None and key in replace_kwargs:
val = replace_kwargs[key]

if val is not None:
if isinstance(val, str):
val = (val,)
if isinstance(val, tuple):
l = list(val)
# Safely insert the scan axis into the logical axes string
if metadata_axis_name not in l:
insert_idx = min(axis, len(l))
l.insert(insert_idx, metadata_axis_name)
replace_kwargs[key] = tuple(l)

return leaf.replace(**replace_kwargs)
return leaf

# We must use a custom is_leaf to catch the VariableState instances
return jax.tree.map(
_update_leaf,
state,
is_leaf=lambda x: hasattr(x, "replace") and hasattr(x, "value"),
)

stacked_params = _add_scan_metadata(stacked_params, scan_axis)
stacked_rest = _add_scan_metadata(stacked_rest, 0)

return nnx.merge(layer_graphdef, stacked_params, stacked_rest)

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."""
Expand All @@ -944,7 +879,17 @@ def pure_layer_fn(state_in, y_in):

return out

def _apply_layers_sequentially(self, layers, x_in, *args, length: int, kv_caches_stacked=None, **kwargs):
def _apply_layers_sequentially(
self,
layers,
x_in,
*args,
length: int,
kv_caches_stacked=None,
skip_block_remat: bool = False,
unroll: int = 1,
**kwargs,
):
"""Runs the layer stack using nnx.scan.

Args:
Expand All @@ -955,6 +900,11 @@ def _apply_layers_sequentially(self, layers, x_in, *args, length: int, kv_caches
kv_caches_stacked: Optional pytree whose leaves have shape [num_layers, ...].
When provided, the i-th slice is passed as `kv_cache=` to layer i and the
updated caches are returned as a third element of the tuple.
skip_block_remat: When True, do not wrap the scanned body in jax.checkpoint.
Used when the scanned module already applies its own (finer-grained,
e.g. per-layer) remat internally, to avoid double rematerialization.
unroll: Number of scan iterations to unroll into straight-line code
(forwarded to jax.lax.scan). unroll >= length fully unrolls the loop.
**kwargs: Keyword args forwarded to the layer (filtered by the layer signature).

Returns:
Expand Down Expand Up @@ -1037,7 +987,12 @@ def layer_fn(carry, scanned_vars):
return new_carry, (new_current_state, updated_kv)
return new_carry, new_current_state

layer_fn_wrapped = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
if skip_block_remat:
# The scanned module applies its own remat internally; wrapping the whole
# body again would double-remat and recompute the entire block.
layer_fn_wrapped = layer_fn
else:
layer_fn_wrapped = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)

if use_kv:
# If kv_caches is provided (e.g., from vLLM), we CANNOT use jax.lax.scan
Expand Down Expand Up @@ -1069,7 +1024,7 @@ def layer_fn(carry, scanned_vars):
params = maxtext_utils_nnx.nnx_ensure_scan_leading_axis(params, length)
state = maxtext_utils_nnx.nnx_ensure_scan_leading_axis(state, length)

final_carry, scanned_state = jax.lax.scan(layer_fn_wrapped, x_in, (params, state))
final_carry, scanned_state = jax.lax.scan(layer_fn_wrapped, x_in, (params, state), unroll=unroll)
returned_kv_stacked = None

# Ensure metadata rank matches the stacked values
Expand Down Expand Up @@ -2032,13 +1987,25 @@ def _apply_gemma4_scanned_blocks(
if attention_metadata is not None:
layer_kwargs["attention_metadata"] = attention_metadata

# Apply the main scan over the full blocks
# Apply the main scan over the full blocks. Gemma4ScannableBlock applies
# per-layer remat internally (local scan + global layer), so skip the
# block-level remat here to avoid double rematerialization. Unrolling the
# block loop (one iteration per repeated block) lets XLA pipeline/free block
# activations across iterations (memory + overlap knob).
block_unroll = max(1, scan_length)
if scan_length > 0:
grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan(
kv_caches, scan_length, attention_pattern_length, stack=False
)
y, self.scanned_blocks, _ = self._apply_layers_sequentially(
self.scanned_blocks, y, *layer_args, length=scan_length, kv_caches_stacked=grouped_kv_caches, **layer_kwargs
self.scanned_blocks,
y,
*layer_args,
length=scan_length,
kv_caches_stacked=grouped_kv_caches,
skip_block_remat=True,
unroll=block_unroll,
**layer_kwargs,
)
maxtext_utils.update_kv_caches_after_scan(
kv_caches, grouped_kv_caches, scan_length, attention_pattern_length, stacked=False
Expand Down
Loading
Loading