Merge pull request #3528 from AI-Hypercomputer:agagik-kimi-checkpoint

PiperOrigin-RevId: 895408378

Author: Google-ML-Automation

src/maxtext/checkpoint_conversion/standalone_scripts/convert_deepseek_family_unscanned_ckpt.py

@@ -70,10 +70,14 @@ def _convert_huggingface_to_jax_weights(base_model_path, model_params, mem_info)
         if key.endswith("_scale_inv"):
           raise ValueError("fp8 checkpoint is not supported.")
         if ds_ckpt.is_key_allowed(key, ds_ckpt.MTP_KEYS_TO_SKIP):
-          mapped_key = ds_ckpt.hf_to_maxtext_mapping(layer, num_experts, first_num_dense_layers, base_num_decoder_layers)[
-              key
-          ]
-          chkpt_vars[mapped_key] = f.get_tensor(key)
+          mapped_key = ds_ckpt.hf_to_maxtext_mapping(
+              layer, num_experts, first_num_dense_layers, base_num_decoder_layers
+          ).get(key)
+          if mapped_key:
+            chkpt_vars[mapped_key] = f.get_tensor(key)
+          else:
+            # This catches keys that are allowed but missing from the mapping dictionary
+            max_logging.log(f"Debug: Allowed key '{key}' (layer {layer}) has no mapping in hf_to_maxtext_mapping.")
 
   logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3))