From 442a8751ef798fb225c3b7a4a9e026380077c507 Mon Sep 17 00:00:00 2001
From: rubik <huajianfeng@zhejianglab.com>
Date: Tue, 30 Jun 2026 15:15:23 +0800
Subject: [PATCH] issue/467 [BUG]: enforce safetensors index-based weight
 loading for robustness Replace fragile glob with model.safetensors.index.json
 lookup to handle non-standard names (e.g., Mistral-Large-Instruct). Includes
 glob fallback

---
 python/infinilm/infer_engine.py   | 26 +++++++++++++++++---------
 python/infinilm/modeling_utils.py | 16 +++++++++++++++-
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py
index 5c4d21484..b45637f75 100644
--- a/python/infinilm/infer_engine.py
+++ b/python/infinilm/infer_engine.py
@@ -12,8 +12,20 @@
 import json
 import os
 
+_MODEL_DEFAULTS = {
+    "gpt2": {"torch_dtype": "float32"},
+    "mistral": {"torch_dtype": "bfloat16"},
+}
+
+def _apply_torch_dtype_defaults(config: dict) -> dict:
+    if config.get("torch_dtype") is None:
+        config["torch_dtype"] = (
+            config.get("dtype") or
+            _MODEL_DEFAULTS.get(config.get("model_type"), {}).get("torch_dtype")
+        )
+    return config
 
-def normalize_hf_config_for_infinilm(config_dict):
+def _normalize_videonsa_config(config_dict):
     model_type = config_dict.get("model_type")
 
     if model_type == "qwen2_5_vl" and config_dict.get("architectures") == [
@@ -36,24 +48,20 @@ def normalize_hf_config_for_infinilm(config_dict):
 
     return config_dict
 
-
 def read_hf_config(model_path):
     config_path = os.path.join(model_path, "config.json")
     with open(config_path, "r") as f:
         config_dict = json.load(f)
 
-    if (
-        config_dict.get("model_type") == "gpt2"
-        and config_dict.get("torch_dtype") is None
-        and config_dict.get("dtype") is None
-    ):
-        config_dict["torch_dtype"] = "float32"
     if "model_type" not in config_dict:
         raise ValueError(
             f"`model_type` is not specified in the config file `{config_path}`."
         )
-    return normalize_hf_config_for_infinilm(config_dict)
 
+    config_dict = _apply_torch_dtype_defaults(config_dict)
+    config_dict = _normalize_videonsa_config(config_dict)
+
+    return config_dict
 
 # config.json (required) defines model architecture, while generation_config.json
 # (optional) defines generation behavior. They are kept as separate readers
diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py
index f4ca09703..33b4dc3c0 100644
--- a/python/infinilm/modeling_utils.py
+++ b/python/infinilm/modeling_utils.py
@@ -183,7 +183,21 @@ def load_model_state_dict_by_file(
     already_loaded_keys = []
     embed_tokens_torch_unscaled = None
 
-    file_list = glob.glob(os.path.join(model_path, "*.safetensors"))
+    index_file_path = os.path.join(model_path, "model.safetensors.index.json")
+    if os.path.exists(index_file_path):
+        # Priority 1: If the index file exists, strictly load exactly what it maps to.
+        # This handles all standard sharded models perfectly, regardless of their actual prefix.
+        print(f"Found index file: {index_file_path}. Loading shards by index.")
+        with open(index_file_path, "r") as f:
+            index_data = json.load(f)
+        weight_map = index_data.get("weight_map", {})
+        unique_filenames = set(weight_map.values())
+        file_list = [os.path.join(model_path, fname) for fname in unique_filenames]
+    else:
+        # Priority 2: If no index file, scan all safetensors files.
+        print("No index file found. Scanning all safetensors files...")
+        file_list = glob.glob(os.path.join(model_path, "*.safetensors"))
+
     if len(file_list) > 0:
         for file_path in tqdm(file_list, desc="Processing files"):
             tqdm.write(f"Processing: {os.path.basename(file_path)}")