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

feat: detect pre-quantized model and skip re-quantization

Browse Source
This commit is contained in:
Christian Medina
2026-07-03 01:19:55 -04:00
parent 4308fd3391
commit b15bc5c975
1 changed files with 33 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
train.py
View File

@@ -33,24 +33,40 @@ def train(config_path):
print(f"Loading model: {config['base_model']}")
# Load bf16 to CPU with BnB 4-bit quantization
print(f"\n[INFO] Loading {config['base_model']} with BnB 4-bit to CPU...")
from transformers import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
)
# Check if model is quantized (has model-XXXXX-of-XXXXX files)
import glob
safetensor_files = glob.glob(f"{config['base_model']}/*.safetensors")
is_quantized = any("of-" in f for f in safetensor_files)
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
quantization_config=bnb_config,
device_map="cpu",
torch_dtype=torch.float16,
trust_remote_code=True,
low_cpu_mem_usage=True,
)
print("✓ Model loaded to CPU with BnB 4-bit (~17.5GB)")
if is_quantized:
print(f"\n[INFO] Loading pre-quantized BnB 4-bit model (no re-quantization)...")
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
device_map="cpu",
torch_dtype=torch.float16,
trust_remote_code=True,
low_cpu_mem_usage=True,
)
print("✓ Model loaded to CPU (BnB 4-bit)")
else:
# Load bf16 with BnB 4-bit quantization
print(f"\n[INFO] Loading {config['base_model']} with BnB 4-bit to CPU...")
from transformers import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
)
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
quantization_config=bnb_config,
device_map="cpu",
torch_dtype=torch.float16,
trust_remote_code=True,
low_cpu_mem_usage=True,
)
print("✓ Model loaded to CPU with BnB 4-bit (~17.5GB)")
# Move to GPU
print(" Moving to GPU 0...")