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

fix: use standard QLoRA pattern with BitsAndBytesConfig

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 07:47:42 -04:00
parent da5eb3abed
commit 357c9fa781
1 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
training/scripts/train.py
View File

@@ -36,17 +36,18 @@ def train(config_path):
print(f"Loading model: {config['base_model']}") print(f"Loading model: {config['base_model']}")
from transformers import BitsAndBytesConfig from transformers import BitsAndBytesConfig
quantization_config = BitsAndBytesConfig( bnb_config = BitsAndBytesConfig(
load_in_4bit=True, load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_quant_type="nf4", bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True, bnb_4bit_use_double_quant=True,
) )
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
quantization_config=quantization_config, quantization_config=bnb_config,
device_map="auto", # Distribute across GPUs dtype=torch.bfloat16,
device_map="auto",
trust_remote_code=True, trust_remote_code=True,
) )
print("Model loaded with QLoRA (4-bit).") print("Model loaded with QLoRA (4-bit).")