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

feat: add BnB 4-bit quantization strategy + use bf16 model

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 21:14:31 -04:00
parent e56ae5e002
commit 815115e97a
2 changed files with 22 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
train.py
View File

@@ -35,8 +35,27 @@ def train(config_path):
# Load model - try multiple strategies # Load model - try multiple strategies
print(f"\n[INFO] Loading {config['base_model']}...") print(f"\n[INFO] Loading {config['base_model']}...")
# Strategy 1: bf16 to CPU with low memory usage # Strategy 1: bf16 model with BnB 4-bit quantization
print("\n[1/4] Trying: bf16 to CPU...") print("\n[1/4] Trying: bf16 with BnB 4-bit...")
try:
from transformers import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True,
)
print("✓ Success: bf16 with BnB 4-bit")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 2: bf16 to CPU
print("\n[2/4] Trying: bf16 to CPU...")
try: try:
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
@@ -49,8 +68,8 @@ def train(config_path):
except Exception as e: except Exception as e:
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# Strategy 2: 4-bit to CPU # Strategy 3: 4-bit to CPU
print("\n[2/4] Trying: 4-bit to CPU...") print("\n[3/4] Trying: 4-bit to CPU...")
try: try:
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
@@ -62,20 +81,6 @@ def train(config_path):
except Exception as e: except Exception as e:
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")
# Strategy 3: bf16 auto
print("\n[3/4] Trying: bf16 auto...")
try:
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
torch_dtype=torch.bfloat16,
device_map="auto",
low_cpu_mem_usage=True,
trust_remote_code=True,
)
print("✓ Success: bf16 auto")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 4: 4-bit auto # Strategy 4: 4-bit auto
print("\n[4/4] Trying: 4-bit auto...") print("\n[4/4] Trying: 4-bit auto...")
try: try:

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

@@ -1,7 +1,7 @@
# LoRA Training Configuration for Llama-2-7b # LoRA Training Configuration for Llama-2-7b
# Dataset: cyron_summary_lora_dataset (20k examples) # Dataset: cyron_summary_lora_dataset (20k examples)
base_model: /data/models/Ornith-1.0-35B-4bit base_model: /data/models/Ornith-1.0-35B
model_type: LlamaForCausalLM model_type: LlamaForCausalLM
tokenizer_type: LlamaTokenizer tokenizer_type: LlamaTokenizer