add glm5.2 indexshare

src/maxtext/checkpoint_conversion/to_maxtext.py

@@ -463,7 +463,30 @@ 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 "indexer" in str(hf_key_single) and str(hf_key_single).startswith("model.layers."):
+          import re
+
+          m = re.match(r"model\.layers\.(\d+)\.(.+)", str(hf_key_single))
+          if m:
+            layer_idx = int(m.group(1))
+            rest = m.group(2)
+            # 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
+        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 +1022,18 @@ 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 and key.startswith("model.layers."):
+            import re
+
+            m = re.match(r"model\.layers\.(\d+)\.(.+)", key)
+            if m:
+              layer_idx = int(m.group(1))
+              rest = m.group(2)
+              # 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.")
         v = hf_state_dict_numpy[key]
         # target dtype is "float32"
@@ -1017,6 +1052,31 @@ 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 and key.startswith("model.layers."):
+            import re
+
+            m = re.match(r"model\.layers\.(\d+)\.(.+)", key)
+            if m:
+              layer_idx = int(m.group(1))
+              rest = m.group(2)
+              # 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
+
     if is_merge_mode:
       tensor_getter = _setup_merge_mode_getter(tensor_getter, config, hf_lora_adapter_path, revision)
 

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,

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,

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):

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

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

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"
 

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: "glm5"
+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

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",
@@ -718,6 +719,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):
@@ -3278,7 +3290,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:
@@ -3497,6 +3509,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 "
@@ -3564,9 +3578,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. "