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

fix: load already-quantized 4-bit model directly (no BnB)

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 18:20:43 -04:00
parent 199176c8d0
commit 1c7e73e927
1 changed files with 3 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
train.py
View File

@@ -33,23 +33,16 @@ def train(config_path):
print(f"Loading model: {config['base_model']}") print(f"Loading model: {config['base_model']}")
# Load model with 4-bit quantization (fit in single GPU) # Load model (already 4-bit on disk, just move to GPU)
print(f"\n[INFO] Loading {config['base_model']} with BnB 4-bit...") print(f"\n[INFO] Loading {config['base_model']} to GPU 0...")
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( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
device_map="cuda:0", device_map="cuda:0",
quantization_config=bnb_config,
torch_dtype=torch.float16, torch_dtype=torch.float16,
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
print("✓ Success: Model loaded to GPU 0 (4-bit)") print("✓ Success: Model loaded to GPU 0")
print(f" GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB") print(f" GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
print(f" Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB") print(f" Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB")