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

fix: use bf16 model with BnB 4-bit (ON-THE-FLY quantization)

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 13:20:26 -04:00
parent 48a2518b4c
commit ffd1d29c35
1 changed files with 27 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
test_model_loading.py
View File

@@ -31,17 +31,22 @@ def check_gpu_memory():
return f"GPU1_ONLY ({gpu1_mem:.1f}GB)" return f"GPU1_ONLY ({gpu1_mem:.1f}GB)"
def test_strategy_1(): def test_strategy_1():
"""Test 1: device_map='auto' (no quantization config)""" """Test 1: bf16 model + BnB 4-bit (ON-THE-FLY quantization)"""
print("\n" + "=" * 80) print("\n" + "=" * 80)
print("TEST 1: device_map='auto' (no BnB)") print("TEST 1: bf16 model + BnB 4-bit (ON-THE-FLY)")
print("=" * 80) print("=" * 80)
try: try:
print(" Loading model...") print(" Loading bf16 model 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( model = AutoModelForCausalLM.from_pretrained(
"/data/models/Ornith-1.0-35B-4bit", "/data/models/Ornith-1.0-35B", # ← bf16 model
quantization_config=bnb_config,
device_map="auto", device_map="auto",
torch_dtype=torch.float16,
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
@@ -55,21 +60,22 @@ def test_strategy_1():
return False, str(e) return False, str(e)
def test_strategy_2(): def test_strategy_2():
"""Test 2: device_map='auto' with BnB 4-bit""" """Test 2: bf16 model + BnB 4-bit (alternative config)"""
print("\n" + "=" * 80) print("\n" + "=" * 80)
print("TEST 2: device_map='auto' + BnB 4-bit") print("TEST 2: bf16 model + BnB 4-bit (alt config)")
print("=" * 80) print("=" * 80)
try: try:
torch.cuda.empty_cache() torch.cuda.empty_cache()
print(" Loading model with BnB 4-bit...") print(" Loading bf16 model with BnB 4-bit...")
bnb_config = BitsAndBytesConfig( bnb_config = BitsAndBytesConfig(
load_in_4bit=True, load_in_4bit=True,
bnb_4bit_quant_type="nf4", bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
) )
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
"/data/models/Ornith-1.0-35B-4bit", "/data/models/Ornith-1.0-35B", # ← bf16 model
quantization_config=bnb_config, quantization_config=bnb_config,
device_map="auto", device_map="auto",
trust_remote_code=True, trust_remote_code=True,
@@ -96,7 +102,7 @@ def test_strategy_3():
# Get model config to determine layers # Get model config to determine layers
from transformers import AutoConfig from transformers import AutoConfig
config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B-4bit", trust_remote_code=True) config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", trust_remote_code=True)
num_layers = config.num_hidden_layers num_layers = config.num_hidden_layers
# Split layers: first half on GPU 0, second half on GPU 1 # Split layers: first half on GPU 0, second half on GPU 1
@@ -114,9 +120,9 @@ def test_strategy_3():
print(f" Created device_map with {len(device_map)} entries") print(f" Created device_map with {len(device_map)} entries")
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
"/data/models/Ornith-1.0-35B-4bit", "/data/models/Ornith-1.0-35B",
device_map=device_map, device_map=device_map,
torch_dtype=torch.float16, torch_dtype=torch.bfloat16,
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
@@ -137,11 +143,11 @@ def test_strategy_4():
try: try:
torch.cuda.empty_cache() torch.cuda.empty_cache()
print(" Loading model to CPU...") print(" Loading bf16 model to CPU...")
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
"/data/models/Ornith-1.0-35B-4bit", "/data/models/Ornith-1.0-35B",
device_map="cpu", device_map="cpu",
torch_dtype=torch.float16, torch_dtype=torch.bfloat16,
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
@@ -180,11 +186,11 @@ def test_strategy_5():
# This is a simplified version - in reality would need more complex logic # This is a simplified version - in reality would need more complex logic
# For now, just test if we can load to one GPU # For now, just test if we can load to one GPU
print(" Loading to GPU 0 only...") print(" Loading bf16 to GPU 0 only...")
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
"/data/models/Ornith-1.0-35B-4bit", "/data/models/Ornith-1.0-35B",
device_map={"": 0}, device_map={"": 0},
torch_dtype=torch.float16, torch_dtype=torch.bfloat16,
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
@@ -195,11 +201,11 @@ def test_strategy_5():
# Now try GPU 1 # Now try GPU 1
torch.cuda.empty_cache() torch.cuda.empty_cache()
print("\n Loading to GPU 1 only...") print("\n Loading bf16 to GPU 1 only...")
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
"/data/models/Ornith-1.0-35B-4bit", "/data/models/Ornith-1.0-35B",
device_map={"": 1}, device_map={"": 1},
torch_dtype=torch.float16, torch_dtype=torch.bfloat16,
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )