From 4c7ca01d7df953af504138cf7f5354f354bf74c8 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Fri, 3 Jul 2026 01:52:07 -0400
Subject: [PATCH] feat: add single shard test script

---
 test_quantize_one_shard.py | 66 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 66 insertions(+)
 create mode 100644 test_quantize_one_shard.py

diff --git a/test_quantize_one_shard.py b/test_quantize_one_shard.py
new file mode 100644
index 0000000..43921e5
--- /dev/null
+++ b/test_quantize_one_shard.py
@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+"""Test quantization on single shard."""
+
+import gc
+import torch
+from pathlib import Path
+from safetensors.torch import load_file
+from bitsandbytes.functional import quantize_nf4
+
+
+def test_one_shard(model_path, output_dir):
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    import glob
+    shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
+    
+    if len(shards) == 0:
+        print("No shards found!")
+        return
+    
+    shard_file = shards[0]  # First shard only
+    print(f"Testing shard: {shard_file}")
+    
+    state_dict = load_file(shard_file, device="cuda:0")
+    
+    weight_keys = [
+        k for k, v in state_dict.items()
+        if "weight" in k and isinstance(v, torch.Tensor) and v.dim() == 2
+    ]
+    
+    print(f"Found {len(weight_keys)} weight tensors\n")
+    
+    quant_states = {}
+    quantized = 0
+    failed = 0
+    
+    for key in weight_keys:
+        try:
+            weight = state_dict[key]
+            qweight, qstate = quantize_nf4(weight, blocksize=64, compress_statistics=True)
+            state_dict[key] = qweight
+            if qstate is not None:
+                quant_states[f"{key}.quant_state"] = qstate
+            quantized += 1
+            print(f"✓ {key} -> {tuple(qweight.shape)}")
+            del weight, qweight, qstate
+            gc.collect()
+            torch.cuda.empty_cache()
+        except Exception as e:
+            failed += 1
+            print(f"✗ {key}: {e}")
+            gc.collect()
+            torch.cuda.empty_cache()
+    
+    print(f"\nResults: {quantized} quantized, {failed} failed, {len(quant_states)} quant states")
+    
+    # Save test output
+    state_dict_cpu = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()}
+    test_output = output_dir / "test_shard.safetensors"
+    torch.save({**state_dict_cpu, **quant_states}, test_output)
+    print(f"Saved to: {test_output}")
+
+
+if __name__ == "__main__":
+    test_one_shard("/data/models/Ornith-1.0-35B", "/data/models/test_quantize")