Default to manual safetensors load to avoid Colab meta-tensor failures

run_standalone_inference.py

@@ -57,6 +57,12 @@ def parse_args() -> argparse.Namespace:
         help="Use transformers/accelerate device_map='auto'. Disabled by default for maximum Colab compatibility.",
     )
     p.add_argument("--prompts-file", default=None, help="Optional .txt file with one prompt per line")
+    p.add_argument(
+        "--load-mode",
+        choices=["manual", "from_pretrained"],
+        default="manual",
+        help="Model loading strategy. `manual` avoids meta-tensor edge cases on Colab.",
+    )
     return p.parse_args()
 
 
@@ -166,6 +172,58 @@ def _has_meta_tensors(model) -> bool:
     return False
 
 
+def _load_model_manual(local_dir: str, torch_dtype, device):
+    from pathlib import Path
+
+    import torch
+    from safetensors.torch import load_file
+    from transformers import AutoConfig, AutoModelForCausalLM
+
+    config = AutoConfig.from_pretrained(local_dir, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
+    weights_path = Path(local_dir) / "model.safetensors"
+    state_dict = load_file(str(weights_path), device="cpu")
+    incompatible = model.load_state_dict(state_dict, strict=False)
+    if incompatible.missing_keys or incompatible.unexpected_keys:
+        raise RuntimeError(
+            "State dict mismatch while manual-loading model.safetensors. "
+            f"missing={len(incompatible.missing_keys)} unexpected={len(incompatible.unexpected_keys)}"
+        )
+    model = model.to(device=device, dtype=torch_dtype)
+    model.eval()
+    return model
+
+
+def _load_model_from_pretrained(local_dir: str, torch_dtype, device, use_device_map: bool):
+    import torch
+    from transformers import AutoModelForCausalLM
+
+    model = AutoModelForCausalLM.from_pretrained(
+        local_dir,
+        trust_remote_code=True,
+        torch_dtype=torch_dtype,
+        device_map="auto" if (device.type == "cuda" and use_device_map) else None,
+        low_cpu_mem_usage=bool(device.type == "cuda" and use_device_map),
+    )
+    if device.type == "cpu" or (device.type == "cuda" and not use_device_map):
+        if _has_meta_tensors(model):
+            print("[warn] Detected meta tensors after load; retrying with device_map='auto' fallback.")
+            del model
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            model = AutoModelForCausalLM.from_pretrained(
+                local_dir,
+                trust_remote_code=True,
+                torch_dtype=torch_dtype,
+                device_map="auto" if device.type == "cuda" else None,
+                low_cpu_mem_usage=bool(device.type == "cuda"),
+            )
+        else:
+            model = model.to(device)
+    model.eval()
+    return model
+
+
 def generate_compat(model, input_ids, max_new_tokens: int, temperature: float, top_k: int, seed: int):
     import torch
     import torch.nn.functional as F
@@ -195,7 +253,7 @@ def main() -> None:
     import torch
     import transformers
     from huggingface_hub import snapshot_download
-    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers import AutoTokenizer
 
     if transformers.__version__ == "4.57.0":
         print(
@@ -216,29 +274,15 @@ def main() -> None:
     sha = Path(local).name
 
     tok = AutoTokenizer.from_pretrained(local, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        local,
-        trust_remote_code=True,
-        torch_dtype=torch_dtype,
-        device_map="auto" if (device.type == "cuda" and args.use_device_map) else None,
-        low_cpu_mem_usage=bool(device.type == "cuda" and args.use_device_map),
-    )
-    if device.type == "cpu" or (device.type == "cuda" and not args.use_device_map):
-        if _has_meta_tensors(model):
-            print("[warn] Detected meta tensors after load; retrying with device_map='auto' fallback.")
-            del model
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-            model = AutoModelForCausalLM.from_pretrained(
-                local,
-                trust_remote_code=True,
-                torch_dtype=torch_dtype,
-                device_map="auto" if device.type == "cuda" else None,
-                low_cpu_mem_usage=bool(device.type == "cuda"),
-            )
-        else:
-            model = model.to(device)
-    model.eval()
+    if args.load_mode == "manual":
+        model = _load_model_manual(local, torch_dtype=torch_dtype, device=device)
+    else:
+        model = _load_model_from_pretrained(
+            local,
+            torch_dtype=torch_dtype,
+            device=device,
+            use_device_map=args.use_device_map,
+        )
 
     cfg = model.config
     vocab_size = int(getattr(cfg, "padded_vocab_size", getattr(cfg, "vocab_size", len(tok))))
@@ -258,7 +302,10 @@ def main() -> None:
     max_prompt_tokens = max(1, context_window - args.max_new_tokens)
 
     print(f"repo={args.repo_id} revision={args.revision} snapshot_sha={sha}")
-    print(f"device={device} dtype={torch_dtype} prompt_style={prompt_style} use_device_map={args.use_device_map}")
+    print(
+        f"device={device} dtype={torch_dtype} prompt_style={prompt_style} "
+        f"use_device_map={args.use_device_map} load_mode={args.load_mode}"
+    )
     print(f"context_window={context_window} max_prompt_tokens={max_prompt_tokens}")
 
     torch.manual_seed(args.seed)