From ffd1d29c358d7e8e878926e8847b03e2ce5abf61 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 13:20:26 -0400
Subject: [PATCH] fix: use bf16 model with BnB 4-bit (ON-THE-FLY quantization)

---
 test_model_loading.py | 48 ++++++++++++++++++++++++-------------------
 1 file changed, 27 insertions(+), 21 deletions(-)

diff --git a/test_model_loading.py b/test_model_loading.py
index 1898243..57061ef 100644
--- a/test_model_loading.py
+++ b/test_model_loading.py
@@ -31,17 +31,22 @@ def check_gpu_memory():
             return f"GPU1_ONLY ({gpu1_mem:.1f}GB)"
 
 def test_strategy_1():
-    """Test 1: device_map='auto' (no quantization config)"""
+    """Test 1: bf16 model + BnB 4-bit (ON-THE-FLY quantization)"""
     print("\n" + "=" * 80)
-    print("TEST 1: device_map='auto' (no BnB)")
+    print("TEST 1: bf16 model + BnB 4-bit (ON-THE-FLY)")
     print("=" * 80)
     
     try:
-        print("   Loading model...")
+        print("   Loading bf16 model with BnB 4-bit...")
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+        )
         model = AutoModelForCausalLM.from_pretrained(
-            "/data/models/Ornith-1.0-35B-4bit",
+            "/data/models/Ornith-1.0-35B",  # ← bf16 model
+            quantization_config=bnb_config,
             device_map="auto",
-            torch_dtype=torch.float16,
             trust_remote_code=True,
             low_cpu_mem_usage=True,
         )
@@ -55,21 +60,22 @@ def test_strategy_1():
         return False, str(e)
 
 def test_strategy_2():
-    """Test 2: device_map='auto' with BnB 4-bit"""
+    """Test 2: bf16 model + BnB 4-bit (alternative config)"""
     print("\n" + "=" * 80)
-    print("TEST 2: device_map='auto' + BnB 4-bit")
+    print("TEST 2: bf16 model + BnB 4-bit (alt config)")
     print("=" * 80)
     
     try:
         torch.cuda.empty_cache()
-        print("   Loading model with BnB 4-bit...")
+        print("   Loading bf16 model with BnB 4-bit...")
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_quant_type="nf4",
             bnb_4bit_compute_dtype=torch.bfloat16,
+            bnb_4bit_use_double_quant=True,
         )
         model = AutoModelForCausalLM.from_pretrained(
-            "/data/models/Ornith-1.0-35B-4bit",
+            "/data/models/Ornith-1.0-35B",  # ← bf16 model
             quantization_config=bnb_config,
             device_map="auto",
             trust_remote_code=True,
@@ -96,7 +102,7 @@ def test_strategy_3():
         
         # Get model config to determine layers
         from transformers import AutoConfig
-        config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B-4bit", trust_remote_code=True)
+        config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", trust_remote_code=True)
         num_layers = config.num_hidden_layers
         
         # Split layers: first half on GPU 0, second half on GPU 1
@@ -114,9 +120,9 @@ def test_strategy_3():
         
         print(f"   Created device_map with {len(device_map)} entries")
         model = AutoModelForCausalLM.from_pretrained(
-            "/data/models/Ornith-1.0-35B-4bit",
+            "/data/models/Ornith-1.0-35B",
             device_map=device_map,
-            torch_dtype=torch.float16,
+            torch_dtype=torch.bfloat16,
             trust_remote_code=True,
             low_cpu_mem_usage=True,
         )
@@ -137,11 +143,11 @@ def test_strategy_4():
     
     try:
         torch.cuda.empty_cache()
-        print("   Loading model to CPU...")
+        print("   Loading bf16 model to CPU...")
         model = AutoModelForCausalLM.from_pretrained(
-            "/data/models/Ornith-1.0-35B-4bit",
+            "/data/models/Ornith-1.0-35B",
             device_map="cpu",
-            torch_dtype=torch.float16,
+            torch_dtype=torch.bfloat16,
             trust_remote_code=True,
             low_cpu_mem_usage=True,
         )
@@ -180,11 +186,11 @@ def test_strategy_5():
         
         # This is a simplified version - in reality would need more complex logic
         # For now, just test if we can load to one GPU
-        print("   Loading to GPU 0 only...")
+        print("   Loading bf16 to GPU 0 only...")
         model = AutoModelForCausalLM.from_pretrained(
-            "/data/models/Ornith-1.0-35B-4bit",
+            "/data/models/Ornith-1.0-35B",
             device_map={"": 0},
-            torch_dtype=torch.float16,
+            torch_dtype=torch.bfloat16,
             trust_remote_code=True,
             low_cpu_mem_usage=True,
         )
@@ -195,11 +201,11 @@ def test_strategy_5():
         
         # Now try GPU 1
         torch.cuda.empty_cache()
-        print("\n   Loading to GPU 1 only...")
+        print("\n   Loading bf16 to GPU 1 only...")
         model = AutoModelForCausalLM.from_pretrained(
-            "/data/models/Ornith-1.0-35B-4bit",
+            "/data/models/Ornith-1.0-35B",
             device_map={"": 1},
-            torch_dtype=torch.float16,
+            torch_dtype=torch.bfloat16,
             trust_remote_code=True,
             low_cpu_mem_usage=True,
         )