diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index 13dd34a1d4..05ee1e2e38 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -63,9 +63,8 @@ DEFAULT_MASK_VALUE, ) -from maxtext.layers import nnx_wrappers from maxtext.layers.attentions import Attention -from maxtext.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned +from maxtext.layers.initializers import nd_dense_init, NdInitializer from maxtext.layers.linears import DenseGeneral from maxtext.layers.normalizations import RMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant @@ -452,141 +451,6 @@ 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 21a4bc6900..01aabb6ad5 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -71,7 +71,6 @@ 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 @@ -400,100 +399,6 @@ 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 fa8dd1f11b..94c4667fb7 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -91,16 +91,6 @@ 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 a9bacf692f..d8465b5e8f 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -13,263 +13,14 @@ # 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. @@ -293,1680 +44,3 @@ 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 309f34e3f4..4c80f63848 100644 --- a/src/maxtext/layers/embeddings.py +++ b/src/maxtext/layers/embeddings.py @@ -26,8 +26,7 @@ from flax import nnx from maxtext.common.common_types import ShardMode, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, Array, Config, DType -from maxtext.layers import nnx_wrappers -from maxtext.layers.initializers import Initializer, default_embed_init, variable_to_logically_partitioned +from maxtext.layers.initializers import Initializer, default_embed_init 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 @@ -43,52 +42,6 @@ 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.""" @@ -240,39 +193,6 @@ 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.""" @@ -375,82 +295,6 @@ 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.""" @@ -727,61 +571,6 @@ 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. @@ -1026,31 +815,6 @@ 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. @@ -1136,43 +900,6 @@ 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. @@ -1543,47 +1270,6 @@ 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. @@ -1952,40 +1638,6 @@ 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 e5e32d044a..cecc68a26c 100644 --- a/src/maxtext/layers/encoders.py +++ b/src/maxtext/layers/encoders.py @@ -19,8 +19,6 @@ 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 @@ -246,35 +244,3 @@ 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 fce3dfb940..9ab52e5fd3 100644 --- a/src/maxtext/layers/multi_token_prediction.py +++ b/src/maxtext/layers/multi_token_prediction.py @@ -17,15 +17,12 @@ 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 @@ -563,37 +560,3 @@ 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 caedc633a2..94c37ae304 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 initializers, linears, mhc, moe, normalizations, quantizations +from maxtext.layers import 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,25 +2725,3 @@ 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 8d2733ee3b..106a48e088 100644 --- a/src/maxtext/models/gemma3.py +++ b/src/maxtext/models/gemma3.py @@ -31,7 +31,6 @@ 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 @@ -606,21 +605,6 @@ 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. @@ -653,11 +637,6 @@ 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""" @@ -744,19 +723,3 @@ 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 72cd841f81..f673488608 100644 --- a/src/maxtext/models/gemma4_vision.py +++ b/src/maxtext/models/gemma4_vision.py @@ -18,15 +18,12 @@ 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 @@ -636,15 +633,3 @@ 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 415aff7ce3..ebbc98d764 100644 --- a/src/maxtext/models/llama4.py +++ b/src/maxtext/models/llama4.py @@ -253,17 +253,6 @@ 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). @@ -812,14 +801,3 @@ 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 90c3d17796..19408299fd 100644 --- a/src/maxtext/models/models.py +++ b/src/maxtext/models/models.py @@ -29,246 +29,16 @@ from maxtext.layers.nnx_decoders import NNXDecoder from maxtext.layers import initializers from maxtext.layers import nnx_wrappers -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.embeddings import Embed +from maxtext.layers.encoders import AudioEncoder, VisionEncoder +from maxtext.layers.multi_token_prediction import MultiTokenPredictionBlock 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, @@ -277,16 +47,13 @@ def transformer_as_linen( model_mode: str = MODEL_MODE_TRAIN, *, name: str | None = None, -) -> 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. +) -> nnx_wrappers.ToLinen: + """Constructs an NNX Transformer wrapped as a Linen module. - 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`). + 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`. Args: config (Config): The configuration object specifying model hyperparameters and options. @@ -294,29 +61,25 @@ 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 Linen/NNX construction. + name (str, optional): Optional module name for construction. Returns: - nnx_wrappers.ToLinen | TransformerLinenPure: - A constructed Transformer model compatible with the specified framework (Linen or NNX). + nnx_wrappers.ToLinen: An NNX Transformer wrapped as a Linen module. """ - 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) + 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, + ) class TransformerLinen(nnx_wrappers.ToLinen): @@ -374,41 +137,7 @@ 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 - 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, - ) + self.decoder = NNXDecoder(config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode, rngs=rngs) # If MTP is enabled via config, set up the MTP block. if getattr(self.config, "mtp_num_layers", 0) > 0: @@ -553,61 +282,27 @@ def __call__( bidirectional_mask_video=bidirectional_mask_video, ) - 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 + 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: - 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 + 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 290847a44c..5e83230932 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -21,7 +21,6 @@ import os from typing import Any, cast -from flax import linen as nn from flax import nnx import jax from jax import lax @@ -42,7 +41,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, variable_to_logically_partitioned +from maxtext.layers.initializers import nd_dense_init from maxtext.layers.linears import DenseGeneral, MlpBlock from maxtext.layers.moe import RoutedMoE from maxtext.layers.normalizations import Qwen3NextRMSNorm, Qwen3NextRMSNormGated, RMSNorm, l2norm @@ -2690,29 +2689,6 @@ 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.""" @@ -3010,29 +2986,6 @@ 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 a9d0b12e92..d12b7d6321 100644 --- a/tests/integration/deepseek_scan_engram_test.py +++ b/tests/integration/deepseek_scan_engram_test.py @@ -14,21 +14,19 @@ """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.""" @@ -90,79 +88,10 @@ def _test_engram_pattern( base_num_decoder_layers=10, ): """Helper method to test different engram layer patterns.""" - - # Setup mock tokenizer - class MockTokenizer: - """Mock tokenizer for testing.""" - - pad_token_id = 0 - - def __len__(self): - return 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 + # 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)") @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 70d59eb6a9..daa6e4567d 100644 --- a/tests/unit/nnx_decoder_test.py +++ b/tests/unit/nnx_decoder_test.py @@ -57,9 +57,6 @@ 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 b0c8b52d19..e83f5be8a8 100644 --- a/tests/unit/router_replay_test.py +++ b/tests/unit/router_replay_test.py @@ -263,78 +263,6 @@ 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 89f04d394f..acf7a5d8f7 100644 --- a/tests/unit/tiling_test.py +++ b/tests/unit/tiling_test.py @@ -21,45 +21,22 @@ 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_linen_loss, vocab_tiling_nnx_loss +from maxtext.utils.vocabulary_tiling import 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: @@ -71,207 +48,18 @@ 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): """ @@ -291,8 +79,6 @@ 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) @@ -317,372 +103,6 @@ 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. @@ -731,9 +151,6 @@ 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)