From 86007c0e0ff590794ac923049af8bcd98585bcd4 Mon Sep 17 00:00:00 2001 From: maxtext authors Date: Thu, 10 Sep 2026 18:57:44 -0700 Subject: [PATCH] PR #4360: [NNX] Delete Linen 4/5: remove the Linen decoder/attention layers and *_as_linen model wrappers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/AI-Hypercomputer/maxtext/pull/4360 ## NNX Migration Route Map 1. ✅ Add NNX scaffolding: `pure_nnx` flag, `init_state_fn`, `TrainStateNNX`, NNX utils. Linen workflow unchanged. (PR #3427) 2. ✅ NNX sharding utilities: `get_abstract_state_nnx`, `get_named_sharding_nnx`, `set_named_sharding_nnx`, `get_partition_spec_nnx`, `get_mesh_from_config`. (PR #3470) 3. ✅ NNX fully supported end-to-end: `TrainStateNNX`, model creation, gradient accumulation, checkpointing, and training loop dispatch. (PR #3500) 4. ✅ Sharding diagnostics on NNX, plus post-training bugfixes that surfaced once the NNX path got exercised end-to-end. (PR #3652) 5. ✅ NNX correctness fixes, feature enablements, and vocab tiling on NNX. 6. ✅ NNX-native DPO. 7. ✅ NNX-native MaxEngine inference. (PR #3821) 8. ✅ NNX-native LoRA + GRPO. (PR #3824) 9. ✅ NNX-aware QK-Clip + remaining checkpoint utilities. (PR #3836) 9.5. ✅ NNX + AQT in MaxEngine + serve-mode reload + gpt3 prefill fix. (PR #3844) 10. ✅ Vocab tiling `custom_vjp` for NNX. 11. ✅ Set NNX defaults to `True`; regenerate sharding goldens; flip back integration-test `pure_nnx=False` annotations. (PR #3526) 12. Delete Linen-specific code paths and NNX compatibility flags — delivered as 5 stacked PRs (12.1–12.5): 12.1. ✅ Collapse Linen dispatch in utils / quantization / checkpointing / sharding. (PR #4357) 12.2. ✅ Collapse Linen dispatch in trainers (pre-train, DiLoCo, GRPO); regenerate the NNX DPO golden. (PR #4358) 12.3. ✅ Collapse Linen dispatch in inference (maxengine, kvcache, vLLM, LoRA). (PR #4359) 12.4. 🔄 **[This PR]** Delete the Linen model code now that 12.1–12.3 removed every reference: the `decoders.py` stack (`Decoder`/`DecoderLayer`/`SequentialBlockDecoderLayers`, only `deepstack_process` kept), the Linen attention/embeddings/encoders/MTP paths, and the `*_as_linen`/`transformer_as_linen` wrappers in the model files. The tests importing those classes are removed in this same PR so the branch stays import-clean. The NNX↔Linen bridge + Linen GRPO reference are kept for the checkpoint converters / torch-gated tests (follow-up). Stacks on 12.3. 12.5. ❌ Remove the `pure_nnx` / `enable_nnx` / `pure_nnx_decoder` config flags. # Description Fourth of the stacked "Delete Linen" series (see 1/5). PRs 1–3 removed every source reference to the Linen model code, so this PR deletes it. That means the Linen decoder stack in `decoders.py` (`Decoder`, `DecoderLayer`, `SequentialBlockDecoderLayers` — only `deepstack_process` stays), the Linen attention implementations, the Linen `embeddings` / `encoders` / `multi_token_prediction` paths, and the `*_as_linen` / `transformer_as_linen` ToLinen wrappers in the model files (the wrapped NNX classes stay). The tests that imported those classes are removed in this same PR so the branch stays import-clean — a repo-wide grep for the deleted symbols now returns only the deletions themselves. Note that the NNX↔Linen bridge (`transformer_as_linen` / `TransformerLinen` / `init_initial_state`) and the Linen GRPO reference are deliberately kept here, since the checkpoint-conversion tools and the torch-gated correctness tests still use them; migrating those to pure NNX is a follow-up. Base branch is `feat/nnx-del-linen-3-inference`. ~16 files, +56 / −3498. # Checklist Before submitting this PR, please make sure (put X in square brackets): - [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label. - [x] I have necessary comments in my code, particularly in hard-to-understand areas. - [x] I have run end-to-end tests tests and provided workload links above if applicable. - [x] I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in [our documentation](https://maxtext.readthedocs.io/en/latest/development.html#adding-new-documentation-files). Copybara import of the project: -- 642e6d079fb1dd546f994b10d3b1421cdc53ef47 by Lance Wang : [NNX] Delete Linen (4/5): remove Linen decoder/attention layers and *_as_linen model wrappers Merging this change closes #4360 Reverts f8d1f3bb1dc99d1439f775ccddc7f21a654b5b00 PiperOrigin-RevId: 979520164 --- src/maxtext/layers/attention_mla.py | 138 +- src/maxtext/layers/attention_op.py | 95 + src/maxtext/layers/attentions.py | 10 + src/maxtext/layers/decoders.py | 1926 +++++++++++++++++ src/maxtext/layers/embeddings.py | 350 ++- src/maxtext/layers/encoders.py | 34 + src/maxtext/layers/multi_token_prediction.py | 37 + src/maxtext/layers/nnx_decoders.py | 24 +- src/maxtext/models/gemma3.py | 37 + src/maxtext/models/gemma4_vision.py | 15 + src/maxtext/models/llama4.py | 22 + src/maxtext/models/models.py | 397 +++- src/maxtext/models/qwen3.py | 49 +- .../integration/deepseek_scan_engram_test.py | 91 +- tests/unit/nnx_decoder_test.py | 3 + tests/unit/router_replay_test.py | 72 + tests/unit/tiling_test.py | 585 ++++- 17 files changed, 3824 insertions(+), 61 deletions(-) diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index 05ee1e2e38..13dd34a1d4 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -63,8 +63,9 @@ DEFAULT_MASK_VALUE, ) +from maxtext.layers import nnx_wrappers from maxtext.layers.attentions import Attention -from maxtext.layers.initializers import nd_dense_init, NdInitializer +from maxtext.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned from maxtext.layers.linears import DenseGeneral from maxtext.layers.normalizations import RMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant @@ -451,6 +452,141 @@ def scan_body_indexer(carry, xs): return indexer_mask, topk_indices, indexer_score +def mla_as_linen( + *, + config: Config, + num_query_heads: int, + num_kv_heads: int, + head_dim: int, + max_target_length: int, + mesh: Mesh, + attention_kernel: str, + inputs_q_shape: Tuple, + inputs_kv_shape: Tuple, + dtype: DType = jnp.float32, + weight_dtype: DType = jnp.float32, + max_prefill_predict_length: int = -1, + dropout_rate: float = 0.0, + kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "normal"), + float32_qk_product: bool = False, # computes logits in float32 for stability. + float32_logits: bool = False, # cast logits in float32 for stability. + quant: Optional[Quant] = None, + kv_quant: Optional[KVQuant] = None, + attention_type: AttentionType = AttentionType.MLA, # Default to MLA attention + attn_logits_soft_cap: float | None = None, + sliding_window_size: int | None = None, + use_ragged_attention: bool = False, + ragged_block_size: int = 256, + use_qk_norm: bool = False, + query_pre_attn_scalar: float | None = None, + use_bias_in_projections: bool = False, # Set to True will enable bias in q, k, v, o projections + # Temperature tuning parameters used for Llama4 + temperature_tuning: bool = False, + temperature_tuning_scale: float = 0.1, + temperature_tuning_floor_scale: float = 8192.0, + # Shard the query activation as the same as the key and value. + # TODO: Find a better sharding axis name. + # TODO: Further break down the Training and Inference axes for the q, k, v. + prefill_query_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, KV_HEAD, KV_HEAD_DIM), + prefill_key_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, KV_HEAD, KV_HEAD_DIM), + prefill_value_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, KV_HEAD, KV_HEAD_DIM), + query_axis_names: AxisNames = (KV_BATCH, LENGTH, KV_HEAD, KV_HEAD_DIM), + key_axis_names: AxisNames = (KV_BATCH, LENGTH, KV_HEAD, KV_HEAD_DIM), + value_axis_names: AxisNames = (KV_BATCH, LENGTH, KV_HEAD, KV_HEAD_DIM), + input_axis_names: AxisNames = (BATCH_ATTN, LENGTH, EMBED), + out_axis_names: AxisNames = (BATCH_ATTN, LENGTH, HEAD, D_KV), + prefill_input_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, EMBED), + decode_input_axis_names: AxisNames = (DECODE_BATCH, DECODE_LENGTH, EMBED), + prefill_out_axis_names: AxisNames = (PREFILL_KV_BATCH, PREFILL_LENGTH, HEAD, D_KV), + decode_out_axis_names: AxisNames = (DECODE_BATCH, DECODE_LENGTH, HEAD, D_KV), + prefill_cache_axis_order: AxisIdxes = (1, 2, 0, 3), + ar_cache_axis_order: AxisIdxes = (1, 2, 0, 3), + compute_axis_order: AxisIdxes = (0, 1, 2, 3), + reshape_q: bool = False, + is_nope_layer: bool = False, + is_vision: bool = False, + model_mode: str = MODEL_MODE_TRAIN, + q_lora_rank: int = 0, + kv_lora_rank: int = 512, + qk_nope_head_dim: int = 128, + qk_rope_head_dim: int = 64, + v_head_dim: int = 128, + max_position_embeddings: int = 4096 * 4, + original_max_position_embeddings: int = 4096, + mscale: float = 1.0, # scaling factor for softmax + rope_factor: float = 40.0, # rotary embedding factor + name: str | None = None, +): + """A factory function to create an MLA as a Linen module. + + This function serves as a bridge to use the NNX-based `MLA` within a + Linen model. + """ + return nnx_wrappers.to_linen( + MLA, + config=config, + num_query_heads=num_query_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + max_target_length=max_target_length, + mesh=mesh, + attention_kernel=attention_kernel, + inputs_q_shape=inputs_q_shape, + inputs_kv_shape=inputs_kv_shape, + dtype=dtype, + weight_dtype=weight_dtype, + max_prefill_predict_length=max_prefill_predict_length, + dropout_rate=dropout_rate, + kernel_init=kernel_init, + float32_qk_product=float32_qk_product, + float32_logits=float32_logits, + quant=quant, + kv_quant=kv_quant, + attention_type=attention_type, + attn_logits_soft_cap=attn_logits_soft_cap, + sliding_window_size=sliding_window_size, + use_ragged_attention=use_ragged_attention, + ragged_block_size=ragged_block_size, + use_qk_norm=use_qk_norm, + query_pre_attn_scalar=query_pre_attn_scalar, + use_bias_in_projections=use_bias_in_projections, + temperature_tuning=temperature_tuning, + temperature_tuning_scale=temperature_tuning_scale, + temperature_tuning_floor_scale=temperature_tuning_floor_scale, + prefill_query_axis_names=prefill_query_axis_names, + prefill_key_axis_names=prefill_key_axis_names, + prefill_value_axis_names=prefill_value_axis_names, + query_axis_names=query_axis_names, + key_axis_names=key_axis_names, + value_axis_names=value_axis_names, + input_axis_names=input_axis_names, + out_axis_names=out_axis_names, + prefill_input_axis_names=prefill_input_axis_names, + decode_input_axis_names=decode_input_axis_names, + prefill_out_axis_names=prefill_out_axis_names, + decode_out_axis_names=decode_out_axis_names, + prefill_cache_axis_order=prefill_cache_axis_order, + ar_cache_axis_order=ar_cache_axis_order, + compute_axis_order=compute_axis_order, + reshape_q=reshape_q, + is_nope_layer=is_nope_layer, + is_vision=is_vision, + model_mode=model_mode, + q_lora_rank=q_lora_rank, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + max_position_embeddings=max_position_embeddings, + original_max_position_embeddings=original_max_position_embeddings, + mscale=mscale, + rope_factor=rope_factor, + name=name, + metadata_fn=variable_to_logically_partitioned, + abstract_init=False, + ) + + class MLA(Attention): """Multi-Head Latent Attention (MLA) layer.""" diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 01aabb6ad5..21a4bc6900 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -71,6 +71,7 @@ from maxtext.kernels.attention.ragged_attention import ragged_gqa from maxtext.kernels.attention.ragged_attention import ragged_mha from maxtext.layers import nnx_wrappers +from maxtext.layers.initializers import variable_to_logically_partitioned from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_utils from maxtext.utils.sharding import logical_to_mesh_axes, maybe_shard_with_pspec, get_logical_axis_rules @@ -399,6 +400,100 @@ def _make_bidirectional_block_mask(bidirectional_mask): return bidirectional_block_mask +def attention_op_as_linen( + *, + config: Config, + mesh: Mesh, + attention_kernel: str, + max_target_length: int, + num_query_heads: int, + num_kv_heads: int, + float32_qk_product: bool = False, + max_prefill_predict_length: int = -1, + float32_logits: bool = False, + flash_axis_names_q: AxisNames = (BATCH_ATTN, HEAD, LENGTH, D_KV), + flash_axis_names_kv: AxisNames = (BATCH_ATTN, HEAD, KV_LENGTH, D_KV), + flash_axis_names_splash_kernel: AxisNames = (HEAD, LENGTH), + prefill_cache_logical_axis_names: AxisNames = ( + CACHE_BATCH_PREFILL, + CACHE_SEQUENCE, + CACHE_HEADS, + CACHE_KV, + ), + cache_logical_axis_names: AxisNames = ( + CACHE_BATCH, + CACHE_SEQUENCE, + CACHE_HEADS, + CACHE_KV, + ), + cache_scale_logical_axis_names: AxisNames = ( + CACHE_SCALE_BATCH, + CACHE_SCALE_SEQUENCE, + CACHE_SCALE_HEADS, + CACHE_SCALE_KV, + ), + ragged_qkv_axis_names: AxisNames = ( + CACHE_BATCH, + CACHE_HEADS, + CACHE_SEQUENCE, + CACHE_KV, + ), + ragged_lengths_names: AxisNames = (CACHE_BATCH,), + compute_axis_order: AxisIdxes = (0, 1, 2, 3), + key_axis_order: AxisIdxes = (2, 0, 1, 3), + reshape_q: bool = False, + dropout_rate: float = 0.0, + dtype: DType = jnp.float32, + quant: Optional[Quant] = None, + kv_quant: Optional[KVQuant] = None, + attention_type: AttentionType = AttentionType.GLOBAL, # Default to global attention + attn_logits_soft_cap: float | None = None, + sliding_window_size: int | None = None, + chunk_attn_window_size: int | None = None, + use_ragged_attention: bool = False, + ragged_block_size: int = 256, +): + """A factory function to create an AttentionOp as a Linen module. + + This function serves as a bridge to use the NNX-based `AttentionOp` within a + Linen model. + """ + return nnx_wrappers.to_linen( + AttentionOp, + config=config, + mesh=mesh, + attention_kernel=attention_kernel, + max_target_length=max_target_length, + num_query_heads=num_query_heads, + num_kv_heads=num_kv_heads, + float32_qk_product=float32_qk_product, + max_prefill_predict_length=max_prefill_predict_length, + float32_logits=float32_logits, + flash_axis_names_q=flash_axis_names_q, + flash_axis_names_kv=flash_axis_names_kv, + flash_axis_names_splash_kernel=flash_axis_names_splash_kernel, + prefill_cache_logical_axis_names=prefill_cache_logical_axis_names, + cache_logical_axis_names=cache_logical_axis_names, + cache_scale_logical_axis_names=cache_scale_logical_axis_names, + ragged_qkv_axis_names=ragged_qkv_axis_names, + ragged_lengths_names=ragged_lengths_names, + compute_axis_order=compute_axis_order, + key_axis_order=key_axis_order, + reshape_q=reshape_q, + dropout_rate=dropout_rate, + dtype=dtype, + quant=quant, + kv_quant=kv_quant, + attention_type=attention_type, + attn_logits_soft_cap=attn_logits_soft_cap, + sliding_window_size=sliding_window_size, + chunk_attn_window_size=chunk_attn_window_size, + use_ragged_attention=use_ragged_attention, + ragged_block_size=ragged_block_size, + metadata_fn=variable_to_logically_partitioned, + ) + + class AttentionOp(nnx.Module): """Attention operation""" diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index 94c4667fb7..fa8dd1f11b 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -91,6 +91,16 @@ def __call__(self, x): return x * jax.lax.rsqrt(jnp.mean(x**2, axis=-1, keepdims=True) + self.eps) +def l2_norm_as_linen(self, eps: float = 1e-6): + """ + Initializes the L2Norm module and returns it as a Linen module. + + Args: + eps: float, epsilon used for numerical stability (default value should be ok for most cases). + """ + return nnx_wrappers.to_linen(L2Norm, eps=eps, metadata_fn=variable_to_logically_partitioned) + + def attention_as_linen( *, config: Config, diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index d8465b5e8f..a9bacf692f 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -13,14 +13,263 @@ # limitations under the License. """Module for decoder layers""" +# pylint: disable=arguments-differ +# pylint: disable=no-name-in-module +import functools +from typing import Any +import warnings + +from flax import linen as nn +from flax import nnx +from flax.linen.partitioning import ScanIn +import jax +from jax.ad_checkpoint import checkpoint_name import jax.numpy as jnp +from jax.sharding import Mesh +from maxtext.common.common_types import Config, DecoderBlockType, ShardMode +from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN +from maxtext.layers import linears +from maxtext.layers import mhc +from maxtext.layers import normalizations +from maxtext.layers import pipeline +from maxtext.layers.nnx_decoders import NNXDecoderLayer, NNXSequentialPipelineStage, NNXScannedPipelineStage +from maxtext.layers import quantizations +from maxtext.layers.attentions import attention_as_linen +from maxtext.layers.embeddings import attend_on_embedding, embed_as_linen, positional_embedding_as_linen +from maxtext.layers.normalizations import rms_norm +from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.models import ( + deepseek, + deepseek4, + deepseek_batchsplit, + deepseek_batchsplit_fp8, + envy, + gemma, + gemma2, + gemma3, + gemma4, + gemma4_small, + gpt3, + gpt_oss, + llama2, + llama4, + mistral, + mixtral, + olmo3, + qwen2, + qwen3, + qwen3_custom, + qwen3_5, + simple_layer, +) +from maxtext.multimodal import utils as mm_utils +from maxtext.utils.sharding import create_sharding +from maxtext.utils import max_logging +from maxtext.utils import max_utils +from maxtext.utils import maxtext_utils +from maxtext.utils import sharding # ------------------------------------------------------------------------------ # The network: Decoder Definitions # ------------------------------------------------------------------------------ +class DecoderLayer(nn.Module): + """ + Transformer decoder layer that attends to the encoder. + This is the core, reusable building block for both the main model's + decoder stack and the auxiliary MTP layers. + """ + + config: Config + mesh: Mesh + model_mode: str + quant: None | Quant = None + + @nn.compact + def __call__( + self, + inputs, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=None, + slot: None | int = None, + kv_cache: jax.Array | None = None, + attention_metadata: dict[str, Any] | None = None, + ): + cfg = self.config + mesh = self.mesh + _maybe_shard_with_logical = functools.partial( + sharding.maybe_shard_with_logical, + mesh=mesh, + shard_mode=cfg.shard_mode, + debug_sharding=cfg.debug_sharding, + ) + + if self.model_mode == MODEL_MODE_PREFILL: + logical_axis_names = ("activation_batch", "prefill_activation_length", "activation_embed") + else: + logical_axis_names = ("activation_batch", "activation_length", "activation_embed") + + if model_mode == MODEL_MODE_PREFILL: + inputs = _maybe_shard_with_logical(inputs, logical_axis_names) + else: + inputs = _maybe_shard_with_logical(inputs, logical_axis_names) + + inputs = checkpoint_name(inputs, "decoder_layer_input") + # inputs: embedded inputs to the decoder with shape [batch, length, emb_dim] + lnx = rms_norm( + num_features=inputs.shape[-1], + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + name="pre_self_attention_norm", + epsilon=cfg.normalization_layer_epsilon, + kernel_axes=("norm",), + )(inputs) + if model_mode == MODEL_MODE_PREFILL: + lnx = _maybe_shard_with_logical(lnx, logical_axis_names) + else: + lnx = _maybe_shard_with_logical(lnx, logical_axis_names) + + attention_layer = attention_as_linen( + config=self.config, + num_query_heads=cfg.num_query_heads, + num_kv_heads=cfg.num_kv_heads, + head_dim=cfg.head_dim, + max_target_length=cfg.max_target_length, + max_prefill_predict_length=cfg.max_prefill_predict_length, + attention_kernel=cfg.attention, + inputs_q_shape=lnx.shape, + inputs_kv_shape=lnx.shape, + mesh=mesh, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + dropout_rate=cfg.dropout_rate, + name="self_attention", + float32_qk_product=cfg.float32_qk_product, + float32_logits=cfg.float32_logits, + quant=self.quant, + kv_quant=quantizations.configure_kv_quant(cfg), + prefill_cache_axis_order=tuple(map(int, cfg.prefill_cache_axis_order.split(","))), + ar_cache_axis_order=tuple(map(int, cfg.ar_cache_axis_order.split(","))), + compute_axis_order=tuple(map(int, cfg.compute_axis_order.split(","))), + reshape_q=cfg.reshape_q, + use_mrope=cfg.use_mrope, + mrope_section=cfg.mrope_section, + share_kv_projections=cfg.share_kv_projections, + model_mode=model_mode, + ) + + attention_lnx, kv_cache = attention_layer( + lnx, + lnx, + decoder_positions, + decoder_segment_ids=decoder_segment_ids, + deterministic=deterministic, + model_mode=model_mode, + kv_cache=kv_cache, + attention_metadata=attention_metadata, + ) + + if model_mode == MODEL_MODE_PREFILL: + attention_lnx = _maybe_shard_with_logical(attention_lnx, logical_axis_names) + else: + attention_lnx = _maybe_shard_with_logical(attention_lnx, logical_axis_names) + + # MLP block. + mlp_lnx = linears.mlp_block( + in_features=lnx.shape[-1], + intermediate_dim=cfg.mlp_dim, + activations=cfg.mlp_activations, + intermediate_dropout_rate=cfg.dropout_rate, + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + name="mlp", + model_mode=model_mode, + config=cfg, + quant=self.quant, + mesh=self.mesh, + )(lnx, deterministic=deterministic) + if model_mode == MODEL_MODE_PREFILL: + mlp_lnx = _maybe_shard_with_logical(mlp_lnx, logical_axis_names) + else: + mlp_lnx = _maybe_shard_with_logical(mlp_lnx, logical_axis_names) + + next_layer_addition = mlp_lnx + attention_lnx + + next_layer_addition_dropped_out = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))( + next_layer_addition, deterministic=deterministic + ) + + layer_output = next_layer_addition_dropped_out + inputs + if model_mode == MODEL_MODE_PREFILL: + layer_output = _maybe_shard_with_logical( + layer_output, + logical_axis_names, + ) + else: + layer_output = _maybe_shard_with_logical( + layer_output, + logical_axis_names, + ) + + if getattr(cfg, "record_internal_nn_metrics", False): + self.sow("intermediates", "activation_mean", jnp.mean(layer_output)) + self.sow("intermediates", "activation_stdev", jnp.std(layer_output)) + self.sow( + "intermediates", + "activation_fraction_zero", + jnp.sum(layer_output == 0) / jnp.size(layer_output), + ) + + if cfg.scan_layers: + return layer_output, None + else: + return layer_output, kv_cache + + +class SequentialBlockDecoderLayers(nn.Module): + """Sequential unscanned series of decoder layers.""" + + decoder_layer: Any + num_decoder_layers: int + config: Config + mesh: Mesh + quant: Quant + model_mode: str + + @nn.compact + def __call__( + self, + inputs: jnp.ndarray, + decoder_segment_ids, + decoder_positions, + deterministic: bool, + model_mode, + slot: None | int = None, + ) -> jnp.ndarray: + for lyr in range(self.num_decoder_layers): + inputs = self.decoder_layer( + config=self.config, mesh=self.mesh, name=f"layers_{lyr}", quant=self.quant, model_mode=model_mode + )( + inputs, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + slot=slot, + ) + if self.config.scan_layers: + inputs = inputs[0] # When scan_layers is True the decoder layers return (outputs, None). + if self.config.scan_layers: + return inputs, None # pytype: disable=bad-return-type + else: + return inputs + + def deepstack_process(hidden_states, bidirectional_mask, visual_embeds): """Process deepstack visual embeddings by adding them to hidden states at visual token positions. @@ -44,3 +293,1680 @@ def deepstack_process(hidden_states, bidirectional_mask, visual_embeds): # Only add where mask is True: hidden_states += visual_embeds * mask hidden_states = hidden_states + visual_embeds_scattered * mask_expanded return hidden_states + + +class Decoder(nn.Module): + """A stack of decoder layers as a part of an encoder-decoder architecture.""" + + config: Config + mesh: Mesh + quant: None | Quant = None + model_mode: str = MODEL_MODE_TRAIN + + def setup(self): + """Initialize decoder layer.""" + self.decoder_layer = self.get_decoder_layers() + self.norm_layer = self.get_norm_layer(num_features=self.config.emb_dim) + if self.config.using_pipeline_parallelism: + remat_policy = self.get_remat_policy() + nnx_blocks = self._get_nnx_decoder_block_classes() + + # Per-stage builder handed to create_pipeline: the pipeline transform invokes it + # once per pipeline stage, passing that stage's rngs, to construct the stage's + # decoder block(s). nnx_blocks (the selected decoder block classes) is captured from + # this setup scope; remat_policy is applied separately by create_pipeline. + def build_pipeline_stage_layers(rngs): + """Builds one pipeline stage module from the selected NNX decoder block classes.""" + return self._build_nnx_pipeline_stage(nnx_blocks, rngs) + + self.pipeline_module = pipeline.create_pipeline( + config=self.config, layers=build_pipeline_stage_layers, mesh=self.mesh, remat_policy=remat_policy + ) + + def minimal_policy(self, with_context=False, with_quantization=False): + """Helper for creating minimal checkpoint policies.""" + names = [ + "query_proj", + "value_proj", + "key_proj", + "kv_proj", + "qkv_proj", + "out_proj", + "mlpwi_0", + "mlpwi_1", + "mlpwi", + "mlpwo", + ] + if with_context: + names.append("context") + if with_quantization: + names.append("quantization") + return jax.checkpoint_policies.save_only_these_names(*names) + + def get_remat_policy(self): + """Get remat policy""" + policy = None + cfg = self.config + 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": + max_logging.log("WARNING: 'minimal_flash' will be deprecated soon, please use 'minimal_with_context' instead.") + max_logging.log("WARNING: 'minimal_flash' will be deprecated soon, please use 'minimal_with_context' instead.") + policy = self.minimal_policy(with_context=True) + elif cfg.remat_policy == "minimal": + # save all except context + policy = self.minimal_policy() + elif cfg.remat_policy == "minimal_with_quantization": + if cfg.scan_layers: + warnings.warn( + "Scan layers can introduce overhead to checkpointed values that in some configurations is slower" + "than not checkpointing at all. If you are using scan layers, benchmark with and without quantization " + "checkpointing in your workflow to see which is faster. Without scan layers, checkpointing quantizations is " + "beneficial for performance." + ) + policy = self.minimal_policy(with_context=False, with_quantization=True) + elif cfg.remat_policy == "minimal_with_context_and_quantization": + if cfg.scan_layers: + warnings.warn( + "Scan layers can introduce overhead to checkpointed values that in some configurations is slower" + "than not checkpointing at all. If you are using scan layers, benchmark with and without quantization " + "checkpointing in your workflow to see which is faster. Without scan layers, checkpointing quantizations is " + "beneficial for performance." + ) + policy = self.minimal_policy(with_context=True, with_quantization=True) + elif cfg.remat_policy == "save_dot_with_context_except_mlp": + policy = jax.checkpoint_policies.save_only_these_names( + "query_proj", + "value_proj", + "key_proj", + "kv_proj", + "qkv_proj", + "context", + "out_proj", + ) + elif cfg.remat_policy == "save_dot_except_mlpwi": + policy = jax.checkpoint_policies.save_only_these_names( + "query_proj", + "value_proj", + "key_proj", + "kv_proj", + "qkv_proj", + "out_proj", + "mlpwo", + ) + elif cfg.remat_policy == "save_dot_except_mlp": + policy = jax.checkpoint_policies.save_only_these_names( + "query_proj", + "value_proj", + "key_proj", + "kv_proj", + "qkv_proj", + "out_proj", + ) + elif cfg.remat_policy == "save_qkv_proj": + policy = jax.checkpoint_policies.save_only_these_names( + "query_proj", + "value_proj", + "key_proj", + "kv_proj", + "qkv_proj", + ) + 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_save_and_offload_names) so that offloading configured via + # `custom` resolves identically to the named presets. + save_names, offload_names = maxtext_utils.get_save_and_offload_names(cfg) + policy = jax.checkpoint_policies.save_and_offload_only_these_names( + names_which_can_be_saved=save_names, + names_which_can_be_offloaded=offload_names, + offload_src="device", + offload_dst="pinned_host", + ) + elif cfg.remat_policy == "save_out_proj": + policy = jax.checkpoint_policies.save_only_these_names( + "out_proj", + ) + else: + assert cfg.remat_policy == "full", "Remat policy needs to be on list of remat policies" + policy = None + return policy + + def get_decoder_layers(self): + """Retrieves a list of decoder layer classes based on the `decoder_block` config. + + Returns: + A list containing one or more `nn.Module` classes for the decoder. + """ + match self.config.decoder_block: + case DecoderBlockType.DEFAULT: + return [DecoderLayer] + case DecoderBlockType.LLAMA2: + return [llama2.LlamaDecoderLayerToLinen] + case DecoderBlockType.MISTRAL: + # TODO(ranran): update to Mistral with sliding window attention + return [mistral.MistralDecoderLayerToLinen] + case DecoderBlockType.MIXTRAL: + return [mixtral.MixtralDecoderLayerToLinen] + case DecoderBlockType.DEEPSEEK: + return [ + deepseek.DeepSeekDenseLayerToLinen, + deepseek.DeepSeekMoELayerToLinen, + ] + case DecoderBlockType.DEEPSEEK4: + return ( + [deepseek4.DeepSeek4ScannableBlockToLinen] if self.config.scan_layers else [deepseek4.DeepSeek4LayerToLinen] + ) + case DecoderBlockType.GEMMA: + return [gemma.GemmaDecoderLayerToLinen] + case DecoderBlockType.GEMMA2: + return [gemma2.Gemma2DecoderLayerToLinen] + case DecoderBlockType.GEMMA3: + return [gemma3.Gemma3DecoderLayerToLinen] + case DecoderBlockType.GEMMA4: + return [gemma4.Gemma4ScannableBlockToLinen] if self.config.scan_layers else [gemma4.Gemma4DecoderLayerToLinen] + case DecoderBlockType.GEMMA4_SMALL: + # PLE input + KV-share donor threading requires per-layer-index state, + # which is not expressible inside ``nn.scan``. + return [gemma4_small.Gemma4SmallDecoderLayerToLinen] + case DecoderBlockType.GPT3: + return [gpt3.Gpt3DecoderLayerToLinen] + case DecoderBlockType.GPT_OSS: + return [gpt_oss.GptOssScannableBlockToLinen] if self.config.scan_layers else [gpt_oss.GptOssDecoderLayerToLinen] + case DecoderBlockType.QWEN2: + return [qwen2.Qwen2DecoderLayerToLinen] + case DecoderBlockType.QWEN3: + return [qwen3.Qwen3DecoderLayerToLinen] + case DecoderBlockType.QWEN3_MOE: + return [qwen3.Qwen3MoeDecoderLayerToLinen] + case DecoderBlockType.QWEN3_CUSTOM_MOE: + return [qwen3_custom.Qwen3CustomMoeDecoderLayerToLinen] + case DecoderBlockType.QWEN3_NEXT: + return [qwen3.Qwen3NextScannableBlockToLinen] if self.config.scan_layers else [qwen3.Qwen3NextDecoderLayerToLinen] + case DecoderBlockType.QWEN3_5: + return [qwen3_5.Qwen3_5ScannableBlockToLinen] if self.config.scan_layers else [qwen3_5.Qwen3_5DecoderLayerToLinen] + case DecoderBlockType.SIMPLE: + return [simple_layer.SimpleDecoderLayerToLinen] + case DecoderBlockType.SIMPLE_MLP: + return [simple_layer.SimpleMlpDecoderLayerToLinen] + case DecoderBlockType.LLAMA4: + return [llama4.Llama4ScannableBlockToLinen] if self.config.scan_layers else [llama4.Llama4DecoderLayerToLinen] + case DecoderBlockType.OLMO3: + return [olmo3.Olmo3ScannableBlockToLinen] if self.config.scan_layers else [olmo3.Olmo3DecoderLayerToLinen] + case DecoderBlockType.ENVY: + return [envy.EnvyScannableBlockToLinen] if self.config.scan_layers else [envy.EnvyDecoderLayerToLinen] + + case _: + # Default case to handle any unknown decoder block types. + raise ValueError(f"Incorrect decoder_block name {self.config.decoder_block.value=}") + + def _get_nnx_decoder_block_classes(self): + """Returns NNX decoder block classes for pipeline stage creation.""" + cfg = self.config + + def get_scannable(normal_cls, scannable_cls): + return [scannable_cls] if cfg.scan_layers else [normal_cls] + + layer_map = { + DecoderBlockType.DEFAULT: [NNXDecoderLayer], + DecoderBlockType.LLAMA2: [llama2.LlamaDecoderLayer], + DecoderBlockType.MISTRAL: [mistral.MistralDecoderLayer], + DecoderBlockType.MIXTRAL: [mixtral.MixtralDecoderLayer], + DecoderBlockType.GEMMA: [gemma.GemmaDecoderLayer], + DecoderBlockType.GEMMA2: [gemma2.Gemma2DecoderLayer], + DecoderBlockType.GEMMA3: [gemma3.Gemma3DecoderLayer], + DecoderBlockType.GEMMA4: get_scannable(gemma4.Gemma4DecoderLayer, gemma4.Gemma4ScannableBlock), + DecoderBlockType.GEMMA4_SMALL: [gemma4_small.Gemma4SmallDecoderLayer], + DecoderBlockType.GPT3: [gpt3.Gpt3DecoderLayer], + DecoderBlockType.GPT_OSS: get_scannable(gpt_oss.GptOssDecoderLayer, gpt_oss.GptOssScannableBlock), + DecoderBlockType.QWEN2: [qwen2.Qwen2DecoderLayer], + DecoderBlockType.QWEN3: [qwen3.Qwen3DecoderLayer], + DecoderBlockType.QWEN3_MOE: [qwen3.Qwen3MoeDecoderLayer], + DecoderBlockType.QWEN3_CUSTOM_MOE: [qwen3_custom.Qwen3CustomMoeDecoderLayer], + DecoderBlockType.QWEN3_5: get_scannable(qwen3_5.Qwen3_5DecoderLayer, qwen3_5.Qwen3_5ScannableBlock), + DecoderBlockType.QWEN3_NEXT: get_scannable(qwen3.Qwen3NextDecoderLayer, qwen3.Qwen3NextScannableBlock), + DecoderBlockType.SIMPLE: [simple_layer.SimpleDecoderLayer], + DecoderBlockType.SIMPLE_MLP: [simple_layer.SimpleMlpDecoderLayer], + DecoderBlockType.DEEPSEEK: [deepseek.DeepSeekDenseLayer, deepseek.DeepSeekMoELayer], + DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock), + DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock), + DecoderBlockType.ENVY: get_scannable(envy.EnvyDecoderLayer, envy.EnvyScannableBlock), + } + + if cfg.decoder_block not in layer_map: + raise ValueError(f"Incorrect decoder_block name {cfg.decoder_block.value=}") + return layer_map[cfg.decoder_block] + + def set_remat_policy(self, block_layers, policy): + """Set remat policy""" + RemattedBlockLayers = [] + for block_layer in block_layers: + if self.config.parameter_memory_host_offload: + # Define parameter movement with mesh-based sharding + def move_to_device(variables): + """Move parameters to device with proper sharding.""" + + def map_fn(path, value): + max_logging.log(f"models.py: Moving parameter {path} to device") + return jax.device_put(value, max_utils.device_space()) + + return jax.tree_util.tree_map_with_path(map_fn, variables) + + # Transform layer class before remat + block_layer = nn.map_variables(block_layer, ["params"], move_to_device, mutable=True) + + # Apply remat policy to layer + layer = nn.remat( + block_layer, + prevent_cse=maxtext_utils.should_prevent_cse_in_remat(self.config), + policy=policy, + static_argnums=(4, 5), # Deterministic and model mode are static arguments. + ) + RemattedBlockLayers.append(layer) + return RemattedBlockLayers + + def _build_nnx_pipeline_stage(self, decoder_blocks, rngs): + """Creates a single NNX pipeline stage module.""" + cfg = self.config + base_stage_cls = decoder_blocks[1] if cfg.decoder_block == DecoderBlockType.DEEPSEEK else decoder_blocks[0] + + if cfg.num_layers_per_pipeline_stage == 1: + return base_stage_cls(config=cfg, mesh=self.mesh, quant=self.quant, model_mode=self.model_mode, rngs=rngs) + elif cfg.scan_layers_per_stage: + return NNXScannedPipelineStage( + base_stage_cls, cfg.num_layers_per_pipeline_stage, cfg, self.mesh, self.quant, self.model_mode, rngs=rngs + ) + return NNXSequentialPipelineStage( + base_stage_cls, cfg.num_layers_per_pipeline_stage, cfg, self.mesh, self.quant, self.model_mode, rngs=rngs + ) + + def get_pipeline_stage_module(self, decoder_blocks): + """get pipeline stage module""" + + def get_layer_to_pipeline(blocks, cfg): + if cfg.decoder_block == DecoderBlockType.DEEPSEEK: + return blocks[1] # return the sparse block + else: + return blocks[0] + + cfg = self.config + base_stage = get_layer_to_pipeline(decoder_blocks, cfg) + if cfg.set_remat_policy_on_layers_per_stage: + policy = self.get_remat_policy() + base_stage = self.set_remat_policy([base_stage], policy)[0] + if cfg.num_layers_per_pipeline_stage == 1: + stage_module = base_stage(config=cfg, mesh=self.mesh, quant=self.quant, model_mode=self.model_mode) + elif cfg.scan_layers_per_stage: + stage_module = self.scan_decoder_layers( + cfg, + base_stage, + cfg.num_layers_per_pipeline_stage, + "layers_per_stage", + self.mesh, + in_axes_tuple=(nn.broadcast,) * 4, + model_mode=self.model_mode, + ) + else: + stage_module = SequentialBlockDecoderLayers( + decoder_layer=base_stage, + num_decoder_layers=cfg.num_layers_per_pipeline_stage, + config=cfg, + mesh=self.mesh, + quant=self.quant, + model_mode=self.model_mode, + ) + return stage_module + + def get_norm_layer(self, num_features: int): + """get normalization layer (return type inherits from nn.Module)""" + if self.config.decoder_block in ( + DecoderBlockType.DEFAULT, + DecoderBlockType.LLAMA2, + DecoderBlockType.MISTRAL, + DecoderBlockType.MIXTRAL, + DecoderBlockType.DEEPSEEK, + DecoderBlockType.DEEPSEEK4, + DecoderBlockType.GEMMA, + DecoderBlockType.GEMMA2, + DecoderBlockType.GEMMA3, + DecoderBlockType.GEMMA4, + DecoderBlockType.GEMMA4_SMALL, + DecoderBlockType.QWEN2, + DecoderBlockType.QWEN3, + DecoderBlockType.QWEN3_MOE, + DecoderBlockType.QWEN3_CUSTOM_MOE, + DecoderBlockType.GPT_OSS, + DecoderBlockType.SIMPLE, + DecoderBlockType.SIMPLE_MLP, + DecoderBlockType.LLAMA4, + DecoderBlockType.OLMO3, + DecoderBlockType.ENVY, + ): + return functools.partial(rms_norm, num_features=num_features, shard_mode=self.config.shard_mode) + elif self.config.decoder_block == DecoderBlockType.GPT3: + return functools.partial(gpt3.gpt3_layer_norm, num_features=num_features, reductions_in_fp32=False, use_bias=True) + elif self.config.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5): + return functools.partial( + normalizations.Qwen3NextRMSNormLinen, num_features=num_features, shard_mode=self.config.shard_mode + ) + else: + raise ValueError(f"Incorrect decoder_block name {self.config.decoder_block.value=}") + + def scan_decoder_layers(self, cfg, decoder_layer, length, metadata_axis_name, mesh, in_axes_tuple, **kwargs): + """scan decoder layers, calls `flax.linen.transforms.scan`""" + initializing = self.is_mutable_collection("params") + params_spec = cfg.param_scan_axis if initializing else ScanIn(cfg.param_scan_axis) + cache_spec = 0 + scan_fn = nn.scan( + decoder_layer, + variable_axes={ + "params": params_spec, + "cache": cache_spec, + "intermediates": 0, + "aqt": 0, + "batch_stats": 0, + "_overwrite_with_gradient": 0, + "indexer_losses": 0, + }, + split_rngs={ + "params": True, + "dropout": cfg.enable_dropout, + }, + in_axes=in_axes_tuple, + length=length, + metadata_params={nn.PARTITION_NAME: metadata_axis_name}, + ) + return scan_fn( + config=cfg, mesh=mesh, name=metadata_axis_name, quant=self.quant, **kwargs # pytype: disable=wrong-keyword-args + ) + + @nn.compact + def _apply_embedding( + self, + shared_embedding: nn.Module | nnx.Module, + decoder_input_tokens, + decoder_positions, + deterministic, + model_mode, + multimodal_input=None, + decoder_input_embeddings=None, + ): + """Applies token and positional embeddings to the input tokens.""" + + cfg = self.config + + # vLLM passes token IDs for text-only requests, but passes complete, premerged + # text and multimodal embeddings for multimodal requests. The latter enter + # through `decoder_input_embeddings` to avoid embedding and merging them again. + y = ( + decoder_input_embeddings + if decoder_input_embeddings is not None + else shared_embedding(decoder_input_tokens.astype("int32"), model_mode=model_mode) + ) + + # Precomputed embeddings are complete (including any multimodal replacements), + # so only merge modality embeddings when token embeddings were created here. + if decoder_input_embeddings is None and multimodal_input is not None: + image_embeddings = multimodal_input.image_embeddings + bidirectional_mask = multimodal_input.bidirectional_mask + image_masks = multimodal_input.image_masks + video_embeddings = getattr(multimodal_input, "video_embeddings", None) + video_masks = getattr(multimodal_input, "video_masks", None) + bidirectional_mask_video = getattr(multimodal_input, "bidirectional_mask_video", None) + audio_embeddings = multimodal_input.audio_embeddings + audio_masks = multimodal_input.audio_masks + + if image_embeddings is not None and cfg.use_multimodal: + if cfg.model_name in [ + "gemma3-4b", + "gemma3-12b", + "gemma3-27b", + "gemma4-26b", + "gemma4-31b", + "gemma4-e2b", + "gemma4-e4b", + "llama4-17b-16e", + "llama4-17b-128e", + "qwen3-omni-30b-a3b", + "qwen3-vl-2b", + "qwen3-vl-4b", + "qwen3-vl-30b-a3b", + "qwen3.5-35b-a3b", + "qwen3.5-397b-a17b", + "maxtext-omni-gemma3-qwen3", + ]: + y = mm_utils.merge_mm_embeddings( + text_embeddings=y, + multimodal_embeddings=image_embeddings, + mask=bidirectional_mask, + token_masks=image_masks, + ) + # TODO(hengtaoguo): Add support for other multimodal models such as Llama4, refactor if needed + else: + raise ValueError(f"Unsupported model_name for multimodal: {cfg.model_name}") + + if video_embeddings is not None and cfg.use_multimodal: + if cfg.model_name in [ + "qwen3-omni-30b-a3b", + "qwen3-vl-2b", + "qwen3-vl-4b", + "qwen3-vl-30b-a3b", + "qwen3.5-35b-a3b", + "qwen3.5-397b-a17b", + ]: + y = mm_utils.merge_mm_embeddings( + text_embeddings=y, + multimodal_embeddings=video_embeddings, + mask=bidirectional_mask_video, + token_masks=video_masks, + ) + else: + raise ValueError(f"Unsupported model_name for video: {cfg.model_name}") + + if audio_embeddings is not None and cfg.use_audio: + if cfg.model_name in ["qwen3-omni-30b-a3b"]: + y = mm_utils.merge_mm_embeddings( + text_embeddings=y, + multimodal_embeddings=audio_embeddings, + mask=audio_masks, + token_masks=None, + ) + else: + raise ValueError(f"Unsupported model_name for audio: {cfg.model_name}") + + y = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))(y, deterministic=deterministic) + y = y.astype(cfg.dtype) + + if cfg.use_untrainable_positional_embedding: + y += positional_embedding_as_linen(embedding_dims=cfg.base_emb_dim)(y.shape[1], decoder_positions) + + if cfg.trainable_position_size > 0: + y += embed_as_linen( + num_embeddings=cfg.trainable_position_size, + num_features=cfg.emb_dim, + dtype=cfg.dtype, + embedding_init=nn.initializers.normal(stddev=1.0), + name="position_embedder", + config=cfg, + mesh=self.mesh, + )(decoder_positions.astype("int32"), model_mode=model_mode) + return y + + @nn.compact + def apply_output_head(self, shared_embedding: nn.Module | nnx.Module, y, deterministic, model_mode): + """Applies final normalization and projects hidden states to logits.""" + + cfg = self.config + if cfg.shard_mode == ShardMode.EXPLICIT: + norm_out_sharding = create_sharding(self.mesh, ("activation_batch", "activation_length", "activation_embed")) + else: + norm_out_sharding = None + + y = self.get_norm_layer(num_features=y.shape[-1])( + dtype=cfg.dtype, + weight_dtype=cfg.weight_dtype, + name="decoder_norm", + epsilon=cfg.normalization_layer_epsilon, + kernel_axes=("norm",), + parameter_memory_host_offload=cfg.parameter_memory_host_offload, + )(y, out_sharding=norm_out_sharding) + y = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))(y, deterministic=deterministic) + + if model_mode in (MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE): + out_sharding = create_sharding(self.mesh, (None, None, "activation_vocab")) + else: + out_sharding = create_sharding( + self.mesh, ("activation_embed_and_logits_batch", "activation_length", "activation_vocab") + ) + + # [batch, length, emb_dim] -> [batch, length, vocab_size] + if cfg.logits_via_embedding: + # Use the transpose of embedding matrix for logit transform. + if isinstance(shared_embedding, nnx.Module): + embedding_table = shared_embedding.embedding.value + else: + embedding_table = shared_embedding.variables["params"]["embedding"] + if isinstance(embedding_table, nn.spmd.LogicallyPartitioned): + embedding_table = embedding_table.unbox() + attend_dtype = jnp.float32 if cfg.logits_dot_in_fp32 else cfg.dtype + logits = attend_on_embedding(y, embedding_table, attend_dtype, self.config, out_sharding) + + if self.config.normalize_embedding_logits: + # Correctly normalize pre-softmax logits for this shared case. + logits = logits / jnp.sqrt(y.shape[-1]) + if cfg.final_logits_soft_cap: + logits = logits / cfg.final_logits_soft_cap + logits = jnp.tanh(logits) * cfg.final_logits_soft_cap + else: + logits = linears.dense_general( + inputs_shape=y.shape, + out_features_shape=cfg.vocab_size, + weight_dtype=cfg.weight_dtype, + dtype=jnp.float32 if cfg.logits_dot_in_fp32 else cfg.dtype, # for logit training stability + kernel_axes=("embed_vocab", "vocab"), + shard_mode=cfg.shard_mode, + name="logits_dense", + matmul_precision=self.config.matmul_precision, + parameter_memory_host_offload=cfg.parameter_memory_host_offload, + )( + y, + out_sharding=out_sharding, + ) # We do not quantize the logits matmul. + + if self.config.cast_logits_to_fp32: + logits = logits.astype(jnp.float32) + + return logits + + @nn.compact + def __call__( + self, + shared_embedding: nn.Module | nnx.Module, + decoder_input_tokens, + decoder_positions, + decoder_segment_ids=None, + deterministic=False, + model_mode=MODEL_MODE_TRAIN, + previous_chunk=None, + slot: None | int = None, + multimodal_input=None, + kv_caches: list[jax.Array] | None = None, + attention_metadata=None, + deepstack_visual_embeds: None | list[jnp.ndarray] = None, + decoder_input_embeddings=None, + ): + cfg = self.config + mesh = self.mesh + assert decoder_input_tokens.ndim == 2 # [batch, len] + + # [batch, length] -> [batch, length, emb_dim] + y = self._apply_embedding( + shared_embedding, + decoder_input_tokens, + decoder_positions, + deterministic, + model_mode, + multimodal_input=multimodal_input, + decoder_input_embeddings=decoder_input_embeddings, + ) + + mhc_expand, mhc_reduce = mhc.get_functions(cfg.mhc_expansion_rate) + if cfg.mhc_expansion_rate > 1: + # (batch, length, emb_dim) --> (batch, length, mhc_expansion_rate, emb_dim) + y = mhc_expand(y) + + policy = self.get_remat_policy() + RemattedBlockLayers = self.set_remat_policy(self.decoder_layer, policy) + # scan does not support kwargs in layer call, passing broadcast_args as positional arg + broadcast_args = ( + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + ) + if cfg.using_pipeline_parallelism: + logical_partition_spec = ( + self.pipeline_module.get_weight_sharding(y, decoder_segment_ids, decoder_positions, deterministic, model_mode) + if cfg.pipeline_fsdp_ag_once or cfg.pipeline_fsdp_ag_per_repeat + else None + ) + if cfg.decoder_block == DecoderBlockType.DEEPSEEK: + assert len(RemattedBlockLayers) == 2, "Scanned layers must have a length of 2 using deepseek." + dense_layer = RemattedBlockLayers[0] + moe_layer = RemattedBlockLayers[1] + num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers + num_moe_layers_outside_pp = num_moe_layers - self.config.pipeline_parallel_layers + logical_axis_rules_pp_as_dp = sharding.logical_axis_rules_pp_act_as_dp(self.config.logical_axis_rules) + # We chose not to pipeline the dense layers, only sparse for SPMD. + with self.mesh, nn.partitioning.axis_rules(logical_axis_rules_pp_as_dp): + 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) + if num_moe_layers_outside_pp > 0: + y, _ = self.scan_decoder_layers( + cfg, + moe_layer, + num_moe_layers_outside_pp, + "moe_layers", + mesh, + in_axes_tuple=(nn.broadcast,) * len(broadcast_args), + model_mode=model_mode, + )(y, *broadcast_args) + y = self.pipeline_module(y, *broadcast_args, logical_partition_spec=logical_partition_spec) + else: # Not DeepSeek + y = self.pipeline_module(y, *broadcast_args, logical_partition_spec=logical_partition_spec) + remaining_layers = self.config.num_decoder_layers - self.config.pipeline_parallel_layers + if remaining_layers > 0: + logical_axis_rules_pp_as_dp = sharding.logical_axis_rules_pp_act_as_dp(self.config.logical_axis_rules) + with self.mesh, nn.partitioning.axis_rules(logical_axis_rules_pp_as_dp): + y, _ = self.scan_decoder_layers( + cfg, + RemattedBlockLayers[0], + remaining_layers, + "layers_outside_pipeline", + mesh, + in_axes_tuple=(nn.broadcast,) * len(broadcast_args), + model_mode=model_mode, + )(y, *broadcast_args) + else: + if cfg.scan_layers: + if cfg.decoder_block == DecoderBlockType.DEEPSEEK: + assert len(RemattedBlockLayers) == 2, "Scanned layers must have a length of 2 using deepseek." + layer_call_kwargs = { + "previous_chunk": previous_chunk, + "slot": slot, + } + dense_layer = RemattedBlockLayers[0] + moe_layer = RemattedBlockLayers[1] + 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, + } + + # Apply Dense Layers + y = self._apply_interleaved_scanned_layers( + y, + 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, + 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"): + # old version of batch-split that fully uses qwix quantization. + if cfg.quantization and cfg.use_qwix_quantization and not cfg.use_manual_quantization: + y = deepseek_batchsplit_fp8.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: + # bf16 and fp8 code path for pure-JAX batch-split. + # fp8 code path supports both manual quantization and qwix + # quantization. + y = deepseek_batchsplit.scan_batch_split_layers( + y, + self.variables["params"]["moe_layers"], + decoder_positions, + mesh=mesh, + cfg=cfg, + num_layers=num_moe_layers, + ) + 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: + bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None + y = self._apply_gemma3_scanned_blocks( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + bidirectional_mask_value, + previous_chunk, + slot, + kv_caches=kv_caches, + attention_metadata=attention_metadata, + ) + elif cfg.decoder_block == DecoderBlockType.GEMMA4: + bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None + y = self._apply_gemma4_scanned_blocks( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + bidirectional_mask_value, + previous_chunk, + slot, + kv_caches=kv_caches, + attention_metadata=attention_metadata, + ) + elif cfg.decoder_block == DecoderBlockType.QWEN3_NEXT: + y = self._apply_qwen3_next_scanned_blocks( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + kv_caches=kv_caches, + attention_metadata=attention_metadata, + ) + elif cfg.decoder_block == DecoderBlockType.DEEPSEEK4: + y = self._apply_deepseek4_scanned_blocks( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + decoder_input_tokens, + ) + else: + RemattedBlockLayer = RemattedBlockLayers[0] + scan_length = int(cfg.num_decoder_layers / cfg.inhomogeneous_layer_cycle_interval) + layer_kwargs = {} + if cfg.decoder_block == DecoderBlockType.LLAMA4: + layer_kwargs = { + "nope_layer_interval": self.config.nope_layer_interval, + "interleave_moe_layer_step": self.config.interleave_moe_layer_step, + } + if cfg.decoder_block == DecoderBlockType.ENVY: + layer_kwargs = { + "interleave_moe_layer_step": self.config.interleave_moe_layer_step, + } + + # Update broadcast_args and in_axes_tuple for vLLM RPA + in_axes_tuple = (nn.broadcast,) * len(broadcast_args) + current_broadcast_args = list(broadcast_args) + current_in_axes_tuple = list(in_axes_tuple) + + if kv_caches is not None: + # Stack kv_caches for scan: [num_layers, ...] + stacked_kv_cache = jnp.stack(kv_caches, axis=0) + + # We pass (y, stacked_kv_cache, 0) as the carry + carry = (y, stacked_kv_cache, 0) + + # We don't pass kv_cache as a scanned argument anymore + + # Pass None for previous_chunk, slot, kv_cache to align with __call__ signature + current_broadcast_args.extend([None, None, None, attention_metadata]) + current_in_axes_tuple.extend([nn.broadcast] * 4) + + max_logging.info(f"DEBUG: len(current_broadcast_args)={len(current_broadcast_args)}") + max_logging.info(f"DEBUG: current_broadcast_args={[type(a) for a in current_broadcast_args]}") + + final_carry, _ = self.scan_decoder_layers( + cfg, + RemattedBlockLayer, + scan_length, + "layers", + mesh, + in_axes_tuple=tuple(current_in_axes_tuple), + model_mode=model_mode, + **layer_kwargs, + )(carry, *current_broadcast_args) + + y, returned_kv_cache, _ = final_carry + + # Update the list of KV caches from the scanned results + for i in range(cfg.num_decoder_layers): + kv_caches[i] = returned_kv_cache[i] + else: + # Fallback to old behavior if kv_caches is None (not vLLM RPA) + current_broadcast_args.append(None) + current_in_axes_tuple.append(nn.broadcast) + + y, _ = self.scan_decoder_layers( + cfg, + RemattedBlockLayer, + scan_length, + "layers", + mesh, + in_axes_tuple=tuple(current_in_axes_tuple), + model_mode=model_mode, + **layer_kwargs, + )(y, *current_broadcast_args) + else: + if cfg.decoder_block == DecoderBlockType.DEEPSEEK: + assert len(RemattedBlockLayers) == 2, "Unscanned layers must have a length of 2 using deepseek." + dense_layer = RemattedBlockLayers[0] + moe_layer = RemattedBlockLayers[1] + + layers = [dense_layer, moe_layer] + layer_prefixes = ["dense_layers", "moe_layers"] + num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers + num_layers_list = [cfg.first_num_dense_layers, num_moe_layers] + # Iterate over the two layer groups (dense and MoE) and apply layer transformation + global_layer_idx_offset = 0 + for layer, num_layers, layer_prefix in zip(layers, num_layers_list, layer_prefixes): + for index in range(num_layers): + global_layer_idx = global_layer_idx_offset + index + kv_cache = kv_caches[index] if kv_caches is not None else None + input_tokens = decoder_input_tokens if cfg.engram_layers else None + y, kv_cache = layer( + config=cfg, + mesh=mesh, + name=f"{layer_prefix}_{index}", + quant=self.quant, + model_mode=self.model_mode, + layer_idx=global_layer_idx, + )( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=previous_chunk, + slot=slot, + kv_cache=kv_cache, + attention_metadata=attention_metadata, + decoder_input_tokens=input_tokens, + ) + if kv_caches is not None and kv_cache is not None: + kv_caches[index] = kv_cache + global_layer_idx_offset += num_layers + elif cfg.decoder_block == DecoderBlockType.GEMMA4_SMALL: + y, kv_caches = self._apply_gemma4_small_layers( + y, + decoder_input_tokens, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + multimodal_input=multimodal_input, + kv_caches=kv_caches, + attention_metadata=attention_metadata, + previous_chunk=previous_chunk, + slot=slot, + ) + else: + for lyr in range(cfg.num_decoder_layers): + RemattedBlockLayer = RemattedBlockLayers[0] + layer_kwargs = {} + layer_call_kwargs = {} + if cfg.decoder_block == DecoderBlockType.GEMMA3: + # Gemma3 uses both global and sliding window attention depending on the layer index. + bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None + layer_kwargs = {"attention_type": gemma3.get_attention_type(layer_id=lyr)} + layer_call_kwargs = {"bidirectional_mask": bidirectional_mask_value} + if cfg.decoder_block == DecoderBlockType.GEMMA4: + # Gemma4 uses both global and sliding window attention depending on the layer index. + bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None + layer_kwargs = {"attention_type": gemma4.get_attention_type(layer_id=lyr)} + layer_call_kwargs = {"bidirectional_mask": bidirectional_mask_value} + if cfg.decoder_block == DecoderBlockType.LLAMA4: + layer_kwargs = { + "is_nope_layer": llama4.determine_is_nope_layer(lyr, self.config.nope_layer_interval), + "is_moe_layer": llama4.determine_is_moe_layer(lyr, self.config.interleave_moe_layer_step), + } + if cfg.decoder_block == DecoderBlockType.ENVY: + layer_kwargs = { + "is_moe_layer": (lyr + 1) % self.config.interleave_moe_layer_step == 0, + } + + if cfg.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, DecoderBlockType.DEEPSEEK4): + layer_kwargs = {"layer_idx": lyr} + if cfg.decoder_block == DecoderBlockType.DEEPSEEK4: + layer_call_kwargs["decoder_input_tokens"] = decoder_input_tokens + kv_cache = None + if kv_caches is not None: + # For all decoder blocks (including QWEN3_NEXT/QWEN3_5 with vLLM flat-list + # kv_caches), pass the per-layer cache directly. For hybrid attention+GDN + # models, kv_caches[lyr] is a regular attention cache for attention layers + # and a (conv_state, recurrent_state) paged-mamba tuple for GDN layers. + kv_cache = kv_caches[lyr] + + if cfg.decoder_block == DecoderBlockType.GPT_OSS: + layer_kwargs = {"attention_type": gpt_oss.get_attention_type(layer_id=lyr)} + if cfg.decoder_block == DecoderBlockType.OLMO3: + layer_kwargs = {"attention_type": olmo3.get_attention_type(layer_id=lyr)} + layer = RemattedBlockLayer( + config=cfg, mesh=mesh, name=f"layers_{lyr}", quant=self.quant, model_mode=self.model_mode, **layer_kwargs + ) + y, returned_cache = layer( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=previous_chunk, + slot=slot, + kv_cache=kv_cache, + attention_metadata=attention_metadata, + **layer_call_kwargs, + ) + if kv_caches is not None and returned_cache is not None: + kv_caches[lyr] = returned_cache + + if deepstack_visual_embeds is not None and lyr < len(deepstack_visual_embeds): + visual_embeds = deepstack_visual_embeds[lyr] + # Use bidirectional_mask to identify visual token positions + bidirectional_mask_value = None + if multimodal_input is not None: + bidirectional_mask_value = ( + multimodal_input.bidirectional_mask + if multimodal_input.bidirectional_mask is not None + else multimodal_input.bidirectional_mask_video + ) + if bidirectional_mask_value is not None and visual_embeds is not None: + y = deepstack_process(y, bidirectional_mask_value, visual_embeds) + + assert isinstance(y, jax.Array) + + # After the final transformer layer, `y` holds the raw, un-normalized hidden state. + if cfg.mhc_expansion_rate > 1: + if cfg.decoder_block == DecoderBlockType.DEEPSEEK4: + hidden_state = mhc.DeepSeek4HyperHeadToLinen( + config=cfg, + mesh=mesh, + name="hc_head", + )(y) + else: + # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim) + hidden_state = mhc_reduce(y) + else: + hidden_state = y + + # When initializing with vLLM RPA attention, we need to run the output head to + # initialize any parameters associated with it. + # Same case applicable to vocab tiling + if self.is_initializing() and (cfg.num_vocab_tiling > 1 or cfg.attention in ("vllm_rpa", "vllm_batched_rpa")): + _ = self.apply_output_head(shared_embedding, hidden_state, deterministic, model_mode) + + # When invoking from vLLM with RPA attention, logit computation is deferred to a later stage. + if cfg.attention in ("vllm_rpa", "vllm_batched_rpa"): + logits = None + # When in the Indexer Dense Warm-up stage, skip the expensive output head projection + # for efficiency, as the main model is frozen and the LM loss is not needed. + # TODO(b/501446870): Investigate model_mode as train at beginning for decoding stage + elif ( + cfg.use_indexer and cfg.indexer_loss_scaling_factor > 0.0 and not cfg.indexer_sparse_training + ) and model_mode == MODEL_MODE_TRAIN: + logits = None + # When vocab tiling is enabled in training mode, full logits won't generate to reduce memory + # Instead, we keep track on the hidden states, which has smaller size compared to full logits + elif cfg.num_vocab_tiling > 1 and model_mode == MODEL_MODE_TRAIN: + logits = None + self.sow("intermediates", "hidden_states", hidden_state) + + else: + logits = self.apply_output_head(shared_embedding, hidden_state, deterministic, model_mode) + + # The API of the Decoder is now a tuple, providing both the main output + # and the raw hidden state needed for auxiliary tasks. + return logits, hidden_state, kv_caches + + def _apply_gemma3_scanned_blocks( + self, + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + bidirectional_mask, + previous_chunk, + slot, + kv_caches=None, + attention_metadata=None, + ): + """Applies Gemma3 scanned decoder blocks, handling main scan and remainders.""" + + cfg = self.config + mesh = self.mesh + + # Define the repeating pattern length and calculate how many full blocks to scan + attention_pattern_length = len(gemma3.GEMMA3_ATTENTION_PATTERN) + scan_length = cfg.num_decoder_layers // attention_pattern_length + + policy = self.get_remat_policy() + RemattedGemma3Block = self.set_remat_policy([gemma3.Gemma3ScannableBlockToLinen], policy)[0] + + layer_kwargs = {"num_of_layers": attention_pattern_length} + + # Apply the main scan over the full blocks + if scan_length > 0: + kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan( + kv_caches, scan_length, attention_pattern_length, stack=True + ) + + broadcast_args_spec = [ + (decoder_segment_ids, nn.broadcast), + (decoder_positions, nn.broadcast), + (deterministic, nn.broadcast), + (model_mode, nn.broadcast), + (slot, nn.broadcast), + (None, nn.broadcast), # page_state + (previous_chunk, nn.broadcast), + (bidirectional_mask, nn.broadcast), + (kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast), + (attention_metadata, nn.broadcast), + ] + broadcast_args = tuple(arg for arg, _ in broadcast_args_spec) + in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec) + + y, returned_kv_cache = self.scan_decoder_layers( + cfg, + RemattedGemma3Block, + scan_length, + "layers", + mesh, + in_axes_tuple=in_axes_tuple, + model_mode=self.model_mode, + **layer_kwargs, + )(y, *broadcast_args) + + maxtext_utils.update_kv_caches_after_scan( + kv_caches, returned_kv_cache, scan_length, attention_pattern_length, stacked=True + ) + + # Apply any remaining layers that did not fit into a full scanned block + num_remaining_layers = cfg.num_decoder_layers % attention_pattern_length + if num_remaining_layers > 0: + # We name the remainder block with a 'remainder' suffix to avoid parameter name collisions + rem_layer_kwargs = {"num_of_layers": num_remaining_layers} + layer = RemattedGemma3Block( + config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode, name="layers_remainder", **rem_layer_kwargs + ) # pytype: disable=wrong-keyword-args + + remainder_kv = None + if kv_caches is not None: + start_idx = scan_length * attention_pattern_length + remainder_kv = tuple(kv_caches[start_idx : start_idx + num_remaining_layers]) + + remainder_args = ( + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + slot, + None, # page_state + previous_chunk, + bidirectional_mask, + remainder_kv, + attention_metadata, + ) + y_and_kv = layer(y, *remainder_args) + + if isinstance(y_and_kv, tuple): + y = y_and_kv[0] + updated_remainder_kv = y_and_kv[1] + else: + y = y_and_kv + updated_remainder_kv = None + + if kv_caches is not None and updated_remainder_kv is not None: + start_idx = scan_length * attention_pattern_length + for offset, updated_item in enumerate(updated_remainder_kv): + kv_caches[start_idx + offset] = updated_item + + return y + + def _apply_qwen3_next_scanned_blocks( + self, + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + kv_caches=None, + attention_metadata=None, + ): + """Applies Qwen3-Next scanned decoder blocks, handling main scan and remainders.""" + + cfg = self.config + mesh = self.mesh + + # Define the repeating pattern length and calculate how many full blocks to scan + block_pattern_len = cfg.inhomogeneous_layer_cycle_interval + num_full_blocks = cfg.num_decoder_layers // block_pattern_len + remainder_layers = cfg.num_decoder_layers % block_pattern_len + + if num_full_blocks > 0: + ScannableBlockToLinen = qwen3.Qwen3NextScannableBlockToLinen + policy = self.get_remat_policy() + + kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan( + kv_caches, num_full_blocks, block_pattern_len, stack=True + ) + + # Positional order must match Qwen3NextScannableBlock.__call__. + broadcast_args_spec = [ + (decoder_segment_ids, nn.broadcast), + (decoder_positions, nn.broadcast), + (deterministic, nn.broadcast), + (model_mode, nn.broadcast), + (previous_chunk, nn.broadcast), + (slot, nn.broadcast), + (kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast), + (attention_metadata, nn.broadcast), + ] + broadcast_args = tuple(arg for arg, _ in broadcast_args_spec) + in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec) + + # For a fully scanned block, apply it inside an nn.scan over the calculated number of full blocks + y, returned_kv_cache = nn.scan( + ScannableBlockToLinen, + variable_axes={ + "params": cfg.param_scan_axis, + "cache": 0, + "intermediates": 0, + "aqt": 0, + "_overwrite_with_gradient": 0, + }, + split_rngs={"params": True, "dropout": cfg.enable_dropout}, + in_axes=in_axes_tuple, + length=num_full_blocks, + unroll=1, + metadata_params={ + nn.PARTITION_NAME: "layers", + "abstract_init": False, + }, + )( + config=cfg, + mesh=mesh, + quant=self.quant, + model_mode=model_mode, + num_of_layers=block_pattern_len, + remat_policy_fn=policy, + apply_internal_remat=True, + # Keep the Linen parameter path identical to the pure-NNX decoder's + # `self.layers`, so a single checkpoint mapping serves both. + name="layers", + )( + y, *broadcast_args + ) + + maxtext_utils.update_kv_caches_after_scan( + kv_caches, returned_kv_cache, num_full_blocks, block_pattern_len, stacked=True + ) + + # Layers past the last whole block go into a short block of their own. It is a + # Qwen3NextScannableBlock like the scanned ones, named `layers_remainder`, so the + # parameter paths still match the pure-NNX decoder's and one mapping serves both. + if remainder_layers > 0: + start_idx = cfg.num_decoder_layers - remainder_layers + remainder_kv = tuple(kv_caches[start_idx:]) if kv_caches is not None else None + y_and_kv = qwen3.Qwen3NextScannableBlockToLinen( + config=cfg, + mesh=mesh, + quant=self.quant, + model_mode=model_mode, + num_of_layers=remainder_layers, + layer_idx_offset=start_idx, + remat_policy_fn=self.get_remat_policy(), + apply_internal_remat=True, + name="layers_remainder", + )( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + remainder_kv, + attention_metadata, + ) + + if isinstance(y_and_kv, tuple): + y = y_and_kv[0] + new_kv = y_and_kv[1] + else: + y = y_and_kv + new_kv = None + + if kv_caches is not None and new_kv is not None: + for offset, updated_item in enumerate(new_kv): + kv_caches[start_idx + offset] = updated_item + + return y + + def _apply_gemma4_scanned_blocks( + self, + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + bidirectional_mask, + previous_chunk, + slot, + kv_caches=None, + attention_metadata=None, + ): + """Applies Gemma4 scanned decoder blocks, handling main scan and remainders.""" + + cfg = self.config + mesh = self.mesh + + # Define the repeating pattern length and calculate how many full blocks to scan + block_pattern_len = len(gemma4.GEMMA4_ATTENTION_PATTERN) + num_full_blocks = cfg.num_decoder_layers // block_pattern_len + remainder_layers = cfg.num_decoder_layers % block_pattern_len + + if num_full_blocks > 0: + ScannableBlockToLinen = gemma4.Gemma4ScannableBlockToLinen + policy = self.get_remat_policy() + # 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 + ) + + broadcast_args_spec = [ + (decoder_segment_ids, nn.broadcast), + (decoder_positions, nn.broadcast), + (deterministic, nn.broadcast), + (model_mode, nn.broadcast), + (slot, nn.broadcast), + (None, nn.broadcast), # page_state + (previous_chunk, nn.broadcast), + (bidirectional_mask, nn.broadcast), + (kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast), + (attention_metadata, nn.broadcast), + ] + broadcast_args = tuple(arg for arg, _ in broadcast_args_spec) + in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec) + + # For a fully scanned block, apply it inside a nn.scan over the calculated number of full blocks + y, returned_kv_cache = nn.scan( + ScannableBlockToLinen, + variable_axes={ + "params": cfg.param_scan_axis, + "cache": 0, + "intermediates": 0, + "aqt": 0, + "_overwrite_with_gradient": 0, + }, + 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, + }, + )( + config=cfg, + mesh=mesh, + 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 + ) + + maxtext_utils.update_kv_caches_after_scan( + kv_caches, returned_kv_cache, num_full_blocks, block_pattern_len, stacked=True + ) + + # Process any remaining layers that don't fit into a full scanned block + for layer_id in range(cfg.num_decoder_layers - remainder_layers, cfg.num_decoder_layers): + attention_type = gemma4.get_attention_type(layer_id) + layer = gemma4.Gemma4DecoderLayerToLinen( + config=cfg, + mesh=mesh, + model_mode=model_mode, + quant=self.quant, + attention_type=attention_type, + layer_idx=layer_id, + ) + kv_cache = kv_caches[layer_id] if kv_caches is not None else None + + # inputs, decoder_segment_ids, decoder_positions, deterministic, model_mode, + # previous_chunk, page_state, slot, bidirectional_mask, kv_cache, attention_metadata + remainder_args = ( + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + None, # page_state + slot, + bidirectional_mask, + kv_cache, + attention_metadata, + ) + + y_and_kv = layer(y, *remainder_args) + if isinstance(y_and_kv, tuple): + y = y_and_kv[0] + new_kv = y_and_kv[1] + else: + y = y_and_kv + new_kv = None + + if kv_caches is not None and new_kv is not None: + kv_caches[layer_id] = new_kv + + return y + + def _apply_deepseek4_scanned_blocks( + self, + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + decoder_input_tokens, + ): + """Applies DeepSeek V4 scanned decoder blocks. + + DeepSeek V4 has some number of prefix layers (defined by `first_num_hash_layers`) + that use static Hash Routing. The remaining layers alternate `compress_ratio=128` (HCA) + and `compress_ratio=4` (CSA) and are evaluated in a single `nn.scan` block. + + For DeepSeek4-Flash (43 hidden layers total): + - 3 Prefix layers (Indices 0, 1, 2) + - 40 Scanned layers: 20 perfectly repeating chunks of [128, 4] + """ + + cfg = self.config + mesh = self.mesh + + broadcast_args = ( + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + slot, + previous_chunk, + ) + + layer_call_kwargs = { + "previous_chunk": previous_chunk, + "slot": slot, + "decoder_input_tokens": decoder_input_tokens, + } + + # 1. Prefix Unrolling + # Prefix layers are unrolled (unscanned) for two architectural reasons: + # 1. Heterogeneous Attention: JAX nn.scan requires identical computation graphs, but the first few layers + # use different attention configurations (e.g., DeepSeek-V4 uses compress_ratios [0, 0, 4] for layers 0, 1, 2). + # 2. Static Hash Routing: The first `first_num_hash_layers` (which is 3 for DeepSeek-V4) use deterministic + # token-to-expert Hash Routing instead of learned top-k routing. + # Therefore, these prefix layers are instantiated individually before we scan the remaining uniform blocks. + num_hash_layers = cfg.first_num_hash_layers + for layer_idx in range(num_hash_layers): + prefix_layer = deepseek4.DeepSeek4LayerToLinen( + config=cfg, + mesh=mesh, + name=f"layers_{layer_idx}", + quant=self.quant, + model_mode=self.model_mode, + layer_idx=layer_idx, + ) + y, _ = prefix_layer( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + **layer_call_kwargs, + ) + + # 2. Chunked Scanning + # The remaining layers perfectly alternate HCA (128) and CSA (4). + num_remaining_layers = cfg.num_decoder_layers - num_hash_layers + num_full_blocks = num_remaining_layers // 2 + + if num_full_blocks > 0: + ScannableBlockToLinen = deepseek4.DeepSeek4ScannableBlockToLinen + policy = self.get_remat_policy() + RemattedDeepSeek4Block = self.set_remat_policy([ScannableBlockToLinen], policy)[0] + + y, _ = nn.scan( + RemattedDeepSeek4Block, + variable_axes={ + "params": cfg.param_scan_axis, + "cache": 0, + "intermediates": 0, + "aqt": 0, + "_overwrite_with_gradient": 0, + }, + split_rngs={"params": True, "dropout": cfg.enable_dropout}, + in_axes=(nn.broadcast,) * len(broadcast_args), + length=num_full_blocks, + metadata_params={ + nn.PARTITION_NAME: "layers", + "abstract_init": False, + }, + )(config=cfg, mesh=mesh, quant=self.quant, model_mode=model_mode, name="scanned_blocks",)(y, *broadcast_args) + + return y + + def _apply_gemma4_small_layers( + self, + y, + decoder_input_tokens, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + multimodal_input=None, + kv_caches=None, + attention_metadata=None, + previous_chunk=None, + slot=None, + ): + """Apply Gemma 4 small (E2B / E4B) decoder layers. + + Threads per-call state through the layer loop: + * ``per_layer_inputs`` from PLE, sliced per layer. + * ``shared_kv_states``: donor-layer-index → (key, value) for + downstream KV-shared layers to consume. + * ``kv_caches``: when running via the vLLM RPA path, the per-layer + cache buffer threaded back from the kernel. KV-shared layers + redirect to the donor's cache slot via ``cache_index_of``. + + Returns ``(y, kv_caches)``. Scan-over-layers and pipeline + parallelism are not supported. + """ + cfg = self.config + mesh = self.mesh + bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None + + per_layer_inputs = None + if cfg.hidden_size_per_layer_input > 0 and cfg.vocab_size_per_layer_input > 0: + per_layer_inputs = gemma4_small.PLEToLinen( + config=cfg, + mesh=mesh, + name="per_layer_embedder", + )(decoder_input_tokens, y) + + layer_types = gemma4_small.build_layer_types(cfg.num_decoder_layers, cfg.model_name) + num_kv_shared = cfg.num_kv_shared_layers + shared_kv_states: dict[int, tuple[jax.Array, jax.Array]] = {} + + # tpu-inference allocates one `kv_caches` slot per non-shared layer; + # KV-shared layers reuse the donor's slot. + cache_index_of = gemma4_small.kv_cache_slot_map(layer_types, num_kv_shared) + + for lyr in range(cfg.num_decoder_layers): + attention_type = layer_types[lyr] + donor_idx = gemma4_small.kv_donor_layer_idx(lyr, layer_types, num_kv_shared) + is_donor = gemma4_small.is_kv_donor_layer(lyr, layer_types, num_kv_shared) + + shared_key = None + shared_value = None + if donor_idx is not None: + if donor_idx not in shared_kv_states: + raise RuntimeError( + f"KV-shared layer {lyr} references donor {donor_idx} but no donor K/V " + f"have been recorded yet. This indicates the layer iteration order is wrong." + ) + shared_key, shared_value = shared_kv_states[donor_idx] + + layer = gemma4_small.Gemma4SmallDecoderLayerToLinen( + config=cfg, + mesh=mesh, + name=f"layers_{lyr}", + quant=self.quant, + model_mode=self.model_mode, + attention_type=attention_type, + layer_idx=lyr, + ) + + # Donor layers expose their rotated, normed K / V to downstream + # shared layers via the decoder layer's compute_shared_kv method. + if is_donor: + donor_k, donor_v = layer(y, decoder_positions, nnx_method="compute_shared_kv") + shared_kv_states[lyr] = (donor_k, donor_v) + # Reuse the just-computed K / V in the layer's own forward pass to + # avoid double-computing the K / V projection / norm / RoPE. + shared_key, shared_value = donor_k, donor_v + + ple_slice = per_layer_inputs[..., lyr, :] if per_layer_inputs is not None else None + + 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, + ) + if kv_caches is not None and kv_cache is not None: + kv_caches[cache_idx] = kv_cache + + return y, kv_caches + + # 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 diff --git a/src/maxtext/layers/embeddings.py b/src/maxtext/layers/embeddings.py index 4c80f63848..309f34e3f4 100644 --- a/src/maxtext/layers/embeddings.py +++ b/src/maxtext/layers/embeddings.py @@ -26,7 +26,8 @@ from flax import nnx from maxtext.common.common_types import ShardMode, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, Array, Config, DType -from maxtext.layers.initializers import Initializer, default_embed_init +from maxtext.layers import nnx_wrappers +from maxtext.layers.initializers import Initializer, default_embed_init, variable_to_logically_partitioned from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils.sharding import logical_to_mesh_axes, create_sharding, truncate_out_sharding @@ -42,6 +43,52 @@ def _maybe_move_embedding_to_device(embedding_table: Array, config: Config) -> A return embedding_table +def embed_as_linen( + *, + num_embeddings: int, + num_features: int, + config: Config, + mesh: Mesh, + cast_input_dtype: None | DType = None, + dtype: DType = jnp.float32, + attend_dtype: None | DType = None, + embedding_init: Initializer = default_embed_init, + name: str | None = None, +): + """Initializes the Embed NNX module and returns it as a Linen module. + + This function serves as a bridge to use the NNX-based `Embed` module within + a Linen model. It wraps the `Embed` module using `nnx.bridge.to_linen`, + making it compatible with the Linen API. + + Args: + num_embeddings: The number of embeddings. + num_features: The number of feature dimensions for each embedding. + config: The model configuration. + cast_input_dtype: The dtype to cast the input to, if any. + dtype: The dtype of the embedding vectors. + attend_dtype: The dtype for the `attend` method. + embedding_init: The initializer for the embedding matrix. + name: The name of the Linen module. + + Returns: + A Linen module that wraps the NNX `Embed` module. + """ + return nnx_wrappers.to_linen( + Embed, + num_embeddings=num_embeddings, + num_features=num_features, + config=config, + mesh=mesh, + cast_input_dtype=cast_input_dtype, + dtype=dtype, + attend_dtype=attend_dtype, + embedding_init=embedding_init, + metadata_fn=variable_to_logically_partitioned, + name=name, + ) + + class Embed(nnx.Module): """A parameterized function from integers [0, n) to d-dimensional vectors.""" @@ -193,6 +240,39 @@ def attend_on_embedding( ) +def rotary_embedding_as_linen( + *, + min_timescale: int, + max_timescale: int, + embedding_dims: int = 0, + cast_as_fprop_dtype: bool = True, + fprop_dtype: DType = jnp.bfloat16, + name: str | None = None, +): + """Initializes the RotaryEmbedding module and returns it as a Linen module. + + Args: + min_timescale: Start of the geometric index. Determines the periodicity of + the added signal. + max_timescale: End of the geometric index. Determines the frequency of the + added signal. + embedding_dims: Dimension of the embedding to be generated. + cast_as_fprop_dtype: Whether to cast the output to the fprop dtype. + fprop_dtype: The dtype of the output. + name: Name of the Linen module. + """ + return nnx_wrappers.to_linen( + RotaryEmbedding, + min_timescale=min_timescale, + max_timescale=max_timescale, + embedding_dims=embedding_dims, + cast_as_fprop_dtype=cast_as_fprop_dtype, + fprop_dtype=fprop_dtype, + metadata_fn=variable_to_logically_partitioned, + name=name, + ) + + class RotaryEmbedding(nnx.Module): """Rotary Position Embedding.""" @@ -295,6 +375,82 @@ def __call__( return x_out +def llama_rotary_embedding_as_linen( + *, + min_timescale: int, + max_timescale: int, + embedding_dims: int = 0, + cast_as_fprop_dtype: bool = True, + fprop_dtype: DType = jnp.bfloat16, + use_scale: bool = True, + name: str | None = None, +): + """Initializes the LLaMARotaryEmbedding module and returns it as a Linen module. + + Args: + min_timescale: Start of the geometric index. Determines the periodicity of + the added signal. + max_timescale: End of the geometric index. Determines the frequency of the + added signal. + embedding_dims: Dimension of the embedding to be generated. + cast_as_fprop_dtype: Whether to cast the output to the fprop dtype. + fprop_dtype: The dtype of the output. + use_scale: Whether to apply LLaMA3.1 scaling factor. + name: Name of the Linen module. + """ + return nnx_wrappers.to_linen( + LLaMARotaryEmbedding, + min_timescale=min_timescale, + max_timescale=max_timescale, + embedding_dims=embedding_dims, + cast_as_fprop_dtype=cast_as_fprop_dtype, + fprop_dtype=fprop_dtype, + use_scale=use_scale, + metadata_fn=variable_to_logically_partitioned, + name=name, + ) + + +def partial_rotary_embedding_as_linen( + *, + min_timescale: int, + max_timescale: int, + mesh: Mesh, + embedding_dims: int = 0, + partial_rotary_factor: float = 0.25, + cast_as_fprop_dtype: bool = True, + fprop_dtype: DType = jnp.bfloat16, + shard_mode: ShardMode = ShardMode.AUTO, + name: str | None = None, +): + """Initializes the PartialRotaryEmbedding module and returns it as a Linen module. + + Args: + min_timescale: Start of the geometric index. Determines the periodicity of + the added signal. + max_timescale: End of the geometric index. Determines the frequency of the + added signal. + embedding_dims: Dimension of the embedding to be generated. + partial_rotary_factor: Ratio of dimensions to apply ROPE to. + cast_as_fprop_dtype: Whether to cast the output to the fprop dtype. + fprop_dtype: The dtype of the output. + name: Name of the Linen module. + """ + return nnx_wrappers.to_linen( + PartialRotaryEmbedding, + min_timescale=min_timescale, + max_timescale=max_timescale, + mesh=mesh, + embedding_dims=embedding_dims, + partial_rotary_factor=partial_rotary_factor, + cast_as_fprop_dtype=cast_as_fprop_dtype, + fprop_dtype=fprop_dtype, + shard_mode=shard_mode, + metadata_fn=variable_to_logically_partitioned, + name=name, + ) + + class PartialRotaryEmbedding(RotaryEmbedding): """Rotary Position Embedding applied to a partial fraction of dimensions.""" @@ -571,6 +727,61 @@ def __call__(self, inputs: jax.Array, position: None | jax.Array = None) -> jax. return outputs +def yarn_rotary_embedding_as_linen( + *, + embedding_dims: int, + mesh: Mesh, + max_position_embeddings: int = 4096 * 4, + original_max_position_embeddings: int = 4096, + beta_fast: float = 32, + beta_slow: float = 1, + rope_theta: float = 10000.0, + rope_factor: float = 40, + cast_as_fprop_dtype: bool = True, + fprop_dtype: DType = jnp.bfloat16, + name: str | None = None, + interleave: bool = True, + truncate: bool = True, + attention_scaling: bool = False, + pairwise: bool = False, + shard_mode: ShardMode = ShardMode.AUTO, +): + """Initializes the YarnRotaryEmbedding module and returns it as a Linen module. + + Args: + embedding_dims: The dimension of the embeddings. + max_position_embeddings: The maximum number of positions. + original_max_position_embeddings: The original maximum number of positions. + beta_fast: The fast beta parameter for YaRN. + beta_slow: The slow beta parameter for YaRN. + rope_theta: The base for the rotary frequencies. + rope_factor: The scaling factor for RoPE. + cast_as_fprop_dtype: Whether to cast the output to `fprop_dtype`. + fprop_dtype: The forward pass dtype. + name: The name of the module. + """ + return nnx_wrappers.to_linen( + YarnRotaryEmbedding, + embedding_dims=embedding_dims, + max_position_embeddings=max_position_embeddings, + mesh=mesh, + original_max_position_embeddings=original_max_position_embeddings, + beta_fast=beta_fast, + beta_slow=beta_slow, + rope_theta=rope_theta, + rope_factor=rope_factor, + cast_as_fprop_dtype=cast_as_fprop_dtype, + fprop_dtype=fprop_dtype, + metadata_fn=variable_to_logically_partitioned, + name=name, + interleave=interleave, + truncate=truncate, + attention_scaling=attention_scaling, + pairwise=pairwise, + shard_mode=shard_mode, + ) + + class YarnRotaryEmbedding(nnx.Module): """Yarn rotary embedding. @@ -815,6 +1026,31 @@ def __call__(self, inputs: Array, position: None | Array = None) -> Array: return output +def positional_embedding_as_linen( + *, + embedding_dims: int, + max_wavelength: int = _MAX_WAVELENGTH, + cast_as_fprop_dtype: bool = False, + fprop_dtype: DType = jnp.bfloat16, +): + """Initializes the PositionalEmbedding module and returns it as a Linen module. + + Args: + embedding_dims: The dimension of the embeddings. + max_wavelength: The maximum wavelength for the sinusoidal positional embeddings. + cast_as_fprop_dtype: Whether to cast output to fprop_dtype. + fprop_dtype: The dtype of the output when cast_as_fprop_dtype is True. + """ + return nnx_wrappers.to_linen( + PositionalEmbedding, + embedding_dims=embedding_dims, + max_wavelength=max_wavelength, + cast_as_fprop_dtype=cast_as_fprop_dtype, + fprop_dtype=fprop_dtype, + metadata_fn=variable_to_logically_partitioned, + ) + + @dataclasses.dataclass(repr=False) class PositionalEmbedding(nnx.Module): """Sinusoidal positional embeddings supporting both uniform and per-batch positions. @@ -900,6 +1136,43 @@ def __call__( return self._compute_embeddings(position) +def llama_vision_rotary_embedding_as_linen( + *, + image_size: int, + patch_size: int, + hidden_size: int, + num_attention_heads: int, + rope_theta: float = 10000.0, + cast_as_fprop_dtype: bool = True, + fprop_dtype: DType = jnp.bfloat16, + name: str | None = None, +): + """Initializes the LlamaVisionRotaryEmbedding module and returns it as a Linen module. + + Args: + image_size: The size of the input image. + patch_size: The size of the image patches. + hidden_size: The size of the hidden dimension. + num_attention_heads: The number of attention heads. + rope_theta: The base theta value for the frequency computation. + cast_as_fprop_dtype: Whether to cast the output to the fprop dtype. + fprop_dtype: The dtype of the output. + name: The name of the Linen module. + """ + return nnx_wrappers.to_linen( + LlamaVisionRotaryEmbedding, + image_size=image_size, + patch_size=patch_size, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + rope_theta=rope_theta, + cast_as_fprop_dtype=cast_as_fprop_dtype, + fprop_dtype=fprop_dtype, + metadata_fn=variable_to_logically_partitioned, + name=name, + ) + + @dataclasses.dataclass(repr=False) class LlamaVisionRotaryEmbedding(nnx.Module): """Rotary position embedding for Llama4 vision encoder. @@ -1270,6 +1543,47 @@ def generate_block_coords(video_grid_thw: Array, max_patches: int, merge_size: i return row, col +def qwen3omnimoe_vision_pos_embed_interpolate_as_linen( + *, + num_position_embeddings: int, + hidden_size: int, + spatial_merge_size: int, + dtype: DType = jnp.float32, + cast_as_fprop_dtype: bool = True, + fprop_dtype: DType = jnp.bfloat16, + name: str | None = None, +): + """Initializes Qwen3OmniMoe bilinear position embedding interpolation as Linen module. + + This implements fast bilinear interpolation of learned 2D positional embeddings + for dynamic input sizes. The embeddings are learned on a fixed grid and interpolated + to match the actual image/video dimensions. + + Args: + num_position_embeddings: Number of position embeddings in the fixed grid (e.g., 1024 for 32x32) + hidden_size: Hidden dimension size + spatial_merge_size: Size of spatial merging blocks + dtype: Data type for embeddings + cast_as_fprop_dtype: Whether to cast the output to the fprop dtype + fprop_dtype: The dtype of the output + name: Module name + + Returns: + A Linen module that wraps the NNX Qwen3OmniMoeVisionPosEmbedInterpolate module. + """ + return nnx_wrappers.to_linen( + Qwen3OmniMoeVisionPosEmbedInterpolate, + num_position_embeddings=num_position_embeddings, + hidden_size=hidden_size, + spatial_merge_size=spatial_merge_size, + dtype=dtype, + cast_as_fprop_dtype=cast_as_fprop_dtype, + fprop_dtype=fprop_dtype, + metadata_fn=variable_to_logically_partitioned, + name=name, + ) + + class Qwen3OmniMoeVisionPosEmbedInterpolate(nnx.Module): """Bilinear interpolation of learned 2D positional embeddings for Qwen3OmniMoe vision. @@ -1638,6 +1952,40 @@ def __call__( return x_out +def qwen3_omni_mrope_embedding_as_linen( + *, + min_timescale: int, + max_timescale: int, + embedding_dims: int = 0, + cast_as_fprop_dtype: bool = True, + fprop_dtype: DType = jnp.bfloat16, + mrope_section: tuple[int, int, int] | None = None, + name: str | None = None, +): + """Initializes Qwen3OmniMoeThinkerTextRotaryEmbedding and returns it as a Linen module. + + Args: + min_timescale: Start of the geometric index. + max_timescale: End of the geometric index (rope_theta). + embedding_dims: Dimension of the embedding (head_dim). + cast_as_fprop_dtype: Whether to cast output to fprop dtype. + fprop_dtype: The dtype of the output. + mrope_section: Tuple of (temporal_dim, height_dim, width_dim) for MRoPE. + name: Name of the Linen module. + """ + return nnx_wrappers.to_linen( + Qwen3OmniMoeThinkerTextRotaryEmbedding, + min_timescale=min_timescale, + max_timescale=max_timescale, + embedding_dims=embedding_dims, + cast_as_fprop_dtype=cast_as_fprop_dtype, + fprop_dtype=fprop_dtype, + mrope_section=mrope_section, + metadata_fn=variable_to_logically_partitioned, + name=name, + ) + + class DeepSeekV4RotaryEmbedding(RotaryEmbedding): """DeepSeek-V4 partial rotary embedding with interleaved frequencies. diff --git a/src/maxtext/layers/encoders.py b/src/maxtext/layers/encoders.py index cecc68a26c..e5e32d044a 100644 --- a/src/maxtext/layers/encoders.py +++ b/src/maxtext/layers/encoders.py @@ -19,6 +19,8 @@ from jax.sharding import Mesh from maxtext.common.common_types import Config, VisionEncoderBlockType +from maxtext.layers import nnx_wrappers +from maxtext.layers import initializers from maxtext.layers import linears @@ -244,3 +246,35 @@ def __call__(self, input_audio, deterministic=False): embeddings = projector(embeddings) return embeddings + + +def vision_encoder_as_linen( + config: Config, + mesh: Mesh, +): + """Creates a VisionEncoder module.""" + module = nnx_wrappers.to_linen( + VisionEncoder, + config=config, + mesh=mesh, + name="vision_encoder", + abstract_init=False, + metadata_fn=initializers.variable_to_logically_partitioned, + ) + return module + + +def audio_encoder_as_linen( + config: Config, + mesh: Mesh, +): + """Creates an AudioEncoder module.""" + module = nnx_wrappers.to_linen( + AudioEncoder, + config=config, + mesh=mesh, + name="audio_encoder", + abstract_init=False, + metadata_fn=initializers.variable_to_logically_partitioned, + ) + return module diff --git a/src/maxtext/layers/multi_token_prediction.py b/src/maxtext/layers/multi_token_prediction.py index 9ab52e5fd3..fce3dfb940 100644 --- a/src/maxtext/layers/multi_token_prediction.py +++ b/src/maxtext/layers/multi_token_prediction.py @@ -17,12 +17,15 @@ from typing import Type +from flax import linen as nn from flax import nnx import jax import jax.numpy as jnp from jax.sharding import Mesh from maxtext.common.common_types import Config, DecoderBlockType, MODEL_MODE_TRAIN, ShardMode from maxtext.layers import moe +from maxtext.layers.decoders import DecoderLayer +from maxtext.layers.initializers import variable_to_logically_partitioned from maxtext.layers.linears import DenseGeneral from maxtext.layers.nnx_decoders import NNXDecoderLayer from maxtext.layers.normalizations import RMSNorm @@ -560,3 +563,37 @@ def calculate_mtp_acceptance_rate(intermediate_outputs, config): total_valid_tokens = jnp.sum(valid_mask) return (correct_predictions / (total_valid_tokens + EPS)) * 100 + + +def multi_token_prediction_block_as_linen( + *, + config: Config, + mesh: Mesh, + transformer_layer_module: Type[DecoderLayer], + decoder: nnx.Module, + rngs: nnx.Rngs, + name: str | None = None, +) -> nn.Module: + """Initializes MultiTokenPredictionBlock as a Linen module. + + Args: + config: Configuration object containing model hyperparameters. + mesh: JAX Mesh for model parallelism. + transformer_layer_module: The Transformer Decoder Layer class to use. + decoder: The decoder module that provides embedding and output head. + rngs: Random number generators for initialization. + name: Optional name for the module. + + Returns: + An instance of MultiTokenPredictionBlock wrapped as a Linen module. + """ + return nnx.bridge.to_linen( + MultiTokenPredictionBlock, + config=config, + mesh=mesh, + transformer_layer_module=transformer_layer_module, + decoder=decoder, + rngs=rngs, + metadata_fn=variable_to_logically_partitioned, + name=name, + ) diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 94c37ae304..caedc633a2 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -38,7 +38,7 @@ ShardMode, ) from maxtext.configs.types import check_forced_routing_support -from maxtext.layers import linears, mhc, moe, normalizations, quantizations +from maxtext.layers import initializers, linears, mhc, moe, normalizations, quantizations from maxtext.layers import nnx_scan, nnx_wrappers from maxtext.layers.attentions import Attention from maxtext.layers.embeddings import Embed, PositionalEmbedding, attend_on_embedding @@ -2725,3 +2725,25 @@ def _append_scanned(name): _add(m) return layers + + +def decoder_as_linen( + config: Config, + mesh: Mesh, + rngs: nnx.Rngs, + model_mode: str, + quant: None | Quant = None, +): + """Creates a Decoder module""" + module = nnx_wrappers.to_linen( + NNXDecoder, + config=config, + mesh=mesh, + model_mode=model_mode, + rngs=rngs, + quant=quant, + name="decoder", + abstract_init=False, + metadata_fn=initializers.variable_to_logically_partitioned, + ) + return module diff --git a/src/maxtext/models/gemma3.py b/src/maxtext/models/gemma3.py index 106a48e088..8d2733ee3b 100644 --- a/src/maxtext/models/gemma3.py +++ b/src/maxtext/models/gemma3.py @@ -31,6 +31,7 @@ from maxtext.layers.linears import DenseGeneral, MlpBlock, Dropout from maxtext.layers.normalizations import RMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.initializers import variable_to_logically_partitioned from maxtext.utils import max_utils from maxtext.utils.sharding import create_sharding, maybe_shard_with_logical @@ -605,6 +606,21 @@ def __call__(self, x: jax.Array, eqn: str = "...tm,md->...td") -> jax.Array: return x +def visionembedder_as_linen( + config: Config, + mesh: Mesh, +): + """Creates a VisionEmbedder module.""" + return nnx_wrappers.to_linen( + VisionEmbedder, + config, + mesh=mesh, + name="VisionEmbedder_0", + abstract_init=False, + metadata_fn=variable_to_logically_partitioned, + ) + + class VisionExit(nnx.Module): """The vision exit layer. @@ -637,6 +653,11 @@ def __call__(self, x): return jnp.reshape(x, (batch_size, height * width, embed_dim)) +def vision_exit_as_linen(x: jax.Array, output_length: int) -> jax.Array: + """A wrapper to use VisionExit as a function.""" + return nnx.bridge.to_linen(VisionExit, output_length=output_length)(x) + + class Gemma3VisionEncoderLayer(nnx.Module): """gemma 3 vision encoder layer""" @@ -723,3 +744,19 @@ def __call__(self, inputs, deterministic, train=False): bn, l, c = x.shape x = jnp.reshape(x, [b, n, l, c]) return x + + +def gemma3visionencoder_as_linen( + config: Config, + mesh: Mesh, +): + """Creates a Gemma3VisionEncoder module.""" + module = nnx_wrappers.to_linen( + Gemma3VisionEncoderLayer, + config=config, + mesh=mesh, + name="Gemma3VisionEncoderLayer_0", + abstract_init=False, + metadata_fn=variable_to_logically_partitioned, + ) + return module diff --git a/src/maxtext/models/gemma4_vision.py b/src/maxtext/models/gemma4_vision.py index f673488608..72cd841f81 100644 --- a/src/maxtext/models/gemma4_vision.py +++ b/src/maxtext/models/gemma4_vision.py @@ -18,12 +18,15 @@ from typing import cast import jax import jax.numpy as jnp +from flax import linen as nn from flax import nnx from jax.sharding import Mesh from maxtext.common.common_types import Config, AttentionType from maxtext.layers import attentions +from maxtext.layers import initializers from maxtext.layers import linears +from maxtext.layers import nnx_wrappers from maxtext.layers import normalizations @@ -633,3 +636,15 @@ def __call__(self, x: jax.Array) -> jax.Array: x_normed = self.norm(x) x_projected = self.projection(x_normed) return x_projected + + +def gemma4_vision_encoder_as_linen(config: Config, mesh: Mesh) -> nn.Module: + """Wraps the Gemma 4 Vision Encoder as a Linen module.""" + return nnx_wrappers.to_linen( + Gemma4VisionEncoderLayer, + config=config, + mesh=mesh, + name="Gemma4VisionEncoderLayer", + abstract_init=False, + metadata_fn=initializers.variable_to_logically_partitioned, + ) diff --git a/src/maxtext/models/llama4.py b/src/maxtext/models/llama4.py index ebbc98d764..415aff7ce3 100644 --- a/src/maxtext/models/llama4.py +++ b/src/maxtext/models/llama4.py @@ -253,6 +253,17 @@ def __call__(self, image_features: Array) -> Array: return hidden_states +def llama4multimodalprojector_as_linen(config: Config, mesh: Mesh): + return nnx_wrappers.to_linen( + Llama4MultiModalProjector, + config=config, + mesh=mesh, + name="Llama4MultiModalProjector_0", + abstract_init=False, + metadata_fn=initializers.variable_to_logically_partitioned, + ) + + def determine_is_nope_layer(layer_id: int, nope_layer_interval: int) -> bool: """ Determines whether the given layer at `layer_id` should use RoPE or not (NoPE). @@ -801,3 +812,14 @@ def __call__( hidden_states = jnp.reshape(hidden_states, [b, t, patch_num, patch_dim]) return hidden_states + + +def llama4visionmodel_as_linen(config: Config, mesh: Mesh) -> nn.Module: + return nnx_wrappers.to_linen( + Llama4VisionModel, + config=config, + mesh=mesh, + name="Llama4VisionModel_0", + abstract_init=False, + metadata_fn=initializers.variable_to_logically_partitioned, + ) diff --git a/src/maxtext/models/models.py b/src/maxtext/models/models.py index 19408299fd..90c3d17796 100644 --- a/src/maxtext/models/models.py +++ b/src/maxtext/models/models.py @@ -29,16 +29,246 @@ from maxtext.layers.nnx_decoders import NNXDecoder from maxtext.layers import initializers from maxtext.layers import nnx_wrappers -from maxtext.layers.embeddings import Embed -from maxtext.layers.encoders import AudioEncoder, VisionEncoder -from maxtext.layers.multi_token_prediction import MultiTokenPredictionBlock +from maxtext.layers.decoders import Decoder +from maxtext.layers.embeddings import Embed, embed_as_linen +from maxtext.layers.encoders import AudioEncoder, VisionEncoder, audio_encoder_as_linen, vision_encoder_as_linen +from maxtext.layers.multi_token_prediction import MultiTokenPredictionBlock, multi_token_prediction_block_as_linen from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.multimodal import processor as mm_processor +from maxtext.utils import max_utils # ------------------------------------------------------------------------------ # The network: Transformer Definitions # ------------------------------------------------------------------------------ +_FORCED_ROUTING_NEEDS_PURE_NNX = ( + "Forced routing (router replay) is only implemented for the pure-NNX" + " decoder (NNXDecoder). Use the NNX model path with pure_nnx_decoder=True" + " (both are the defaults)." +) + + +class TransformerLinenPure(nn.Module): + """An autoregressive transformer model.""" + + # Make new attributes required, so that all Transformer dependencies (train, decode, + # compile, etc) will error instead of silently use defaults. + # pylint: disable=attribute-defined-outside-init + config: Config + mesh: Mesh + quant: Quant + # Possible model_mode values can be found in maxtext.common.common_types. + # We generally use maxtext.common.common_types.MODEL_MODE_TRAIN or + # maxtext.common.common_types.MODEL_MODE_PREFILL for initializations here. + # TODO: Make model_mode required after confirming no users are affected. + model_mode: str = MODEL_MODE_TRAIN # May be different than the model_mode passed to __call__ + # pylint: enable=attribute-defined-outside-init + + def init(self, *args, model_mode: str = MODEL_MODE_TRAIN, **kwargs): + """Initializes the model.""" + module = self.clone(model_mode=model_mode) + kwargs["model_mode"] = model_mode + return nn.Module.init(module, *args, **kwargs) + + def apply(self, *args, model_mode: str = MODEL_MODE_TRAIN, **kwargs): + """Applies the model.""" + module = self.clone(model_mode=model_mode) + kwargs["model_mode"] = model_mode + return nn.Module.apply(module, *args, **kwargs) + + def setup(self): + """Initialize shared_embedding & decoder layers.""" + + cfg = self.config + mesh = self.mesh + self.shared_embedding = embed_as_linen( + num_embeddings=cfg.vocab_size, + num_features=cfg.emb_dim, + dtype=cfg.dtype, + attend_dtype=jnp.float32 if cfg.logits_dot_in_fp32 else cfg.dtype, # for logit training stability + embedding_init=nn.initializers.normal(stddev=1.0), + name="token_embedder", + config=cfg, + mesh=self.mesh, + ) + self.vision_encoder = ( + vision_encoder_as_linen(config=cfg, mesh=mesh) if getattr(self.config, "use_multimodal", False) else None + ) + self.audio_encoder = audio_encoder_as_linen(config=cfg, mesh=mesh) if getattr(cfg, "use_audio", False) else None + self.decoder = Decoder(config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode) + + # If MTP is enabled via config, set up the MTP block. + if getattr(self.config, "mtp_num_layers", 0) > 0: + # Get the list of layer blueprints for the current model. + # For MTP, we use the DecoderLayer blueprint to ensure architectural consistency. + # By convention, this is the last layer in the list. + layer_types = self.decoder.get_decoder_layers() + mtp_layer_linen = layer_types[-1] + # UNWRAP: The MTP block is pure NNX. If the decoder returned a Linen wrapper, + # extract the native NNX class to preserve parameter tracing/scoping. + mtp_layer_nnx = getattr(mtp_layer_linen, "module_class", mtp_layer_linen) + self.mtp_block = multi_token_prediction_block_as_linen( + config=self.config, + mesh=self.mesh, + transformer_layer_module=mtp_layer_nnx, # pyrefly: ignore[bad-argument-type] + decoder=self.decoder, # pyrefly: ignore[bad-argument-type] + rngs=self.make_rng("mtp_block"), # pyrefly: ignore[bad-argument-type] + ) + + def logits_from_hidden_states_for_vocab_tiling(self, hidden_states, deterministic, model_mode): + """ + Compute logits from hidden states (wrapping decoder.apply_output_head). + This function is only used for vocabulary tiling. + """ + logits = self.decoder.apply_output_head( + shared_embedding=self.shared_embedding, + y=hidden_states, + deterministic=deterministic, + model_mode=model_mode, + ) + return logits + + def __call__( + self, + decoder_input_tokens: jnp.ndarray, + decoder_positions: jnp.ndarray, + decoder_segment_ids=None, + encoder_images: None | jnp.ndarray = None, + encoder_image_masks: None | jnp.ndarray = None, + encoder_videos: None | jnp.ndarray = None, + encoder_video_masks: None | jnp.ndarray = None, + encoder_video_grid_thw: None | jnp.ndarray = None, + encoder_audios: None | jnp.ndarray = None, + enable_dropout=True, + model_mode=MODEL_MODE_TRAIN, + previous_chunk=None, + true_length: None | int = None, + slot: None | int = None, + decoder_target_tokens: None | jnp.ndarray = None, + decoder_target_mask: None | jnp.ndarray = None, + nnx_method=None, + kv_caches: list[jax.Array] | None = None, + attention_metadata: dict[str, Any] | None = None, + forced_routed_experts: jnp.ndarray | None = None, + ): + """Applies Transformer decoder-branch on encoded-input and target. + + Args: + true_length: (Optional) Prompt length before padding + slot: (Optional) An integer representing the decode batch index selected + for this request. + """ + + if decoder_segment_ids is not None and model_mode == MODEL_MODE_AUTOREGRESSIVE: + raise ValueError( + f"During autoregressive decoding we assume the tokens are in the active sequence" + f" which is always {DECODING_ACTIVE_SEQUENCE_INDICATOR}." + ) + + bidirectional_mask_image = None + bidirectional_mask_video = None + image_embeddings = None + video_embeddings = None + audio_embeddings = None + deepstack_visual_embeds = None + + if getattr(self.config, "use_multimodal", False) and encoder_images is not None: + image_embeddings, deepstack_visual_embeds = self.vision_encoder( # pyrefly: ignore[not-callable] + input_images=encoder_images, deterministic=not enable_dropout + ) + bidirectional_mask_image = mm_processor.get_bidirectional_mask_vision( + self.config, decoder_input_tokens, is_video=False + ) + + if getattr(self.config, "use_multimodal", False) and encoder_videos is not None: + video_embeddings, deepstack_visual_embeds = self.vision_encoder( # pyrefly: ignore[not-callable] + input_images=encoder_videos, + input_masks=encoder_video_masks, + video_grid_thw=encoder_video_grid_thw, + deterministic=not enable_dropout, + ) + encoder_video_masks = mm_processor.downsample_video_mask_to_tokens(encoder_video_masks, self.config) + bidirectional_mask_video = mm_processor.get_bidirectional_mask_vision( + self.config, decoder_input_tokens, is_video=True + ) + + if getattr(self.config, "use_multimodal", False) and encoder_audios is not None and self.audio_encoder is not None: + audio_embeddings = self.audio_encoder(input_audio=encoder_audios, deterministic=not enable_dropout) + + # Create audio mask for placeholder tokens (qwen3-omni models) + audio_masks = None + if audio_embeddings is not None: + audio_masks = mm_processor.get_bidirectional_mask_audio(self.config, decoder_input_tokens) + + multimodal_input = None + if image_embeddings is not None or video_embeddings is not None or audio_embeddings is not None: + multimodal_input = MultimodalInput( + image_embeddings=image_embeddings, + image_masks=encoder_image_masks, + video_embeddings=video_embeddings, + video_masks=encoder_video_masks, + audio_embeddings=audio_embeddings, + audio_masks=audio_masks, + bidirectional_mask=bidirectional_mask_image, + bidirectional_mask_video=bidirectional_mask_video, + ) + + if forced_routed_experts is not None: + raise NotImplementedError(_FORCED_ROUTING_NEEDS_PURE_NNX) + + logits, hidden_state, kv_caches = self.decoder( + shared_embedding=self.shared_embedding, + decoder_input_tokens=decoder_input_tokens, + decoder_positions=decoder_positions, + decoder_segment_ids=decoder_segment_ids, + deterministic=not enable_dropout, + model_mode=model_mode, + previous_chunk=previous_chunk, + slot=slot, + multimodal_input=multimodal_input, + kv_caches=kv_caches, + attention_metadata=attention_metadata, + deepstack_visual_embeds=deepstack_visual_embeds, + ) # pytype: disable=wrong-keyword-args + + # If we are initializing the model AND MTP is enabled, we must create + # dummy target tensors. This allows Flax to trace the MTPBlock and create + # all its necessary parameters, without requiring the main training pipeline + # to be aware of this initialization detail. + if self.is_initializing() and getattr(self.config, "mtp_num_layers", 0) > 0: + if decoder_target_tokens is None: + dummy_shape = decoder_input_tokens.shape + decoder_target_tokens = jnp.ones(dummy_shape, dtype=jnp.int32) + decoder_target_mask = jnp.ones(dummy_shape, dtype=jnp.int32) + decoder_segment_ids = jnp.ones(dummy_shape, dtype=jnp.int32) + + # The Multi-Token Prediction (MTP) block functions as a "side-car" to the main + # model, active only during training. It computes an auxiliary loss based on + # predicting multiple future tokens, as described in the DeepSeek-V3 paper. + # To ensure architectural consistency, it uses two key components from the parent Transformer: + # 1. The same `DecoderLayer` blueprint for its internal transformer blocks. + # 2. The `shared_embedding` for both embedding future tokens and for its final + # logit projection. + # Its only effect is to "sow" these losses; it does not alter the primary logits output. + if getattr(self.config, "mtp_num_layers", 0) > 0: + self.mtp_block( + shared_embedding=self.shared_embedding, + main_hidden_state=hidden_state, + input_ids=decoder_input_tokens, + target_ids=decoder_target_tokens, + target_mask=decoder_target_mask, + position_ids=decoder_positions, + decoder_segment_ids=decoder_segment_ids, + deterministic=not enable_dropout, + model_mode=model_mode, + ) + + if self.config.attention in ("vllm_rpa", "vllm_batched_rpa"): + # In vLLM, logits are computed separately after updating the KV cache. + return hidden_state, kv_caches + + return logits + def transformer_as_linen( config: Config, @@ -47,13 +277,16 @@ def transformer_as_linen( model_mode: str = MODEL_MODE_TRAIN, *, name: str | None = None, -) -> nnx_wrappers.ToLinen: - """Constructs an NNX Transformer wrapped as a Linen module. +) -> nnx_wrappers.ToLinen | TransformerLinenPure: + """Constructs a Transformer model as a Linen or NNX module. + + This function returns an autoregressive Transformer model as either a Linen module + or an NNX-wrapped module, depending on the `config.enable_nnx` flag. The returned module + is suitable for training, evaluation, or decoding. - Returns a `TransformerLinen` that wraps the NNX-style Transformer so it can be - driven through the Linen init/apply API (checkpoint conversion, AOT compile, - the inference engine). Pure-NNX call sites build `Transformer` directly via - `model_creation_utils.from_config`. + If `config.enable_nnx` is True, returns a `TransformerLinen` that wraps the NNX-style + Transformer for integration with NNX-specific APIs and workflows. + Otherwise, returns a pure Flax Linen implementation (`TransformerLinenPure`). Args: config (Config): The configuration object specifying model hyperparameters and options. @@ -61,25 +294,29 @@ def transformer_as_linen( quant (Quant): The quantization module or configuration to use. model_mode (str, optional): The operational mode for the model, e.g. training, prefill, or autoregressive. Defaults to `MODEL_MODE_TRAIN`. - name (str, optional): Optional module name for construction. + name (str, optional): Optional module name for Linen/NNX construction. Returns: - nnx_wrappers.ToLinen: An NNX Transformer wrapped as a Linen module. + nnx_wrappers.ToLinen | TransformerLinenPure: + A constructed Transformer model compatible with the specified framework (Linen or NNX). """ - return TransformerLinen( - Transformer, - args=(), - kwargs=nn.FrozenDict( - { - "mesh": mesh, - "config": config, - "quant": quant, - "model_mode": model_mode, - } - ), - metadata_fn=initializers.variable_to_logically_partitioned, - name=name, - ) + if config.enable_nnx: + return TransformerLinen( + Transformer, + args=(), + kwargs=nn.FrozenDict( + { + "mesh": mesh, + "config": config, + "quant": quant, + "model_mode": model_mode, + } + ), + metadata_fn=initializers.variable_to_logically_partitioned, + name=name, + ) + else: + return TransformerLinenPure(config, mesh, quant, model_mode=model_mode, name=name) class TransformerLinen(nnx_wrappers.ToLinen): @@ -137,7 +374,41 @@ def __init__( VisionEncoder(config=cfg, mesh=mesh, rngs=rngs) if getattr(cfg, "use_multimodal", False) else None ) self.audio_encoder = AudioEncoder(config=cfg, mesh=mesh, rngs=rngs) if getattr(cfg, "use_audio", False) else None - self.decoder = NNXDecoder(config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode, rngs=rngs) + if cfg.pure_nnx_decoder: + self.decoder = NNXDecoder(config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode, rngs=rngs) + else: + decoder_linen = Decoder(config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode) + self.decoder = nnx_wrappers.ToNNX(decoder_linen, rngs=rngs) + + batch_size, seq_len = max_utils.get_batch_seq_len_for_mode(config=cfg, model_mode=model_mode) + dummy_decoder_input_tokens = jnp.ones((batch_size, seq_len), dtype=jnp.int32) + dummy_decoder_positions = jnp.ones((batch_size, seq_len), dtype=jnp.int32) + + if self.config.attention in ("vllm_rpa", "vllm_batched_rpa"): + try: + # pylint: disable=import-outside-toplevel + from tpu_inference.layers.common.attention_metadata import AttentionMetadata # pytype: disable=import-error + except ImportError as e: + raise ImportError( + "vLLM RPA attention requires the vllm-tpu package. Please install it with `pip install vllm-tpu`." + ) from e + dummy_attention_metadata = AttentionMetadata( + input_positions=jnp.ones((batch_size * seq_len,), dtype=jnp.int32), + block_tables=jnp.ones((seq_len,), dtype=jnp.int32), + seq_lens=jnp.ones((1), dtype=jnp.int32), + query_start_loc=jnp.ones((2), dtype=jnp.int32), + request_distribution=jnp.ones((3), dtype=jnp.int32), + ) + else: + dummy_attention_metadata = None + + if not cfg.pure_nnx_decoder: + self.decoder.lazy_init( # pyrefly: ignore[missing-attribute] + shared_embedding=self.token_embedder, + decoder_input_tokens=dummy_decoder_input_tokens, + decoder_positions=dummy_decoder_positions, + attention_metadata=dummy_attention_metadata, + ) # If MTP is enabled via config, set up the MTP block. if getattr(self.config, "mtp_num_layers", 0) > 0: @@ -282,27 +553,61 @@ def __call__( bidirectional_mask_video=bidirectional_mask_video, ) - res = self.decoder( - shared_embedding=self.token_embedder, - decoder_input_tokens=decoder_input_tokens, - decoder_input_embeddings=decoder_input_embeddings, - decoder_positions=decoder_positions, - decoder_segment_ids=decoder_segment_ids, - deterministic=not enable_dropout, - model_mode=model_mode, - previous_chunk=previous_chunk, - slot=slot, - multimodal_input=multimodal_input, - kv_caches=kv_caches, - attention_metadata=attention_metadata, - deepstack_visual_embeds=deepstack_visual_embeds, - forced_routed_experts=forced_routed_experts, - ) # pytype: disable=wrong-keyword-args - if isinstance(res, tuple) and len(res) == 4: - logits, hidden_state, kv_caches, expert_indices = res + mutable_collections = [] + if getattr(self.config, "record_internal_nn_metrics", False): + mutable_collections.append("intermediates") + if getattr(self.config, "distill_beta", 0.0) > 0.0 and "intermediates" not in mutable_collections: + mutable_collections.append("intermediates") + if getattr(self.config, "load_balance_loss_weight", 0.0) > 0.0 and "intermediates" not in mutable_collections: + mutable_collections.append("intermediates") + + if self.config.pure_nnx_decoder: + res = self.decoder( + shared_embedding=self.token_embedder, + decoder_input_tokens=decoder_input_tokens, + decoder_input_embeddings=decoder_input_embeddings, + decoder_positions=decoder_positions, + decoder_segment_ids=decoder_segment_ids, + deterministic=not enable_dropout, + model_mode=model_mode, + previous_chunk=previous_chunk, + slot=slot, + multimodal_input=multimodal_input, + kv_caches=kv_caches, + attention_metadata=attention_metadata, + deepstack_visual_embeds=deepstack_visual_embeds, + forced_routed_experts=forced_routed_experts, + ) # pytype: disable=wrong-keyword-args + if isinstance(res, tuple) and len(res) == 4: + logits, hidden_state, kv_caches, expert_indices = res + else: + logits, hidden_state, kv_caches = res + expert_indices = None else: - logits, hidden_state, kv_caches = res - expert_indices = None + if forced_routed_experts is not None: + raise NotImplementedError(_FORCED_ROUTING_NEEDS_PURE_NNX) + + res = self.decoder( + shared_embedding=self.token_embedder, + decoder_input_tokens=decoder_input_tokens, + decoder_input_embeddings=decoder_input_embeddings, + decoder_positions=decoder_positions, + decoder_segment_ids=decoder_segment_ids, + deterministic=not enable_dropout, + model_mode=model_mode, + previous_chunk=previous_chunk, + slot=slot, + multimodal_input=multimodal_input, + kv_caches=kv_caches, + attention_metadata=attention_metadata, + deepstack_visual_embeds=deepstack_visual_embeds, + mutable=mutable_collections, # pyrefly: ignore[unexpected-keyword] + ) # pytype: disable=wrong-keyword-args + if isinstance(res, tuple) and len(res) == 4: + logits, hidden_state, kv_caches, expert_indices = res + else: + logits, hidden_state, kv_caches = res + expert_indices = None # If we are initializing the model AND MTP is enabled, we must create # dummy target tensors. This allows Flax to trace the MTPBlock and create diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 5e83230932..290847a44c 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -21,6 +21,7 @@ import os from typing import Any, cast +from flax import linen as nn from flax import nnx import jax from jax import lax @@ -41,7 +42,7 @@ from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.embeddings import PositionalEmbedding, Qwen3OmniMoeVisionPosEmbedInterpolate -from maxtext.layers.initializers import nd_dense_init +from maxtext.layers.initializers import nd_dense_init, variable_to_logically_partitioned from maxtext.layers.linears import DenseGeneral, MlpBlock from maxtext.layers.moe import RoutedMoE from maxtext.layers.normalizations import Qwen3NextRMSNorm, Qwen3NextRMSNormGated, RMSNorm, l2norm @@ -2689,6 +2690,29 @@ def __call__(self, hidden_states: Array) -> Array: return output +def qwen3omni_visionencoder_as_linen(config: Config, mesh: Mesh) -> nn.Module: + """Convert Qwen3OmniMoeVisionEncoder to Linen module.""" + return nnx_wrappers.to_linen( + Qwen3OmniMoeVisionEncoder, + config=config, + mesh=mesh, + name="Qwen3OmniMoeVisionEncoder_0", + abstract_init=False, + metadata_fn=max_initializers.variable_to_logically_partitioned, + ) + + +def qwen3omni_visionprojector_as_linen(config: Config, mesh: Mesh) -> nn.Module: + """Convert Qwen3OmniMoeVisionProjector to Linen module.""" + return nnx_wrappers.to_linen( + Qwen3OmniMoeVisionProjector, + config=config, + name="Qwen3OmniMoeVisionProjector_0", + abstract_init=False, + metadata_fn=max_initializers.variable_to_logically_partitioned, + ) + + class Qwen3OmniAudioEncoderLayer(nnx.Module): """Transformer encoder layer for audio model.""" @@ -2986,6 +3010,29 @@ def __call__(self, hidden_states: Array) -> Array: return hidden_states +def qwen3omni_audioencoder_as_linen(config: Config, mesh: Mesh): + """Convert AudioEncoder (convs + transformer layers, no projector) to Linen module.""" + return nnx_wrappers.to_linen( + Qwen3OmniAudioEncoder, + config=config, + mesh=mesh, + name="Qwen3OmniAudioEncoder_0", + abstract_init=False, + metadata_fn=variable_to_logically_partitioned, + ) + + +def qwen3omni_audioprojector_as_linen(config: Config, mesh: Mesh): + """Convert AudioProjector to Linen module.""" + return nnx_wrappers.to_linen( + Qwen3OmniAudioProjector, + config=config, + name="Qwen3OmniAudioProjector_0", + abstract_init=False, + metadata_fn=variable_to_logically_partitioned, + ) + + # Vision encoder Linen wrappers Qwen3OmniMoeVisionPatchMergerToLinen = nnx_wrappers.to_linen_class( Qwen3OmniMoeVisionPatchMerger, diff --git a/tests/integration/deepseek_scan_engram_test.py b/tests/integration/deepseek_scan_engram_test.py index d12b7d6321..a9d0b12e92 100644 --- a/tests/integration/deepseek_scan_engram_test.py +++ b/tests/integration/deepseek_scan_engram_test.py @@ -14,19 +14,21 @@ """Unit tests for DeepSeek Engram across scanned decoder layers.""" +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 -# The Linen Decoder this test exercised was removed in PR12 (Delete Linen). -# NNX decoder coverage is in tests/unit/nnx_decoders_test.py. -pytestmark = pytest.mark.skip( - reason="Linen Decoder removed in PR12 (Delete Linen); NNX decoder coverage is in tests/unit/nnx_decoders_test.py" -) - class DummyEmbedding: """Dummy embedding layer for testing.""" @@ -88,10 +90,79 @@ def _test_engram_pattern( base_num_decoder_layers=10, ): """Helper method to test different engram layer patterns.""" - # The Linen Decoder this exercised was removed in PR12 (Delete Linen); - # NNX decoder coverage lives in tests/unit/nnx_decoders_test.py. - del mock_from_pretrained, engram_layers_str, expected_keys, first_num_dense_layers, base_num_decoder_layers - raise unittest.SkipTest("Linen Decoder removed in PR12 (Delete Linen)") + + # Setup mock tokenizer + class MockTokenizer: + """Mock tokenizer for testing.""" + + pad_token_id = 0 + + def __len__(self): + return 128 + + 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}]", + f"first_num_dense_layers={first_num_dense_layers}", + f"base_num_decoder_layers={base_num_decoder_layers}", + f"num_decoder_layers={base_num_decoder_layers}", + ] + ) + + 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 jax.set_mesh(mesh), jax.disable_jit(): + 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 @pytest.mark.tpu_only @patch("transformers.AutoTokenizer.from_pretrained") diff --git a/tests/unit/nnx_decoder_test.py b/tests/unit/nnx_decoder_test.py index daa6e4567d..70d59eb6a9 100644 --- a/tests/unit/nnx_decoder_test.py +++ b/tests/unit/nnx_decoder_test.py @@ -57,6 +57,9 @@ def _model(self): enable_checkpointing=False, scan_layers=False, num_vocab_tiling=2, + pure_nnx=True, + enable_nnx=True, + pure_nnx_decoder=True, ) model = model_creation_utils.from_config(cfg, devices=jax.devices(), model_mode=MODEL_MODE_PREFILL, rngs=nnx.Rngs(0)) return cfg, model diff --git a/tests/unit/router_replay_test.py b/tests/unit/router_replay_test.py index e83f5be8a8..b0c8b52d19 100644 --- a/tests/unit/router_replay_test.py +++ b/tests/unit/router_replay_test.py @@ -263,6 +263,78 @@ def setUp(self): os.environ["SKIP_JAX_PRECOMPILE"] = "1" self.seq_len, self.batch_size, self.top_k = 8, 1, 2 + def _run(self, **overrides): + """Builds a tiny model and applies it with forced routing enabled.""" + cfg = _init_test_cfg( + extra_args=["attention=dot_product"], + **_tiny_qwen35_kwargs( + self.seq_len, + self.batch_size, + 4, + self.top_k, + run_name="test_router_replay_guard", + base_num_decoder_layers=1, + num_decoder_layers=1, + **overrides, + ), + ) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + ids = jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32) + positions = jnp.arange(self.seq_len, dtype=jnp.int32)[None, :] + segmentation = jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32) + forced = jnp.zeros((self.batch_size, self.seq_len, self.top_k), dtype=jnp.int32) + + model = models.transformer_as_linen(config=cfg, mesh=mesh, quant=None, model_mode="train") + params = model.init( + {"params": jax.random.PRNGKey(0), "dropout": jax.random.PRNGKey(1)}, + ids, + positions, + segmentation, + enable_dropout=False, + ) + return model.apply( + params, + ids, + positions, + segmentation, + enable_dropout=False, + forced_routed_experts=forced, + ) + + def test_linen_decoder_rejects_forced_routing(self): + with self.assertRaisesRegex(NotImplementedError, "pure-NNX decoder"): + self._run(scan_layers=False, pure_nnx_decoder=False) + + def test_linen_decoder_without_forced_routing_still_works(self): + """The guard must not break the Linen path when replay is unused.""" + cfg = _init_test_cfg( + extra_args=["attention=dot_product"], + **_tiny_qwen35_kwargs( + self.seq_len, + self.batch_size, + 4, + self.top_k, + run_name="test_router_replay_guard_linen_ok", + base_num_decoder_layers=1, + num_decoder_layers=1, + scan_layers=False, + pure_nnx_decoder=False, + ), + ) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + ids = jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32) + positions = jnp.arange(self.seq_len, dtype=jnp.int32)[None, :] + segmentation = jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32) + model = models.transformer_as_linen(config=cfg, mesh=mesh, quant=None, model_mode="train") + params = model.init( + {"params": jax.random.PRNGKey(0), "dropout": jax.random.PRNGKey(1)}, + ids, + positions, + segmentation, + enable_dropout=False, + ) + self.assertIsNotNone(params) + def test_gemma4_scanned_rejects_forced_routing(self): cfg = _init_test_cfg( extra_args=["attention=dot_product"], diff --git a/tests/unit/tiling_test.py b/tests/unit/tiling_test.py index acf7a5d8f7..89f04d394f 100644 --- a/tests/unit/tiling_test.py +++ b/tests/unit/tiling_test.py @@ -21,22 +21,45 @@ import unittest import pytest +from flax import linen as nn from flax import nnx from flax.linen import partitioning as nn_partitioning import jax import jax.numpy as jnp +from jax.sharding import Mesh from maxtext.configs import pyconfig +from maxtext.common.common_types import Config from maxtext.common.common_types import MODEL_MODE_TRAIN +from maxtext.layers import quantizations +from maxtext.models import models from maxtext.utils import max_utils from maxtext.utils import maxtext_utils from maxtext.utils import maxtext_utils_nnx from maxtext.utils import model_creation_utils -from maxtext.utils.vocabulary_tiling import vocab_tiling_nnx_loss +from maxtext.utils.vocabulary_tiling import vocab_tiling_linen_loss, vocab_tiling_nnx_loss from tests.utils.test_helpers import get_test_config_path +def compute_loss_linen(intermediate_outputs, logits, data, config, model, params, is_train): + """ + A loss function wrapper that deals with both vocab tiling or non-vocab tiling cases + """ + if config.num_vocab_tiling > 1: + hidden_state_key = ("intermediates", "decoder", "hidden_states") + hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0] + total_loss, _ = vocab_tiling_linen_loss(hidden_states, data, config, model, params, is_train) + else: + one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) + xent, _ = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) + xent = nn.with_logical_constraint(xent, ("activation_embed_and_logits_batch", "activation_length")) + # Mask out paddings at the end of each example. + xent = xent * (data["targets_segmentation"] != 0) + total_loss = jnp.sum(xent) + return total_loss + + class LossAndGradientCorrectnessTest(unittest.TestCase): """ Unit tests for verifying loss and gradient correctness of: @@ -48,18 +71,207 @@ def setUp(self): """ Set up common configurations and dummy data for the tests. """ + # vocab_tiling on the Linen path uses transformer_as_linen + model.apply, + # so this class must stay on Linen even when NNX defaults are flipped to + # True. The NNX-side equivalents live in VocabTilingNNXTest below. self.base_config = [ None, get_test_config_path(), "base_emb_dim=32", "vocab_size=128", + "enable_nnx=False", + "pure_nnx=False", + "pure_nnx_decoder=False", ] self.rng = jax.random.PRNGKey(1234) self.batch_size = 1 self.seq_len = 64 + self.dummy_inputs = jnp.ones((self.batch_size, self.seq_len), dtype=jnp.int32) self.rtol = 1e-2 self.atol = 1e-2 + def get_grads(self, cfg: Config, params, data): + """ + Computes and returns the gradients for a given configuration and set of parameters. + """ + quant = quantizations.configure_quantization(cfg) + devices_array = maxtext_utils.create_device_mesh(cfg) + mesh = Mesh(devices_array, cfg.mesh_axes) + model = models.transformer_as_linen(cfg, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) + + @jax.jit + def grad_fn(p, d): + def loss_fn(train_params): + logits, intermediate_outputs = model.apply( + train_params, + decoder_input_tokens=self.dummy_inputs, + decoder_positions=self.dummy_inputs, + mutable=["intermediates"], + ) + return compute_loss_linen(intermediate_outputs, logits, d, cfg, model, train_params, is_train=True) + + return jax.value_and_grad(loss_fn)(p) + + return grad_fn(params, data) + + def assert_pytrees_all_close(self, tree1, tree2, error_message=None): + """Manually assert that two pytrees are all-close.""" + + # Map jnp.allclose to every leaf + # This creates a new pytree of the same structure, but with boolean leaves + leaves_are_close = jax.tree_util.tree_map( + lambda x, y: jnp.allclose(x, y, rtol=self.rtol, atol=self.atol), + tree1, + tree2, + ) + + # Get a flat list of all boolean leaves + all_leaves = jax.tree_util.tree_leaves(leaves_are_close) + + # Check if every single leaf was True + all_are_close = all(all_leaves) + + # Assert the final result + if not all_are_close: + # For a better error, you could find the first False leaf + print("Pytrees are not all-close. Mismatched leaves:") + # This part is more complex, which is why chex is preferred + mismatches = jax.tree_util.tree_map( + lambda x, y, z: "MISMATCH" if not z else "MATCH", tree1, tree2, leaves_are_close + ) + print(mismatches) + raise AssertionError(error_message) + + @pytest.mark.tpu_only + def test_gradient_accumulation(self): + """ + Tests GA loss and gradient correctness. + """ + cfg_non_ga = pyconfig.initialize( + self.base_config, + run_name="non_GA_grad_test", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=4, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + gradient_accumulation_steps=1, + ) + quant_non_ga = quantizations.configure_quantization(cfg_non_ga) + devices_array_non_ga = maxtext_utils.create_device_mesh(cfg_non_ga) + mesh_non_ga = Mesh(devices_array_non_ga, cfg_non_ga.mesh_axes) + model_non_ga = models.transformer_as_linen( + cfg_non_ga, mesh=mesh_non_ga, quant=quant_non_ga, model_mode=MODEL_MODE_TRAIN + ) + + rng_model, rng_targets = jax.random.split(self.rng) + + params = model_non_ga.init( + {"params": rng_model, "dropout": rng_model}, + self.dummy_inputs, + self.dummy_inputs, + ) + + data = { + "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_ga.vocab_size), + "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), + } + + loss_non_ga, grads_non_ga = self.get_grads(cfg_non_ga, params, data) + + cfg_ga = pyconfig.initialize( + self.base_config, + run_name="GA_grad_test", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=1, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + gradient_accumulation_steps=4, + ) + loss_ga, grads_ga = self.get_grads(cfg_ga, params, data) + # Loss correctness test + assert jnp.allclose(loss_non_ga, loss_ga, rtol=self.rtol), "Losses do not match for gradient accumulation test." + + # Gradient correctness test + self.assert_pytrees_all_close( + grads_non_ga, + grads_ga, + "Gradients of embedding table do not match for GA.", + ) + + @pytest.mark.tpu_only + def test_vocab_tiling_gradient_with_z_loss(self): + """ + Tests loss and gradient correctness when z-loss is enabled, comparing + standard computation vs. vocabulary tiling computation. + """ + cfg_non_tiling = pyconfig.initialize( + self.base_config, + run_name="grad_test_z_loss_no_tiling", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + logits_via_embedding=False, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=1, + z_loss_multiplier=1e-4, # Enable z-loss + ) + quant_non_tiling = quantizations.configure_quantization(cfg_non_tiling) + devices_array_non_tiling = maxtext_utils.create_device_mesh(cfg_non_tiling) + mesh_non_tiling = Mesh(devices_array_non_tiling, cfg_non_tiling.mesh_axes) + model_non_tiling = models.transformer_as_linen( + cfg_non_tiling, mesh=mesh_non_tiling, quant=quant_non_tiling, model_mode=MODEL_MODE_TRAIN + ) + + rng_model, rng_targets = jax.random.split(self.rng) + + params = model_non_tiling.init( + {"params": rng_model, "dropout": rng_model}, + self.dummy_inputs, + self.dummy_inputs, + ) + + data = { + "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_tiling.vocab_size), + "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), + } + + loss_non_tiling, grads_non_tiling = self.get_grads(cfg_non_tiling, params, data) + + cfg_tiling = pyconfig.initialize( + self.base_config, + run_name="grad_test_z_loss_with_tiling", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + logits_via_embedding=False, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=4, + z_loss_multiplier=1e-4, # Enable z-loss + ) + loss_tiling, grads_tiling = self.get_grads(cfg_tiling, params, data) + + # Loss correctness test + assert jnp.allclose(loss_non_tiling, loss_tiling, rtol=self.rtol), "Losses do not match when z-loss is enabled." + + # Gradient correctness test + self.assert_pytrees_all_close( + grads_non_tiling, + grads_tiling, + "Gradients do not match for vocab tiling when z-loss is enabled.", + ) + @pytest.mark.tpu_only def test_vocab_tiling_nnx_loss(self): """ @@ -79,6 +291,8 @@ def test_vocab_tiling_nnx_loss(self): matmul_precision="high", num_vocab_tiling=4, z_loss_multiplier=1e-4, + enable_nnx=True, + pure_nnx=True, ) rng_model, rng_hidden, rng_targets = jax.random.split(self.rng, 3) rngs = maxtext_utils_nnx.create_nnx_rngs(cfg, rng_key=rng_model) @@ -103,6 +317,372 @@ def test_vocab_tiling_nnx_loss(self): xent_sum_tiled, xent_sum_ref, rtol=self.rtol, atol=self.atol ), f"NNX vocab tiling loss {xent_sum_tiled} does not match non-tiled reference {xent_sum_ref}." + @pytest.mark.cpu_only + def test_linen_vocab_tiling_hidden_gradient_respects_outer_loss_scale(self): + """The tiled loss's hidden-state gradient obeys the chain rule.""" + cfg = pyconfig.initialize( + self.base_config, + run_name="linen_vocab_tiling_outer_loss_scale", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + logits_via_embedding=False, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=4, + ) + quant = quantizations.configure_quantization(cfg) + devices_array = maxtext_utils.create_device_mesh(cfg) + mesh = Mesh(devices_array, cfg.mesh_axes) + model = models.transformer_as_linen(cfg, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) + rng_model, rng_hidden, rng_targets = jax.random.split(self.rng, 3) + params = model.init( + {"params": rng_model, "dropout": rng_model}, + self.dummy_inputs, + self.dummy_inputs, + ) + hidden_states = jax.random.normal(rng_hidden, (self.batch_size, self.seq_len, cfg.emb_dim), dtype=jnp.float32) + data = { + "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg.vocab_size), + "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), + } + + def scaled_loss(h, outer_loss_scale): + total_loss, _ = vocab_tiling_linen_loss(h, data, cfg, model, params, is_train=True) + return outer_loss_scale * total_loss + + grad_with_scale = jax.jit(jax.grad(scaled_loss, argnums=0)) + base_grad = grad_with_scale(hidden_states, 1.0) + doubled_grad = grad_with_scale(hidden_states, 2.0) + + assert jnp.all(jnp.isfinite(base_grad)) + assert jnp.any(base_grad != 0) + assert jnp.allclose(doubled_grad, 2.0 * base_grad, rtol=1e-5, atol=1e-6) + + @pytest.mark.tpu_only + def test_vocab_tiling_gradient_non_tied_embedding(self): + """ + Tests loss and gradient correctness for a model with non-tied embeddings (FSDP). + """ + cfg_non_tiling = pyconfig.initialize( + self.base_config, + run_name="grad_test_non_tied_no_tiling", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + logits_via_embedding=False, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=1, + ) + quant_non_tiling = quantizations.configure_quantization(cfg_non_tiling) + devices_array_non_tiling = maxtext_utils.create_device_mesh(cfg_non_tiling) + mesh_non_tiling = Mesh(devices_array_non_tiling, cfg_non_tiling.mesh_axes) + model_non_tiling = models.transformer_as_linen( + cfg_non_tiling, mesh=mesh_non_tiling, quant=quant_non_tiling, model_mode=MODEL_MODE_TRAIN + ) + + rng_model, rng_targets = jax.random.split(self.rng) + + params = model_non_tiling.init( + {"params": rng_model, "dropout": rng_model}, + self.dummy_inputs, + self.dummy_inputs, + ) + + data = { + "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_tiling.vocab_size), + "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), + } + + loss_non_tiling, grads_non_tiling = self.get_grads(cfg_non_tiling, params, data) + + cfg_tiling = pyconfig.initialize( + self.base_config, + run_name="value_and_grad_test_non_tied_with_tiling", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + logits_via_embedding=False, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=4, + ) + loss_tiling, grads_tiling = self.get_grads(cfg_tiling, params, data) + # Loss correctness test + assert jnp.allclose(loss_non_tiling, loss_tiling, rtol=self.rtol), "Losses do not match for non-tied embeddings." + + # Gradient correctness test + self.assert_pytrees_all_close( + grads_non_tiling, + grads_tiling, + "Gradients of embedding table do not match for non-tied embeddings.", + ) + + @pytest.mark.tpu_only + @pytest.mark.scheduled_only + def test_vocab_tiling_gradient_tied_embedding(self): + """ + Tests loss and gradient correctness for a model with tied embeddings (FSDP). + """ + cfg_non_tiling = pyconfig.initialize( + self.base_config, + run_name="value_and_grad_test_tied_no_tiling", + enable_checkpointing=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + logits_via_embedding=True, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=1, + ) + + quant_non_tiling = quantizations.configure_quantization(cfg_non_tiling) + devices_array_non_tiling = maxtext_utils.create_device_mesh(cfg_non_tiling) + mesh_non_tiling = Mesh(devices_array_non_tiling, cfg_non_tiling.mesh_axes) + model_non_tiling = models.transformer_as_linen( + cfg_non_tiling, mesh=mesh_non_tiling, quant=quant_non_tiling, model_mode=MODEL_MODE_TRAIN + ) + + rng_model, rng_targets = jax.random.split(self.rng) + + params = model_non_tiling.init( + {"params": rng_model, "dropout": rng_model}, + self.dummy_inputs, + self.dummy_inputs, + ) + + data = { + "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_tiling.vocab_size), + "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), + } + + loss_non_tiling, grads_non_tiling = self.get_grads(cfg_non_tiling, params, data) + + cfg_tiling = pyconfig.initialize( + self.base_config, + run_name="grad_test_tied_with_tiling", + enable_checkpointing=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + logits_via_embedding=True, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=4, + ) + loss_tiling, grads_tiling = self.get_grads(cfg_tiling, params, data) + + assert jnp.allclose(loss_non_tiling, loss_tiling, rtol=self.rtol), "Losses do not match for tied embeddings." + + self.assert_pytrees_all_close( + grads_non_tiling, grads_tiling, "Gradients of embedding table do not match for tied embeddings." + ) + + @pytest.mark.tpu_only + @pytest.mark.scheduled_only + def test_vocab_tiling_gradient_data_parallelism(self): + """ + Tests loss and gradient correctness for data parallelism sharding. + """ + cfg_non_tiling = pyconfig.initialize( + self.base_config, + run_name="value_and_grad_test_dp_non_tiling", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + ici_data_parallelism=4, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=1, + ) + quant_non_tiling = quantizations.configure_quantization(cfg_non_tiling) + devices_array_non_tiling = maxtext_utils.create_device_mesh(cfg_non_tiling) + mesh_non_tiling = Mesh(devices_array_non_tiling, cfg_non_tiling.mesh_axes) + model_non_tiling = models.transformer_as_linen( + cfg_non_tiling, mesh=mesh_non_tiling, quant=quant_non_tiling, model_mode=MODEL_MODE_TRAIN + ) + + rng_model, rng_targets = jax.random.split(self.rng) + + params = model_non_tiling.init( + {"params": rng_model, "dropout": rng_model}, + self.dummy_inputs, + self.dummy_inputs, + ) + + data = { + "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_tiling.vocab_size), + "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), + } + + loss_non_tiling, grads_non_tiling = self.get_grads(cfg_non_tiling, params, data) + + cfg_tiling = pyconfig.initialize( + self.base_config, + run_name="value_and_grad_test_dp_tiling", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + logits_via_embedding=False, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + ici_data_parallelism=4, + num_vocab_tiling=4, + ) + loss_tiling, grads_tiling = self.get_grads(cfg_tiling, params, data) + # Loss correctness test + assert jnp.allclose(loss_non_tiling, loss_tiling, rtol=self.rtol), "Losses do not match for data parallelism." + + # Gradient correctness test + self.assert_pytrees_all_close( + grads_non_tiling, grads_tiling, "Gradients of embedding table do not match for data parallelism." + ) + + @pytest.mark.tpu_only + @pytest.mark.scheduled_only + def test_vocab_tiling_gradient_tensor_parallelism(self): + """ + Tests loss and gradient correctness for tensor parallelism sharding. + """ + cfg_non_tiling = pyconfig.initialize( + self.base_config, + run_name="value_and_grad_test_tp_non_tiling", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + ici_tensor_parallelism=4, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=1, + ) + quant_non_tiling = quantizations.configure_quantization(cfg_non_tiling) + devices_array_non_tiling = maxtext_utils.create_device_mesh(cfg_non_tiling) + mesh_non_tiling = Mesh(devices_array_non_tiling, cfg_non_tiling.mesh_axes) + model_non_tiling = models.transformer_as_linen( + cfg_non_tiling, mesh=mesh_non_tiling, quant=quant_non_tiling, model_mode=MODEL_MODE_TRAIN + ) + + rng_model, rng_targets = jax.random.split(self.rng) + + params = model_non_tiling.init( + {"params": rng_model, "dropout": rng_model}, + self.dummy_inputs, + self.dummy_inputs, + ) + + data = { + "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_tiling.vocab_size), + "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), + } + + loss_non_tiling, grads_non_tiling = self.get_grads(cfg_non_tiling, params, data) + + cfg_tiling = pyconfig.initialize( + self.base_config, + run_name="value_and_grad_test_tp_tiling", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + logits_via_embedding=False, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + ici_tensor_parallelism=4, + num_vocab_tiling=4, + ) + loss_tiling, grads_tiling = self.get_grads(cfg_tiling, params, data) + # Loss correctness test + assert jnp.allclose(loss_non_tiling, loss_tiling, rtol=self.rtol), "Losses do not match for tensor parallelism." + + # Gradient correctness test + self.assert_pytrees_all_close( + grads_non_tiling, grads_tiling, "Gradients of embedding table do not match for tensor parallelism." + ) + + @pytest.mark.tpu_only + @pytest.mark.scheduled_only + def test_vocab_tiling_gradient_context_parallelism(self): + """ + Tests loss and gradient correctness for context parallelism sharding. + """ + cfg_non_tiling = pyconfig.initialize( + self.base_config, + run_name="value_and_grad_test_cp_non_tiling", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + ici_context_parallelism=4, + base_num_decoder_layers=0, + dataset_type="synthetic", + packing=False, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=1, + ) + quant_non_tiling = quantizations.configure_quantization(cfg_non_tiling) + devices_array_non_tiling = maxtext_utils.create_device_mesh(cfg_non_tiling) + mesh_non_tiling = Mesh(devices_array_non_tiling, cfg_non_tiling.mesh_axes) + model_non_tiling = models.transformer_as_linen( + cfg_non_tiling, mesh=mesh_non_tiling, quant=quant_non_tiling, model_mode=MODEL_MODE_TRAIN + ) + + rng_model, rng_targets = jax.random.split(self.rng) + + params = model_non_tiling.init( + {"params": rng_model, "dropout": rng_model}, + self.dummy_inputs, + self.dummy_inputs, + ) + + data = { + "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_tiling.vocab_size), + "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), + } + + loss_non_tiling, grads_non_tiling = self.get_grads(cfg_non_tiling, params, data) + + cfg_tiling = pyconfig.initialize( + self.base_config, + run_name="value_and_grad_test_cp_tiling", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + logits_via_embedding=False, + base_num_decoder_layers=0, + ici_context_parallelism=4, + dataset_type="synthetic", + packing=False, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=4, + ) + loss_tiling, grads_tiling = self.get_grads(cfg_tiling, params, data) + + # Loss correctness test + assert jnp.allclose(loss_non_tiling, loss_tiling, rtol=self.rtol), "Losses do not match for context parallelism." + + # Gradient correctness test + self.assert_pytrees_all_close( + grads_non_tiling, grads_tiling, "Gradients of embedding table do not match for context parallelism." + ) + class VocabTilingNNXTest(unittest.TestCase): """Loss + gradient parity for the NNX vocab-tiling `custom_vjp` path. @@ -151,6 +731,9 @@ def _build_cfg_and_model( z_loss_multiplier=z_loss_multiplier, ici_context_parallelism=ici_context_parallelism, vocab_tiling_ag_once=vocab_tiling_ag_once, + pure_nnx=True, + enable_nnx=True, + pure_nnx_decoder=True, ) mesh = maxtext_utils.get_mesh_from_config(cfg) rngs = maxtext_utils_nnx.create_nnx_rngs(cfg)