fix: copy config.json for load test
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user