Support scan version of Engram layers

Author: RissyRan

src/maxtext/configs/types.py

@@ -2346,8 +2346,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         raise ValueError(
             "Engram requires both 'hf_access_token' and 'tokenizer_path' " "to load the Hugging Face tokenizer."
         )
-      if self.scan_layers:
-        raise NotImplementedError("Currently Engram only supports unscanned version. Please set scan_layers=False.")
       if len(self.engram_vocab_bases) != (self.engram_max_ngram_size - 1):
         raise ValueError(
             f"Engram vocab size mismatch: expected {self.engram_max_ngram_size - 1} (max_ngram_size - 1), "

src/maxtext/layers/decoders.py

@@ -854,45 +854,85 @@ def __call__(
               "slot": slot,
           }
           dense_layer = RemattedBlockLayers[0]
-          dense_layer.__call__ = functools.partial(dense_layer.__call__, **layer_call_kwargs)
-          y, _ = self.scan_decoder_layers(
-              cfg,
-              dense_layer,
-              cfg.first_num_dense_layers,
-              "dense_layers",
-              mesh,
-              in_axes_tuple=(nn.broadcast,) * len(broadcast_args),
-              model_mode=model_mode,
-          )(y, *broadcast_args)
           moe_layer = RemattedBlockLayers[1]
-          moe_layer.__call__ = functools.partial(moe_layer.__call__, **layer_call_kwargs)
-          num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers
+          if cfg.engram_layers:
+            original_dense_call = dense_layer.__call__
+            original_moe_call = moe_layer.__call__
+            dense_layer.__call__ = functools.partial(dense_layer.__call__, **layer_call_kwargs)
+            moe_layer.__call__ = functools.partial(moe_layer.__call__, **layer_call_kwargs)
+
+            common_kwargs = {
+                "dense_layer": dense_layer,
+                "moe_layer": moe_layer,
+                "original_dense_call": original_dense_call,
+                "original_moe_call": original_moe_call,
+                "layer_call_kwargs": layer_call_kwargs,
+                "decoder_segment_ids": decoder_segment_ids,
+                "decoder_positions": decoder_positions,
+                "deterministic": deterministic,
+                "model_mode": model_mode,
+                "decoder_input_tokens": decoder_input_tokens,
+                "broadcast_args": broadcast_args,
+            }
 
-          # If batch-split schedule is used and initialization is complete,
-          # as detected by immutable params, use deepseek_batchsplit custom
-          # scan with initialized parameters.
-          if cfg.use_batch_split_schedule and not self.is_mutable_collection("params"):
-            y = deepseek_batchsplit.scan_batch_split_layers(
+            # Apply Dense Layers
+            y = self._apply_interleaved_scanned_layers(
                 y,
-                self.variables["params"]["moe_layers"],
-                decoder_positions,
-                decoder_segment_ids,
-                model_mode=model_mode,
-                mesh=mesh,
-                quant=self.quant,
-                cfg=cfg,
-                policy=policy,
+                layer_type="dense",
+                start_idx=0,
+                end_idx=cfg.first_num_dense_layers,
+                engram_indices=cfg.engram_layers,
+                **common_kwargs,
+            )
+
+            # Apply MoE Layers
+            y = self._apply_interleaved_scanned_layers(
+                y,
+                layer_type="moe",
+                start_idx=cfg.first_num_dense_layers,
+                end_idx=cfg.num_decoder_layers,
+                engram_indices=cfg.engram_layers,
+                **common_kwargs,
             )
           else:
+            dense_layer.__call__ = functools.partial(dense_layer.__call__, **layer_call_kwargs)
             y, _ = self.scan_decoder_layers(
                 cfg,
-                moe_layer,
-                num_moe_layers,
-                "moe_layers",
+                dense_layer,
+                cfg.first_num_dense_layers,
+                "dense_layers",
                 mesh,
                 in_axes_tuple=(nn.broadcast,) * len(broadcast_args),
                 model_mode=model_mode,
             )(y, *broadcast_args)
+            moe_layer.__call__ = functools.partial(moe_layer.__call__, **layer_call_kwargs)
+            num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers
+
+            # If batch-split schedule is used and initialization is complete,
+            # as detected by immutable params, use deepseek_batchsplit custom
+            # scan with initialized parameters.
+            if cfg.use_batch_split_schedule and not self.is_mutable_collection("params"):
+              y = deepseek_batchsplit.scan_batch_split_layers(
+                  y,
+                  self.variables["params"]["moe_layers"],
+                  decoder_positions,
+                  decoder_segment_ids,
+                  model_mode=model_mode,
+                  mesh=mesh,
+                  quant=self.quant,
+                  cfg=cfg,
+                  policy=policy,
+              )
+            else:
+              y, _ = self.scan_decoder_layers(
+                  cfg,
+                  moe_layer,
+                  num_moe_layers,
+                  "moe_layers",
+                  mesh,
+                  in_axes_tuple=(nn.broadcast,) * len(broadcast_args),
+                  model_mode=model_mode,
+              )(y, *broadcast_args)
         elif cfg.decoder_block == DecoderBlockType.GEMMA3:
           y = self._apply_gemma3_scanned_blocks(
               y,
@@ -1107,3 +1147,74 @@ def _apply_gemma3_scanned_blocks(
           **layer_call_kwargs,
       )
     return y
+
+  # TODO(b/490118813): Relocate the following functions to their designated directories
+  # once the plug-in strategy is implemented: _find_next_boundary(), _apply_single_engram_layer()
+  # _apply_scanned_chunk() and _apply_interleaved_scanned_layers().
+  def _find_next_boundary(self, current_idx, end_idx, engram_indices):
+    """Finds the next index boundary, either the next Engram layer index or the overall end index."""
+    next_engrams = [l for l in engram_indices if l > current_idx]
+    if next_engrams:
+      return min(end_idx, *next_engrams)
+    return end_idx
+
+  def _apply_single_engram_layer(self, y, current_idx, layer_type, **kwargs):
+    """Applies a single, unscanned Engram layer."""
+    layer = kwargs["dense_layer"] if layer_type == "dense" else kwargs["moe_layer"]
+    layer_prefix = "dense_layers" if layer_type == "dense" else "moe_layers"
+    original_call = kwargs["original_dense_call"] if layer_type == "dense" else kwargs["original_moe_call"]
+    layer_call_kwargs = kwargs["layer_call_kwargs"]
+
+    layer.__call__ = original_call
+    y, _ = layer(
+        config=self.config,
+        mesh=self.mesh,
+        name=f"{layer_prefix}_engram_{current_idx}",
+        quant=self.quant,
+        model_mode=self.model_mode,
+        layer_idx=current_idx,
+    )(
+        y,
+        kwargs["decoder_segment_ids"],
+        kwargs["decoder_positions"],
+        kwargs["deterministic"],
+        kwargs["model_mode"],
+        decoder_input_tokens=kwargs["decoder_input_tokens"],
+        **layer_call_kwargs,
+    )
+    layer.__call__ = functools.partial(original_call, **layer_call_kwargs)
+    return y
+
+  def _apply_scanned_chunk(self, y, current_idx, next_boundary, layer_type, **kwargs):
+    """Applies a contiguous chunk of layers using the scan operation."""
+    layer = kwargs["dense_layer"] if layer_type == "dense" else kwargs["moe_layer"]
+    layer_prefix = "dense_layers" if layer_type == "dense" else "moe_layers"
+    broadcast_args = kwargs["broadcast_args"]
+    scan_length = next_boundary - current_idx
+
+    if scan_length > 0:
+      y, _ = self.scan_decoder_layers(
+          self.config,
+          layer,
+          scan_length,
+          f"{layer_prefix}_{current_idx}_{next_boundary - 1}",
+          self.mesh,
+          in_axes_tuple=(nn.broadcast,) * len(broadcast_args),
+          model_mode=kwargs["model_mode"],
+      )(y, *broadcast_args)
+    return y
+
+  def _apply_interleaved_scanned_layers(self, y, layer_type, start_idx, end_idx, engram_indices, **kwargs):
+    """Applies a mix of scanned standard layers and unscanned Engram layers."""
+    current_idx = start_idx
+    while current_idx < end_idx:
+      if current_idx in engram_indices:
+        # Handle individual unscanned Engram layer
+        y = self._apply_single_engram_layer(y, current_idx, layer_type, **kwargs)
+        current_idx += 1
+      else:
+        # Find next boundary and scan the chunk
+        next_boundary = self._find_next_boundary(current_idx, end_idx, engram_indices)
+        y = self._apply_scanned_chunk(y, current_idx, next_boundary, layer_type, **kwargs)
+        current_idx = next_boundary
+    return y

tests/unit/deepseek_scan_engram_test.py

@@ -0,0 +1,175 @@
+# Copyright 2023-2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for DeepSeek Engram across scanned decoder layers."""
+
+import gc
+import os
+import unittest
+from unittest.mock import patch
+
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+
+from maxtext.configs import pyconfig
+from maxtext.utils.globals import MAXTEXT_PKG_DIR
+from maxtext.common.common_types import MODEL_MODE_TRAIN
+from maxtext.layers.decoders import Decoder
+from maxtext.utils import maxtext_utils
+import pytest
+
+
+class DummyEmbedding:
+  """Dummy embedding layer for testing."""
+
+  def __init__(self, emb_dim: int):
+    self.emb_dim = emb_dim
+
+  def __call__(self, x, model_mode):
+    return jnp.ones((x.shape[0], x.shape[1], self.emb_dim))
+
+
+class TestDeepSeekScanEngram(unittest.TestCase):
+  """Test DeepSeek decoder block with scan_layers=True and engram_layers."""
+
+  _COMMON_CONFIG = [
+      "run_name=test_deepseek_scan_engram",
+      "model_name=deepseek-custom",
+      "override_model_config=True",
+      "decoder_block=deepseek",
+      "scan_layers=True",
+      "first_num_dense_layers=5",
+      "base_num_decoder_layers=10",
+      "num_decoder_layers=10",
+      "mhc_expansion_rate=4",
+      "attention=dot_product",
+      "per_device_batch_size=2",
+      "max_target_length=8",
+      "max_prefill_predict_length=8",
+      "enable_checkpointing=False",
+      "engram_num_heads=1",
+      "engram_head_dim=32",
+      "engram_vocab_bases=[226240,226240]",
+      "engram_max_ngram_size=3",
+      "engram_kernel_size=4",
+      "hf_access_token=dummy",
+      "tokenizer_path=dummy",
+  ]
+
+  def _test_engram_pattern(self, mock_from_pretrained, engram_layers_str, expected_keys):
+    """Helper method to test different engram layer patterns."""
+
+    # Setup mock tokenizer
+    class MockTokenizer:
+      """Mock tokenizer for testing."""
+
+      pad_token_id = 0
+
+      def __len__(self):
+        return 1000
+
+      def __call__(self, x):
+        return jnp.ones_like(x)
+
+      def convert_ids_to_tokens(self, *args, **kwargs):
+        return "a"
+
+      def decode(self, *args, **kwargs):
+        return "a"
+
+      def batch_decode(self, token_ids, *args, **kwargs):
+        return ["a" for _ in token_ids]
+
+    mock_from_pretrained.return_value = MockTokenizer()
+
+    config_path = os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")
+    config = pyconfig.initialize([None, config_path] + self._COMMON_CONFIG + [f"engram_layers=[{engram_layers_str}]"])
+
+    devices_array = maxtext_utils.create_device_mesh(config)
+    mesh = Mesh(devices_array, config.mesh_axes)
+
+    decoder = Decoder(
+        config=config,
+        mesh=mesh,
+        model_mode=MODEL_MODE_TRAIN,
+    )
+
+    batch_size = config.global_batch_size_to_load
+    seq_len = config.max_target_length
+
+    decoder_input_tokens = jnp.ones((batch_size, seq_len), dtype=jnp.int32)
+    decoder_positions = jnp.ones((batch_size, seq_len), dtype=jnp.int32)
+    decoder_segment_ids = jnp.ones((batch_size, seq_len), dtype=jnp.int32)
+
+    shared_embedding = DummyEmbedding(emb_dim=config.emb_dim)
+
+    with mesh:
+      variables = decoder.init(
+          {"params": jax.random.PRNGKey(0), "dropout": jax.random.PRNGKey(1), "aqt": jax.random.PRNGKey(2)},
+          shared_embedding=shared_embedding,
+          decoder_input_tokens=decoder_input_tokens,
+          decoder_positions=decoder_positions,
+          decoder_segment_ids=decoder_segment_ids,
+          deterministic=True,
+          model_mode=MODEL_MODE_TRAIN,
+      )
+
+    self.assertIn("params", variables)
+    params = variables["params"]
+    for key in expected_keys:
+      self.assertIn(key, params)
+
+    del variables
+    del params
+    del decoder
+    jax.clear_caches()
+    gc.collect()
+
+  @pytest.mark.tpu_only
+  @patch("transformers.AutoTokenizer.from_pretrained")
+  def test_decoder_init_engram_2_8(self, mock_from_pretrained):
+    """Test engram layers at indices 2 and 8."""
+    self._test_engram_pattern(
+        mock_from_pretrained,
+        "2,8",
+        [
+            "dense_layers_0_1",
+            "dense_layers_engram_2",
+            "dense_layers_3_4",
+            "moe_layers_5_7",
+            "moe_layers_engram_8",
+            "moe_layers_9_9",
+        ],
+    )
+
+  @pytest.mark.tpu_only
+  @patch("transformers.AutoTokenizer.from_pretrained")
+  def test_decoder_init_engram_0_5(self, mock_from_pretrained):
+    """Test engram layers at indices 0 and 5 - first engram layer of block."""
+    self._test_engram_pattern(
+        mock_from_pretrained,
+        "0,5",
+        ["dense_layers_engram_0", "dense_layers_1_4", "moe_layers_engram_5", "moe_layers_6_9"],
+    )
+
+  @pytest.mark.tpu_only
+  @patch("transformers.AutoTokenizer.from_pretrained")
+  def test_decoder_init_engram_4_9(self, mock_from_pretrained):
+    """Test engram layers at indices 4 and 9 - last engram layer of block."""
+    self._test_engram_pattern(
+        mock_from_pretrained,
+        "4,9",
+        ["dense_layers_0_3", "dense_layers_engram_4", "moe_layers_5_8", "moe_layers_engram_9"],
+    )

tests/unit/train_compile_test.py

@@ -830,7 +830,7 @@ def test_engram_integration(self):
             "compile_topology_num_slices=1",
             "model_name=deepseek-custom",
             "per_device_batch_size=4",
-            "scan_layers=False",  # TODO(ranran): update to scan_layers=True after support
+            "scan_layers=True",
             "max_target_length=1024",
             "attention=flash",
             "use_tokamax_splash=True",