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

fix: switch to CompressedTensors 4-bit checkpoint (BnB on CPU not working)

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 14:17:52 -04:00
parent 81dea9b9c8
commit da8658ae73
1 changed files with 46 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

122
test_model_loading.py
View File

@@ -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: