AI-Hypercomputer · entrpn · Aug 21, 2026 · Aug 21, 2026 · Aug 21, 2026 · Aug 21, 2026
@@ -52,10 +52,12 @@
 import argparse
 from functools import partial
 import json
+import logging
 import os
 import sys
 import threading
 import time
+
 from typing import Any, Callable, List, Sequence
 import absl
 import ml_dtypes
@@ -474,10 +476,23 @@ def _get_hf_loading_function(hf_source_keys_or_key, tensor_getter, hook_fn, mt_t
   if not isinstance(hf_source_keys_or_key, list):
     # Case 1: Single hf key (str)
     def _loader(getter, key, shape, hook):
+      if key is None:
+        return np.zeros(shape, dtype=np.float32)
       if isinstance(key, (list, tuple)):
         tensors = tuple(getter(k) for k in key)
         return apply_hook_fns(tensors, shape, hook)
-      return apply_hook_fns(getter(key), shape, hook)
+      try:
+        tensor = getter(key)
+      except ValueError as e:
+        if "not found in HF checkpoint index" in str(e):
+          max_logging.log(f"Warning: Key {key} not found in HF checkpoint index; falling back to zeros with shape {shape}.")
+          return np.zeros(shape, dtype=np.float32)
+        raise e
-      except ValueError as e:
-        if "not found in HF checkpoint index" in str(e):
-          return np.zeros(shape, dtype=np.float32)
-        raise e
+      except ValueError as e:
+        if "not found in HF checkpoint index" in str(e):
+          max_logging.log(f"Warning: key {key} not found in HF checkpoint index. Initializing with zeros.")
+          return np.zeros(shape, dtype=np.float32)
+        raise e
-      except ValueError as e:
-        if "not found in HF checkpoint index" in str(e):
-          return np.zeros(shape, dtype=np.float32)
-        raise e
+      except ValueError as e:
+        if "not found in HF checkpoint index" in str(e):
+          max_logging.log(f"Warning: key {key} not found in HF checkpoint index. Initializing with zeros.")
+          return np.zeros(shape, dtype=np.float32)
+        raise e
+
+
+      return apply_hook_fns(tensor, shape, hook)
+
+
 
     load_fn = partial(
         _loader,
@@ -727,7 +742,7 @@ def convert_lora_to_maxtext_adapter(
     max_logging.log("Warning: You want an Instruct version, so we are using the base model architecture instead.")
     model_key = model_key.replace("-Instruct", "")
   hf_config_obj = HF_MODEL_CONFIGS[model_key]
-  hf_config_dict = hf_config_obj.to_dict()
+  hf_config_dict = hf_config_obj.to_dict() if hasattr(hf_config_obj, "to_dict") else hf_config_obj
   param_map_mt_to_hf = PARAM_MAPPING[model_key](hf_config_dict, config, config.scan_layers)
 
   mt_adapter_tree = {}
@@ -868,6 +883,7 @@ def main(
 ) -> None:
   overall_start = time.time()
   # Check if the user is using an Instruct version. If so, use the base model architecture
+  model_name_original = None
   for i, arg in enumerate(args):
     if arg.startswith("model_name="):
       model_name_arg = args[i].split("=")[1]
@@ -878,23 +894,29 @@ def main(
         args[i] = f"model_name={model_name_arg}"
       break
 
+  # Initialize maxtext config
+  config = pyconfig.initialize(args)
+  max_utils.print_system_information()
+
+  if model_name_original is None:
+    model_name_original = config.model_name
+
   # check the supported model ids
   if model_name_original not in HF_IDS:
     raise ValueError(
-        f"Unsupported model name: {model_name_original}.\
-                      Supported models are: {list(HF_IDS.keys())}"
+        f"Unsupported model name: {model_name_original}."
+        f" Supported models are: {list(HF_IDS.keys())}"
     )
 
   model_id = hf_model_path or HF_IDS[model_name_original]
 
-  # Initialize maxtext config
-  config = pyconfig.initialize(args)
-  max_utils.print_system_information()
 
   if not config.base_output_directory:
     output_directory = f"tmp/{config.run_name}"
   else:
     output_directory = config.base_output_directory
+  output_directory = os.path.abspath(output_directory)
+
 
   hf_token = config.hf_access_token
 
@@ -1009,16 +1031,20 @@ def _eager_getter(key):
     model_key = config.model_name
     # load config
     hf_config_obj = HF_MODEL_CONFIGS[model_key]
-    hf_config_dict = hf_config_obj.to_dict()
+    hf_config_dict = hf_config_obj.to_dict() if hasattr(hf_config_obj, "to_dict") else hf_config_obj
     # example of param mapping (gemma2, maxtext:huggingface):
     # "params-decoder-layers_{maxtext_layer_idx}-pre_self_attention_norm_global-scale":
     #   f"model.layers.{global_layer_idx}.input_layernorm.weight",
     param_map_mt_to_hf = PARAM_MAPPING[model_key](hf_config_dict, config, config.scan_layers)
     # Example of Hook FN mapping, to perform reshape:
     # f"params-decoder-layers_{maxtext_layer_idx}-self_attention_global-key-kernel": reshape_kernel,
     hook_fn_map_mt = HOOK_FNS[model_key](hf_config_dict, config, config.scan_layers, saving_to_hf=False)
+    for h in hook_fn_map_mt.values():
+      if hasattr(h, "set_getter"):
+        h.set_getter(tensor_getter)
     max_logging.log("Parameter mappings and hooks obtained.")
 
+
     maxtext_abstract_dict, abstract_params_treedef = get_maxtext_model_info(config)
 
     # Weight transformation

@@ -1936,4 +1936,27 @@ def __init__(self, **kwargs):
     "olmo3-7b": olmo3_7b_config,
     "olmo3-7b-pt": olmo3_7b_config,
     "olmo3-32b": olmo3_32b_config,
+    "kimi-k3": {
+        "model_type": "kimi_k3",
+        "architectures": ["KimiK3ForConditionalGeneration"],
+        "text_config": {
+            "hidden_size": 7168,
+            "num_hidden_layers": 93,
+            "num_attention_heads": 64,
+            "num_key_value_heads": 64,
+            "vocab_size": 163840,
+            "n_routed_experts": 896,
+            "num_experts_per_tok": 16,
+            "n_shared_experts": 2,
+        },
+        "hidden_size": 7168,
+        "num_hidden_layers": 93,
+        "num_attention_heads": 64,
+        "num_key_value_heads": 64,
+        "vocab_size": 163840,
+        "n_routed_experts": 896,
+        "num_experts_per_tok": 16,
+        "n_shared_experts": 2,
+    },
 }
+