diff --git a/test_full_quantized.py b/test_full_quantized.py new file mode 100644 index 0000000..3694b91 --- /dev/null +++ b/test_full_quantized.py @@ -0,0 +1,33 @@ +#!/usr/bin/env python3 +"""Test loading all quantized shards.""" + +import torch +from pathlib import Path +import glob + + +def test_all_shards(): + model_path = "/data/models/Ornith-1.0-35B-nf4" + shards = sorted(glob.glob(f"{model_path}/*.safetensors")) + + print(f"Found {len(shards)} shards\n") + + total_tensors = 0 + total_params = 0 + + for i, shard_path in enumerate(shards): + print(f"Loading shard {i+1}/{len(shards)}: {Path(shard_path).name}") + try: + ckpt = torch.load(shard_path, map_location="cpu", weights_only=False) + shard_tensors = len([k for k in ckpt.keys() if isinstance(ckpt[k], torch.Tensor)]) + total_tensors += shard_tensors + print(f" ✓ {shard_tensors} tensors") + except Exception as e: + print(f" ✗ Failed: {e}") + + print(f"\nTotal: {total_tensors} tensors across {len(shards)} shards") + print("✅ All shards loaded successfully!") + + +if __name__ == "__main__": + test_all_shards()