From 48e44cdf1c1445519db4bb5f25eeeb0f88b837cd Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 14:30:19 -0400
Subject: [PATCH] feat: add Test 11 - PEFT prepare + manual 4-bit quantization

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

diff --git a/test_model_loading.py b/test_model_loading.py
index 7758250..5577a92 100644
--- a/test_model_loading.py
+++ b/test_model_loading.py
@@ -462,6 +462,82 @@ def test_strategy_10():
         traceback.print_exc()
         return False, str(e)
 
+def test_strategy_11():
+    """Test 11: PEFT prepare_model_for_kbit_training + manual quantization"""
+    print("\n" + "=" * 80)
+    print("TEST 11: PEFT prepare + manual 4-bit quantization")
+    print("=" * 80)
+    
+    try:
+        torch.cuda.empty_cache()
+        print("   Step 1: Load bf16 model to CPU...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B",
+            device_map="cpu",
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print(f"   ✓ Model loaded: {type(model).__name__}")
+        print(f"   ✓ Model class: {model.__class__.__name__}")
+        print(f"   ✓ Model loaded to CPU (~70GB)")
+        
+        # Check CPU memory
+        import psutil
+        mem = psutil.virtual_memory()
+        print(f"   CPU RAM: {mem.used / 1e9:.2f}GB / {mem.total / 1e9:.2f}GB")
+        
+        print("\n   Step 2: Apply PEFT prepare_model_for_kbit_training...")
+        from peft import prepare_model_for_kbit_training
+        model = prepare_model_for_kbit_training(
+            model,
+            use_gradient_checkpointing=False,
+        )
+        print("   ✓ Model prepared for k-bit training")
+        
+        print("\n   Step 3: Manually quantize Linear layers to 4-bit...")
+        from bitsandbytes.nn import Linear4bit
+        from torch import nn
+        
+        # Count and replace Linear layers
+        linear_count = 0
+        for name, module in model.named_modules():
+            if isinstance(module, nn.Linear) and 'lm_head' not in name:
+                # Replace with 4-bit version
+                new_module = Linear4bit(
+                    module.in_features,
+                    module.out_features,
+                    bias=module.bias is not None,
+                )
+                # Copy weights
+                new_module.weight = nn.Parameter(
+                    module.weight.data.clone()
+                )
+                if module.bias is not None:
+                    new_module.bias = nn.Parameter(
+                        module.bias.data.clone()
+                    )
+                # Replace in model
+                layers = name.split('.')
+                parent = model
+                for layer in layers[:-1]:
+                    parent = getattr(parent, layer)
+                setattr(parent, layers[-1], new_module)
+                linear_count += 1
+        
+        print(f"   ✓ Replaced {linear_count} Linear layers with 4-bit")
+        
+        print("\n   Step 4: Move to GPU...")
+        model = model.to("cuda:0")
+        pattern = check_gpu_memory()
+        print(f"   Pattern: {pattern}")
+        return True, pattern
+    except Exception as e:
+        print(f"\n   ✗ FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+        return False, str(e)
+
 if __name__ == "__main__":
     print("=" * 80)
     print("Testing multiple model loading strategies")
@@ -484,6 +560,7 @@ if __name__ == "__main__":
         ("Test 8: CompressedTensors 4-bit → GPU 0 only", test_strategy_8),
         ("Test 9: CompressedTensors 4-bit → GPU (test)", test_strategy_9),
         ("Test 10: CompressedTensors 4-bit → FSDP", test_strategy_10),
+        ("Test 11: PEFT prepare + manual 4-bit", test_strategy_11),
     ]
     
     for name, test_func in tests: