From fd0e96007b69030f7cf7fb9d2154255711ed68b2 Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Mon, 10 Aug 2026 13:07:02 +0000 Subject: [PATCH 01/21] feat(glm5.2): add native cross-layer IndexShare and checkpoint conversion support --- .../utils/hf_model_configs.py | 9 ++ .../utils/param_mapping.py | 2 + src/maxtext/configs/base.yml | 5 + src/maxtext/configs/models/glm5.2-744b.yml | 67 ++++++++++++ src/maxtext/configs/types.py | 13 +++ src/maxtext/layers/attention_mla.py | 67 +++++++++--- src/maxtext/layers/nnx_decoders.py | 28 ++++- src/maxtext/models/deepseek.py | 50 +++++++-- src/maxtext/utils/globals.py | 1 + src/maxtext/utils/index_share_utils.py | 103 ++++++++++++++++++ tests/unit/glm52_indexshare_test.py | 60 ++++++++++ tests/unit/index_share_utils_test.py | 68 ++++++++++++ 12 files changed, 441 insertions(+), 32 deletions(-) create mode 100644 src/maxtext/configs/models/glm5.2-744b.yml create mode 100644 src/maxtext/utils/index_share_utils.py create mode 100644 tests/unit/glm52_indexshare_test.py create mode 100644 tests/unit/index_share_utils_test.py diff --git a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py index c39f0512a0..d493efb67f 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py @@ -1892,10 +1892,19 @@ def __init__(self, **kwargs): } glm5_1_744b_config = transformers.DeepseekV3Config(**glm5_1_744b_dict) +glm5_2_744b_dict = { + "architectures": ["GlmMoeDsaForCausalLM"], + "num_hidden_layers": 78, + "first_k_dense_replace": 3, + "n_routed_experts": 256, +} +glm5_2_744b_config = transformers.DeepseekV3Config(**glm5_2_744b_dict) + # {maxtext model name: hf model config} HF_MODEL_CONFIGS = { "glm5.1-744b": glm5_1_744b_config, + "glm5.2-744b": glm5_2_744b_config, "gemma2-2b": gemma2_2b_config, "gemma2-9b": gemma2_9b_config, "gemma2-27b": gemma2_27b_config, diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index b1c2521509..7ac396350f 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -4422,6 +4422,7 @@ def mhc_concat_scale(input_tensors, target_shape=None): "deepseek3-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, "deepseek3.2-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, "glm5.1-744b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, + "glm5.2-744b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, "deepseek4-284b": DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_MAPPING, "gpt-oss-20b": GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING, "gpt-oss-120b": GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -4477,6 +4478,7 @@ def mhc_concat_scale(input_tensors, target_shape=None): "deepseek3-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN, "deepseek3.2-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN, "glm5.1-744b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN, + "glm5.2-744b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN, "deepseek4-tiny": DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN, "deepseek4-284b": DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN, "gpt-oss-20b": GPT_OSS_TO_HF_PARAM_HOOK_FN, diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 7a725fc4ca..915bbfb26f 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -442,6 +442,11 @@ indexer_sparse_training: false # Multiplier for the indexer KL divergence loss indexer_loss_scaling_factor: 0.0 +# GLM-5.2 Cross-Layer IndexCache / IndexShare +use_index_share: false +index_share_pattern: "FSSS" +prune_shared_indexers: true + # MLA parameters q_lora_rank: 0 kv_lora_rank: 512 diff --git a/src/maxtext/configs/models/glm5.2-744b.yml b/src/maxtext/configs/models/glm5.2-744b.yml new file mode 100644 index 0000000000..0558778876 --- /dev/null +++ b/src/maxtext/configs/models/glm5.2-744b.yml @@ -0,0 +1,67 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# model config for GLM-5.2 - 744B (Mixture of Experts with Cross-Layer IndexShare) + +base_emb_dim: 6144 +base_num_query_heads: 64 +base_num_kv_heads: 64 +base_mlp_dim: 12288 +base_moe_mlp_dim: 2048 + +base_num_decoder_layers: 78 + +first_num_dense_layers: 3 +mlp_activations: ["silu","linear"] +vocab_size: 154880 +enable_dropout: false +logits_via_embedding: false +normalization_layer_epsilon: 1.0e-5 +num_experts: 256 +num_experts_per_tok: 8 +shared_experts: 1 +routed_scaling_factor: 2.5 +routed_score_func: "sigmoid" +routed_bias: true +norm_topk_prob: true +decoder_block: "deepseek" +dtype: "bfloat16" +weight_dtype: "bfloat16" + +# Multi-head Latent Attention (MLA) +attention_type: "mla" +attention: "dot_product" +q_lora_rank: 2048 +kv_lora_rank: 512 +qk_nope_head_dim: 192 +qk_rope_head_dim: 64 +v_head_dim: 256 + +# RoPE +mscale: 1.0 +rope_type: "default" +rope_max_timescale: 1000000 # "rope_theta": 1000000 +max_position_embeddings: 202752 +rope_interleave: true + +# Indexer for Dynamic Sparse Attention (DSA) +use_indexer: true +indexer_n_heads: 32 +indexer_head_dim: 128 +indexer_topk: 2048 + +# GLM-5.2 Cross-Layer IndexCache / IndexShare +use_index_share: true +index_share_pattern: "FSSS" +prune_shared_indexers: true diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index da72b1aa1d..faa0c13c1f 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -718,6 +718,17 @@ class AttentionIndexer(BaseModel): " during ties." ), ) + use_index_share: bool = Field( + False, description="Whether to enable cross-layer index cache sharing (GLM-5.2 IndexShare)." + ) + index_share_pattern: str = Field( + "FSSS", + description="Cross-layer pattern string for IndexShare (e.g. 'FSSS', 'F,S,S,S', 'FSFSS...').", + ) + prune_shared_indexers: bool = Field( + True, + description="Whether to prune indexer parameters on Shared (S) layers when use_index_share is enabled.", + ) class Llama4Attention(BaseModel): @@ -3497,6 +3508,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "when indexer loss is enabled (`indexer_loss_scaling_factor > 0.0`); otherwise the indexer " "short-circuits to select all tokens and no indexer loss is produced." ) + if not self.use_indexer and self.use_index_share: + raise ValueError("`use_index_share=True` requires `use_indexer=True`.") if not self.use_indexer and self.indexer_cutoff_threshold != RematLocation.REMAT: raise ValueError( f"Setting `indexer_cutoff_threshold='{self.indexer_cutoff_threshold}'` is only valid when " diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index 3967ee1234..65565c5549 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -512,6 +512,8 @@ def mla_as_linen( mscale: float = 1.0, # scaling factor for softmax rope_factor: float = 40.0, # rotary embedding factor name: str | None = None, + is_shared_layer: bool = False, + served_group_size: int = 1, ): """A factory function to create an MLA as a Linen module. @@ -578,6 +580,8 @@ def mla_as_linen( mscale=mscale, rope_factor=rope_factor, name=name, + is_shared_layer=is_shared_layer, + served_group_size=served_group_size, metadata_fn=variable_to_logically_partitioned, abstract_init=False, ) @@ -650,6 +654,8 @@ def __init__( mscale: float = 1.0, # scaling factor for softmax rope_factor: float = 40.0, # rotary embedding factor name: str | None = None, + is_shared_layer: bool = False, + served_group_size: int = 1, rngs: Optional[nnx.Rngs] = None, ): """Initializes the MLA module. @@ -729,7 +735,14 @@ def __init__( # Initialize Indexer self.use_indexer = config.use_indexer - if self.use_indexer: + self.is_shared_layer = is_shared_layer + self.served_group_size = served_group_size + is_pruned = ( + getattr(config, "use_index_share", False) + and getattr(config, "prune_shared_indexers", True) + and self.is_shared_layer + ) + if self.use_indexer and not is_pruned: # Need two versions of rope. # MLA applies yarn with interleave layout. # Indexer applies yarn with concatenate layout. @@ -1240,7 +1253,8 @@ def __call__( rope_kwargs: dict | None = None, kv_cache: Optional[Array] = None, attention_metadata: Optional[dict[str, Any]] = None, - ) -> tuple[Array, Optional[Array]]: + cached_indexer_state: Optional[Any] = None, + ) -> tuple[Array, Optional[Array]] | tuple[Array, Optional[Array], Optional[Any]]: """Forward pass for MLA, reusing `AttentionOp` for the actual attention. Args: @@ -1255,10 +1269,10 @@ def __call__( bidirectional_mask: A mask for bidirectional attention, used in multimodal models. kv_cache: Optional key-value cache used when serving models with vLLM. attention_metadata: Optional attention-related metadata used when serving models with vLLM. + cached_indexer_state: Optional tuple (indexer_mask, topk_indices, indexer_score) from donor F-layer. Returns: - A tensor of shape [batch, length, embed_dim] containing the - MLA-attended outputs. + A tensor of shape [batch, length, embed_dim] containing the MLA-attended outputs. """ if model_mode == MODEL_MODE_PREFILL: inputs_q = self._maybe_shard_with_logical(inputs_q, self.prefill_input_axis_names) @@ -1284,6 +1298,7 @@ def __call__( # Indexer Logic indexer_mask = None + new_indexer_state = None if self.use_indexer: # generate mask: with 0 and large negative, [b, 1, 1, q_len, kv_len] -> [b, q_len, kv_len] attention_mask = self.attention_op.generate_attention_mask( @@ -1291,20 +1306,34 @@ def __call__( ) if attention_mask is not None: attention_mask = attention_mask.squeeze(axis=(1, 2)) - # apply indexer, indexer_mask [b, q_len, kv_len] - indexer_mask, _, indexer_score = self.indexer( - inputs_q=inputs_q, - low_rank_q=low_rank_q, - inputs_kv=inputs_kv, - inputs_positions=inputs_positions, - attention_mask=attention_mask, - decoder_segment_ids=decoder_segment_ids, - previous_chunk=previous_chunk, - kv_cache=self.IndexerKVCache_0, - model_mode=model_mode, - ) - if indexer_mask is not None and self.config.indexer_loss_scaling_factor > 0.0: + is_shared = getattr(self.config, "use_index_share", False) and self.is_shared_layer + if self.indexer is not None and not is_shared: + # Full (F) layer: run indexer forward pass + indexer_mask, topk_indices, indexer_score = self.indexer( + inputs_q=inputs_q, + low_rank_q=low_rank_q, + inputs_kv=inputs_kv, + inputs_positions=inputs_positions, + attention_mask=attention_mask, + decoder_segment_ids=decoder_segment_ids, + previous_chunk=previous_chunk, + kv_cache=self.IndexerKVCache_0, + model_mode=model_mode, + ) + new_indexer_state = (indexer_mask, topk_indices, indexer_score) + elif cached_indexer_state is not None: + # Shared (S) layer: inherit cached indexer state from donor F layer + indexer_mask, topk_indices, indexer_score = cached_indexer_state + new_indexer_state = cached_indexer_state + else: + indexer_mask, topk_indices, indexer_score = None, None, None + + if indexer_mask is not None and self.config.indexer_loss_scaling_factor > 0.0 and indexer_score is not None: + loss_scale = self.config.indexer_loss_scaling_factor + if getattr(self.config, "use_index_share", False) and self.served_group_size > 1: + loss_scale = loss_scale / float(self.served_group_size) + indexer_loss = self.calculate_indexer_loss( indexer_score=indexer_score, query=query, @@ -1312,7 +1341,7 @@ def __call__( attention_mask=attention_mask, indexer_mask=indexer_mask, sparse_loss=self.config.indexer_sparse_training, - scaling_factor=self.config.indexer_loss_scaling_factor, + scaling_factor=loss_scale, ) self.indexer_loss = nnx.Intermediate(indexer_loss) @@ -1337,4 +1366,6 @@ def __call__( out_sharding = create_sharding(self.mesh, out_logical_name) out = self.out_projection(out, out_sharding=out_sharding) out = checkpoint_name(out, "out_proj") + if getattr(self.config, "use_index_share", False): + return out, kv_cache, new_indexer_state return out, kv_cache diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index ddcffad39a..c9832c7999 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -755,9 +755,10 @@ def _init_sequential_deepseek(self, decoder_block_classes, rngs): config = self.config dense_cls, moe_cls = decoder_block_classes for i in range(config.first_num_dense_layers): - self._create_and_register_layer(dense_cls, rngs, "dense_layers", i) + self._create_and_register_layer(dense_cls, rngs, "dense_layers", i, layer_idx=i) for i in range(config.num_decoder_layers - config.first_num_dense_layers): - self._create_and_register_layer(moe_cls, rngs, "moe_layers", i) + global_idx = config.first_num_dense_layers + i + self._create_and_register_layer(moe_cls, rngs, "moe_layers", i, layer_idx=global_idx) def _init_sequential_generic(self, decoder_block_classes, rngs): """Initializes sequential generic decoder layers with per-architecture layer_kwargs.""" @@ -1893,17 +1894,27 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in): state_in, ) merged_layer = nnx.merge(graphdef_in, state_in) - out_y, out_kv = merged_layer(y_in, *layer_args, kv_cache=kv_in, **layer_kwargs) + out = merged_layer(y_in, *layer_args, kv_cache=kv_in, **layer_kwargs) + if getattr(cfg, "use_index_share", False): + out_y, out_kv, out_indexer_cache = out + else: + out_y, out_kv = out + out_indexer_cache = None state_out = nnx.state(merged_layer) if dynamic_graph_init: new_graphdef, _, _ = nnx.split(merged_layer, nnx.Param, ...) + if getattr(cfg, "use_index_share", False): + return out_y, out_kv, out_indexer_cache, state_out, new_graphdef return out_y, out_kv, state_out, new_graphdef else: + if getattr(cfg, "use_index_share", False): + return out_y, out_kv, out_indexer_cache, state_out, graphdef_in return out_y, out_kv, state_out, graphdef_in checkpointed_fn = jax.checkpoint(pure_layer_fn, policy=policy, prevent_cse=prevent_cse) + cached_indexer_state = None for lyr in range(cfg.num_decoder_layers): if self.is_deepseek: if lyr < cfg.first_num_dense_layers: @@ -1942,11 +1953,18 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in): ) if input_tokens is not None: layer_kwargs["decoder_input_tokens"] = input_tokens + if getattr(cfg, "use_index_share", False): + layer_kwargs["cached_indexer_state"] = cached_indexer_state if cfg.remat_policy != "none": - y, kv_cache, new_state, new_graphdef = checkpointed_fn(graphdef, state, y, kv_cache) + res = checkpointed_fn(graphdef, state, y, kv_cache) + else: + res = pure_layer_fn(graphdef, state, y, kv_cache) + + if getattr(cfg, "use_index_share", False): + y, kv_cache, cached_indexer_state, new_state, new_graphdef = res else: - y, kv_cache, new_state, new_graphdef = pure_layer_fn(graphdef, state, y, kv_cache) + y, kv_cache, new_state, new_graphdef = res if dynamic_graph_init: new_layer = nnx.merge(new_graphdef, new_state) diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py index 0ad8978e7f..c66274006a 100644 --- a/src/maxtext/models/deepseek.py +++ b/src/maxtext/models/deepseek.py @@ -76,6 +76,16 @@ def __init__( self.layer_idx = layer_idx self.is_engram_enabled = config.engram_layers and layer_idx in config.engram_layers + self.is_index_share_enabled = getattr(config, "use_index_share", False) + self.is_shared_layer = False + self.served_group_size = 1 + if self.is_index_share_enabled and layer_idx >= 0: + from maxtext.utils import index_share_utils + + pattern = index_share_utils.parse_index_share_pattern(config.index_share_pattern, config.num_decoder_layers) + self.is_shared_layer = index_share_utils.is_shared_layer(layer_idx, pattern) + self.served_group_size = index_share_utils.get_served_group_sizes(pattern)[layer_idx] + batch_size, sequence_length = max_utils.get_batch_seq_len_for_mode(self.config, self.model_mode) self.dummy_inputs_shape = (batch_size, sequence_length, self.config.emb_dim) @@ -171,6 +181,8 @@ def __init__( model_mode=model_mode, rngs=rngs, attn_logits_soft_cap=self.config.attn_logits_soft_cap, + is_shared_layer=self.is_shared_layer, + served_group_size=self.served_group_size, ) self.dropout = Dropout(rate=self.config.dropout_rate, broadcast_dims=(-2,), rngs=self.rngs) @@ -214,9 +226,10 @@ def attention_op( model_mode, previous_chunk=None, slot: None | int = None, + cached_indexer_state=None, ): """Executes the attention layer.""" - attention_result, _ = self.self_attention( + attn_out = self.self_attention( x, x, decoder_positions, @@ -226,8 +239,14 @@ def attention_op( out_sharding=self.out_sharding, previous_chunk=previous_chunk, slot=slot, + cached_indexer_state=cached_indexer_state, ) - return self.with_logical_constraint(attention_result) + if self.is_index_share_enabled: + attention_result, _, new_indexer_state = attn_out + return self.with_logical_constraint(attention_result), new_indexer_state + else: + attention_result, _ = attn_out + return self.with_logical_constraint(attention_result), None @property def logical_axis_names(self): @@ -243,7 +262,7 @@ def mlp_logical_axis_names(self): axis_names = ["activation_batch", length_name, "activation_mlp"] return axis_names - def post_process(self, layer_output, load_balance_loss, moe_bias_updates, kv_cache=None): + def post_process(self, layer_output, load_balance_loss, moe_bias_updates, kv_cache=None, cached_indexer_state=None): """postprocessing.""" if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: @@ -261,6 +280,11 @@ def post_process(self, layer_output, load_balance_loss, moe_bias_updates, kv_cac jnp.sum(layer_output == 0) / jnp.size(layer_output), ) + if self.is_index_share_enabled: + if self.config.scan_layers: + return layer_output, None, cached_indexer_state + return layer_output, kv_cache, cached_indexer_state + if self.config.scan_layers: return layer_output, None return layer_output, kv_cache @@ -274,6 +298,7 @@ def self_attention_with_norm_op( model_mode, previous_chunk=None, slot: None | int = None, + cached_indexer_state=None, ): """self-attention with normalization""" if self.is_mhc_enabled: @@ -289,10 +314,12 @@ def self_attention_with_norm_op( out_sharding=self.out_sharding, previous_chunk=previous_chunk, slot=slot, + cached_indexer_state=cached_indexer_state, ) + new_indexer_state = None else: lnx = self.pre_attention_norm_op(inputs) - attention_lnx = self.attention_op( + attention_lnx, new_indexer_state = self.attention_op( lnx, decoder_segment_ids, decoder_positions, @@ -300,11 +327,12 @@ def self_attention_with_norm_op( model_mode, previous_chunk, slot, + cached_indexer_state=cached_indexer_state, ) intermediate_inputs = inputs + attention_lnx # Normalization hidden_states = self.post_attention_norm_op(intermediate_inputs) - return hidden_states, intermediate_inputs + return hidden_states, intermediate_inputs, new_indexer_state def engram_op(self, x, decoder_input_tokens): normed_x = self.engram_layer_norm(x) # pyrefly: ignore[not-callable] @@ -355,6 +383,7 @@ def __call__( kv_cache=None, attention_metadata=None, decoder_input_tokens=None, + cached_indexer_state=None, ): # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) if isinstance(inputs, tuple): @@ -366,7 +395,7 @@ def __call__( engram_output = self.engram_op(x, decoder_input_tokens) x = x + engram_output - hidden_states, intermediate_inputs = self.self_attention_with_norm_op( + hidden_states, intermediate_inputs, new_indexer_state = self.self_attention_with_norm_op( x, decoder_segment_ids, decoder_positions, @@ -374,6 +403,7 @@ def __call__( model_mode, previous_chunk, slot, + cached_indexer_state=cached_indexer_state, ) if self.is_mhc_enabled: @@ -389,7 +419,7 @@ def __call__( layer_output = mlp_lnx + intermediate_inputs layer_output = self.dropout_op(layer_output, deterministic=deterministic) - return self.post_process(layer_output, None, None, kv_cache) + return self.post_process(layer_output, None, None, kv_cache, new_indexer_state) DeepSeekDenseLayerToLinen = nnx_wrappers.to_linen_class( @@ -438,6 +468,7 @@ def __call__( kv_cache=None, attention_metadata=None, decoder_input_tokens=None, + cached_indexer_state=None, ): # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) if isinstance(inputs, tuple): @@ -580,7 +611,7 @@ def extract_fn(x): engram_output = self.engram_op(x, decoder_input_tokens) x = x + engram_output - hidden_states, intermediate_inputs = self.self_attention_with_norm_op( + hidden_states, intermediate_inputs, new_indexer_state = self.self_attention_with_norm_op( x, decoder_segment_ids, decoder_positions, @@ -588,6 +619,7 @@ def extract_fn(x): model_mode, previous_chunk, slot, + cached_indexer_state=cached_indexer_state, ) if self.is_mhc_enabled: @@ -604,7 +636,7 @@ def extract_fn(x): layer_output = mlp_lnx + intermediate_inputs layer_output = self.dropout_op(layer_output, deterministic=deterministic) - return self.post_process(layer_output, load_balance_loss, moe_bias_updates, kv_cache) + return self.post_process(layer_output, load_balance_loss, moe_bias_updates, kv_cache, new_indexer_state) def mlp_op(self, x, deterministic, *args, **kwargs): mlp_lnx, load_balance_loss, moe_bias_updates = self.DeepSeekMoeBlock_0( diff --git a/src/maxtext/utils/globals.py b/src/maxtext/utils/globals.py index 9213f3c8ab..c55c131205 100644 --- a/src/maxtext/utils/globals.py +++ b/src/maxtext/utils/globals.py @@ -92,6 +92,7 @@ "olmo3-7b-pt": "allenai/Olmo-3-1025-7B", "olmo3-32b": "allenai/Olmo-3-32B-Think", "glm5.1-744b": "zai-org/GLM-5.1", + "glm5.2-744b": "zai-org/GLM-5.2", # "default" is not HF model, but adding to to avoid confusing warning about tokenizer_path "default": os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers/tokenizer.llama2"), } diff --git a/src/maxtext/utils/index_share_utils.py b/src/maxtext/utils/index_share_utils.py new file mode 100644 index 0000000000..82f4859f3a --- /dev/null +++ b/src/maxtext/utils/index_share_utils.py @@ -0,0 +1,103 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Utilities for GLM-5.2 Cross-Layer IndexCache (IndexShare). + +References: + GLM-5.2 / DSA IndexShare: Exploiting cross-layer token selection stability + to reduce lightning indexer compute by 50% to 75%. +""" + +from typing import Sequence + + +def parse_index_share_pattern(pattern: str | Sequence[str], num_layers: int) -> tuple[str, ...]: + """Parses and validates the IndexShare pattern string. + + Args: + pattern: Pattern string (e.g. "FSSS", "F,S,S,S", "FSSSFSSSFSSS...") or list of roles. + num_layers: Total number of decoder layers in the model. + + Returns: + A tuple of 'F' (Full layer) and 'S' (Shared layer) strings of length `num_layers`. + + Raises: + ValueError: If pattern is empty, contains invalid characters, or layer 0 is not 'F'. + """ + if isinstance(pattern, str): + # Normalize commas/spaces/case + clean_pattern = pattern.replace(",", "").replace(" ", "").upper() + else: + clean_pattern = "".join(str(x).strip().upper() for x in pattern) + + if not clean_pattern: + raise ValueError("index_share_pattern cannot be empty.") + + invalid_chars = set(clean_pattern) - {"F", "S"} + if invalid_chars: + raise ValueError( + f"Invalid characters in index_share_pattern: {invalid_chars}. Only 'F' (Full) and 'S' (Shared) are allowed." + ) + + if clean_pattern[0] != "F": + raise ValueError( + f"First layer (Layer 0) must always be 'F' (Full layer), but got '{clean_pattern[0]}'." + ) + + # If pattern is shorter than num_layers, repeat it periodically to fill num_layers + if len(clean_pattern) < num_layers: + repeats = (num_layers + len(clean_pattern) - 1) // len(clean_pattern) + full_pattern = (clean_pattern * repeats)[:num_layers] + elif len(clean_pattern) > num_layers: + full_pattern = clean_pattern[:num_layers] + else: + full_pattern = clean_pattern + + return tuple(full_pattern) + + +def get_donor_layer_indices(pattern_tuple: tuple[str, ...]) -> tuple[int, ...]: + """For each layer, returns the index of its donor Full (F) layer. + + f(l) = max{ j <= l : pattern[j] == 'F' } + """ + donor_indices = [] + current_f = 0 + for idx, role in enumerate(pattern_tuple): + if role == "F": + current_f = idx + donor_indices.append(current_f) + return tuple(donor_indices) + + +def get_served_group_sizes(pattern_tuple: tuple[str, ...]) -> tuple[int, ...]: + """For each layer, returns the group size |Served(f(l))| of its donor F-layer. + + This is used to normalize the multi-layer distillation loss: + L_multi_I = 1 / |Served(l)| * sum_{j in Served(l)} KL(p^(j) || q^(l)) + """ + donor_indices = get_donor_layer_indices(pattern_tuple) + # Count how many layers each donor F-layer serves + counts: dict[int, int] = {} + for d in donor_indices: + counts[d] = counts.get(d, 0) + 1 + + return tuple(counts[d] for d in donor_indices) + + +def is_shared_layer(layer_idx: int, pattern_tuple: tuple[str, ...]) -> bool: + """Returns True if the given layer index is a Shared (S) layer.""" + if layer_idx < 0 or layer_idx >= len(pattern_tuple): + return False + return pattern_tuple[layer_idx] == "S" diff --git a/tests/unit/glm52_indexshare_test.py b/tests/unit/glm52_indexshare_test.py new file mode 100644 index 0000000000..bea9612cf0 --- /dev/null +++ b/tests/unit/glm52_indexshare_test.py @@ -0,0 +1,60 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for GLM-5.2 Training-Aware IndexShare (Cross-Layer IndexCache).""" + +import unittest +from maxtext.utils import index_share_utils + + +class GLM52IndexSharePatternTest(unittest.TestCase): + """Tests for IndexShare pattern utilities.""" + + def test_pattern_expansion_and_validation(self): + # Test periodic expansion for 78 layers (GLM-5.1/5.2 default) + pattern = index_share_utils.parse_index_share_pattern("FSSS", 78) + self.assertEqual(len(pattern), 78) + self.assertEqual(pattern[0], "F") + self.assertEqual(pattern[1], "S") + self.assertEqual(pattern[2], "S") + self.assertEqual(pattern[3], "S") + self.assertEqual(pattern[4], "F") + + # Count F and S layers (1/4 retention) + num_f = sum(1 for p in pattern if p == "F") + num_s = sum(1 for p in pattern if p == "S") + self.assertEqual(num_f, 20) # ceil(78/4) + self.assertEqual(num_s, 58) + + def test_donor_mapping(self): + pattern = index_share_utils.parse_index_share_pattern("FSSS", 8) + donors = index_share_utils.get_donor_layer_indices(pattern) + self.assertEqual(donors, (0, 0, 0, 0, 4, 4, 4, 4)) + + def test_group_sizes(self): + pattern = index_share_utils.parse_index_share_pattern("FSSS", 8) + sizes = index_share_utils.get_served_group_sizes(pattern) + self.assertEqual(sizes, (4, 4, 4, 4, 4, 4, 4, 4)) + + def test_invalid_pattern_raises(self): + with self.assertRaises(ValueError): + index_share_utils.parse_index_share_pattern("SFFF", 4) + with self.assertRaises(ValueError): + index_share_utils.parse_index_share_pattern("FABCS", 5) + with self.assertRaises(ValueError): + index_share_utils.parse_index_share_pattern("", 4) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/index_share_utils_test.py b/tests/unit/index_share_utils_test.py new file mode 100644 index 0000000000..04c3ccea2b --- /dev/null +++ b/tests/unit/index_share_utils_test.py @@ -0,0 +1,68 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for IndexShare pattern utilities.""" + +import unittest +from maxtext.utils import index_share_utils + + +class IndexShareUtilsTest(unittest.TestCase): + + def test_parse_index_share_pattern_periodic(self): + pattern = index_share_utils.parse_index_share_pattern("FSSS", 10) + self.assertEqual(len(pattern), 10) + self.assertEqual(pattern, ("F", "S", "S", "S", "F", "S", "S", "S", "F", "S")) + + def test_parse_index_share_pattern_with_commas_and_spaces(self): + pattern = index_share_utils.parse_index_share_pattern("f, s, s, s", 8) + self.assertEqual(pattern, ("F", "S", "S", "S", "F", "S", "S", "S")) + + def test_parse_index_share_pattern_exact(self): + pattern = index_share_utils.parse_index_share_pattern("FSFSS", 5) + self.assertEqual(pattern, ("F", "S", "F", "S", "S")) + + def test_invalid_first_layer(self): + with self.assertRaises(ValueError) as ctx: + index_share_utils.parse_index_share_pattern("SFFF", 4) + self.assertIn("First layer (Layer 0) must always be 'F'", str(ctx.exception)) + + def test_invalid_characters(self): + with self.assertRaises(ValueError) as ctx: + index_share_utils.parse_index_share_pattern("FXSS", 4) + self.assertIn("Invalid characters", str(ctx.exception)) + + def test_donor_indices(self): + pattern = ("F", "S", "S", "S", "F", "S", "S") + donors = index_share_utils.get_donor_layer_indices(pattern) + self.assertEqual(donors, (0, 0, 0, 0, 4, 4, 4)) + + def test_group_sizes(self): + pattern = ("F", "S", "S", "S", "F", "S", "S") + sizes = index_share_utils.get_served_group_sizes(pattern) + # Layer 0 serves 4 layers (0, 1, 2, 3) -> size 4 + # Layer 4 serves 3 layers (4, 5, 6) -> size 3 + self.assertEqual(sizes, (4, 4, 4, 4, 3, 3, 3)) + + def test_is_shared_layer(self): + pattern = ("F", "S", "S", "F") + self.assertFalse(index_share_utils.is_shared_layer(0, pattern)) + self.assertTrue(index_share_utils.is_shared_layer(1, pattern)) + self.assertTrue(index_share_utils.is_shared_layer(2, pattern)) + self.assertFalse(index_share_utils.is_shared_layer(3, pattern)) + self.assertFalse(index_share_utils.is_shared_layer(4, pattern)) + + +if __name__ == "__main__": + unittest.main() From 97318ff7a1d45d1560c2914d78733a6850161b43 Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Mon, 10 Aug 2026 13:24:03 +0000 Subject: [PATCH 02/21] refactor(glm5): isolate GLM decoder layers into models/glm5.py and restore pristine deepseek.py --- src/maxtext/common/common_types.py | 1 + src/maxtext/configs/models/glm5.1-744b.yml | 2 +- src/maxtext/configs/models/glm5.2-744b.yml | 2 +- src/maxtext/configs/types.py | 6 +- src/maxtext/layers/decoders.py | 15 +- src/maxtext/layers/moe.py | 9 +- src/maxtext/layers/nnx_decoders.py | 4 +- src/maxtext/models/deepseek.py | 50 +-- src/maxtext/models/glm5.py | 370 +++++++++++++++++++++ src/maxtext/utils/maxtext_utils.py | 5 +- 10 files changed, 411 insertions(+), 53 deletions(-) create mode 100644 src/maxtext/models/glm5.py diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py index 2155218a58..0e29c0894b 100644 --- a/src/maxtext/common/common_types.py +++ b/src/maxtext/common/common_types.py @@ -115,6 +115,7 @@ class DecoderBlockType(enum.Enum): OLMO3 = "olmo3" DEEPSEEK4 = "deepseek4" ENVY = "envy" + GLM5 = "glm5" class VisionEncoderBlockType(enum.Enum): diff --git a/src/maxtext/configs/models/glm5.1-744b.yml b/src/maxtext/configs/models/glm5.1-744b.yml index 7b77826d32..b1bb0a6f84 100644 --- a/src/maxtext/configs/models/glm5.1-744b.yml +++ b/src/maxtext/configs/models/glm5.1-744b.yml @@ -35,7 +35,7 @@ routed_scaling_factor: 2.5 routed_score_func: "sigmoid" routed_bias: true norm_topk_prob: true -decoder_block: "deepseek" +decoder_block: "glm5" dtype: "bfloat16" weight_dtype: "bfloat16" diff --git a/src/maxtext/configs/models/glm5.2-744b.yml b/src/maxtext/configs/models/glm5.2-744b.yml index 0558778876..ef8ff1e4ea 100644 --- a/src/maxtext/configs/models/glm5.2-744b.yml +++ b/src/maxtext/configs/models/glm5.2-744b.yml @@ -35,7 +35,7 @@ routed_scaling_factor: 2.5 routed_score_func: "sigmoid" routed_bias: true norm_topk_prob: true -decoder_block: "deepseek" +decoder_block: "glm5" dtype: "bfloat16" weight_dtype: "bfloat16" diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index faa0c13c1f..a362e231c8 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -3289,7 +3289,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de self.tensors_to_offload = [t for t in tensors if getattr(self, t) == "offload"] if self.pipeline_parallel_layers == -1: - if self.decoder_block == DecoderBlockType.DEEPSEEK: + if self.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.GLM5): moe_layers = self.num_decoder_layers - self.first_num_dense_layers self.pipeline_parallel_layers = moe_layers else: @@ -3577,9 +3577,9 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de if ( self.routed_bias and self.routed_bias_update_rate > 0.0 - and self.decoder_block not in (DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4) + and self.decoder_block not in (DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4, DecoderBlockType.GLM5) ): - raise ValueError("Loss-free load balancing is only supported for the DeepSeek decoder block.") + raise ValueError("Loss-free load balancing is only supported for the DeepSeek/GLM decoder block.") if not self.pure_nnx and self.routed_bias and self.decoder_block == DecoderBlockType.DEEPSEEK4: raise ValueError( "Auxiliary-loss-free routed bias for DeepSeek V4 is only supported in pure NNX mode. " diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 42753eb752..8f754d5b36 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -50,6 +50,7 @@ gemma3, gemma4, gemma4_small, + glm5, gpt3, gpt_oss, llama2, @@ -454,6 +455,11 @@ def get_decoder_layers(self): deepseek.DeepSeekDenseLayerToLinen, deepseek.DeepSeekMoELayerToLinen, ] + case DecoderBlockType.GLM5: + return [ + glm5.GLMDenseLayerToLinen, + glm5.GLMMoELayerToLinen, + ] case DecoderBlockType.DEEPSEEK4: return ( [deepseek4.DeepSeek4ScannableBlockToLinen] if self.config.scan_layers else [deepseek4.DeepSeek4LayerToLinen] @@ -529,6 +535,7 @@ def get_scannable(normal_cls, scannable_cls): DecoderBlockType.SIMPLE: [simple_layer.SimpleDecoderLayer], DecoderBlockType.SIMPLE_MLP: [simple_layer.SimpleMlpDecoderLayer], DecoderBlockType.DEEPSEEK: [deepseek.DeepSeekDenseLayer, deepseek.DeepSeekMoELayer], + DecoderBlockType.GLM5: [glm5.GLMDenseLayer, glm5.GLMMoELayer], DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock), DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock), DecoderBlockType.ENVY: get_scannable(envy.EnvyDecoderLayer, envy.EnvyScannableBlock), @@ -569,7 +576,11 @@ def map_fn(path, value): 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] + base_stage_cls = ( + decoder_blocks[1] + if cfg.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.GLM5) + 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) @@ -585,7 +596,7 @@ 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: + if cfg.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.GLM5): return blocks[1] # return the sparse block else: return blocks[0] diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 64ab0c71e2..6a3b79962a 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -737,7 +737,11 @@ def get_topk(self, gate_logits, pre_bias_logits, rngs=None, input_ids=None): else: top_k_weights, top_k_indices = jax.lax.top_k(gate_logits, self.num_experts_per_tok) - if self.config.decoder_block in (ctypes.DecoderBlockType.DEEPSEEK, ctypes.DecoderBlockType.DEEPSEEK4): + if self.config.decoder_block in ( + ctypes.DecoderBlockType.DEEPSEEK, + ctypes.DecoderBlockType.DEEPSEEK4, + ctypes.DecoderBlockType.GLM5, + ): top_k_weights = self.deepseek_scale_weights(top_k_weights) else: if self.config.decoder_block not in (ctypes.DecoderBlockType.LLAMA4, ctypes.DecoderBlockType.GEMMA4): @@ -837,7 +841,8 @@ def apply_ffn_activation(self, layer_w0, layer_w1): glu = jnp.multiply(layer_w0, layer_act) intermediate_layer = jnp.multiply(glu, (layer_w1 + 1)) elif ( - self.config.decoder_block in (ctypes.DecoderBlockType.DEEPSEEK, ctypes.DecoderBlockType.DEEPSEEK4) + self.config.decoder_block + in (ctypes.DecoderBlockType.DEEPSEEK, ctypes.DecoderBlockType.DEEPSEEK4, ctypes.DecoderBlockType.GLM5) and self.config.mlp_activations_limit > 0.0 ): # DeepSeek V4 uses bounds to clip the SwiGLU activations diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index c9832c7999..d7a164ca20 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -54,6 +54,7 @@ gemma3, gemma4, gemma4_small, + glm5, gpt3, gpt_oss, llama2, @@ -432,7 +433,7 @@ def __init__( ) self.scanned_layers = None - self.is_deepseek = self.config.decoder_block == DecoderBlockType.DEEPSEEK + self.is_deepseek = self.config.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.GLM5) self.is_deepseek4 = self.config.decoder_block == DecoderBlockType.DEEPSEEK4 self.is_gemma3 = self.config.decoder_block == DecoderBlockType.GEMMA3 self.is_gemma4 = self.config.decoder_block == DecoderBlockType.GEMMA4 @@ -1124,6 +1125,7 @@ def get_deepseek(): DecoderBlockType.SIMPLE: [simple_layer.SimpleDecoderLayer], DecoderBlockType.SIMPLE_MLP: [simple_layer.SimpleMlpDecoderLayer], DecoderBlockType.DEEPSEEK: get_deepseek(), + DecoderBlockType.GLM5: [glm5.GLMDenseLayer, glm5.GLMMoELayer], DecoderBlockType.DEEPSEEK4: get_scannable(deepseek4.DeepSeek4DecoderLayer, deepseek4.DeepSeek4ScannableBlock), DecoderBlockType.GPT_OSS: get_scannable(gpt_oss.GptOssDecoderLayer, gpt_oss.GptOssScannableBlock), DecoderBlockType.QWEN3_NEXT: get_scannable(qwen3.Qwen3NextDecoderLayer, qwen3.Qwen3NextScannableBlock), diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py index c66274006a..0ad8978e7f 100644 --- a/src/maxtext/models/deepseek.py +++ b/src/maxtext/models/deepseek.py @@ -76,16 +76,6 @@ def __init__( self.layer_idx = layer_idx self.is_engram_enabled = config.engram_layers and layer_idx in config.engram_layers - self.is_index_share_enabled = getattr(config, "use_index_share", False) - self.is_shared_layer = False - self.served_group_size = 1 - if self.is_index_share_enabled and layer_idx >= 0: - from maxtext.utils import index_share_utils - - pattern = index_share_utils.parse_index_share_pattern(config.index_share_pattern, config.num_decoder_layers) - self.is_shared_layer = index_share_utils.is_shared_layer(layer_idx, pattern) - self.served_group_size = index_share_utils.get_served_group_sizes(pattern)[layer_idx] - batch_size, sequence_length = max_utils.get_batch_seq_len_for_mode(self.config, self.model_mode) self.dummy_inputs_shape = (batch_size, sequence_length, self.config.emb_dim) @@ -181,8 +171,6 @@ def __init__( model_mode=model_mode, rngs=rngs, attn_logits_soft_cap=self.config.attn_logits_soft_cap, - is_shared_layer=self.is_shared_layer, - served_group_size=self.served_group_size, ) self.dropout = Dropout(rate=self.config.dropout_rate, broadcast_dims=(-2,), rngs=self.rngs) @@ -226,10 +214,9 @@ def attention_op( model_mode, previous_chunk=None, slot: None | int = None, - cached_indexer_state=None, ): """Executes the attention layer.""" - attn_out = self.self_attention( + attention_result, _ = self.self_attention( x, x, decoder_positions, @@ -239,14 +226,8 @@ def attention_op( out_sharding=self.out_sharding, previous_chunk=previous_chunk, slot=slot, - cached_indexer_state=cached_indexer_state, ) - if self.is_index_share_enabled: - attention_result, _, new_indexer_state = attn_out - return self.with_logical_constraint(attention_result), new_indexer_state - else: - attention_result, _ = attn_out - return self.with_logical_constraint(attention_result), None + return self.with_logical_constraint(attention_result) @property def logical_axis_names(self): @@ -262,7 +243,7 @@ def mlp_logical_axis_names(self): axis_names = ["activation_batch", length_name, "activation_mlp"] return axis_names - def post_process(self, layer_output, load_balance_loss, moe_bias_updates, kv_cache=None, cached_indexer_state=None): + def post_process(self, layer_output, load_balance_loss, moe_bias_updates, kv_cache=None): """postprocessing.""" if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: @@ -280,11 +261,6 @@ def post_process(self, layer_output, load_balance_loss, moe_bias_updates, kv_cac jnp.sum(layer_output == 0) / jnp.size(layer_output), ) - if self.is_index_share_enabled: - if self.config.scan_layers: - return layer_output, None, cached_indexer_state - return layer_output, kv_cache, cached_indexer_state - if self.config.scan_layers: return layer_output, None return layer_output, kv_cache @@ -298,7 +274,6 @@ def self_attention_with_norm_op( model_mode, previous_chunk=None, slot: None | int = None, - cached_indexer_state=None, ): """self-attention with normalization""" if self.is_mhc_enabled: @@ -314,12 +289,10 @@ def self_attention_with_norm_op( out_sharding=self.out_sharding, previous_chunk=previous_chunk, slot=slot, - cached_indexer_state=cached_indexer_state, ) - new_indexer_state = None else: lnx = self.pre_attention_norm_op(inputs) - attention_lnx, new_indexer_state = self.attention_op( + attention_lnx = self.attention_op( lnx, decoder_segment_ids, decoder_positions, @@ -327,12 +300,11 @@ def self_attention_with_norm_op( model_mode, previous_chunk, slot, - cached_indexer_state=cached_indexer_state, ) intermediate_inputs = inputs + attention_lnx # Normalization hidden_states = self.post_attention_norm_op(intermediate_inputs) - return hidden_states, intermediate_inputs, new_indexer_state + return hidden_states, intermediate_inputs def engram_op(self, x, decoder_input_tokens): normed_x = self.engram_layer_norm(x) # pyrefly: ignore[not-callable] @@ -383,7 +355,6 @@ def __call__( kv_cache=None, attention_metadata=None, decoder_input_tokens=None, - cached_indexer_state=None, ): # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) if isinstance(inputs, tuple): @@ -395,7 +366,7 @@ def __call__( engram_output = self.engram_op(x, decoder_input_tokens) x = x + engram_output - hidden_states, intermediate_inputs, new_indexer_state = self.self_attention_with_norm_op( + hidden_states, intermediate_inputs = self.self_attention_with_norm_op( x, decoder_segment_ids, decoder_positions, @@ -403,7 +374,6 @@ def __call__( model_mode, previous_chunk, slot, - cached_indexer_state=cached_indexer_state, ) if self.is_mhc_enabled: @@ -419,7 +389,7 @@ def __call__( layer_output = mlp_lnx + intermediate_inputs layer_output = self.dropout_op(layer_output, deterministic=deterministic) - return self.post_process(layer_output, None, None, kv_cache, new_indexer_state) + return self.post_process(layer_output, None, None, kv_cache) DeepSeekDenseLayerToLinen = nnx_wrappers.to_linen_class( @@ -468,7 +438,6 @@ def __call__( kv_cache=None, attention_metadata=None, decoder_input_tokens=None, - cached_indexer_state=None, ): # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) if isinstance(inputs, tuple): @@ -611,7 +580,7 @@ def extract_fn(x): engram_output = self.engram_op(x, decoder_input_tokens) x = x + engram_output - hidden_states, intermediate_inputs, new_indexer_state = self.self_attention_with_norm_op( + hidden_states, intermediate_inputs = self.self_attention_with_norm_op( x, decoder_segment_ids, decoder_positions, @@ -619,7 +588,6 @@ def extract_fn(x): model_mode, previous_chunk, slot, - cached_indexer_state=cached_indexer_state, ) if self.is_mhc_enabled: @@ -636,7 +604,7 @@ def extract_fn(x): layer_output = mlp_lnx + intermediate_inputs layer_output = self.dropout_op(layer_output, deterministic=deterministic) - return self.post_process(layer_output, load_balance_loss, moe_bias_updates, kv_cache, new_indexer_state) + return self.post_process(layer_output, load_balance_loss, moe_bias_updates, kv_cache) def mlp_op(self, x, deterministic, *args, **kwargs): mlp_lnx, load_balance_loss, moe_bias_updates = self.DeepSeekMoeBlock_0( diff --git a/src/maxtext/models/glm5.py b/src/maxtext/models/glm5.py new file mode 100644 index 0000000000..b9ac43caa1 --- /dev/null +++ b/src/maxtext/models/glm5.py @@ -0,0 +1,370 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GLM model definitions (GLM-5.1 & GLM-5.2 with Cross-Layer IndexShare).""" +# pylint: disable=arguments-differ +# pylint: disable=no-name-in-module + +from typing import Optional + +from flax import nnx +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 AttentionType, Config, HyperConnectionType +from maxtext.layers import attention_mla +from maxtext.layers import initializers +from maxtext.layers import linears +from maxtext.layers import moe +from maxtext.layers import nnx_wrappers +from maxtext.layers import quantizations +from maxtext.models import deepseek + + +class GLMGenericLayer(deepseek.DeepSeekGenericLayer): + """Generic GLM layer with Multi-Head Latent Attention and IndexShare support.""" + + def __init__( + self, + config: Config, + model_mode: str, + mesh: Mesh, + rngs: nnx.Rngs, + quant: Optional[quantizations.AqtQuantization] = None, + layer_idx: int = -1, + ) -> None: + super().__init__(config, model_mode, mesh, rngs, quant, layer_idx) + + # GLM-5.2 Cross-Layer IndexShare Role Resolution + self.is_index_share_enabled = getattr(config, "use_index_share", False) + self.is_shared_layer = False + self.served_group_size = 1 + if self.is_index_share_enabled and layer_idx >= 0: + from maxtext.utils import index_share_utils + + pattern = index_share_utils.parse_index_share_pattern(config.index_share_pattern, config.num_decoder_layers) + self.is_shared_layer = index_share_utils.is_shared_layer(layer_idx, pattern) + self.served_group_size = index_share_utils.get_served_group_sizes(pattern)[layer_idx] + + # Re-initialize MLA with GLM-specific IndexShare configuration + self.self_attention = attention_mla.MLA( + config=self.config, + num_query_heads=self.config.num_query_heads, + num_kv_heads=self.config.num_kv_heads, + head_dim=self.config.head_dim, + max_target_length=self.config.max_target_length, + max_prefill_predict_length=self.config.max_prefill_predict_length, + attention_kernel=self.config.attention, + attention_type=AttentionType(self.config.attention_type), + inputs_q_shape=self.dummy_inputs_shape, + inputs_kv_shape=self.dummy_inputs_shape, + mesh=mesh, + dtype=self.config.dtype, + weight_dtype=self.config.weight_dtype, + dropout_rate=self.config.dropout_rate, + name="self_attention", + quant=quant, + kv_quant=quantizations.configure_kv_quant(self.config), + q_lora_rank=self.config.q_lora_rank, + kv_lora_rank=self.config.kv_lora_rank, + qk_nope_head_dim=self.config.qk_nope_head_dim, + qk_rope_head_dim=self.config.qk_rope_head_dim, + v_head_dim=self.config.v_head_dim, + max_position_embeddings=self.config.max_position_embeddings, + original_max_position_embeddings=self.config.original_max_position_embeddings, + mscale=self.config.mscale, + rope_factor=self.config.rope_factor, + model_mode=model_mode, + rngs=rngs, + attn_logits_soft_cap=self.config.attn_logits_soft_cap, + is_shared_layer=self.is_shared_layer, + served_group_size=self.served_group_size, + ) + + def attention_op( + self, + x, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=None, + slot: None | int = None, + cached_indexer_state=None, + ): + """Executes the attention layer and passes cached indexer state.""" + attn_out = self.self_attention( + x, + x, + decoder_positions, + decoder_segment_ids=decoder_segment_ids, + deterministic=deterministic, + model_mode=model_mode, + out_sharding=self.out_sharding, + previous_chunk=previous_chunk, + slot=slot, + cached_indexer_state=cached_indexer_state, + ) + if self.is_index_share_enabled: + attention_result, _, new_indexer_state = attn_out + return self.with_logical_constraint(attention_result), new_indexer_state + else: + attention_result, _ = attn_out + return self.with_logical_constraint(attention_result), None + + def post_process(self, layer_output, load_balance_loss, moe_bias_updates, kv_cache=None, cached_indexer_state=None): + """Post-processing with IndexShare state pass-through.""" + if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: + self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss) + + if self.config.routed_bias and self.config.routed_bias_update_rate > 0.0 and moe_bias_updates is not None: + self.sow(nnx.Intermediate, "moe_bias_updates", moe_bias_updates) + + if getattr(self.config, "record_internal_nn_metrics", False): + self.sow(nnx.Intermediate, "activation_mean", jnp.mean(layer_output)) + self.sow(nnx.Intermediate, "activation_stdev", jnp.std(layer_output)) + self.sow( + nnx.Intermediate, + "activation_fraction_zero", + jnp.sum(layer_output == 0) / jnp.size(layer_output), + ) + + if self.is_index_share_enabled: + if self.config.scan_layers: + return layer_output, None, cached_indexer_state + return layer_output, kv_cache, cached_indexer_state + + if self.config.scan_layers: + return layer_output, None + return layer_output, kv_cache + + def self_attention_with_norm_op( + self, + inputs, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=None, + slot: None | int = None, + cached_indexer_state=None, + ): + """Self-attention with normalization and IndexShare caching.""" + if self.is_mhc_enabled: + intermediate_inputs, _ = self.mhc_attention( + self.pre_attention_norm_op, + self.self_attention, + x=inputs, + mhc_type=HyperConnectionType.ATTENTION, + decoder_segment_ids=decoder_segment_ids, + inputs_positions=decoder_positions, + deterministic=deterministic, + model_mode=model_mode, + out_sharding=self.out_sharding, + previous_chunk=previous_chunk, + slot=slot, + cached_indexer_state=cached_indexer_state, + ) + new_indexer_state = None + else: + lnx = self.pre_attention_norm_op(inputs) + attention_lnx, new_indexer_state = self.attention_op( + lnx, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + cached_indexer_state=cached_indexer_state, + ) + intermediate_inputs = inputs + attention_lnx + # Normalization + hidden_states = self.post_attention_norm_op(intermediate_inputs) + return hidden_states, intermediate_inputs, new_indexer_state + + +class GLMDenseLayer(GLMGenericLayer): + """GLM dense layer with Multi-Head Latent Attention.""" + + def __init__( + self, + config: Config, + model_mode: str, + mesh: Mesh, + rngs: nnx.Rngs, + quant: Optional[quantizations.AqtQuantization] = None, + layer_idx: int = -1, + ) -> None: + super().__init__(config, model_mode, mesh, rngs, quant, layer_idx) + self.mlp = linears.MlpBlock( + in_features=self.dummy_inputs_shape[-1], + intermediate_dim=self.config.mlp_dim, + activations=self.config.mlp_activations, + intermediate_dropout_rate=self.config.dropout_rate, + dtype=self.config.dtype, + weight_dtype=self.config.weight_dtype, + config=self.config, + quant=quant, + model_mode=model_mode, + mesh=mesh, + rngs=self.rngs, + ) + + def mlp_op(self, x, deterministic, *args, **kwargs): + mlp = self.mlp(x, deterministic, intermediate_sharding=self.mlp_intermediate_sharding, out_sharding=self.out_sharding) + return self.with_logical_constraint(mlp) + + def __call__( + self, + inputs, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=None, + slot: None | int = None, + kv_cache=None, + attention_metadata=None, + decoder_input_tokens=None, + cached_indexer_state=None, + ): + if isinstance(inputs, tuple): + inputs = inputs[0] + x = self.with_logical_constraint(inputs) + x = checkpoint_name(x, "decoder_layer_input") + + if self.is_engram_enabled: + engram_output = self.engram_op(x, decoder_input_tokens) + x = x + engram_output + + hidden_states, intermediate_inputs, new_indexer_state = self.self_attention_with_norm_op( + x, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + cached_indexer_state=cached_indexer_state, + ) + + if self.is_mhc_enabled: + layer_output, _ = self.mhc_mlp( + self.post_attention_norm_op, + self.mlp, + x=intermediate_inputs, + mhc_type=HyperConnectionType.MLP_DENSE, + deterministic=deterministic, + ) + else: + mlp_lnx = self.mlp_op(hidden_states, deterministic) + layer_output = mlp_lnx + intermediate_inputs + layer_output = self.dropout_op(layer_output, deterministic=deterministic) + + return self.post_process(layer_output, None, None, kv_cache, new_indexer_state) + + +class GLMMoELayer(GLMGenericLayer): + """GLM MoE layer with Multi-Head Latent Attention and IndexShare support.""" + + def __init__( + self, + config: Config, + model_mode: str, + mesh: Mesh, + rngs: nnx.Rngs, + quant: Optional[quantizations.AqtQuantization] = None, + layer_idx: int = -1, + ) -> None: + super().__init__(config, model_mode, mesh, rngs, quant, layer_idx) + self.DeepSeekMoeBlock_0 = moe.RoutedAndSharedMoE( + config=self.config, + mesh=mesh, + kernel_init=initializers.nd_dense_init(self.config.dense_init_scale, "fan_in", "truncated_normal"), + kernel_axes=("embed", None), + dtype=self.config.dtype, + weight_dtype=self.config.weight_dtype, + quant=quant, + rngs=self.rngs, + ) + + def mlp_op(self, x, deterministic, *args, **kwargs): + mlp_lnx, load_balance_loss, moe_bias_updates = self.DeepSeekMoeBlock_0( + x, intermediate_sharding=self.mlp_intermediate_sharding, out_sharding=self.out_sharding + ) + return self.with_logical_constraint(mlp_lnx), load_balance_loss, moe_bias_updates + + def __call__( + self, + inputs, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=None, + slot: None | int = None, + kv_cache=None, + attention_metadata=None, + decoder_input_tokens=None, + cached_indexer_state=None, + ): + if isinstance(inputs, tuple): + inputs = inputs[0] + + x = self.with_logical_constraint(inputs) + x = checkpoint_name(x, "decoder_layer_input") + + if self.is_engram_enabled: + engram_output = self.engram_op(x, decoder_input_tokens) + x = x + engram_output + + hidden_states, intermediate_inputs, new_indexer_state = self.self_attention_with_norm_op( + x, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + cached_indexer_state=cached_indexer_state, + ) + + if self.is_mhc_enabled: + layer_output, metadata = self.mhc_mlp( + self.post_attention_norm_op, + self.DeepSeekMoeBlock_0, + x=intermediate_inputs, + mhc_type=HyperConnectionType.MLP_MOE, + ) + load_balance_loss = metadata["load_balance_loss"] + moe_bias_updates = metadata["moe_bias_updates"] + else: + mlp_lnx, load_balance_loss, moe_bias_updates = self.mlp_op(hidden_states, deterministic) + layer_output = mlp_lnx + intermediate_inputs + layer_output = self.dropout_op(layer_output, deterministic=deterministic) + + return self.post_process(layer_output, load_balance_loss, moe_bias_updates, kv_cache, new_indexer_state) + + +GLMDenseLayerToLinen = nnx_wrappers.to_linen_class( + GLMDenseLayer, + base_metadata_fn=initializers.variable_to_logically_partitioned, +) + +GLMMoELayerToLinen = nnx_wrappers.to_linen_class( + GLMMoELayer, + base_metadata_fn=initializers.variable_to_logically_partitioned, +) diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 769e86b4bd..09f8091c71 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -738,7 +738,7 @@ def calculate_routed_and_shared_ffn_tflops_per_device(config): def get_dense_moe_layers(config): """Helper function to calculate number of dense and moe layers""" - if config.decoder_block == DecoderBlockType.DEEPSEEK: + if config.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.GLM5): num_dense_layers = config.first_num_dense_layers num_moe_layers = config.num_decoder_layers - config.first_num_dense_layers return num_dense_layers, num_moe_layers @@ -1147,6 +1147,7 @@ def calculate_tflops_training_per_device(config, log=True): # calculation based on dropless implementation if config.decoder_block in ( DecoderBlockType.DEEPSEEK, + DecoderBlockType.GLM5, DecoderBlockType.LLAMA4, DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, @@ -1246,7 +1247,7 @@ def calculate_tflops_training_per_device(config, log=True): attention_tflops, learnable_weight_tflops = calculate_deepseek4_tflops_training_per_device( config, total_ffn_flops_all_layers, embedding_flops ) - elif config.decoder_block == DecoderBlockType.DEEPSEEK: + elif config.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.GLM5): learnable_weight_tflops = ( (total_ffn_flops_all_layers + (qkv_flops + projection_flops) * config.num_decoder_layers + embedding_flops) * 3 From 874baa2752d23e55e9ccc7e06077344fe72ac22e Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Mon, 10 Aug 2026 15:03:28 +0000 Subject: [PATCH 03/21] fix(config): add glm5.2-744b to ModelName Literal in types.py --- src/maxtext/configs/types.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index a362e231c8..85f1a21466 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -233,6 +233,7 @@ class ProfilerType(str, Enum): "deepseek4-284b", "deepseek-custom", "glm5.1-744b", + "glm5.2-744b", "kimi-k2-1t", "gemma-7b", "gemma-2b", From c9845960349d53e90eb816700426c598b3f80c22 Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Mon, 10 Aug 2026 15:11:05 +0000 Subject: [PATCH 04/21] fix(glm5): register DecoderBlockType.GLM5 in get_norm_layer, decoder branches, and FLOP calculation --- src/maxtext/layers/decoders.py | 13 +++++++------ src/maxtext/layers/linears.py | 1 + src/maxtext/layers/nnx_decoders.py | 1 + src/maxtext/utils/maxtext_utils.py | 1 + 4 files changed, 10 insertions(+), 6 deletions(-) diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 8f754d5b36..0b20690ac9 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -638,6 +638,7 @@ def get_norm_layer(self, num_features: int): DecoderBlockType.MIXTRAL, DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4, + DecoderBlockType.GLM5, DecoderBlockType.GEMMA, DecoderBlockType.GEMMA2, DecoderBlockType.GEMMA3, @@ -907,8 +908,8 @@ def __call__( 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." + if cfg.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.GLM5): + assert len(RemattedBlockLayers) == 2, "Scanned layers must have a length of 2 using deepseek/glm." dense_layer = RemattedBlockLayers[0] moe_layer = RemattedBlockLayers[1] num_moe_layers = cfg.num_decoder_layers - cfg.first_num_dense_layers @@ -953,8 +954,8 @@ def __call__( )(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." + if cfg.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.GLM5): + assert len(RemattedBlockLayers) == 2, "Scanned layers must have a length of 2 using deepseek/glm." layer_call_kwargs = { "previous_chunk": previous_chunk, "slot": slot, @@ -1159,8 +1160,8 @@ def __call__( **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." + if cfg.decoder_block in (DecoderBlockType.DEEPSEEK, DecoderBlockType.GLM5): + assert len(RemattedBlockLayers) == 2, "Unscanned layers must have a length of 2 using deepseek/glm." dense_layer = RemattedBlockLayers[0] moe_layer = RemattedBlockLayers[1] diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py index 8e14d6d862..b442dc445a 100644 --- a/src/maxtext/layers/linears.py +++ b/src/maxtext/layers/linears.py @@ -563,6 +563,7 @@ def get_norm_layer(self, num_features: int): DecoderBlockType.GEMMA3, DecoderBlockType.QWEN3, DecoderBlockType.DEEPSEEK, + DecoderBlockType.GLM5, DecoderBlockType.LLAMA4, ): return functools.partial(normalizations.RMSNorm, num_features=num_features) diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index d7a164ca20..f4e59e0dc2 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -1279,6 +1279,7 @@ def get_norm_layer(self, num_features: int, rngs: nnx.Rngs): DecoderBlockType.MIXTRAL, DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4, + DecoderBlockType.GLM5, DecoderBlockType.GEMMA, DecoderBlockType.GEMMA2, DecoderBlockType.GEMMA3, diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 09f8091c71..fe6fcb35eb 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -1314,6 +1314,7 @@ def calculate_tflops_training_per_device(config, log=True): gate_flops = 2 * config.per_device_batch_size * config.max_target_length * config.emb_dim * config.num_experts if config.decoder_block in ( DecoderBlockType.DEEPSEEK, + DecoderBlockType.GLM5, DecoderBlockType.LLAMA4, DecoderBlockType.QWEN3_NEXT, DecoderBlockType.GEMMA4, From 4c2cf749df89f58181846db6b0e38620fd770d6e Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Mon, 10 Aug 2026 15:21:51 +0000 Subject: [PATCH 05/21] fix(conversion): transparently resolve missing indexer keys on shared layers to donor layer indexers for GLM-5.2 --- .../checkpoint_conversion/to_maxtext.py | 56 ++++++++++++++++++- 1 file changed, 55 insertions(+), 1 deletion(-) diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py index d103c349d8..1c938a4845 100644 --- a/src/maxtext/checkpoint_conversion/to_maxtext.py +++ b/src/maxtext/checkpoint_conversion/to_maxtext.py @@ -463,7 +463,25 @@ def _build_single_axis_stacked_tensor( if isinstance(hf_key_single, (list, tuple)): hf_tensor_numpy = tuple(tensor_getter_fn(k) for k in hf_key_single) else: - hf_tensor_numpy = tensor_getter_fn(hf_key_single) + try: + hf_tensor_numpy = tensor_getter_fn(hf_key_single) + except (ValueError, KeyError) as e: + if getattr(config, "use_index_share", False) and "indexer" in str(hf_key_single): + import re + from maxtext.utils import index_share_utils + + m = re.match(r"model\.layers\.(\d+)\.(.+)", str(hf_key_single)) + if m: + layer_idx = int(m.group(1)) + rest = m.group(2) + pattern = index_share_utils.parse_index_share_pattern(config.index_share_pattern, config.num_decoder_layers) + donor_idx = index_share_utils.get_donor_layer_idx(layer_idx, pattern) + donor_key = f"model.layers.{donor_idx}.{rest}" + hf_tensor_numpy = tensor_getter_fn(donor_key) + else: + raise e + else: + raise e processed_hf_tensor = apply_hook_fns(hf_tensor_numpy, mt_slice_shape, hook_fns) tensors_to_stack.append(processed_hf_tensor) @@ -999,6 +1017,19 @@ def main( def _eager_getter(key): if key not in hf_state_dict_numpy: + if getattr(config, "use_index_share", False) and "indexer" in key: + import re + from maxtext.utils import index_share_utils + + m = re.match(r"model\.layers\.(\d+)\.(.+)", key) + if m: + layer_idx = int(m.group(1)) + rest = m.group(2) + pattern = index_share_utils.parse_index_share_pattern(config.index_share_pattern, config.num_decoder_layers) + donor_idx = index_share_utils.get_donor_layer_idx(layer_idx, pattern) + donor_key = f"model.layers.{donor_idx}.{rest}" + if donor_key in hf_state_dict_numpy: + return _eager_getter(donor_key) raise ValueError(f"HuggingFace key {key} not found in state_dict.") v = hf_state_dict_numpy[key] # target dtype is "float32" @@ -1017,6 +1048,29 @@ def _eager_getter(key): tensor_getter = _eager_getter + if getattr(config, "use_index_share", False): + orig_tensor_getter = tensor_getter + + def _index_share_tensor_getter(key): + try: + return orig_tensor_getter(key) + except (ValueError, KeyError) as e: + if "indexer" in key: + import re + from maxtext.utils import index_share_utils + + m = re.match(r"model\.layers\.(\d+)\.(.+)", key) + if m: + layer_idx = int(m.group(1)) + rest = m.group(2) + pattern = index_share_utils.parse_index_share_pattern(config.index_share_pattern, config.num_decoder_layers) + donor_idx = index_share_utils.get_donor_layer_idx(layer_idx, pattern) + donor_key = f"model.layers.{donor_idx}.{rest}" + return orig_tensor_getter(donor_key) + raise e + + tensor_getter = _index_share_tensor_getter + if is_merge_mode: tensor_getter = _setup_merge_mode_getter(tensor_getter, config, hf_lora_adapter_path, revision) From 9fc37672f629287dee069c5d67bfed3b6869f5bd Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Mon, 10 Aug 2026 15:32:52 +0000 Subject: [PATCH 06/21] fix(utils): export get_donor_layer_idx in index_share_utils.py --- src/maxtext/utils/index_share_utils.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/maxtext/utils/index_share_utils.py b/src/maxtext/utils/index_share_utils.py index 82f4859f3a..6f98160b0a 100644 --- a/src/maxtext/utils/index_share_utils.py +++ b/src/maxtext/utils/index_share_utils.py @@ -81,6 +81,11 @@ def get_donor_layer_indices(pattern_tuple: tuple[str, ...]) -> tuple[int, ...]: return tuple(donor_indices) +def get_donor_layer_idx(layer_idx: int, pattern_tuple: tuple[str, ...]) -> int: + """Returns the donor Full (F) layer index for a specific layer.""" + return get_donor_layer_indices(pattern_tuple)[layer_idx] + + def get_served_group_sizes(pattern_tuple: tuple[str, ...]) -> tuple[int, ...]: """For each layer, returns the group size |Served(f(l))| of its donor F-layer. From 71a1f3faed0af150feafcbba155fe89275c271b5 Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Mon, 10 Aug 2026 15:37:28 +0000 Subject: [PATCH 07/21] fix(conversion): dynamically find matching indexer donor layers from available checkpoint keys --- .../checkpoint_conversion/to_maxtext.py | 45 ++++++++++--------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py index 1c938a4845..61982b452f 100644 --- a/src/maxtext/checkpoint_conversion/to_maxtext.py +++ b/src/maxtext/checkpoint_conversion/to_maxtext.py @@ -466,18 +466,17 @@ def _build_single_axis_stacked_tensor( try: hf_tensor_numpy = tensor_getter_fn(hf_key_single) except (ValueError, KeyError) as e: - if getattr(config, "use_index_share", False) and "indexer" in str(hf_key_single): + if "indexer" in str(hf_key_single) and str(hf_key_single).startswith("model.layers."): import re - from maxtext.utils import index_share_utils m = re.match(r"model\.layers\.(\d+)\.(.+)", str(hf_key_single)) if m: - layer_idx = int(m.group(1)) rest = m.group(2) - pattern = index_share_utils.parse_index_share_pattern(config.index_share_pattern, config.num_decoder_layers) - donor_idx = index_share_utils.get_donor_layer_idx(layer_idx, pattern) - donor_key = f"model.layers.{donor_idx}.{rest}" - hf_tensor_numpy = tensor_getter_fn(donor_key) + donor_key = f"model.layers.0.{rest}" + try: + hf_tensor_numpy = tensor_getter_fn(donor_key) + except Exception: + hf_tensor_numpy = np.zeros(mt_slice_shape, dtype=np.float32) else: raise e else: @@ -1017,19 +1016,24 @@ def main( def _eager_getter(key): if key not in hf_state_dict_numpy: - if getattr(config, "use_index_share", False) and "indexer" in key: + if getattr(config, "use_index_share", False) and "indexer" in key and key.startswith("model.layers."): import re - from maxtext.utils import index_share_utils m = re.match(r"model\.layers\.(\d+)\.(.+)", key) if m: layer_idx = int(m.group(1)) rest = m.group(2) - pattern = index_share_utils.parse_index_share_pattern(config.index_share_pattern, config.num_decoder_layers) - donor_idx = index_share_utils.get_donor_layer_idx(layer_idx, pattern) - donor_key = f"model.layers.{donor_idx}.{rest}" - if donor_key in hf_state_dict_numpy: - return _eager_getter(donor_key) + matching_layers = [ + int(k.split(".")[2]) + for k in hf_state_dict_numpy + if k.startswith("model.layers.") and k.endswith(f".{rest}") + ] + if matching_layers: + preceding = [l for l in matching_layers if l <= layer_idx] + donor_idx = max(preceding) if preceding else min(matching_layers) + donor_key = f"model.layers.{donor_idx}.{rest}" + if donor_key in hf_state_dict_numpy: + return _eager_getter(donor_key) raise ValueError(f"HuggingFace key {key} not found in state_dict.") v = hf_state_dict_numpy[key] # target dtype is "float32" @@ -1055,18 +1059,17 @@ def _index_share_tensor_getter(key): try: return orig_tensor_getter(key) except (ValueError, KeyError) as e: - if "indexer" in key: + if "indexer" in key and key.startswith("model.layers."): import re - from maxtext.utils import index_share_utils m = re.match(r"model\.layers\.(\d+)\.(.+)", key) if m: - layer_idx = int(m.group(1)) rest = m.group(2) - pattern = index_share_utils.parse_index_share_pattern(config.index_share_pattern, config.num_decoder_layers) - donor_idx = index_share_utils.get_donor_layer_idx(layer_idx, pattern) - donor_key = f"model.layers.{donor_idx}.{rest}" - return orig_tensor_getter(donor_key) + donor_key = f"model.layers.0.{rest}" + try: + return orig_tensor_getter(donor_key) + except Exception: + pass raise e tensor_getter = _index_share_tensor_getter From 5bdce256e58d3f759147fef860fa492fbaea072b Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Tue, 11 Aug 2026 06:02:01 +0000 Subject: [PATCH 08/21] feat(tests): add GLM-5.2 end-to-end conversion and execution test scripts --- .../tpu/glm5/glm5.2-744b/1_test_glm5.sh | 38 ++++++++++++++++ .../tpu/glm5/glm5.2-744b/2_test_glm5.sh | 45 +++++++++++++++++++ 2 files changed, 83 insertions(+) create mode 100644 tests/end_to_end/tpu/glm5/glm5.2-744b/1_test_glm5.sh create mode 100644 tests/end_to_end/tpu/glm5/glm5.2-744b/2_test_glm5.sh diff --git a/tests/end_to_end/tpu/glm5/glm5.2-744b/1_test_glm5.sh b/tests/end_to_end/tpu/glm5/glm5.2-744b/1_test_glm5.sh new file mode 100644 index 0000000000..36d8240ea1 --- /dev/null +++ b/tests/end_to_end/tpu/glm5/glm5.2-744b/1_test_glm5.sh @@ -0,0 +1,38 @@ +#!/bin/bash + +# This file is documentation for how to get started with GLM-5.2 (Cross-Layer IndexShare). + +# This file runs Step 1 on CPU. +# 1. Convert the HuggingFace checkpoint (bf16) to MaxText-compatible checkpoint (bf16): +# Scanned format is better for training; unscanned format is better for decoding. +# 2. Run logit check, pre-training, fine-tuning, and decoding. + +set -ex + +export MODEL_NAME='glm5.2-744b' +export TOKENIZER_PATH='zai-org/GLM-5.2' + +# Installing torch for checkpoint conversion and forward_pass_logit_checker.py +python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu + +if [ -z "${BASE_OUTPUT_PATH}" ]; then + export BASE_OUTPUT_PATH=gs://runner-maxtext-logs/$(date +%Y-%m-%d-%H-%M) + echo "BASE_OUTPUT_PATH is not set" +fi +BASE_OUTPUT_PATH=${BASE_OUTPUT_PATH%/} +echo using BASE_OUTPUT_PATH = ${BASE_OUTPUT_PATH} + +# Step 1: Checkpoint conversion +# HF checkpoint: https://huggingface.co/zai-org/GLM-5.2 +BF16_LOCAL_PATH=${BF16_LOCAL_PATH:-/home/rishabhbaghel_google_com/glm5.2_raw} + +# scanned +python3 -m maxtext.checkpoint_conversion.to_maxtext src/maxtext/configs/base.yml \ + model_name=${MODEL_NAME} scan_layers=true \ + base_output_directory=${BASE_OUTPUT_PATH}/scanned hf_access_token=$HF_TOKEN \ + hardware=cpu skip_jax_distributed_system=True \ + checkpoint_storage_concurrent_gb=1024 \ + --hf_model_path=$BF16_LOCAL_PATH \ + --lazy_load_tensors=False \ + --eager_load_method=safetensors \ + --save_dtype=bfloat16 diff --git a/tests/end_to_end/tpu/glm5/glm5.2-744b/2_test_glm5.sh b/tests/end_to_end/tpu/glm5/glm5.2-744b/2_test_glm5.sh new file mode 100644 index 0000000000..1f3f556c7b --- /dev/null +++ b/tests/end_to_end/tpu/glm5/glm5.2-744b/2_test_glm5.sh @@ -0,0 +1,45 @@ +#!/bin/bash + +# This file runs Step 2 on TPU cluster for GLM-5.2 (Cross-Layer IndexShare). +# 1. Forward pass logit check against golden logits. +# 2. High-throughput distributed pre-training with IndexShare. +# 3. Decoding & sanity prompt generation. + +set -ex + +export MODEL_NAME='glm5.2-744b' +export TOKENIZER_PATH='zai-org/GLM-5.2' + +# Installing torch CPU for tokenizer / evaluation helpers +python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu + +if [ -z "${BASE_OUTPUT_PATH}" ]; then + export BASE_OUTPUT_PATH=gs://runner-maxtext-logs/$(date +%Y-%m-%d-%H-%M) + echo "BASE_OUTPUT_PATH is not set" +fi +BASE_OUTPUT_PATH=${BASE_OUTPUT_PATH%/} +echo using BASE_OUTPUT_PATH = ${BASE_OUTPUT_PATH} + +SCANNED_CKPT_PATH=${SCANNED_CKPT_PATH:-gs://maxtext-glm5-europe-west4/maxtext-glm-5.2-bf16-converted-final-78l/0/items} +export DATASET_PATH=gs://maxtext-dataset + +# 1. Forward Logit & Generation Test +python3 -m maxtext.inference.decode src/maxtext/configs/base.yml \ + base_output_directory=${BASE_OUTPUT_PATH} \ + run_name=decode_glm52 \ + model_name=${MODEL_NAME} \ + tokenizer_type=huggingface \ + tokenizer_path=${TOKENIZER_PATH} \ + load_parameters_path=${SCANNED_CKPT_PATH} \ + scan_layers=true \ + attention=dot_product \ + sparse_matmul=false \ + dtype=bfloat16 \ + weight_dtype=bfloat16 \ + per_device_batch_size=1 \ + max_prefill_predict_length=64 \ + max_target_length=128 \ + ici_fsdp_parallelism=16 \ + ici_expert_parallelism=4 \ + checkpoint_storage_concurrent_gb=1024 \ + prompt="The capital of France is" From 3bc1393e18e4ccc13e62c5311a5fc9cafd73a8c8 Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Tue, 11 Aug 2026 06:03:14 +0000 Subject: [PATCH 09/21] feat(glm5.2): explicitly pass IndexShare configuration in test script --- tests/end_to_end/tpu/glm5/glm5.2-744b/2_test_glm5.sh | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/end_to_end/tpu/glm5/glm5.2-744b/2_test_glm5.sh b/tests/end_to_end/tpu/glm5/glm5.2-744b/2_test_glm5.sh index 1f3f556c7b..11a031da5e 100644 --- a/tests/end_to_end/tpu/glm5/glm5.2-744b/2_test_glm5.sh +++ b/tests/end_to_end/tpu/glm5/glm5.2-744b/2_test_glm5.sh @@ -23,7 +23,7 @@ echo using BASE_OUTPUT_PATH = ${BASE_OUTPUT_PATH} SCANNED_CKPT_PATH=${SCANNED_CKPT_PATH:-gs://maxtext-glm5-europe-west4/maxtext-glm-5.2-bf16-converted-final-78l/0/items} export DATASET_PATH=gs://maxtext-dataset -# 1. Forward Logit & Generation Test +# 1. Forward Logit & Generation Test with GLM-5.2 Cross-Layer IndexShare python3 -m maxtext.inference.decode src/maxtext/configs/base.yml \ base_output_directory=${BASE_OUTPUT_PATH} \ run_name=decode_glm52 \ @@ -42,4 +42,9 @@ python3 -m maxtext.inference.decode src/maxtext/configs/base.yml \ ici_fsdp_parallelism=16 \ ici_expert_parallelism=4 \ checkpoint_storage_concurrent_gb=1024 \ + use_indexer=true \ + use_index_share=true \ + index_share_pattern="FSSS" \ + prune_shared_indexers=true \ prompt="The capital of France is" + From 325e42fe9e5c49f710625caf4439a1eb56a94dca Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Tue, 11 Aug 2026 06:06:12 +0000 Subject: [PATCH 10/21] feat(eval): add GLM-5.2 sanity evaluation and prompt generation script --- scratch/predict_glm52_prompts.py | 153 +++++++++++++++++++++++++++++++ 1 file changed, 153 insertions(+) create mode 100644 scratch/predict_glm52_prompts.py diff --git a/scratch/predict_glm52_prompts.py b/scratch/predict_glm52_prompts.py new file mode 100644 index 0000000000..929b3c6725 --- /dev/null +++ b/scratch/predict_glm52_prompts.py @@ -0,0 +1,153 @@ +import functools +import os +import sys +from typing import Sequence +import numpy as np +import jax +import jax.numpy as jnp +from transformers import AutoTokenizer + +from maxtext.configs import pyconfig +from maxtext.layers import quantizations +from maxtext.models import models +from maxtext.utils import max_logging +from maxtext.utils import max_utils +from maxtext.utils import maxtext_utils +from maxtext.common.common_types import DECODING_ACTIVE_SEQUENCE_INDICATOR, MODEL_MODE_TRAIN +from maxtext.utils import model_creation_utils + + +def get_top_k(logits_1d, tokenizer, k=10): + probs = jax.nn.softmax(logits_1d, axis=-1) + top_indices = np.argsort(np.asarray(logits_1d))[-k:][::-1] + results = [] + for idx in top_indices: + try: + tok_str = tokenizer.decode([int(idx)]) + except Exception: + tok_str = f"" + results.append((int(idx), tok_str, float(logits_1d[idx]), float(probs[idx]))) + return results + + +def main(argv: Sequence[str]): + import absl.logging + absl.logging.set_verbosity(absl.logging.INFO) + config = pyconfig.initialize(argv) + print("Initializing JAX distributed system for GLM-5.2...", flush=True) + jax.config.update("jax_default_prng_impl", "unsafe_rbg") + devices_array = maxtext_utils.create_device_mesh(config) + mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) + + print(f"JAX Process {jax.process_index()}/{jax.process_count()} initialized. Mesh shape: {mesh.shape}", flush=True) + tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_path, trust_remote_code=True) + print(f"Loaded tokenizer from {config.tokenizer_path}", flush=True) + + print(f"Building GLM-5.2 model from checkpoint: {config.load_parameters_path}...", flush=True) + model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN) + print("GLM-5.2 model created and checkpoint restored successfully!", flush=True) + + test_cases = [ + ("Raw Prompt 1", "The capital of France is"), + ("Raw Prompt 2", "The largest planet in our solar system is"), + ("Raw Prompt 3", "Deep learning is a subset of machine learning that focuses on"), + ("GLM Tagged Math", "<|user|>\nWhat is 25 * 4? Give only the number.\n<|assistant|>\n"), + ("GLM Tagged Code", "<|user|>\nWrite a Python function to check if a number is prime.\n<|assistant|>\n"), + ("GLM Tagged QA", "<|user|>\nWhat is the boiling point of water in Celsius?\n<|assistant|>\n"), + ] + + from flax import nnx + + @nnx.jit + def forward_step(model, tokens, positions, segment_ids): + return model( + decoder_input_tokens=tokens, + decoder_positions=positions, + decoder_segment_ids=segment_ids, + enable_dropout=False, + ) + + max_len = config.max_target_length + output_log_path = "/tmp/glm52_predictions_output.txt" + out_file = open(output_log_path, "w") + + def log_out(msg): + print(msg, flush=True) + out_file.write(msg + "\n") + out_file.flush() + + if jax.process_index() == 0: + log_out("=" * 80) + log_out("GLM-5.2 Cross-Layer IndexShare 744B Model Sanity Evaluation") + log_out(f"Model: {config.model_name} | Checkpoint: {config.load_parameters_path}") + log_out(f"IndexShare Pattern: {config.index_share_pattern} | Use Index Share: {config.use_index_share}") + log_out("=" * 80 + "\n") + + for label, prompt_str in test_cases: + token_ids = tokenizer.encode(prompt_str) + seq_len = len(token_ids) + if jax.process_index() == 0: + log_out("\n" + "=" * 80) + log_out(f"Test Case: [{label}]") + log_out(f"Prompt: {repr(prompt_str)}") + log_out(f"Prompt Tokens ({seq_len} tokens): {token_ids}") + log_out("-" * 80) + + current_tokens = np.zeros((config.global_batch_size_to_train_on, max_len), dtype=np.int32) + current_tokens[:, :seq_len] = np.array(token_ids, dtype=np.int32) + positions = np.stack([np.arange(max_len, dtype=np.int32) for _ in range(config.global_batch_size_to_train_on)]) + segment_ids = np.zeros((config.global_batch_size_to_train_on, max_len), dtype=np.int32) + segment_ids[:, :seq_len] = DECODING_ACTIVE_SEQUENCE_INDICATOR + + # Step 1: Top Next-Token Prediction + logits = forward_step(model, current_tokens, positions, segment_ids) + gathered_logits = jax.experimental.multihost_utils.process_allgather(logits, tiled=True) + if gathered_logits.ndim == 4: + gathered_logits = jnp.reshape(gathered_logits, (-1, max_len, config.vocab_size)) + + last_logits = np.asarray(gathered_logits[0, seq_len - 1, :]) + top_tokens = get_top_k(last_logits, tokenizer, k=10) + + if jax.process_index() == 0: + log_out("\nTop 10 Predictions for Next Token:") + log_out(f"{'Rank':<5} | {'Token ID':<10} | {'Token':<22} | {'Logit':<10} | {'Probability':<12}") + log_out("-" * 68) + for rank, (t_id, t_str, logit_val, prob_val) in enumerate(top_tokens, 1): + log_out(f"{rank:<5} | {t_id:<10} | {repr(t_str):<22} | {logit_val:<10.4f} | {prob_val:<12.6f}") + + # Step 2: Greedy Autoregressive Generation + gen_tokens = list(token_ids) + curr_len = seq_len + max_gen_tokens = min(40, max_len - seq_len) + for _ in range(max_gen_tokens): + if curr_len >= max_len: + break + segment_ids[:, :curr_len] = DECODING_ACTIVE_SEQUENCE_INDICATOR + logits = forward_step(model, current_tokens, positions, segment_ids) + gathered_logits = jax.experimental.multihost_utils.process_allgather(logits, tiled=True) + if gathered_logits.ndim == 4: + gathered_logits = jnp.reshape(gathered_logits, (-1, max_len, config.vocab_size)) + + next_tok = int(np.argmax(np.asarray(gathered_logits[0, curr_len - 1, :]))) + gen_tokens.append(next_tok) + current_tokens[:, curr_len] = next_tok + curr_len += 1 + + if next_tok in [tokenizer.eos_token_id, 154820]: + break + + if jax.process_index() == 0: + continuation_text = tokenizer.decode(gen_tokens[seq_len:]) + full_text = tokenizer.decode(gen_tokens) + log_out(f"\n[Generated Continuation]:\n{repr(continuation_text)}") + log_out(f"\n[Full Generated Text]:\n{repr(full_text)}\n") + + out_file.close() + if jax.process_index() == 0: + gcs_dest = "gs://maxtext-glm5-europe-west4/predictions_glm52_78l.txt" + os.system(f"gcloud storage cp {output_log_path} {gcs_dest} || true") + log_out(f"\nSaved full predictions log to: {gcs_dest}") + + +if __name__ == "__main__": + main(sys.argv[1:]) From 5d47740f0ed0154105498d1e3dd8098f8a328d41 Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Tue, 11 Aug 2026 07:49:55 +0000 Subject: [PATCH 11/21] feat(xprof): add named scopes glm_full_layer_indexer and glm_shared_layer_index_reuse for XProf/XPlane profiling --- src/maxtext/layers/attention_mla.py | 36 +++++++++++++++++------------ src/maxtext/models/glm5.py | 11 +++++++-- 2 files changed, 30 insertions(+), 17 deletions(-) diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index 65565c5549..294229a44b 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -1310,22 +1310,28 @@ def __call__( is_shared = getattr(self.config, "use_index_share", False) and self.is_shared_layer if self.indexer is not None and not is_shared: # Full (F) layer: run indexer forward pass - indexer_mask, topk_indices, indexer_score = self.indexer( - inputs_q=inputs_q, - low_rank_q=low_rank_q, - inputs_kv=inputs_kv, - inputs_positions=inputs_positions, - attention_mask=attention_mask, - decoder_segment_ids=decoder_segment_ids, - previous_chunk=previous_chunk, - kv_cache=self.IndexerKVCache_0, - model_mode=model_mode, - ) - new_indexer_state = (indexer_mask, topk_indices, indexer_score) + with jax.named_scope("glm_full_layer_indexer"): + indexer_mask, topk_indices, indexer_score = self.indexer( + inputs_q=inputs_q, + low_rank_q=low_rank_q, + inputs_kv=inputs_kv, + inputs_positions=inputs_positions, + attention_mask=attention_mask, + decoder_segment_ids=decoder_segment_ids, + previous_chunk=previous_chunk, + kv_cache=self.IndexerKVCache_0, + model_mode=model_mode, + ) + indexer_mask = checkpoint_name(indexer_mask, "full_layer_indexer_mask") + topk_indices = checkpoint_name(topk_indices, "full_layer_topk_indices") + new_indexer_state = (indexer_mask, topk_indices, indexer_score) elif cached_indexer_state is not None: - # Shared (S) layer: inherit cached indexer state from donor F layer - indexer_mask, topk_indices, indexer_score = cached_indexer_state - new_indexer_state = cached_indexer_state + # Shared (S) layer: inherit cached indexer state from donor F layer (zero indexer GEMMs) + with jax.named_scope("glm_shared_layer_index_reuse"): + indexer_mask, topk_indices, indexer_score = cached_indexer_state + indexer_mask = checkpoint_name(indexer_mask, "shared_layer_reused_mask") + topk_indices = checkpoint_name(topk_indices, "shared_layer_reused_indices") + new_indexer_state = cached_indexer_state else: indexer_mask, topk_indices, indexer_score = None, None, None diff --git a/src/maxtext/models/glm5.py b/src/maxtext/models/glm5.py index b9ac43caa1..dc3db3fdb6 100644 --- a/src/maxtext/models/glm5.py +++ b/src/maxtext/models/glm5.py @@ -49,14 +49,21 @@ def __init__( # GLM-5.2 Cross-Layer IndexShare Role Resolution self.is_index_share_enabled = getattr(config, "use_index_share", False) - self.is_shared_layer = False - self.served_group_size = 1 if self.is_index_share_enabled and layer_idx >= 0: from maxtext.utils import index_share_utils + import absl.logging pattern = index_share_utils.parse_index_share_pattern(config.index_share_pattern, config.num_decoder_layers) self.is_shared_layer = index_share_utils.is_shared_layer(layer_idx, pattern) self.served_group_size = index_share_utils.get_served_group_sizes(pattern)[layer_idx] + if layer_idx == 0: + num_f = pattern.count("F") + num_s = pattern.count("S") + absl.logging.info( + f"[GLM-5.2 IndexShare Active] Total layers: {config.num_decoder_layers} | " + f"Pattern: {config.index_share_pattern} | Full (F) layers with active indexers: {num_f} | " + f"Shared (S) layers with pruned indexers: {num_s} (Pruned {num_s / config.num_decoder_layers * 100:.1f}% indexer compute/parameters)" + ) # Re-initialize MLA with GLM-specific IndexShare configuration self.self_attention = attention_mla.MLA( From edd46883b1ccd7de11efde890817ea568c7b5fc9 Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Tue, 11 Aug 2026 07:53:37 +0000 Subject: [PATCH 12/21] fix(glm5): initialize default is_shared_layer and served_group_size for abstract scanned layers --- src/maxtext/models/glm5.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/maxtext/models/glm5.py b/src/maxtext/models/glm5.py index dc3db3fdb6..2ddaf2be2f 100644 --- a/src/maxtext/models/glm5.py +++ b/src/maxtext/models/glm5.py @@ -49,6 +49,8 @@ def __init__( # GLM-5.2 Cross-Layer IndexShare Role Resolution self.is_index_share_enabled = getattr(config, "use_index_share", False) + self.is_shared_layer = False + self.served_group_size = 1 if self.is_index_share_enabled and layer_idx >= 0: from maxtext.utils import index_share_utils import absl.logging From a5b64e6d511d9002f37b9daa77a6fb4f97eff74c Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Tue, 11 Aug 2026 09:24:54 +0000 Subject: [PATCH 13/21] feat(glm5.2): enable IndexShare carry in scanned layers execution --- src/maxtext/layers/attention_mla.py | 14 ++- src/maxtext/layers/nnx_decoders.py | 144 +++++++++++++++++++--------- src/maxtext/models/glm5.py | 8 ++ 3 files changed, 121 insertions(+), 45 deletions(-) diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index 294229a44b..086cafd272 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -1254,6 +1254,7 @@ def __call__( kv_cache: Optional[Array] = None, attention_metadata: Optional[dict[str, Any]] = None, cached_indexer_state: Optional[Any] = None, + layer_idx: Optional[Any] = None, ) -> tuple[Array, Optional[Array]] | tuple[Array, Optional[Array], Optional[Any]]: """Forward pass for MLA, reusing `AttentionOp` for the actual attention. @@ -1286,7 +1287,9 @@ def __call__( if model_mode != MODEL_MODE_TRAIN and decoder_segment_ids is None: decoder_segment_ids = jnp.ones(inputs_q.shape[:2], dtype=jnp.int32) - query, low_rank_q = self.mla_query_projection(inputs_q, inputs_positions, model_mode) + query, low_rank_q = self.mla_q_projection( + inputs_q, inputs_positions, decoder_segment_ids, model_mode, previous_chunk, rope_kwargs + ) if self.config.force_q_layout: query = layout.with_layout_constraint(query, DLL(major_to_minor=(0, 2, 3, 1))) key, value, cached_values = self.mla_kv_projection( @@ -1307,7 +1310,14 @@ def __call__( if attention_mask is not None: attention_mask = attention_mask.squeeze(axis=(1, 2)) - is_shared = getattr(self.config, "use_index_share", False) and self.is_shared_layer + if getattr(self.config, "use_index_share", False): + if layer_idx is not None: + is_shared = (layer_idx % 4 != 0) + else: + is_shared = self.is_shared_layer + else: + is_shared = False + if self.indexer is not None and not is_shared: # Full (F) layer: run indexer forward pass with jax.named_scope("glm_full_layer_indexer"): diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index f4e59e0dc2..703ea355dc 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -989,6 +989,14 @@ def _extract_matching_state(template, full): updated_graphdef = [graphdef] use_kv = kv_caches_stacked is not None + is_index_share = getattr(self.config, "use_index_share", False) + cached_indexer_state = kwargs.get("cached_indexer_state", None) + start_layer_idx = kwargs.get("start_layer_idx", 0) + + if is_index_share: + init_scan_carry = (x_in, cached_indexer_state, start_layer_idx) + else: + init_scan_carry = x_in def layer_fn(carry, scanned_vars): # Ensure metadata rank matches the sliced values @@ -1014,14 +1022,28 @@ def layer_fn(carry, scanned_vars): if kv_cache_layer is not None: call_kwargs["kv_cache"] = kv_cache_layer - layer_out = layer(carry, *args, **call_kwargs) + if is_index_share: + y_in, current_cached_indexer, lyr_idx = carry + call_kwargs["cached_indexer_state"] = current_cached_indexer + call_kwargs["layer_idx"] = lyr_idx + else: + y_in = carry + + layer_out = layer(y_in, *args, **call_kwargs) if isinstance(layer_out, tuple): - new_carry = layer_out[0] + new_carry_y = layer_out[0] updated_kv = layer_out[1] if len(layer_out) > 1 else None + new_indexer_state = layer_out[2] if len(layer_out) > 2 else None else: - new_carry = layer_out + new_carry_y = layer_out updated_kv = None + new_indexer_state = None + + if is_index_share: + new_carry = (new_carry_y, new_indexer_state, lyr_idx + 1) + else: + new_carry = new_carry_y # Extract the updated state to return it if dynamic_graph_init: @@ -1054,7 +1076,7 @@ def layer_fn(carry, scanned_vars): # kv_caches_stacked is actually the original kv_caches list in this new flow kv_caches_list = kv_caches_stacked - current_carry = x_in + current_carry = init_scan_carry for i in range(length): # Statically slice the parameters and state for this layer @@ -1069,16 +1091,28 @@ def layer_fn(carry, scanned_vars): # Update the list in-place (mutates the list passed by reference) kv_caches_list[i] = updated_kv + if is_index_share: + final_carry, out_indexer_state, _ = current_carry + else: + final_carry = current_carry + out_indexer_state = None + # We don't need to rebuild scanned_state or return it because during # inference with vLLM, parameters do not change and we don't need intermediates. - return current_carry, layers, None + return final_carry, layers, None, out_indexer_state else: params = maxtext_utils_nnx.nnx_ensure_scan_leading_axis(params, length) state = maxtext_utils_nnx.nnx_ensure_scan_leading_axis(state, length) - final_carry, scanned_state = jax.lax.scan(layer_fn_wrapped, x_in, (params, state), unroll=unroll) + scan_res_carry, scanned_state = jax.lax.scan(layer_fn_wrapped, init_scan_carry, (params, state), unroll=unroll) returned_kv_stacked = None + if is_index_share: + final_carry, out_indexer_state, _ = scan_res_carry + else: + final_carry = scan_res_carry + out_indexer_state = None + # Move the scan axis to each variable's param_scan_axis and restore its name # in the sharding metadata. jax.lax.scan emits it at position 0. scanned_state = maxtext_utils_nnx.nnx_add_and_sync_scan_axis(scanned_state, metadata_axis_name) @@ -1095,6 +1129,8 @@ def layer_fn(carry, scanned_vars): nnx.update(layers, clean_state) out_layers = layers + if is_index_share: + return final_carry, out_layers, returned_kv_stacked if use_kv else None, out_indexer_state return final_carry, out_layers, returned_kv_stacked if use_kv else None def get_decoder_layers(self): @@ -1792,52 +1828,74 @@ def __call__( *layer_args, **common_kwargs, ) - else: - y, self.dense_layers, _ = self._apply_layers_sequentially( - self.dense_layers, - y, - *layer_args, - length=cfg.first_num_dense_layers, - **layer_kwargs, - ) - - num_moe = cfg.num_decoder_layers - cfg.first_num_dense_layers + if getattr(cfg, "use_index_share", False): + y, self.dense_layers, _, cached_indexer_state = self._apply_layers_sequentially( + self.dense_layers, + y, + *layer_args, + length=cfg.first_num_dense_layers, + start_layer_idx=0, + cached_indexer_state=None, + **layer_kwargs, + ) - if cfg.use_batch_split_schedule: - policy = self.get_remat_policy() - mock_params = self._build_linen_params(self.moe_layers) + num_moe = cfg.num_decoder_layers - cfg.first_num_dense_layers - if cfg.quantization and cfg.use_qwix_quantization and not cfg.use_manual_quantization: - y = deepseek_batchsplit_fp8.scan_batch_split_layers( - y, - mock_params, - decoder_positions, - decoder_segment_ids, - model_mode=model_mode, - mesh=self.mesh, - quant=self.quant, - cfg=cfg, - policy=policy, - ) - else: - # bf16 code path - y = deepseek_batchsplit.scan_batch_split_layers( - y, - mock_params, - decoder_positions, - mesh=self.mesh, - cfg=cfg, - num_layers=num_moe, - ) - else: - y, self.moe_layers, _ = self._apply_layers_sequentially( + y, self.moe_layers, _, _ = self._apply_layers_sequentially( self.moe_layers, y, *layer_args, length=num_moe, + start_layer_idx=cfg.first_num_dense_layers, + cached_indexer_state=cached_indexer_state, + **layer_kwargs, + ) + else: + y, self.dense_layers, _ = self._apply_layers_sequentially( + self.dense_layers, + y, + *layer_args, + length=cfg.first_num_dense_layers, **layer_kwargs, ) + num_moe = cfg.num_decoder_layers - cfg.first_num_dense_layers + + if cfg.use_batch_split_schedule: + policy = self.get_remat_policy() + mock_params = self._build_linen_params(self.moe_layers) + + if cfg.quantization and cfg.use_qwix_quantization and not cfg.use_manual_quantization: + y = deepseek_batchsplit_fp8.scan_batch_split_layers( + y, + mock_params, + decoder_positions, + decoder_segment_ids, + model_mode=model_mode, + mesh=self.mesh, + quant=self.quant, + cfg=cfg, + policy=policy, + ) + else: + # bf16 code path + y = deepseek_batchsplit.scan_batch_split_layers( + y, + mock_params, + decoder_positions, + mesh=self.mesh, + cfg=cfg, + num_layers=num_moe, + ) + else: + y, self.moe_layers, _ = self._apply_layers_sequentially( + self.moe_layers, + y, + *layer_args, + length=num_moe, + **layer_kwargs, + ) + elif self.is_deepseek4: y = self._apply_deepseek4_scanned_blocks( y, diff --git a/src/maxtext/models/glm5.py b/src/maxtext/models/glm5.py index 2ddaf2be2f..c1c7c6c9cf 100644 --- a/src/maxtext/models/glm5.py +++ b/src/maxtext/models/glm5.py @@ -112,6 +112,7 @@ def attention_op( previous_chunk=None, slot: None | int = None, cached_indexer_state=None, + layer_idx=None, ): """Executes the attention layer and passes cached indexer state.""" attn_out = self.self_attention( @@ -125,6 +126,7 @@ def attention_op( previous_chunk=previous_chunk, slot=slot, cached_indexer_state=cached_indexer_state, + layer_idx=layer_idx, ) if self.is_index_share_enabled: attention_result, _, new_indexer_state = attn_out @@ -169,6 +171,7 @@ def self_attention_with_norm_op( previous_chunk=None, slot: None | int = None, cached_indexer_state=None, + layer_idx=None, ): """Self-attention with normalization and IndexShare caching.""" if self.is_mhc_enabled: @@ -198,6 +201,7 @@ def self_attention_with_norm_op( previous_chunk, slot, cached_indexer_state=cached_indexer_state, + layer_idx=layer_idx, ) intermediate_inputs = inputs + attention_lnx # Normalization @@ -249,6 +253,7 @@ def __call__( attention_metadata=None, decoder_input_tokens=None, cached_indexer_state=None, + layer_idx=None, ): if isinstance(inputs, tuple): inputs = inputs[0] @@ -268,6 +273,7 @@ def __call__( previous_chunk, slot, cached_indexer_state=cached_indexer_state, + layer_idx=layer_idx, ) if self.is_mhc_enabled: @@ -329,6 +335,7 @@ def __call__( attention_metadata=None, decoder_input_tokens=None, cached_indexer_state=None, + layer_idx=None, ): if isinstance(inputs, tuple): inputs = inputs[0] @@ -349,6 +356,7 @@ def __call__( previous_chunk, slot, cached_indexer_state=cached_indexer_state, + layer_idx=layer_idx, ) if self.is_mhc_enabled: From 00eb67295fb4f7421af5744ec2f64c6e93c441ce Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Tue, 11 Aug 2026 10:29:24 +0000 Subject: [PATCH 14/21] fix(profiler): block until ready on active profiled steps to capture full 78-layer forward and backward passes --- src/maxtext/common/profiler.py | 13 +++++++++++-- src/maxtext/trainers/pre_train/train.py | 2 ++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/src/maxtext/common/profiler.py b/src/maxtext/common/profiler.py index f6f56c372e..e633ab2b4a 100644 --- a/src/maxtext/common/profiler.py +++ b/src/maxtext/common/profiler.py @@ -75,6 +75,13 @@ def __init__(self, config, offset_step=0): if advanced_config: self.profiling_options.advanced_configuration = advanced_config + def is_active(self, step=None): + if self.mode == "": + return False + if step is not None: + return self.start_initial_profile_step <= step <= self.finished_initial_profile_step + return getattr(self, "_active", False) + def maybe_activate_profiler(self, step, state): """Conditionally activates the profiler based on the current step. This method checks if the current training step matches the step designated @@ -83,13 +90,14 @@ def maybe_activate_profiler(self, step, state): """ if self.mode != "" and (step == self.start_initial_profile_step or self.should_activate_periodic_profile(step)): optional_postfix = f"step_{step}" if self.profile_period > 0 else "" + self._active = True self.activate(blocking_object=state, optional_postfix=optional_postfix) def activate(self, blocking_object=None, optional_postfix=""): """Start the profiler. nsys profiler becomes no-op when libcudart.so is not available on the system.""" if self.profile_cleanly and blocking_object is not None: - jax.block_until_ready(blocking_object) + jax.tree_util.tree_map(lambda x: x.block_until_ready() if hasattr(x, "block_until_ready") else x, blocking_object) if self.managed_mldiagnostics and self.mode == "xplane": # Handle the special profiling logic for managed_mldiagnostics @@ -121,13 +129,14 @@ def maybe_deactivate_profiler(self, step, state): deactivating a periodic profile. """ if self.mode != "" and (step == self.finished_initial_profile_step or self.should_deactivate_periodic_profile(step)): + self._active = False self.deactivate(blocking_object=state) def deactivate(self, blocking_object=None): """End the profiler. The result is uploaded to the output bucket.""" if self.profile_cleanly and blocking_object is not None: - jax.block_until_ready(blocking_object) + jax.tree_util.tree_map(lambda x: x.block_until_ready() if hasattr(x, "block_until_ready") else x, blocking_object) if self.managed_mldiagnostics and self.mode == "xplane": # Handle the special profileing logic for managed_mldiagnostics diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index a992ca6fab..ead22423e6 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -755,6 +755,8 @@ def training_loop_iteration( if shard_optimizer_over_data and isinstance(model, nn.Module): state = sharding.maybe_shard_with_name(state, state_mesh_shardings, shard_mode) state, metrics = p_train_step(state, example_batch, *step_rng_args) + if prof.is_active(step): + jax.tree_util.tree_map(lambda x: x.block_until_ready() if hasattr(x, "block_until_ready") else x, metrics) step_time_delta = datetime.datetime.now() - last_step_completion last_step_completion = datetime.datetime.now() From ecc8152b6d9f8d35e2e26aba1a629a0d68d0c34a Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Tue, 11 Aug 2026 10:38:05 +0000 Subject: [PATCH 15/21] fix(indexshare): provide invariant concrete dummy tensor structure for scan carry --- src/maxtext/layers/nnx_decoders.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 703ea355dc..ef3a78b2a2 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -994,6 +994,15 @@ def _extract_matching_state(template, full): start_layer_idx = kwargs.get("start_layer_idx", 0) if is_index_share: + cached_indexer_state = kwargs.get("cached_indexer_state", None) + if cached_indexer_state is None: + batch, seq_len = x_in.shape[0], x_in.shape[1] + topk = getattr(self.config, "indexer_topk", 2048) + n_heads = getattr(self.config, "indexer_n_heads", 32) + dummy_mask = jnp.zeros((batch, seq_len, seq_len), dtype=jnp.bool_) + dummy_indices = jnp.zeros((batch, seq_len, topk), dtype=jnp.int32) + dummy_score = jnp.zeros((batch, n_heads, seq_len, seq_len), dtype=jnp.float32) + cached_indexer_state = (dummy_mask, dummy_indices, dummy_score) init_scan_carry = (x_in, cached_indexer_state, start_layer_idx) else: init_scan_carry = x_in From 64bc8e5c9730757eaa513ee50270e2e3593beec3 Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Tue, 11 Aug 2026 10:50:36 +0000 Subject: [PATCH 16/21] fix(decoder): fix indentation of scanned use_index_share execution block --- src/maxtext/layers/nnx_decoders.py | 118 ++++++++++++++--------------- 1 file changed, 59 insertions(+), 59 deletions(-) diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index ef3a78b2a2..1d11fe5b8e 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -1837,74 +1837,74 @@ def __call__( *layer_args, **common_kwargs, ) - if getattr(cfg, "use_index_share", False): - y, self.dense_layers, _, cached_indexer_state = self._apply_layers_sequentially( - self.dense_layers, - y, - *layer_args, - length=cfg.first_num_dense_layers, - start_layer_idx=0, - cached_indexer_state=None, - **layer_kwargs, - ) + elif getattr(cfg, "use_index_share", False): + y, self.dense_layers, _, cached_indexer_state = self._apply_layers_sequentially( + self.dense_layers, + y, + *layer_args, + length=cfg.first_num_dense_layers, + start_layer_idx=0, + cached_indexer_state=None, + **layer_kwargs, + ) - num_moe = cfg.num_decoder_layers - cfg.first_num_dense_layers + num_moe = cfg.num_decoder_layers - cfg.first_num_dense_layers + + y, self.moe_layers, _, _ = self._apply_layers_sequentially( + self.moe_layers, + y, + *layer_args, + length=num_moe, + start_layer_idx=cfg.first_num_dense_layers, + cached_indexer_state=cached_indexer_state, + **layer_kwargs, + ) + else: + y, self.dense_layers, _ = self._apply_layers_sequentially( + self.dense_layers, + y, + *layer_args, + length=cfg.first_num_dense_layers, + **layer_kwargs, + ) - y, self.moe_layers, _, _ = self._apply_layers_sequentially( + num_moe = cfg.num_decoder_layers - cfg.first_num_dense_layers + + if cfg.use_batch_split_schedule: + policy = self.get_remat_policy() + mock_params = self._build_linen_params(self.moe_layers) + + if cfg.quantization and cfg.use_qwix_quantization and not cfg.use_manual_quantization: + y = deepseek_batchsplit_fp8.scan_batch_split_layers( + y, + mock_params, + decoder_positions, + decoder_segment_ids, + model_mode=model_mode, + mesh=self.mesh, + quant=self.quant, + cfg=cfg, + policy=policy, + ) + else: + # bf16 code path + y = deepseek_batchsplit.scan_batch_split_layers( + y, + mock_params, + decoder_positions, + mesh=self.mesh, + cfg=cfg, + num_layers=num_moe, + ) + else: + y, self.moe_layers, _ = self._apply_layers_sequentially( self.moe_layers, y, *layer_args, length=num_moe, - start_layer_idx=cfg.first_num_dense_layers, - cached_indexer_state=cached_indexer_state, - **layer_kwargs, - ) - else: - y, self.dense_layers, _ = self._apply_layers_sequentially( - self.dense_layers, - y, - *layer_args, - length=cfg.first_num_dense_layers, **layer_kwargs, ) - num_moe = cfg.num_decoder_layers - cfg.first_num_dense_layers - - if cfg.use_batch_split_schedule: - policy = self.get_remat_policy() - mock_params = self._build_linen_params(self.moe_layers) - - if cfg.quantization and cfg.use_qwix_quantization and not cfg.use_manual_quantization: - y = deepseek_batchsplit_fp8.scan_batch_split_layers( - y, - mock_params, - decoder_positions, - decoder_segment_ids, - model_mode=model_mode, - mesh=self.mesh, - quant=self.quant, - cfg=cfg, - policy=policy, - ) - else: - # bf16 code path - y = deepseek_batchsplit.scan_batch_split_layers( - y, - mock_params, - decoder_positions, - mesh=self.mesh, - cfg=cfg, - num_layers=num_moe, - ) - else: - y, self.moe_layers, _ = self._apply_layers_sequentially( - self.moe_layers, - y, - *layer_args, - length=num_moe, - **layer_kwargs, - ) - elif self.is_deepseek4: y = self._apply_deepseek4_scanned_blocks( y, From 171a6ba2734f260237aff77e77bbf290ef356595 Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Tue, 11 Aug 2026 10:52:34 +0000 Subject: [PATCH 17/21] fix(mla): call correct mla_query_projection method --- src/maxtext/layers/attention_mla.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index 086cafd272..e126100ad3 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -1287,9 +1287,7 @@ def __call__( if model_mode != MODEL_MODE_TRAIN and decoder_segment_ids is None: decoder_segment_ids = jnp.ones(inputs_q.shape[:2], dtype=jnp.int32) - query, low_rank_q = self.mla_q_projection( - inputs_q, inputs_positions, decoder_segment_ids, model_mode, previous_chunk, rope_kwargs - ) + query, low_rank_q = self.mla_query_projection(inputs_q, inputs_positions, model_mode) if self.config.force_q_layout: query = layout.with_layout_constraint(query, DLL(major_to_minor=(0, 2, 3, 1))) key, value, cached_values = self.mla_kv_projection( From aae73ee20d80ab9e5f5620b623989f5dfd3945c4 Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Tue, 11 Aug 2026 10:54:02 +0000 Subject: [PATCH 18/21] fix(mla): use jax.lax.cond for scanned indexer conditional execution --- src/maxtext/layers/attention_mla.py | 72 ++++++++++++++++------------- 1 file changed, 41 insertions(+), 31 deletions(-) diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index e126100ad3..99e97ffbaa 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -1308,38 +1308,48 @@ def __call__( if attention_mask is not None: attention_mask = attention_mask.squeeze(axis=(1, 2)) - if getattr(self.config, "use_index_share", False): - if layer_idx is not None: - is_shared = (layer_idx % 4 != 0) + if self.indexer is not None: + def _run_full(_): + with jax.named_scope("glm_full_layer_indexer"): + mask, indices, score = self.indexer( + inputs_q=inputs_q, + low_rank_q=low_rank_q, + inputs_kv=inputs_kv, + inputs_positions=inputs_positions, + attention_mask=attention_mask, + decoder_segment_ids=decoder_segment_ids, + previous_chunk=previous_chunk, + kv_cache=self.IndexerKVCache_0, + model_mode=model_mode, + ) + mask = checkpoint_name(mask, "full_layer_indexer_mask") + indices = checkpoint_name(indices, "full_layer_topk_indices") + return mask, indices, score + + def _run_shared(_): + with jax.named_scope("glm_shared_layer_index_reuse"): + mask, indices, score = cached_indexer_state + mask = checkpoint_name(mask, "shared_layer_reused_mask") + indices = checkpoint_name(indices, "shared_layer_reused_indices") + return mask, indices, score + + if getattr(self.config, "use_index_share", False) and cached_indexer_state is not None: + if layer_idx is not None: + is_full = (layer_idx % 4 == 0) + indexer_mask, topk_indices, indexer_score = jax.lax.cond( + is_full, + _run_full, + _run_shared, + operand=None, + ) + elif self.is_shared_layer: + indexer_mask, topk_indices, indexer_score = _run_shared(None) + else: + indexer_mask, topk_indices, indexer_score = _run_full(None) else: - is_shared = self.is_shared_layer - else: - is_shared = False - - if self.indexer is not None and not is_shared: - # Full (F) layer: run indexer forward pass - with jax.named_scope("glm_full_layer_indexer"): - indexer_mask, topk_indices, indexer_score = self.indexer( - inputs_q=inputs_q, - low_rank_q=low_rank_q, - inputs_kv=inputs_kv, - inputs_positions=inputs_positions, - attention_mask=attention_mask, - decoder_segment_ids=decoder_segment_ids, - previous_chunk=previous_chunk, - kv_cache=self.IndexerKVCache_0, - model_mode=model_mode, - ) - indexer_mask = checkpoint_name(indexer_mask, "full_layer_indexer_mask") - topk_indices = checkpoint_name(topk_indices, "full_layer_topk_indices") - new_indexer_state = (indexer_mask, topk_indices, indexer_score) - elif cached_indexer_state is not None: - # Shared (S) layer: inherit cached indexer state from donor F layer (zero indexer GEMMs) - with jax.named_scope("glm_shared_layer_index_reuse"): - indexer_mask, topk_indices, indexer_score = cached_indexer_state - indexer_mask = checkpoint_name(indexer_mask, "shared_layer_reused_mask") - topk_indices = checkpoint_name(topk_indices, "shared_layer_reused_indices") - new_indexer_state = cached_indexer_state + indexer_mask, topk_indices, indexer_score = _run_full(None) + + new_indexer_state = (indexer_mask, topk_indices, indexer_score) else: indexer_mask, topk_indices, indexer_score = None, None, None From e55e0a7fd47556fe4b20313f47cc54f19b6357d9 Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Tue, 11 Aug 2026 10:55:30 +0000 Subject: [PATCH 19/21] fix(indexshare): match exact dummy indexer mask and score shapes and dtypes --- src/maxtext/layers/nnx_decoders.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 1d11fe5b8e..add7ea07c9 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -998,10 +998,9 @@ def _extract_matching_state(template, full): if cached_indexer_state is None: batch, seq_len = x_in.shape[0], x_in.shape[1] topk = getattr(self.config, "indexer_topk", 2048) - n_heads = getattr(self.config, "indexer_n_heads", 32) - dummy_mask = jnp.zeros((batch, seq_len, seq_len), dtype=jnp.bool_) + dummy_mask = jnp.zeros((batch, seq_len, seq_len), dtype=jnp.bfloat16) dummy_indices = jnp.zeros((batch, seq_len, topk), dtype=jnp.int32) - dummy_score = jnp.zeros((batch, n_heads, seq_len, seq_len), dtype=jnp.float32) + dummy_score = jnp.zeros((batch, seq_len, seq_len), dtype=jnp.float32) cached_indexer_state = (dummy_mask, dummy_indices, dummy_score) init_scan_carry = (x_in, cached_indexer_state, start_layer_idx) else: From 6a07b6eb23979c921ba3e635ae1550d5e8d789f3 Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Tue, 11 Aug 2026 13:09:16 +0000 Subject: [PATCH 20/21] chore: remove scratch script from repository --- scratch/predict_glm52_prompts.py | 153 ------------------------------- 1 file changed, 153 deletions(-) delete mode 100644 scratch/predict_glm52_prompts.py diff --git a/scratch/predict_glm52_prompts.py b/scratch/predict_glm52_prompts.py deleted file mode 100644 index 929b3c6725..0000000000 --- a/scratch/predict_glm52_prompts.py +++ /dev/null @@ -1,153 +0,0 @@ -import functools -import os -import sys -from typing import Sequence -import numpy as np -import jax -import jax.numpy as jnp -from transformers import AutoTokenizer - -from maxtext.configs import pyconfig -from maxtext.layers import quantizations -from maxtext.models import models -from maxtext.utils import max_logging -from maxtext.utils import max_utils -from maxtext.utils import maxtext_utils -from maxtext.common.common_types import DECODING_ACTIVE_SEQUENCE_INDICATOR, MODEL_MODE_TRAIN -from maxtext.utils import model_creation_utils - - -def get_top_k(logits_1d, tokenizer, k=10): - probs = jax.nn.softmax(logits_1d, axis=-1) - top_indices = np.argsort(np.asarray(logits_1d))[-k:][::-1] - results = [] - for idx in top_indices: - try: - tok_str = tokenizer.decode([int(idx)]) - except Exception: - tok_str = f"" - results.append((int(idx), tok_str, float(logits_1d[idx]), float(probs[idx]))) - return results - - -def main(argv: Sequence[str]): - import absl.logging - absl.logging.set_verbosity(absl.logging.INFO) - config = pyconfig.initialize(argv) - print("Initializing JAX distributed system for GLM-5.2...", flush=True) - jax.config.update("jax_default_prng_impl", "unsafe_rbg") - devices_array = maxtext_utils.create_device_mesh(config) - mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) - - print(f"JAX Process {jax.process_index()}/{jax.process_count()} initialized. Mesh shape: {mesh.shape}", flush=True) - tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_path, trust_remote_code=True) - print(f"Loaded tokenizer from {config.tokenizer_path}", flush=True) - - print(f"Building GLM-5.2 model from checkpoint: {config.load_parameters_path}...", flush=True) - model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN) - print("GLM-5.2 model created and checkpoint restored successfully!", flush=True) - - test_cases = [ - ("Raw Prompt 1", "The capital of France is"), - ("Raw Prompt 2", "The largest planet in our solar system is"), - ("Raw Prompt 3", "Deep learning is a subset of machine learning that focuses on"), - ("GLM Tagged Math", "<|user|>\nWhat is 25 * 4? Give only the number.\n<|assistant|>\n"), - ("GLM Tagged Code", "<|user|>\nWrite a Python function to check if a number is prime.\n<|assistant|>\n"), - ("GLM Tagged QA", "<|user|>\nWhat is the boiling point of water in Celsius?\n<|assistant|>\n"), - ] - - from flax import nnx - - @nnx.jit - def forward_step(model, tokens, positions, segment_ids): - return model( - decoder_input_tokens=tokens, - decoder_positions=positions, - decoder_segment_ids=segment_ids, - enable_dropout=False, - ) - - max_len = config.max_target_length - output_log_path = "/tmp/glm52_predictions_output.txt" - out_file = open(output_log_path, "w") - - def log_out(msg): - print(msg, flush=True) - out_file.write(msg + "\n") - out_file.flush() - - if jax.process_index() == 0: - log_out("=" * 80) - log_out("GLM-5.2 Cross-Layer IndexShare 744B Model Sanity Evaluation") - log_out(f"Model: {config.model_name} | Checkpoint: {config.load_parameters_path}") - log_out(f"IndexShare Pattern: {config.index_share_pattern} | Use Index Share: {config.use_index_share}") - log_out("=" * 80 + "\n") - - for label, prompt_str in test_cases: - token_ids = tokenizer.encode(prompt_str) - seq_len = len(token_ids) - if jax.process_index() == 0: - log_out("\n" + "=" * 80) - log_out(f"Test Case: [{label}]") - log_out(f"Prompt: {repr(prompt_str)}") - log_out(f"Prompt Tokens ({seq_len} tokens): {token_ids}") - log_out("-" * 80) - - current_tokens = np.zeros((config.global_batch_size_to_train_on, max_len), dtype=np.int32) - current_tokens[:, :seq_len] = np.array(token_ids, dtype=np.int32) - positions = np.stack([np.arange(max_len, dtype=np.int32) for _ in range(config.global_batch_size_to_train_on)]) - segment_ids = np.zeros((config.global_batch_size_to_train_on, max_len), dtype=np.int32) - segment_ids[:, :seq_len] = DECODING_ACTIVE_SEQUENCE_INDICATOR - - # Step 1: Top Next-Token Prediction - logits = forward_step(model, current_tokens, positions, segment_ids) - gathered_logits = jax.experimental.multihost_utils.process_allgather(logits, tiled=True) - if gathered_logits.ndim == 4: - gathered_logits = jnp.reshape(gathered_logits, (-1, max_len, config.vocab_size)) - - last_logits = np.asarray(gathered_logits[0, seq_len - 1, :]) - top_tokens = get_top_k(last_logits, tokenizer, k=10) - - if jax.process_index() == 0: - log_out("\nTop 10 Predictions for Next Token:") - log_out(f"{'Rank':<5} | {'Token ID':<10} | {'Token':<22} | {'Logit':<10} | {'Probability':<12}") - log_out("-" * 68) - for rank, (t_id, t_str, logit_val, prob_val) in enumerate(top_tokens, 1): - log_out(f"{rank:<5} | {t_id:<10} | {repr(t_str):<22} | {logit_val:<10.4f} | {prob_val:<12.6f}") - - # Step 2: Greedy Autoregressive Generation - gen_tokens = list(token_ids) - curr_len = seq_len - max_gen_tokens = min(40, max_len - seq_len) - for _ in range(max_gen_tokens): - if curr_len >= max_len: - break - segment_ids[:, :curr_len] = DECODING_ACTIVE_SEQUENCE_INDICATOR - logits = forward_step(model, current_tokens, positions, segment_ids) - gathered_logits = jax.experimental.multihost_utils.process_allgather(logits, tiled=True) - if gathered_logits.ndim == 4: - gathered_logits = jnp.reshape(gathered_logits, (-1, max_len, config.vocab_size)) - - next_tok = int(np.argmax(np.asarray(gathered_logits[0, curr_len - 1, :]))) - gen_tokens.append(next_tok) - current_tokens[:, curr_len] = next_tok - curr_len += 1 - - if next_tok in [tokenizer.eos_token_id, 154820]: - break - - if jax.process_index() == 0: - continuation_text = tokenizer.decode(gen_tokens[seq_len:]) - full_text = tokenizer.decode(gen_tokens) - log_out(f"\n[Generated Continuation]:\n{repr(continuation_text)}") - log_out(f"\n[Full Generated Text]:\n{repr(full_text)}\n") - - out_file.close() - if jax.process_index() == 0: - gcs_dest = "gs://maxtext-glm5-europe-west4/predictions_glm52_78l.txt" - os.system(f"gcloud storage cp {output_log_path} {gcs_dest} || true") - log_out(f"\nSaved full predictions log to: {gcs_dest}") - - -if __name__ == "__main__": - main(sys.argv[1:]) From b8b51836430b01d9cb06bf44e0b6ce5e8b7ddba5 Mon Sep 17 00:00:00 2001 From: Rishabh Baghel Date: Thu, 13 Aug 2026 11:28:23 +0000 Subject: [PATCH 21/21] fix(glm5.2): address review comments for indexshare and checkpoint conversion - Dynamically search backwards for preceding donor indexer layers during checkpoint conversion - Precompute is_full_array and served_group_size_array for dynamic JAX tracing in scanned MLA execution - Guard IndexShare initialization log with process_index == 0 to prevent multi-host log spam - Add test coverage for multi-group donor index resolution --- .../checkpoint_conversion/to_maxtext.py | 39 ++++++++++--------- src/maxtext/layers/attention_mla.py | 24 ++++++++++-- src/maxtext/models/glm5.py | 2 +- tests/unit/glm52_indexshare_test.py | 12 ++++++ 4 files changed, 55 insertions(+), 22 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py index 61982b452f..4149baa6fb 100644 --- a/src/maxtext/checkpoint_conversion/to_maxtext.py +++ b/src/maxtext/checkpoint_conversion/to_maxtext.py @@ -471,11 +471,17 @@ def _build_single_axis_stacked_tensor( m = re.match(r"model\.layers\.(\d+)\.(.+)", str(hf_key_single)) if m: + layer_idx = int(m.group(1)) rest = m.group(2) - donor_key = f"model.layers.0.{rest}" - try: - hf_tensor_numpy = tensor_getter_fn(donor_key) - except Exception: + # Search backwards for the closest preceding donor layer containing the key + for candidate_idx in range(layer_idx - 1, -1, -1): + donor_key = f"model.layers.{candidate_idx}.{rest}" + try: + hf_tensor_numpy = tensor_getter_fn(donor_key) + break + except Exception: + continue + else: hf_tensor_numpy = np.zeros(mt_slice_shape, dtype=np.float32) else: raise e @@ -1023,15 +1029,9 @@ def _eager_getter(key): if m: layer_idx = int(m.group(1)) rest = m.group(2) - matching_layers = [ - int(k.split(".")[2]) - for k in hf_state_dict_numpy - if k.startswith("model.layers.") and k.endswith(f".{rest}") - ] - if matching_layers: - preceding = [l for l in matching_layers if l <= layer_idx] - donor_idx = max(preceding) if preceding else min(matching_layers) - donor_key = f"model.layers.{donor_idx}.{rest}" + # Search backwards for the closest preceding donor layer containing the key + for candidate_idx in range(layer_idx - 1, -1, -1): + donor_key = f"model.layers.{candidate_idx}.{rest}" if donor_key in hf_state_dict_numpy: return _eager_getter(donor_key) raise ValueError(f"HuggingFace key {key} not found in state_dict.") @@ -1064,12 +1064,15 @@ def _index_share_tensor_getter(key): m = re.match(r"model\.layers\.(\d+)\.(.+)", key) if m: + layer_idx = int(m.group(1)) rest = m.group(2) - donor_key = f"model.layers.0.{rest}" - try: - return orig_tensor_getter(donor_key) - except Exception: - pass + # Search backwards for the closest preceding donor layer containing the key + for candidate_idx in range(layer_idx - 1, -1, -1): + donor_key = f"model.layers.{candidate_idx}.{rest}" + try: + return orig_tensor_getter(donor_key) + except (ValueError, KeyError): + continue raise e tensor_getter = _index_share_tensor_getter diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index 99e97ffbaa..60646e9cff 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -737,6 +737,18 @@ def __init__( self.use_indexer = config.use_indexer self.is_shared_layer = is_shared_layer self.served_group_size = served_group_size + if getattr(config, "use_index_share", False): + from maxtext.utils import index_share_utils + + pattern = index_share_utils.parse_index_share_pattern( + config.index_share_pattern, config.num_decoder_layers + ) + self.is_full_array = jnp.array([role == "F" for role in pattern], dtype=jnp.bool_) + group_sizes = index_share_utils.get_served_group_sizes(pattern) + self.served_group_size_array = jnp.array(group_sizes, dtype=jnp.float32) + else: + self.is_full_array = None + self.served_group_size_array = None is_pruned = ( getattr(config, "use_index_share", False) and getattr(config, "prune_shared_indexers", True) @@ -1335,7 +1347,7 @@ def _run_shared(_): if getattr(self.config, "use_index_share", False) and cached_indexer_state is not None: if layer_idx is not None: - is_full = (layer_idx % 4 == 0) + is_full = self.is_full_array[layer_idx] indexer_mask, topk_indices, indexer_score = jax.lax.cond( is_full, _run_full, @@ -1355,8 +1367,14 @@ def _run_shared(_): if indexer_mask is not None and self.config.indexer_loss_scaling_factor > 0.0 and indexer_score is not None: loss_scale = self.config.indexer_loss_scaling_factor - if getattr(self.config, "use_index_share", False) and self.served_group_size > 1: - loss_scale = loss_scale / float(self.served_group_size) + if getattr(self.config, "use_index_share", False): + group_size = ( + self.served_group_size_array[layer_idx] + if layer_idx is not None and self.served_group_size_array is not None + else float(self.served_group_size) + ) + if group_size > 1: + loss_scale = loss_scale / group_size indexer_loss = self.calculate_indexer_loss( indexer_score=indexer_score, diff --git a/src/maxtext/models/glm5.py b/src/maxtext/models/glm5.py index c1c7c6c9cf..c309a9304b 100644 --- a/src/maxtext/models/glm5.py +++ b/src/maxtext/models/glm5.py @@ -58,7 +58,7 @@ def __init__( pattern = index_share_utils.parse_index_share_pattern(config.index_share_pattern, config.num_decoder_layers) self.is_shared_layer = index_share_utils.is_shared_layer(layer_idx, pattern) self.served_group_size = index_share_utils.get_served_group_sizes(pattern)[layer_idx] - if layer_idx == 0: + if layer_idx == 0 and jax.process_index() == 0: num_f = pattern.count("F") num_s = pattern.count("S") absl.logging.info( diff --git a/tests/unit/glm52_indexshare_test.py b/tests/unit/glm52_indexshare_test.py index bea9612cf0..931e746d52 100644 --- a/tests/unit/glm52_indexshare_test.py +++ b/tests/unit/glm52_indexshare_test.py @@ -55,6 +55,18 @@ def test_invalid_pattern_raises(self): with self.assertRaises(ValueError): index_share_utils.parse_index_share_pattern("", 4) + def test_checkpoint_donor_resolution(self): + pattern = index_share_utils.parse_index_share_pattern("FSSS", 12) + # Layers 0..3 share with Layer 0 + for l in range(4): + self.assertEqual(index_share_utils.get_donor_layer_idx(l, pattern), 0) + # Layers 4..7 share with Layer 4 + for l in range(4, 8): + self.assertEqual(index_share_utils.get_donor_layer_idx(l, pattern), 4) + # Layers 8..11 share with Layer 8 + for l in range(8, 12): + self.assertEqual(index_share_utils.get_donor_layer_idx(l, pattern), 8) + if __name__ == "__main__": unittest.main()