49 lines
1.7 KiB
Python
49 lines
1.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Test loading quantized shard to check for MISMATCH."""
|
|
|
|
import torch
|
|
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)
|
|
|
|
print("\nLoading quantized test shard...")
|
|
try:
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
"/data/models/test_quantize",
|
|
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()
|