diff --git a/test_model_loading.py b/test_model_loading.py index 5bb60bb..7758250 100644 --- a/test_model_loading.py +++ b/test_model_loading.py @@ -283,34 +283,23 @@ def quantize_model_bnb(model, quant_type="4bit"): # return False, str(e) def test_strategy_6(): - """Test 6: Load bf16 to CPU with BnB 4-bit, then move to GPU""" + """Test 6: Use CompressedTensors 4-bit checkpoint (pre-quantized)""" print("\n" + "=" * 80) - print("TEST 6: bf16 to CPU → BnB 4-bit (device_map=cpu)") + print("TEST 6: CompressedTensors 4-bit checkpoint") print("=" * 80) try: torch.cuda.empty_cache() - print(" Step 1: Load bf16 model to CPU with BnB 4-bit...") - bnb_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - ) + print(" Step 1: Load CompressedTensors 4-bit checkpoint...") model = AutoModelForCausalLM.from_pretrained( - "/data/models/Ornith-1.0-35B", - quantization_config=bnb_config, - device_map="cpu", # ← Quantize on CPU, not GPU! + "/data/models/Ornith-1.0-35B-4bit", # ← Pre-quantized checkpoint + torch_dtype=torch.float16, 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 with BnB 4-bit (~17.5GB)") - - # Check CPU memory - import psutil - mem = psutil.virtual_memory() - print(f" CPU RAM: {mem.used / 1e9:.2f}GB / {mem.total / 1e9:.2f}GB") + print(f" ✓ Model loaded (~18GB on disk)") print("\n Step 2: Move to GPU 0...") model = model.to("cuda:0") @@ -330,9 +319,9 @@ def test_strategy_6(): return False, str(e) def test_strategy_7(): - """Test 7: bf16 to CPU with BnB 4-bit → accelerate distribute""" + """Test 7: CompressedTensors 4-bit → accelerate distribute""" print("\n" + "=" * 80) - print("TEST 7: bf16 to CPU → BnB 4-bit → accelerate") + print("TEST 7: CompressedTensors 4-bit → accelerate") print("=" * 80) try: @@ -340,22 +329,18 @@ def test_strategy_7(): # Detect layer names dynamically print(" Detecting layer names...") - layer_names = get_layer_names("/data/models/Ornith-1.0-35B") + layer_names = get_layer_names("/data/models/Ornith-1.0-35B-4bit") - print("\n Step 1: Load bf16 model to CPU with BnB 4-bit...") - bnb_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - ) + print("\n Step 1: Load CompressedTensors 4-bit checkpoint...") model = AutoModelForCausalLM.from_pretrained( - "/data/models/Ornith-1.0-35B", - quantization_config=bnb_config, - device_map="cpu", + "/data/models/Ornith-1.0-35B-4bit", + torch_dtype=torch.float16, trust_remote_code=True, low_cpu_mem_usage=True, ) - print(" ✓ Model loaded to CPU with BnB 4-bit (~17.5GB)") + print(f" ✓ Model loaded: {type(model).__name__}") + print(f" ✓ Model class: {model.__class__.__name__}") + print(f" ✓ Model loaded (~18GB on disk)") print("\n Step 2: Use accelerate to distribute across GPUs...") from accelerate import infer_auto_device_map @@ -370,8 +355,8 @@ def test_strategy_7(): # Reload with device_map model = AutoModelForCausalLM.from_pretrained( - "/data/models/Ornith-1.0-35B", - quantization_config=bnb_config, + "/data/models/Ornith-1.0-35B-4bit", + torch_dtype=torch.float16, device_map=device_map, trust_remote_code=True, low_cpu_mem_usage=True, @@ -388,27 +373,23 @@ def test_strategy_7(): return False, str(e) def test_strategy_8(): - """Test 8: bf16 to CPU with BnB 4-bit → GPU 0 only""" + """Test 8: CompressedTensors 4-bit → GPU 0 only""" print("\n" + "=" * 80) - print("TEST 8: bf16 to CPU → BnB 4-bit → GPU 0 only") + print("TEST 8: CompressedTensors 4-bit → GPU 0 only") print("=" * 80) try: torch.cuda.empty_cache() - print(" Step 1: Load bf16 model to CPU with BnB 4-bit...") - bnb_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - ) + print(" Step 1: Load CompressedTensors 4-bit checkpoint...") model = AutoModelForCausalLM.from_pretrained( - "/data/models/Ornith-1.0-35B", - quantization_config=bnb_config, - device_map="cpu", + "/data/models/Ornith-1.0-35B-4bit", + torch_dtype=torch.float16, trust_remote_code=True, low_cpu_mem_usage=True, ) - print(" ✓ Model loaded to CPU with BnB 4-bit (~17.5GB)") + print(f" ✓ Model loaded: {type(model).__name__}") + print(f" ✓ Model class: {model.__class__.__name__}") + print(f" ✓ Model loaded (~18GB on disk)") print("\n Step 2: Move to GPU 0 only...") model = model.to("cuda:0") @@ -422,25 +403,23 @@ def test_strategy_8(): return False, str(e) def test_strategy_9(): - """Test 9: bf16 to CPU with BnB 8-bit → GPU""" + """Test 9: CompressedTensors 4-bit → GPU (test distribution)""" print("\n" + "=" * 80) - print("TEST 9: bf16 to CPU → BnB 8-bit → GPU") + print("TEST 9: CompressedTensors 4-bit → GPU (test)") print("=" * 80) try: torch.cuda.empty_cache() - print(" Step 1: Load bf16 model to CPU with BnB 8-bit...") - bnb_config = BitsAndBytesConfig( - load_in_8bit=True, - ) + print(" Step 1: Load CompressedTensors 4-bit checkpoint...") model = AutoModelForCausalLM.from_pretrained( - "/data/models/Ornith-1.0-35B", - quantization_config=bnb_config, - device_map="cpu", + "/data/models/Ornith-1.0-35B-4bit", + torch_dtype=torch.float16, trust_remote_code=True, low_cpu_mem_usage=True, ) - print(" ✓ Model loaded to CPU with BnB 8-bit (~35GB)") + print(f" ✓ Model loaded: {type(model).__name__}") + print(f" ✓ Model class: {model.__class__.__name__}") + print(f" ✓ Model loaded (~18GB on disk)") print("\n Step 2: Move to GPU...") model = model.to("cuda:0") @@ -454,27 +433,23 @@ def test_strategy_9(): return False, str(e) def test_strategy_10(): - """Test 10: bf16 to CPU with BnB 4-bit → FSDP""" + """Test 10: CompressedTensors 4-bit → FSDP""" print("\n" + "=" * 80) - print("TEST 10: bf16 to CPU → BnB 4-bit → FSDP") + print("TEST 10: CompressedTensors 4-bit → FSDP") print("=" * 80) try: torch.cuda.empty_cache() - print(" Step 1: Load bf16 model to CPU with BnB 4-bit...") - bnb_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - ) + print(" Step 1: Load CompressedTensors 4-bit checkpoint...") model = AutoModelForCausalLM.from_pretrained( - "/data/models/Ornith-1.0-35B", - quantization_config=bnb_config, - device_map="cpu", + "/data/models/Ornith-1.0-35B-4bit", + torch_dtype=torch.float16, trust_remote_code=True, low_cpu_mem_usage=True, ) - print(" ✓ Model loaded to CPU with BnB 4-bit (~17.5GB)") + print(f" ✓ Model loaded: {type(model).__name__}") + print(f" ✓ Model class: {model.__class__.__name__}") + print(f" ✓ Model loaded (~18GB on disk)") print("\n Step 2: Move to GPU...") model = model.to("cuda:0") @@ -504,16 +479,11 @@ if __name__ == "__main__": results = [] tests = [ - # ("Test 1: device_map=auto (no BnB)", test_strategy_1), - # ("Test 2: device_map=auto + BnB 4-bit", test_strategy_2), - # ("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), - ("Test 7: bf16 to CPU → BnB 4-bit → accelerate", test_strategy_7), - ("Test 8: bf16 to CPU → BnB 4-bit → GPU 0 only", test_strategy_8), - ("Test 9: bf16 to CPU → BnB 8-bit → GPU", test_strategy_9), - ("Test 10: bf16 to CPU → FSDP → GPU", test_strategy_10), + ("Test 6: CompressedTensors 4-bit → GPU", test_strategy_6), + ("Test 7: CompressedTensors 4-bit → accelerate", test_strategy_7), + ("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), ] for name, test_func in tests: