Support -Instruct models on to_huggingface

Author: hengtaoguo

src/maxtext/checkpoint_conversion/to_huggingface.py

@@ -118,11 +118,20 @@ def _get_model_mappings(
   Raises:
     ValueError: If mappings for the specified `model_name` are not found.
   """
-  if model_name not in PARAM_MAPPING or model_name not in HF_SHAPE or model_name not in HOOK_FNS:
-    raise ValueError(f"Mappings not found for model: {model_name}. Available PARAM_MAPPING keys: {PARAM_MAPPING.keys()}")
+  target_model_name = model_name
+  if target_model_name not in PARAM_MAPPING and target_model_name.endswith("-Instruct"):
+    base_name = target_model_name[: -len("-Instruct")]
+    if base_name in HF_IDS and base_name in PARAM_MAPPING:
+      target_model_name = base_name
+      max_logging.log(f"Mappings for '{model_name}' not found, falling back to base model '{target_model_name}'.")
+
+  if target_model_name not in PARAM_MAPPING or target_model_name not in HF_SHAPE or target_model_name not in HOOK_FNS:
+    raise ValueError(
+        f"Mappings not found for model: {model_name} (resolved to {target_model_name}). Available PARAM_MAPPING keys: {PARAM_MAPPING.keys()}"
+    )
 
-  param_mapping = PARAM_MAPPING[model_name](hf_config_dict, maxtext_config, scan_layers)
-  hook_fn_mapping = HOOK_FNS[model_name](hf_config_dict, maxtext_config, scan_layers, saving_to_hf=True)
+  param_mapping = PARAM_MAPPING[target_model_name](hf_config_dict, maxtext_config, scan_layers)
+  hook_fn_mapping = HOOK_FNS[target_model_name](hf_config_dict, maxtext_config, scan_layers, saving_to_hf=True)
 
   # Promote composite hook keys into param_mapping.
   # If HOOK_FNS defines a composite tuple key (e.g., (wi_0, wi_1) for MoE gate_up_proj),
@@ -138,7 +147,7 @@ def _get_model_mappings(
 
   return {
       "param_mapping": param_mapping,
-      "shape_mapping": HF_SHAPE[model_name](hf_config_dict),
+      "shape_mapping": HF_SHAPE[target_model_name](hf_config_dict),
       "hook_fn_mapping": hook_fn_mapping,
   }
 
@@ -228,7 +237,7 @@ def _validate_or_update_architecture(hf_config, max_config, override: bool):
 
     if not is_match:
       if override:
-        max_logging.log(f"⚠️ Overwriting HF Config '{hf_attr}': {hf_value} -> {mt_value} (from MaxText '{mt_attr}')")
+        max_logging.log(f"Overwriting HF Config '{hf_attr}': {hf_value} -> {mt_value} (from MaxText '{mt_attr}')")
         setattr(hf_config, hf_attr, mt_value)
       else:
         mismatches.append(f"{hf_attr} (HF={hf_value} vs MaxText={mt_value})")
@@ -275,9 +284,18 @@ def main(argv: Sequence[str]) -> None:
 
   # 1. Get HuggingFace Model Configuration
   model_key = config.model_name
+
   if model_key not in HF_IDS:
     raise ValueError(f"Unsupported model name: {config.model_name}. Supported models are: {list(HF_IDS.keys())}")
-  hf_config_obj = HF_MODEL_CONFIGS[model_key]
+
+  config_key = model_key
+  if config_key not in HF_MODEL_CONFIGS and config_key.endswith("-Instruct"):
+    base_key = config_key[: -len("-Instruct")]
+    if base_key in HF_MODEL_CONFIGS:
+      max_logging.log(f"⚠️ Config for '{config_key}' not found, falling back to base '{base_key}'.")
+      config_key = base_key
+
+  hf_config_obj = HF_MODEL_CONFIGS[config_key]
 
   # Validate architecture consistency (raising ValueError on mismatch) or override HF config if specified.
   _validate_or_update_architecture(hf_config_obj, config, override=FLAGS.override_model_architecture)