From e0c10eb653849568d3ccf624fe3569b0af017663 Mon Sep 17 00:00:00 2001 From: maxtext authors Date: Mon, 10 Aug 2026 10:42:09 -0700 Subject: [PATCH] Support activation rematerialization with small Gemma models. PiperOrigin-RevId: 962240704 --- src/maxtext/configs/types.py | 4 +- src/maxtext/layers/decoders.py | 2 +- src/maxtext/layers/nnx_decoders.py | 58 ++++++++++---- src/maxtext/models/gemma4.py | 2 +- tests/unit/gemma4_small_test.py | 124 ++++++++++++++++++++++++++++- tests/unit/nnx_decoders_test.py | 1 + 6 files changed, 170 insertions(+), 21 deletions(-) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 68684cb38a..d827bd5295 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1254,11 +1254,11 @@ class PipelineParallelism(BaseModel): class RematAndOffload(BaseModel): """Configuration for gradient checkpointing (rematerialization) and offloading.""" - remat_policy: str = Field( + remat_policy: str | None = Field( RematPolicy.FULL.value, description="The rematerialization policy, trading off speed and memory.", ) - remat_policy_for_vit: str = Field("minimal", description="Remat policy for multimodal model's vision encoder.") + remat_policy_for_vit: str | None = Field("minimal", description="Remat policy for multimodal model's vision encoder.") decoder_layer_input: RematLocation = Field( RematLocation.DEVICE, description="Remat policy for the decoder layer's input." ) diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 42753eb752..be026bc630 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -347,7 +347,7 @@ def get_remat_policy(self): """Get remat policy""" policy = None cfg = self.config - if cfg.remat_policy != "none": + if cfg.remat_policy and cfg.remat_policy != "none": if cfg.remat_policy in ("minimal_with_context", "minimal_flash"): # save all if cfg.remat_policy == "minimal_flash": diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index ee02407502..8a24334ce4 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -1161,7 +1161,7 @@ def get_remat_policy(self): """Get remat policy for jax.checkpoint.""" policy = None cfg = self.config - if cfg.remat_policy != "none": + if cfg.remat_policy and cfg.remat_policy != "none": if cfg.remat_policy in {"minimal_with_context", "minimal_flash"}: if cfg.remat_policy == "minimal_flash": max_logging.log("WARNING: 'minimal_flash' will be deprecated soon, please use 'minimal_with_context' instead.") @@ -2274,21 +2274,47 @@ def _apply_gemma4_small_layers( cache_idx = cache_index_of[lyr] kv_cache = kv_caches[cache_idx] if kv_caches is not None else None - y, kv_cache = layer( - y, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - previous_chunk=previous_chunk, - slot=slot, - bidirectional_mask=bidirectional_mask_value, - kv_cache=kv_cache, - attention_metadata=attention_metadata, - per_layer_input=ple_slice, - shared_key=shared_key, - shared_value=shared_value, - ) + policy = self.get_remat_policy() + prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg) + + # When activation rematerialization is enabled (e.g. remat_policy='full' or 'minimal' + # during training), wrap each unscanned layer in jax.checkpoint to prevent OOMs. When + # remat is disabled (remat_policy='none' or None), call the layer directly. + if cfg.remat_policy and cfg.remat_policy != "none": + y, kv_cache = self._apply_layer_with_remat( + layer, + y, + policy, + prevent_cse, + decoder_segment_ids=decoder_segment_ids, + decoder_positions=decoder_positions, + deterministic=deterministic, + model_mode=model_mode, + previous_chunk=previous_chunk, + slot=slot, + bidirectional_mask=bidirectional_mask_value, + kv_cache=kv_cache, + attention_metadata=attention_metadata, + per_layer_input=ple_slice, + shared_key=shared_key, + shared_value=shared_value, + ) + else: + y, kv_cache = layer( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=previous_chunk, + slot=slot, + bidirectional_mask=bidirectional_mask_value, + kv_cache=kv_cache, + attention_metadata=attention_metadata, + per_layer_input=ple_slice, + shared_key=shared_key, + shared_value=shared_value, + ) if kv_caches is not None and kv_cache is not None: kv_caches[cache_idx] = kv_cache diff --git a/src/maxtext/models/gemma4.py b/src/maxtext/models/gemma4.py index 14fc68cebf..fe5328a3f1 100644 --- a/src/maxtext/models/gemma4.py +++ b/src/maxtext/models/gemma4.py @@ -520,7 +520,7 @@ def _remat_enabled(self): is ``None`` for both ``"none"`` and ``"full"``, so it cannot distinguish "no remat" from "full remat" on its own. """ - return self.apply_internal_remat and self.config.remat_policy != "none" + return self.apply_internal_remat and bool(self.config.remat_policy) and self.config.remat_policy != "none" def _scan_local_layers(self, y, layer_kwargs): """Runs the local (sliding-window) layers via a per-layer rematerialized ``jax.lax.scan``.""" diff --git a/tests/unit/gemma4_small_test.py b/tests/unit/gemma4_small_test.py index 880be61487..0de35e25ad 100644 --- a/tests/unit/gemma4_small_test.py +++ b/tests/unit/gemma4_small_test.py @@ -12,12 +12,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for Gemma 4 small (E2B / E4B) layer-pattern helpers.""" +"""Unit tests for Gemma 4 small (E2B / E4B) layer-pattern helpers and rematerialization.""" +import os import unittest +from unittest import mock +from flax import nnx +import jax +import jax.numpy as jnp +from jax.sharding import Mesh +import numpy as np + +from maxtext.common import common_types from maxtext.common.common_types import AttentionType +from maxtext.configs import pyconfig +from maxtext.layers import embeddings +from maxtext.layers import nnx_decoders from maxtext.models import gemma4_small +from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR L = AttentionType.LOCAL_SLIDING @@ -132,5 +145,114 @@ def test_slot_map_without_sharing_is_identity(self): self.assertEqual(slot_map, {i: i for i in range(10)}) +class Gemma4SmallDecoderRematTest(unittest.TestCase): + """Verify that NNXDecoder applies rematerialization per layer for gemma4-small.""" + + _BASE_CONFIG_PATH = os.path.join(MAXTEXT_CONFIGS_DIR, "base.yml") + _NUM_LAYERS = 10 + _NUM_KV_SHARED = 5 # Last 5 of 10 layers share K/V (1 full period of pattern). + _NUM_Q_HEADS = 4 + _NUM_KV_HEADS = 1 + _HEAD_DIM = 32 + _GLOBAL_HEAD_DIM = 64 + _HIDDEN_SIZE = 128 + _PLE_DIM = 32 + _VOCAB = 256 + + def _build_jax_config(self, remat_policy="none"): + """Builds a small E2B-shaped MaxText config with trimmed dimensions for fast tests.""" + return pyconfig.initialize( + ["", self._BASE_CONFIG_PATH], + model_name="gemma4-e2b", + remat_policy=remat_policy, + scan_layers=False, + use_multimodal=False, + override_model_config=True, + # Override shapes for fast tests: + base_num_decoder_layers=self._NUM_LAYERS, + base_num_query_heads=self._NUM_Q_HEADS, + base_num_kv_heads=self._NUM_KV_HEADS, + base_emb_dim=self._HIDDEN_SIZE, + base_mlp_dim=4 * self._HIDDEN_SIZE, + head_dim=self._HEAD_DIM, + global_head_dim=self._GLOBAL_HEAD_DIM, + vocab_size=self._VOCAB, + vocab_size_per_layer_input=self._VOCAB, + hidden_size_per_layer_input=self._PLE_DIM, + num_kv_shared_layers=self._NUM_KV_SHARED, + max_target_length=64, + max_prefill_predict_length=8, + attention="dot_product", # avoid splash on CPU + dtype="float32", + weight_dtype="float32", + float32_qk_product=True, + float32_logits=True, + matmul_precision="highest", + dropout_rate=0.0, + ) + + def test_gemma4_small_decoder_remat(self): + cfg_no_remat = self._build_jax_config(remat_policy="none") + cfg_remat = self._build_jax_config(remat_policy="full") + + mesh = Mesh(np.array(jax.devices()), axis_names=("x",)) + rngs = nnx.Rngs(0) + + decoder_no_remat = nnx_decoders.NNXDecoder(config=cfg_no_remat, mesh=mesh, rngs=rngs) + decoder_remat = nnx_decoders.NNXDecoder(config=cfg_remat, mesh=mesh, rngs=rngs) + nnx.update(decoder_remat, nnx.state(decoder_no_remat)) + + embed = embeddings.Embed( + num_embeddings=cfg_no_remat.vocab_size, + num_features=cfg_no_remat.emb_dim, + dtype=cfg_no_remat.dtype, + config=cfg_no_remat, + mesh=mesh, + rngs=rngs, + ) + + tokens = jax.random.randint(jax.random.key(1), (2, 8), 0, self._VOCAB) + positions = jnp.arange(8)[None, :] + + def loss_fn(model): + out, *_ = model( + embed, + tokens, + decoder_positions=positions, + deterministic=True, + model_mode=common_types.MODEL_MODE_TRAIN, + ) + return jnp.sum(out) + + # Spy on _apply_layer_with_remat (which wraps each layer in jax.checkpoint). + # When rematerialization is disabled ('none'), layers are called directly + # without checkpointing (call count is 0). + with mock.patch.object( + decoder_no_remat, + "_apply_layer_with_remat", + wraps=decoder_no_remat._apply_layer_with_remat, # pylint: disable=protected-access) + ) as spy_no_remat: + loss_no_remat, grad_no_remat = nnx.value_and_grad(loss_fn)(decoder_no_remat) + self.assertEqual(spy_no_remat.call_count, 0) + + # When rematerialization is enabled ('full'), every unscanned decoder layer + # must be wrapped in jax.checkpoint via _apply_layer_with_remat once per layer. + with mock.patch.object( + decoder_remat, + "_apply_layer_with_remat", + wraps=decoder_remat._apply_layer_with_remat, # pylint: disable=protected-access) + ) as spy_remat: + loss_remat, grad_remat = nnx.value_and_grad(loss_fn)(decoder_remat) + self.assertEqual(spy_remat.call_count, self._NUM_LAYERS) + + # Rematerialization should not alter outputs or gradients. + np.testing.assert_allclose(loss_no_remat, loss_remat, rtol=1e-4, atol=1e-4) + jax.tree.map( + lambda g1, g2: np.testing.assert_allclose(g1, g2, rtol=1e-4, atol=1e-4), + nnx.state(grad_no_remat), + nnx.state(grad_remat), + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index abdfb30495..c2832f4ce9 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -1163,6 +1163,7 @@ def test_gemma4_small_decoder_with_mock_cache_and_ple(self): "model_name=gemma4-e2b", "scan_layers=False", "attention=dot_product", + "remat_policy=none", "num_decoder_layers=3", "num_kv_shared_layers=1", "base_emb_dim=128",