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

feat: use BnB 4-bit model for QLoRA training

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 20:04:55 -04:00
parent 76ff09b140
commit 57b169f202
2 changed files with 7 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
train.py
View File

@@ -33,49 +33,26 @@ def train(config_path):
print(f"Loading model: {config['base_model']}") print(f"Loading model: {config['base_model']}")
# Load bf16 model to CPU, then quantize with BnB # Load BnB 4-bit model (already quantized)
print(f"\n[INFO] Loading {config['base_model']} bf16 to CPU...") print(f"\n[INFO] Loading {config['base_model']} (BnB 4-bit)...")
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
device_map="cpu",
torch_dtype=torch.bfloat16,
trust_remote_code=True,
low_cpu_mem_usage=True,
)
print("✓ Model loaded to CPU (~70GB bf16)")
# Apply BnB 4-bit quantization on CPU
print(" Applying BnB 4-bit quantization on CPU...")
from transformers import BitsAndBytesConfig from transformers import BitsAndBytesConfig
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.float16, bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
) )
# Reload with quantization config
del model
import gc
gc.collect()
torch.cuda.empty_cache()
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
device_map="cpu",
quantization_config=bnb_config, quantization_config=bnb_config,
device_map="auto",
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("✓ Model quantized to 4-bit on CPU (~17.5GB)") print("✓ Success: Model loaded (BnB 4-bit)")
# Move to GPU
print(" Moving to GPU 0...")
model = model.to("cuda:0")
print("✓ Success: Model loaded to GPU 0 (4-bit)")
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") if torch.cuda.device_count() > 1:
print(f" GPU 1: {torch.cuda.memory_allocated(1) / 1e9:.2f} GB")
# Add LoRA # Add LoRA
lora_config = LoraConfig( lora_config = LoraConfig(

2
training/configs/ornith-35b-lora.yaml
View File

@@ -1,7 +1,7 @@
# LoRA Training Configuration for Ornith-1.0-35B # LoRA Training Configuration for Ornith-1.0-35B
# Dataset: cyron_summary_lora_dataset (20k examples) # Dataset: cyron_summary_lora_dataset (20k examples)
base_model: /data/models/Ornith-1.0-35B base_model: /data/models/Ornith-1.0-35B-bnb-4bit
model_type: Qwen3_5MoeForCausalLM model_type: Qwen3_5MoeForCausalLM
tokenizer_type: AutoTokenizer tokenizer_type: AutoTokenizer