feat: add Test 6 - bf16 to CPU → BnB 4-bit quantize → GPU

test_model_loading.py

@@ -219,6 +219,54 @@ def test_strategy_5():
         print(f"\n   ✗ FAILED: {e}")
         return False, str(e)
 
+def test_strategy_6():
+    """Test 6: Load bf16 to CPU, quantize with BnB, then move to GPU"""
+    print("\n" + "=" * 80)
+    print("TEST 6: bf16 to CPU → BnB 4-bit quantize → GPU")
+    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("   ✓ 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 BnB 4-bit quantization...")
+        from peft import prepare_model_for_kbit_training
+        model = prepare_model_for_kbit_training(
+            model,
+            use_gradient_checkpointing=False,
+        )
+        print("   ✓ Model quantized to 4-bit (~17.5GB)")
+        
+        print("\n   Step 3: Move to GPU 0...")
+        model = model.to("cuda:0")
+        pattern = check_gpu_memory()
+        print(f"   Pattern after move to GPU 0: {pattern}")
+        
+        print("\n   Step 4: Move to GPU 1...")
+        model = model.to("cuda:1")
+        pattern = check_gpu_memory()
+        print(f"   Pattern after move to GPU 1: {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")
@@ -241,6 +289,7 @@ if __name__ == "__main__":
         ("Test 3: Explicit device_map", test_strategy_3),
         ("Test 4: Load to CPU then GPU", test_strategy_4),
         ("Test 5: Sequential layer loading", test_strategy_5),
+        ("Test 6: bf16 to CPU → BnB 4-bit → GPU", test_strategy_6),
     ]
     
     for name, test_func in tests: