diff --git a/test_model_loading.py b/test_model_loading.py index 57061ef..be80840 100644 --- a/test_model_loading.py +++ b/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: