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

fix: load 4-bit model AS-IS without BnB (already quantized)

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 15:56:00 -04:00
parent 1197235b4d
commit 3089f27901
1 changed files with 14 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files

117
training/scripts/train.py Executable file → Normal file
View File

@@ -32,106 +32,16 @@ def train(config_path):
print(f"Loading model: {config['base_model']}") print(f"Loading model: {config['base_model']}")
# Load model - try 4-bit first, fall back to bf16 with CPU offload # Load model AS-IS (already 4-bit quantized with CompressedTensors)
print(f"Loading model: {config['base_model']}") print(f"\n[INFO] Loading {config['base_model']}...")
from transformers import BitsAndBytesConfig, AutoConfig
# Try multiple loading strategies in order model = AutoModelForCausalLM.from_pretrained(
print(f"Loading model: {config['base_model']}") config["base_model"],
torch_dtype=torch.float16,
# Strategy 1: Load 4-bit with BitsAndBytes device_map="auto",
print("\n[1/4] Trying: 4-bit with BitsAndBytes...") trust_remote_code=True,
try: )
from transformers import BitsAndBytesConfig print("✓ Success: Model loaded successfully")
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True,
)
print("✓ Success: 4-bit model loaded")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 1: 4-bit AS-IS with auto device_map
print("\n[1/6] Trying: 4-bit AS-IS...")
try:
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
torch_dtype=torch.float16,
device_map="auto",
trust_remote_code=True,
)
print("✓ Success: 4-bit AS-IS")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 2: 4-bit with bf16 compute
print("\n[2/6] Trying: 4-bit with bf16 compute...")
try:
from transformers import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
model = AutoModelForCausalLM.from_pretrained(config["base_model"], quantization_config=bnb_config, device_map="auto", trust_remote_code=True)
print("✓ Success: 4-bit with bf16 compute")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 3: 4-bit to CPU
print("\n[3/6] Trying: 4-bit to CPU...")
try:
model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.float16, device_map="cpu", trust_remote_code=True)
print("✓ Success: 4-bit to CPU")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 4: 4-bit fp32
print("\n[4/6] Trying: 4-bit fp32...")
try:
model = AutoModelForCausalLM.from_pretrained(config["base_model"], torch_dtype=torch.float32, device_map="auto", trust_remote_code=True)
print("✓ Success: 4-bit fp32")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 5: bf16 with accelerate CPU offload
print("\n[5/6] Trying: bf16 with accelerate CPU offload...")
try:
from accelerate import Accelerator
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
torch_dtype=torch.bfloat16,
device_map="cpu",
low_cpu_mem_usage=True,
trust_remote_code=True,
)
accelerator = Accelerator(cpu_offload=True)
model = accelerator.prepare(model)
print("✓ Success: bf16 with accelerate CPU offload")
except Exception as e:
print(f"✗ Failed: {e}")
# Strategy 6: bf16 to CPU with gradient checkpointing
print("\n[6/6] Trying: bf16 to CPU with gradient checkpointing...")
try:
model = AutoModelForCausalLM.from_pretrained(
config["base_model"],
torch_dtype=torch.bfloat16,
device_map="cpu",
low_cpu_mem_usage=True,
trust_remote_code=True,
)
model.gradient_checkpointing_enable()
print("✓ Success: bf16 to CPU with gradient checkpointing")
except Exception as e:
print(f"✗ Failed: {e}")
raise RuntimeError("All loading strategies failed!")
# Note: Skip prepare_model_for_kbit_training to avoid OOM
# Add LoRA # Add LoRA
lora_config = LoraConfig( lora_config = LoraConfig(
@@ -153,10 +63,11 @@ def train(config_path):
from datasets import load_dataset from datasets import load_dataset
import os import os
# Get dataset paths - dataset is in training/data/ relative to repo root # Get dataset paths - assuming run from /home/workdir/artifacts
repo_root = Path(__file__).parent.parent.parent repo_root = Path(__file__).parent
train_path = str(repo_root / "training" / "data" / "train.jsonl") train_path = str(repo_root / "training" / "data" / "train.jsonl")
test_path = str(repo_root / "training" / "data" / "test.jsonl") test_path = str(repo_root / "training" / "data" / "test.jsonl")
print(f"Looking for dataset at: {train_path}")
dataset = load_dataset( dataset = load_dataset(
"json", "json",
@@ -210,7 +121,7 @@ def train(config_path):
def main(): def main():
parser = argparse.ArgumentParser(description="Train LoRA adapter") parser = argparse.ArgumentParser(description="Train LoRA adapter")
parser.add_argument("--config", type=str, default="configs/llama2-7b-lora.yaml", parser.add_argument("--config", type=str, default="training/configs/ornith-35b-lora.yaml",
help="Training configuration file") help="Training configuration file")
parser.add_argument("--check-only", action="store_true", parser.add_argument("--check-only", action="store_true",
help="Validate config and dependencies without training") help="Validate config and dependencies without training")
@@ -248,7 +159,7 @@ def check_setup(config_path):
# Check dataset files # Check dataset files
print("\n3. Dataset files:") print("\n3. Dataset files:")
repo_root = Path(__file__).parent.parent.parent repo_root = Path(__file__).parent
train_path = repo_root / "training" / "data" / "train.jsonl" train_path = repo_root / "training" / "data" / "train.jsonl"
test_path = repo_root / "training" / "data" / "test.jsonl" test_path = repo_root / "training" / "data" / "test.jsonl"