Onboard GLM-5.2

src/dependencies/requirements/requirements.txt

@@ -33,13 +33,13 @@ pylint
 pytest
 pytype
 qwix>=0.1.6
-sentencepiece
+sentencepiece>=0.2.0
 tensorboard-plugin-profile
 tensorboardx
 tensorflow-datasets
 tensorflow-text
 tensorflow
-tiktoken
+tiktoken>=0.5.0
 tokamax>=0.0.4
 transformers
 google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip

src/maxtext/checkpoint_conversion/to_maxtext.py

@@ -53,6 +53,7 @@
 from functools import partial
 import json
 import os
+import re
 import sys
 import threading
 import time
@@ -119,20 +120,25 @@ def __init__(self, model_id, token, revision=None):
     self.current_shard_content = {}
     # Cache for resolved local shard paths
     self._local_shard_paths = {}
+    # Cache for open safetensors file readers
+    self._open_files = {}
     # Use a lock to serialize heavy RAM operations, but NOT downloads
     self._ram_lock = threading.Lock()
     self._initialize_index()
 
   def __getstate__(self):
-    """Allows pickling/copying by excluding the non-pickleable lock."""
+    """Allows pickling/copying by excluding the non-pickleable lock and files."""
     state = self.__dict__.copy()
     del state["_ram_lock"]
+    if "_open_files" in state:
+      del state["_open_files"]
     return state
 
   def __setstate__(self, state):
-    """Restores state after pickling/copying and recreates a new lock."""
+    """Restores state after pickling/copying and recreates a new lock and file cache."""
     self.__dict__.update(state)
     self._ram_lock = threading.Lock()
+    self._open_files = {}
 
   def _initialize_index(self):
     """Fetches and parses the Hugging Face model index file to build a shard map."""
@@ -206,7 +212,17 @@ def get_tensor(self, key: str) -> np.ndarray:
     # This prevents multiple threads from simultaneously allocating large chunks of RAM.
     with self._ram_lock:
       with safe_open(local_path, framework="np", device="cpu") as f:
-        return f.get_tensor(key)
+        tensor = f.get_tensor(key)
+      # Prune older downloaded shards if downloading remotely to prevent disk/RAM exhaustion
+      if not self.is_local and len(self._local_shard_paths) > 3:
+        oldest_shard = next(iter(self._local_shard_paths))
+        oldest_path = self._local_shard_paths.pop(oldest_shard)
+        if os.path.exists(oldest_path):
+          try:
+            os.remove(oldest_path)
+          except OSError:
+            pass
+      return tensor
 
 
 class LazyTensor:
@@ -448,7 +464,28 @@ 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."):
+          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:  # pylint: disable=broad-exception-caught
+                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)
 
@@ -984,6 +1021,16 @@ 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."):
+            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"
@@ -1002,6 +1049,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 and key.startswith("model.layers."):
+            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

@@ -1884,8 +1884,27 @@ def __init__(self, **kwargs):
 qwen3_vl_30b_a3b_config = PTConfig(**qwen3_vl_30b_a3b_dict)
 
 
+glm5_1_744b_dict = {
+    "architectures": ["GlmMoeDsaForCausalLM"],
+    "num_hidden_layers": 78,
+    "first_k_dense_replace": 3,
+    "n_routed_experts": 256,
+}
+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/hf_shape.py

@@ -1312,6 +1312,8 @@ def DEEPSEEKV4_HF_WEIGHTS_TO_SHAPE(config):
     "deepseek3-671b": DEEPSEEK_HF_WEIGHTS_TO_SHAPE,
     "deepseek3.2-671b": DEEPSEEK_HF_WEIGHTS_TO_SHAPE,
     "deepseek4-284b": DEEPSEEKV4_HF_WEIGHTS_TO_SHAPE,
+    "glm5.1-744b": DEEPSEEK_HF_WEIGHTS_TO_SHAPE,
+    "glm5.2-744b": DEEPSEEK_HF_WEIGHTS_TO_SHAPE,
     "gpt-oss-20b": GPT_OSS_HF_WEIGHTS_TO_SHAPE,
     "gpt-oss-120b": GPT_OSS_HF_WEIGHTS_TO_SHAPE,
     "mixtral-8x7b": MIXTRAL_HF_WEIGHTS_TO_SHAPE,