From 3892ea7fec052543feea6e76270a52b5e32fa17a Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 12:23:34 -0400
Subject: [PATCH] feat: test with BnB 4-bit quantization to force proper
 loading

---
 test_model_loading.py | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/test_model_loading.py b/test_model_loading.py
index d08fbf9..e853f0a 100644
--- a/test_model_loading.py
+++ b/test_model_loading.py
@@ -24,6 +24,47 @@ def test_strategy_1():
         )
         print("   ✓ Model loaded successfully")
         
+        # Test 1b: Try with load_in_4bit=True (force quantization)
+        print("\n   Testing with load_in_4bit=True (force quantization)...")
+        torch.cuda.empty_cache()
+        
+        from transformers import BitsAndBytesConfig
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+        )
+        
+        model2 = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B-4bit",
+            quantization_config=bnb_config,
+            device_map="auto",
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print("   ✓ Model loaded with BnB 4-bit")
+        
+        print("\n   Memory with BnB 4-bit:")
+        for i in range(torch.cuda.device_count()):
+            mem = torch.cuda.memory_allocated(i) / 1e9
+            print(f"   GPU {i}: {mem:.2f} GB")
+        
+        gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
+        gpu1_mem = torch.cuda.memory_allocated(1) / 1e9
+        
+        print("\n   Distribution Pattern:")
+        if abs(gpu0_mem - gpu1_mem) < 1.0:
+            if gpu0_mem < 10.0:
+                print("   ✓ DISTRIBUTED: Model split across both GPUs")
+                return True
+            else:
+                print("   ⚠ DUPLICATE: Same model on both GPUs")
+                print(f"      Each GPU has ~{gpu0_mem:.2f}GB")
+                return False
+        else:
+            print("   ✗ NOT DISTRIBUTED: Model on one GPU only")
+            return False
+        
         # Check memory usage
         print("\n   Memory Usage:")
         for i in range(torch.cuda.device_count()):