57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
#!/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():
|
|
# 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(
|
|
str(test_dir),
|
|
device_map="cpu",
|
|
torch_dtype=torch.float16,
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
print("✓ Model loaded successfully!")
|
|
|
|
# Check for mismatched parameters
|
|
print("\nChecking parameter shapes...")
|
|
mismatch_count = 0
|
|
for name, param in model.named_parameters():
|
|
if hasattr(param, 'quant_state') and param.quant_state is not None:
|
|
# Quantized parameter - check if it loaded correctly
|
|
expected_shape = config.get_shape_for_parameter(name) if hasattr(config, 'get_shape_for_parameter') else None
|
|
if expected_shape and tuple(param.shape) != expected_shape:
|
|
print(f"✗ MISMATCH: {name} - expected {expected_shape}, got {tuple(param.shape)}")
|
|
mismatch_count += 1
|
|
else:
|
|
print(f"✓ {name} - {tuple(param.shape)}")
|
|
|
|
if mismatch_count == 0:
|
|
print("\n✅ No MISMATCH errors!")
|
|
else:
|
|
print(f"\n❌ {mismatch_count} MISMATCH errors found!")
|
|
|
|
except Exception as e:
|
|
print(f"✗ Failed to load: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_load()
|