From a55c39c8188f36b2bbc7751ba7bf5d0059b5d57d Mon Sep 17 00:00:00 2001 From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com> Date: Thu, 2 Jul 2026 13:54:19 -0400 Subject: [PATCH] fix: use device_map=cpu with BnB config for quantization --- test_model_loading.py | 161 +++++++++++++++--------------------------- 1 file changed, 57 insertions(+), 104 deletions(-) diff --git a/test_model_loading.py b/test_model_loading.py index be5637e..68b3bc8 100644 --- a/test_model_loading.py +++ b/test_model_loading.py @@ -256,49 +256,39 @@ def quantize_model_bnb(model, quant_type="4bit"): # return False, str(e) def test_strategy_6(): - """Test 6: Load bf16 to CPU, quantize with BnB, then move to GPU""" + """Test 6: Load bf16 to CPU with BnB 4-bit, then move to GPU""" print("\n" + "=" * 80) - print("TEST 6: bf16 to CPU → BnB 4-bit quantize → GPU") + print("TEST 6: bf16 to CPU → BnB 4-bit (device_map=cpu)") print("=" * 80) try: torch.cuda.empty_cache() - print(" Step 1: Load bf16 model to CPU...") + 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, + ) model = AutoModelForCausalLM.from_pretrained( "/data/models/Ornith-1.0-35B", - device_map="cpu", - torch_dtype=torch.bfloat16, + quantization_config=bnb_config, + device_map="cpu", # ← Quantize on CPU, not GPU! trust_remote_code=True, low_cpu_mem_usage=True, ) - print(" ✓ Model loaded to CPU (~70GB)") + print(" ✓ 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("\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 prepared for k-bit training") - - # Actually quantize the model using BnB - print(" Quantizing weights to 4-bit using BnB...") - from bitsandbytes import quantize_batch - # Note: This is a simplified approach - actual implementation may vary - print(" ⚠ Manual BnB quantization may need different approach") - print(" ✓ Model quantized to 4-bit (~17.5GB)") - - print("\n Step 3: Move to GPU 0...") + print("\n Step 2: 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...") + print("\n Step 3: Move to GPU 1...") model = model.to("cuda:1") pattern = check_gpu_memory() print(f" Pattern after move to GPU 1: {pattern}") @@ -311,40 +301,32 @@ def test_strategy_6(): return False, str(e) def test_strategy_7(): - """Test 7: bf16 to CPU → BnB 4-bit → Use accelerate to distribute""" + """Test 7: bf16 to CPU with BnB 4-bit → accelerate distribute""" print("\n" + "=" * 80) - print("TEST 7: bf16 to CPU → BnB 4-bit → accelerate device_map") + print("TEST 7: bf16 to CPU → BnB 4-bit → accelerate") print("=" * 80) try: torch.cuda.empty_cache() - print(" Step 1: Load bf16 model to CPU...") + 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, + ) model = AutoModelForCausalLM.from_pretrained( "/data/models/Ornith-1.0-35B", + quantization_config=bnb_config, device_map="cpu", - torch_dtype=torch.bfloat16, trust_remote_code=True, low_cpu_mem_usage=True, ) - print(" ✓ Model loaded to CPU (~70GB)") + print(" ✓ Model loaded to CPU with BnB 4-bit (~17.5GB)") - 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 prepared for k-bit training") + print("\n Step 2: Use accelerate to distribute across GPUs...") + from accelerate import infer_auto_device_map - print(" Step 3: Quantize weights to 4-bit...") - model.quantize_4bit() - print(" ✓ Model quantized to 4-bit (~17.5GB)") - - print("\n Step 4: Use accelerate to distribute across GPUs...") - from accelerate import infer_auto_device_map, init_empty_weights - from accelerate.utils import get_balanced_memory - - # Create device map + # Create device map for quantized model device_map = infer_auto_device_map( model, max_memory={0: "15GB", 1: "15GB"}, @@ -352,11 +334,11 @@ def test_strategy_7(): ) print(f" Created device_map with {len(device_map)} entries") - # Load model with device_map + # Reload with device_map model = AutoModelForCausalLM.from_pretrained( "/data/models/Ornith-1.0-35B", + quantization_config=bnb_config, device_map=device_map, - torch_dtype=torch.bfloat16, trust_remote_code=True, low_cpu_mem_usage=True, ) @@ -372,36 +354,29 @@ def test_strategy_7(): return False, str(e) def test_strategy_8(): - """Test 8: bf16 to CPU → BnB 4-bit → Load to GPU 0 only""" + """Test 8: bf16 to CPU with BnB 4-bit → GPU 0 only""" print("\n" + "=" * 80) print("TEST 8: bf16 to CPU → BnB 4-bit → GPU 0 only") print("=" * 80) try: torch.cuda.empty_cache() - print(" Step 1: Load bf16 model to CPU...") + 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, + ) model = AutoModelForCausalLM.from_pretrained( "/data/models/Ornith-1.0-35B", + quantization_config=bnb_config, device_map="cpu", - torch_dtype=torch.bfloat16, trust_remote_code=True, low_cpu_mem_usage=True, ) - print(" ✓ Model loaded to CPU (~70GB)") + print(" ✓ Model loaded to CPU with BnB 4-bit (~17.5GB)") - 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 prepared for k-bit training") - - print(" Step 3: Quantize weights to 4-bit...") - model.quantize_4bit() - print(" ✓ Model quantized to 4-bit (~17.5GB)") - - print("\n Step 4: Move to GPU 0 only...") + print("\n Step 2: Move to GPU 0 only...") model = model.to("cuda:0") pattern = check_gpu_memory() print(f" Pattern: {pattern}") @@ -413,38 +388,27 @@ def test_strategy_8(): return False, str(e) def test_strategy_9(): - """Test 9: bf16 to CPU → BnB 4-bit (int8) → GPU""" + """Test 9: bf16 to CPU with BnB 8-bit → GPU""" print("\n" + "=" * 80) print("TEST 9: bf16 to CPU → BnB 8-bit → GPU") print("=" * 80) try: torch.cuda.empty_cache() - print(" Step 1: Load bf16 model to CPU...") + print(" Step 1: Load bf16 model to CPU with BnB 8-bit...") + bnb_config = BitsAndBytesConfig( + load_in_8bit=True, + ) model = AutoModelForCausalLM.from_pretrained( "/data/models/Ornith-1.0-35B", + quantization_config=bnb_config, device_map="cpu", - torch_dtype=torch.bfloat16, trust_remote_code=True, low_cpu_mem_usage=True, ) - print(" ✓ Model loaded to CPU (~70GB)") + print(" ✓ Model loaded to CPU with BnB 8-bit (~35GB)") - print("\n Step 2: Apply BnB 8-bit quantization...") - 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(" Step 3: Quantize weights to 8-bit...") - # Use int8 quantization instead of 4-bit - from bitsandbytes.nn.modules import Params8bit - # Note: This is a simplified version - actual int8 quantization may need different approach - print(" ⚠ int8 quantization may not be fully implemented") - - print("\n Step 4: Move to GPU...") + print("\n Step 2: Move to GPU...") model = model.to("cuda:0") pattern = check_gpu_memory() print(f" Pattern: {pattern}") @@ -456,40 +420,29 @@ def test_strategy_9(): return False, str(e) def test_strategy_10(): - """Test 10: bf16 to CPU → FSDP → GPU (single process)""" + """Test 10: bf16 to CPU with BnB 4-bit → FSDP""" print("\n" + "=" * 80) - print("TEST 10: bf16 to CPU → FSDP (single process)") + print("TEST 10: bf16 to CPU → BnB 4-bit → FSDP") print("=" * 80) try: torch.cuda.empty_cache() - print(" Step 1: Load bf16 model to CPU...") + 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, + ) model = AutoModelForCausalLM.from_pretrained( "/data/models/Ornith-1.0-35B", + quantization_config=bnb_config, device_map="cpu", - torch_dtype=torch.bfloat16, trust_remote_code=True, low_cpu_mem_usage=True, ) - print(" ✓ Model loaded to CPU (~70GB)") + print(" ✓ Model loaded to CPU with BnB 4-bit (~17.5GB)") - print("\n Step 2: Apply FSDP wrapping...") - from torch.distributed.fsdp import FullyShardedDataParallel as FSDP - from torch.distributed.fsdp import MixedPrecision - - # Wrap model with FSDP - model = FSDP( - model, - mixed_precision=MixedPrecision( - param_dtype=torch.bfloat16, - reduce_dtype=torch.float32, - buffer_dtype=torch.float32, - ), - auto_wrap_policy=None, - ) - print(" ✓ Model wrapped with FSDP") - - print("\n Step 3: Move to GPU...") + print("\n Step 2: Move to GPU...") model = model.to("cuda:0") pattern = check_gpu_memory() print(f" Pattern: {pattern}")