Enable selective parameter training strategy

Author: RissyRan

src/maxtext/configs/base.yml

@@ -367,6 +367,8 @@ index_topk: 2048
 sparse_indexer_loss: False
 # Multiplier for the indexer KL divergence loss
 indexer_loss_scaling_factor: 0.0
+# Whether to enable sparse training for indexer by detaching its input from the computational graph
+indexer_sparse_training: False
 
 # MLA parameters
 q_lora_rank: 0
@@ -797,6 +799,10 @@ adam_eps: 1.e-8 # A small constant applied to denominator outside of the square
 adam_eps_root: 0. # A small constant applied to denominator inside the square root.
 adam_weight_decay: 0.1 # AdamW Weight decay
 adamw_mask: [] # List of parameter names/patterns to exclude from weight decay in AdamW, like ['bias', '.*norm', '.*ln.*'].
+# List of parameter names/patterns to train.
+# If non-empty, all other parameters will be frozen. Example: ['.*indexer.*'].
+# If empty (default), all parameters are trained.
+trainable_parameters_mask: []
 mu_dtype: "" # data type to store "mu" of AdamW tracking the first moment. Inherits from  weight_dtype if unset.
 # Setting nu_dtype is not yet supported by optax, instead nu_dtype is always inherited from weights.
 # See b/399961932 for more.

src/maxtext/configs/types.py

@@ -537,6 +537,10 @@ class AttentionIndexer(BaseModel):
   index_topk: NonNegativeInt = Field(2048, description="Number of tokens selected by the query token in indexer.")
   sparse_indexer_loss: bool = Field(False, description="Determines the token selection strategy for indexer loss.")
   indexer_loss_scaling_factor: float = Field(0.0, description="Multiplier for the indexer KL divergence loss.")
+  indexer_sparse_training: bool = Field(
+      False,
+      description="Whether to enable sparse training for indexer by detaching its input from the computational graph.",
+  )
 
 
 class Llama4Attention(BaseModel):
@@ -1195,6 +1199,13 @@ class AdamW(BaseModel):
           "List of parameter names/patterns to exclude from weight decay in AdamW," " like ['bias', '.*norm', '.*ln.*']"
       ),
   )
+  trainable_parameters_mask: list[str] = Field(
+      default_factory=list,
+      description=(
+          "List of parameter names/patterns to train. If non-empty, all other parameters will be frozen, "
+          "example: ['.*indexer.*']. If empty (default), all parameters are trained."
+      ),
+  )
   mu_dtype: str = Field(
       "",
       description="Data type for 'mu' (first moment) in AdamW. Inherits from weight_dtype if empty.",

src/maxtext/layers/attention_mla.py

@@ -266,12 +266,23 @@ def __call__(
     bsz, seqlen, _ = inputs_q.shape  # s = t = seqlen
 
     # Query Processing: Project from Latent low_rank_q
-    q = self.wq_b(low_rank_q)  # [b, t, q_lora_rank] -> [b, t, h * d]
+    if self.config.indexer_sparse_training:
+      # Detach indexer input from the computational graph so main loss doesn't backprop through indexer,
+      # and indexer loss doesn't backprop into main model embeddings/latent variables.
+      inputs_q_for_indexer = jax.lax.stop_gradient(inputs_q)
+      low_rank_q_for_indexer = jax.lax.stop_gradient(low_rank_q)
+      inputs_kv_for_indexer = jax.lax.stop_gradient(inputs_kv)
+    else:
+      inputs_q_for_indexer = inputs_q
+      low_rank_q_for_indexer = low_rank_q
+      inputs_kv_for_indexer = inputs_kv
+
+    q = self.wq_b(low_rank_q_for_indexer)  # [b, t, q_lora_rank] -> [b, t, h * d]
     q = q.reshape(bsz, seqlen, self.n_heads, self.head_dim)  # [b, t, h, d]
     q = self.apply_partial_rope(q, inputs_positions=inputs_positions)
 
     # Key Processing: Project from Input
-    k = self.wk(inputs_kv)  # [b, s, embed_dim] -> [b, s, d]
+    k = self.wk(inputs_kv_for_indexer)  # [b, s, embed_dim] -> [b, s, d]
     k = self.k_norm(k)
     k = k[:, :, None, :]  # [b, s, d] -> [b, s, 1, d]
     k = self.apply_partial_rope(k, inputs_positions=inputs_positions)
@@ -283,7 +294,7 @@ def __call__(
     logits = jnp.einsum("bthd, bsd -> btsh", q, k, precision=self.config.matmul_precision)
     logits = jax.nn.relu(logits)
     # Compute head weights: project from input, [b, t, embed_dim] -> [b, t, h]
-    weights = self.weights_proj(inputs_q)
+    weights = self.weights_proj(inputs_q_for_indexer)
     # Weights scaling affect indexer_score, but does not affect topk_indices. Keep scaling for numerical stability.
     # https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/87e509a2e5a100d221c97df52c6e8be7835f0057/inference/model.py#L478-L480
     weights = weights * (self.n_heads**-0.5) * self.softmax_scale

src/maxtext/optimizers/optimizers.py

@@ -24,31 +24,35 @@
 from maxtext.utils.muon_utils import get_muon_weight_dimension_numbers
 
 
-def get_adamw_mask(config):
-  """Create a mask function for AdamW optimizer to exclude certain parameters from weight decay."""
-  if not getattr(config, "adamw_mask", None):
+def _get_path_mask_fn(patterns, match_returns_true=True):
+  """Helper to create a mask function from a list of regex patterns."""
+  if not patterns:
     return None
 
-  compiled_patterns = [re.compile(pattern) for pattern in config.adamw_mask]
+  compiled_patterns = [re.compile(pattern) for pattern in patterns]
 
   def mask_fn(params):
-    def _is_decayed(path, _):
+    def _is_masked(path, _):
       # Join path keys into a single string for pattern matching (e.g., "layer1/bias")
       path_str = "/".join(str(getattr(p, "key", getattr(p, "idx", getattr(p, "name", p)))) for p in path)
-      # If any pattern in adamw_mask matches the path, exclude from weight decay (return False).
-      # Otherwise, apply weight decay (return True).
-      return not any(pattern.search(path_str) for pattern in compiled_patterns)
+      matched = any(pattern.search(path_str) for pattern in compiled_patterns)
+      return matched if match_returns_true else not matched
 
-    return jax.tree_util.tree_map_with_path(_is_decayed, params)
+    return jax.tree_util.tree_map_with_path(_is_masked, params)
 
   return mask_fn
 
 
+def get_adamw_mask(config):
+  """Create a mask function for AdamW optimizer to exclude certain parameters from weight decay."""
+  return _get_path_mask_fn(getattr(config, "adamw_mask", None), match_returns_true=False)
+
+
 def get_optimizer(config, learning_rate_schedule, model=None):
   """Create optimizer."""
   if config.opt_type == "adamw":
     # Create AdamW Optimizer following Llama2's training details, see https://arxiv.org/pdf/2307.09288.pdf section 2.2
-    return optax.adamw(
+    base_opt = optax.adamw(
         learning_rate_schedule,
         b1=config.adam_b1,
         b2=config.adam_b2,
@@ -59,7 +63,7 @@ def get_optimizer(config, learning_rate_schedule, model=None):
         mask=get_adamw_mask(config),
     )
   elif config.opt_type == "adam_pax":
-    return adam_pax(
+    base_opt = adam_pax(
         learning_rate_schedule,
         beta1=config.adam_b1,
         beta2=config.adam_b2,
@@ -69,7 +73,7 @@ def get_optimizer(config, learning_rate_schedule, model=None):
         mask=get_adamw_mask(config),
     )
   elif config.opt_type == "sgd":
-    return optax.sgd(learning_rate_schedule)
+    base_opt = optax.sgd(learning_rate_schedule)
   elif config.opt_type == "muon":
     # extract muon dimension number from model structure
     if model is not None:
@@ -92,10 +96,19 @@ def get_optimizer(config, learning_rate_schedule, model=None):
         "adam_eps_root": config.adam_eps_root,
         "adam_weight_decay": config.adam_weight_decay,
     }
-    return muon(**muon_kwargs)
+    base_opt = muon(**muon_kwargs)
   else:
     raise ValueError(f"{config.opt_type=} is not a supported.")
 
+  # If a whitelist of trainable parameters is provided, freeze everything else.
+  # When trainable_parameters_mask is empty, freeze_mask_fn is None and all parameters are trained.
+  trainable_patterns = getattr(config, "trainable_parameters_mask", None)
+  freeze_mask_fn = _get_path_mask_fn(trainable_patterns, match_returns_true=False)
+  if freeze_mask_fn is not None:
+    return optax.chain(base_opt, optax.masked(optax.set_to_zero(), freeze_mask_fn))
+
+  return base_opt
+
 
 def adam_pax(
     learning_rate_fn: optax.Schedule,

src/maxtext/trainers/pre_train/train.py

@@ -232,6 +232,13 @@ def loss_fn(model, config, data, dropout_rng, params, is_train=True):
 
     if indexer_losses:
       indexer_loss = jnp.mean(jnp.concatenate(indexer_losses))
+      # DeepSeek V3.2: When `indexer_sparse_training` is true, we optimize the indexer
+      # using ONLY the indexer loss, and the main model using ONLY the language modeling loss.
+      # To do this, we decouple the gradients. We detach the indexer input from the
+      # computational graph inside the Indexer module itself (in attention_mla.py)
+      # by stopping gradients on its inputs.
+      # So here, we just add the indexer loss to the total loss. The gradients will
+      # naturally separate because the inputs to the indexer were stopped.
       loss += indexer_loss
     else:
       max_logging.debug("No indexer loss found.")

tests/unit/attention_test.py

@@ -37,8 +37,10 @@
     DEFAULT_MASK_VALUE,
 )
 from maxtext.layers.attention_mla import MLA
+from maxtext.layers.attention_mla import Indexer
 from maxtext.layers.attention_op import ChunkedCausalMask, _generate_chunk_attention_mask, _make_bidirectional_block_mask
 from maxtext.layers.attentions import Attention
+from maxtext.layers import embeddings
 from maxtext.configs import pyconfig
 from maxtext.models.qwen3 import Qwen3NextGatedDeltaNet
 import numpy as np
@@ -1693,6 +1695,78 @@ def test_indexer_loss_kl_divergence_zero(self):
 
     np.testing.assert_allclose(loss, 0.0, atol=1e-5)
 
+  def test_indexer_gradients(self):
+    # Test that gradients flow back to inputs when indexer_sparse_training=False
+    # but do NOT flow back when indexer_sparse_training=True
+    bsz, seqlen = 2, 8
+    inputs_positions = jnp.broadcast_to(jnp.arange(seqlen)[None, :], (bsz, seqlen))
+
+    for sparse_training in [False, True]:
+      with self.subTest(indexer_sparse_training=sparse_training):
+        argv = [
+            "",
+            get_test_config_path(),
+            "run_name=test",
+            "attention_type=mla",
+            f"indexer_sparse_training={sparse_training}",
+            "max_target_length=16",
+            "index_topk=4",
+            "index_n_heads=2",
+            "index_head_dim=8",
+            "emb_dim=16",
+            "qk_rope_head_dim=4",
+            "q_lora_rank=16",
+        ]
+        config = pyconfig.initialize(argv)
+        rngs = nnx.Rngs(0)
+        mesh = jax.sharding.Mesh(jax.devices(), ("data",))
+        rope = embeddings.RotaryEmbedding(
+            min_timescale=1,
+            max_timescale=10000,
+            mesh=mesh,
+            embedding_dims=config.qk_rope_head_dim,
+            fprop_dtype=jnp.float32,
+            rngs=rngs,
+        )
+        rope.interleave = False
+
+        indexer = Indexer(
+            config=config,
+            rotary_embedding=rope,
+            rngs=rngs,
+        )
+
+        inputs_q = jnp.ones((bsz, seqlen, config.emb_dim))
+        low_rank_q = jnp.ones((bsz, seqlen, config.q_lora_rank))
+        inputs_kv = jnp.ones((bsz, seqlen, config.emb_dim))
+
+        def loss_fn(inputs_q, low_rank_q, inputs_kv, indexer):
+          _, _, indexer_score = indexer(
+              inputs_q=inputs_q,
+              low_rank_q=low_rank_q,
+              inputs_kv=inputs_kv,
+              inputs_positions=inputs_positions,
+          )
+          # A dummy loss function (e.g., sum of scores)
+          return jnp.sum(indexer_score)
+
+        # Calculate gradients with respect to the 3 inputs
+        grad_fn = nnx.grad(loss_fn, argnums=(0, 1, 2))
+        grads = grad_fn(inputs_q, low_rank_q, inputs_kv, indexer)
+
+        grad_q, grad_low_rank, grad_kv = grads
+
+        if sparse_training:
+          # Gradients should be exactly zero
+          self.assertTrue(jnp.all(grad_q == 0.0))
+          self.assertTrue(jnp.all(grad_low_rank == 0.0))
+          self.assertTrue(jnp.all(grad_kv == 0.0))
+        else:
+          # Gradients should be non-zero
+          self.assertFalse(jnp.all(grad_q == 0.0))
+          self.assertFalse(jnp.all(grad_low_rank == 0.0))
+          self.assertFalse(jnp.all(grad_kv == 0.0))
+
 
 class Qwen3NextGatedDeltaNetTest(unittest.TestCase):
   """Test for the Gated Delta Net in Qwen3-Next"""

tests/unit/optimizers_test.py

@@ -362,5 +362,71 @@ def test_optimizer_without_mask(self, opt_type, mock_path):
       self.assertIsNone(kwargs["mask"])
 
 
+class TrainableParametersMaskTest(parameterized.TestCase):
+  """Tests for the trainable parameters mask functionality via get_optimizer"""
+
+  def test_get_optimizer_with_trainable_mask(self):
+    """Test get_optimizer with a valid trainable_parameters_mask."""
+    argv = [
+        "",
+        get_test_config_path(),
+        "run_name=test_with_trainable_mask",
+        "trainable_parameters_mask=['.*indexer.*', 'layer_norm']",
+    ]
+    config = pyconfig.initialize(argv)
+
+    # Use a constant learning rate > 0 to ensure non-zero updates
+    def learning_rate_schedule(step):
+      return 1.0
+
+    opt = optimizers.get_optimizer(config, learning_rate_schedule)
+
+    # We can test the optimizer by creating some dummy params and gradients
+    # and checking if the updates are zeroed out for non-trainable parameters.
+    params = {
+        "layer1": {"kernel": jax.numpy.ones((2, 2)), "indexer": jax.numpy.ones((2, 2))},
+        "layer2": {"layer_norm": {"scale": jax.numpy.ones((2, 2))}},
+        "layer3": {"ln": {"scale": jax.numpy.ones((2, 2))}},
+    }
+
+    # Give some non-zero gradients
+    grads = jax.tree_util.tree_map(lambda x: jax.numpy.ones_like(x) * 0.5, params)
+
+    # Initialize optimizer state
+    opt_state = opt.init(params)
+
+    # Compute updates
+    updates, _ = opt.update(grads, opt_state, params)
+
+    # 'layer1/kernel' doesn't match the trainable mask, so it should be frozen (update == 0)
+    self.assertTrue(jax.numpy.all(updates["layer1"]["kernel"] == 0))
+    # 'layer3/ln/scale' doesn't match the trainable mask, so it should be frozen (update == 0)
+    self.assertTrue(jax.numpy.all(updates["layer3"]["ln"]["scale"] == 0))
+    # 'layer1/indexer' matches, so it should be trained (update != 0)
+    self.assertFalse(jax.numpy.all(updates["layer1"]["indexer"] == 0))
+    # 'layer2/layer_norm/scale' matches, so it should be trained (update != 0)
+    self.assertFalse(jax.numpy.all(updates["layer2"]["layer_norm"]["scale"] == 0))
+
+  def test_get_optimizer_without_trainable_mask(self):
+    """Test get_optimizer when trainable_parameters_mask is empty."""
+    argv = ["", get_test_config_path(), "run_name=test", "trainable_parameters_mask=[]"]
+    config = pyconfig.initialize(argv)
+
+    # Use a constant learning rate > 0 to ensure non-zero updates
+    def learning_rate_schedule(step):
+      return 1.0
+
+    opt = optimizers.get_optimizer(config, learning_rate_schedule)
+
+    params = {"layer1": {"kernel": jax.numpy.ones((2, 2))}}
+    grads = {"layer1": {"kernel": jax.numpy.ones((2, 2)) * 0.5}}
+
+    opt_state = opt.init(params)
+    updates, _ = opt.update(grads, opt_state, params)
+
+    # When no trainable mask is provided, nothing is frozen by this mechanism
+    self.assertFalse(jax.numpy.all(updates["layer1"]["kernel"] == 0))
+
+
 if __name__ == "__main__":
   unittest.main()