cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: use device_map=cpu with BnB config for quantization

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 13:54:19 -04:00
parent 3480dd9fbd
commit a55c39c818
1 changed files with 57 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files

161
test_model_loading.py
View File

@@ -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}")