From a5490840132f655ca5d825c0402910f6b8372a46 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Fri, 3 Jul 2026 01:52:46 -0400
Subject: [PATCH] feat: add test for quantized shard loading

---
 test_load_quantized.py | 48 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)
 create mode 100644 test_load_quantized.py

diff --git a/test_load_quantized.py b/test_load_quantized.py
new file mode 100644
index 0000000..3203d42
--- /dev/null
+++ b/test_load_quantized.py
@@ -0,0 +1,48 @@
+#!/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()