fix: copy config.json for load test

This commit is contained in:
Christian Medina
2026-07-03 01:57:29 -04:00
parent a549084013
commit 7115d56723

View File

@@ -1,18 +1,26 @@
#!/usr/bin/env python3
"""Test loading quantized shard to check for MISMATCH."""
import gc
import torch
import shutil
from pathlib import Path
from transformers import AutoModelForCausalLM, AutoConfig
def test_load():
print("Loading original model config...")
config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", trust_remote_code=True)
# Copy config.json from original model
test_dir = Path("/data/models/test_quantize")
config_src = Path("/data/models/Ornith-1.0-35B") / "config.json"
config_dst = test_dir / "config.json"
if not config_dst.exists():
shutil.copy2(config_src, config_dst)
print(f"Copied config.json to {test_dir}")
print("\nLoading quantized test shard...")
try:
model = AutoModelForCausalLM.from_pretrained(
"/data/models/test_quantize",
str(test_dir),
device_map="cpu",
torch_dtype=torch.float16,
trust_remote_code=True,