30 lines
847 B
Python
30 lines
847 B
Python
#!/usr/bin/env python3
|
|
"""Test quantized shard shapes."""
|
|
|
|
import torch
|
|
from pathlib import Path
|
|
|
|
|
|
def test_shapes():
|
|
print("Loading quantized test shard...")
|
|
ckpt = torch.load("/data/models/test_quantize/model.safetensors", map_location="cpu")
|
|
|
|
print(f"Loaded {len(ckpt)} tensors\n")
|
|
|
|
# Check for mismatched shapes
|
|
mismatch_count = 0
|
|
for key, tensor in list(ckpt.items())[:20]: # Check first 20
|
|
if isinstance(tensor, torch.Tensor):
|
|
print(f"✓ {key} -> {tuple(tensor.shape)}")
|
|
elif hasattr(tensor, 'shape'):
|
|
print(f"✓ {key} -> {tensor.shape}")
|
|
else:
|
|
print(f"? {key} -> {type(tensor)}")
|
|
|
|
print(f"\nTotal tensors: {len(ckpt)}")
|
|
print("If all shapes look correct, the quantization is working!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_shapes()
|