From 3089f27901f403886a6564305cc4da6da198fe28 Mon Sep 17 00:00:00 2001 From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com> Date: Wed, 1 Jul 2026 15:56:00 -0400 Subject: [PATCH] fix: load 4-bit model AS-IS without BnB (already quantized) --- training/scripts/train.py | 117 +++++--------------------------------- 1 file changed, 14 insertions(+), 103 deletions(-) mode change 100755 => 100644 training/scripts/train.py diff --git a/training/scripts/train.py b/training/scripts/train.py old mode 100755 new mode 100644 index 85f8712..0e25155 --- a/training/scripts/train.py +++ b/training/scripts/train.py @@ -32,106 +32,16 @@ def train(config_path): print(f"Loading model: {config['base_model']}") - # Load model - try 4-bit first, fall back to bf16 with CPU offload - print(f"Loading model: {config['base_model']}") - from transformers import BitsAndBytesConfig, AutoConfig + # Load model AS-IS (already 4-bit quantized with CompressedTensors) + print(f"\n[INFO] Loading {config['base_model']}...") - # Try multiple loading strategies in order - print(f"Loading model: {config['base_model']}") - - # Strategy 1: Load 4-bit with BitsAndBytes - print("\n[1/4] Trying: 4-bit with BitsAndBytes...") - 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 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 + model = AutoModelForCausalLM.from_pretrained( + config["base_model"], + torch_dtype=torch.float16, + device_map="auto", + trust_remote_code=True, + ) + print("✓ Success: Model loaded successfully") # Add LoRA lora_config = LoraConfig( @@ -153,10 +63,11 @@ def train(config_path): from datasets import load_dataset import os - # Get dataset paths - dataset is in training/data/ relative to repo root - repo_root = Path(__file__).parent.parent.parent + # Get dataset paths - assuming run from /home/workdir/artifacts + repo_root = Path(__file__).parent train_path = str(repo_root / "training" / "data" / "train.jsonl") test_path = str(repo_root / "training" / "data" / "test.jsonl") + print(f"Looking for dataset at: {train_path}") dataset = load_dataset( "json", @@ -210,7 +121,7 @@ def train(config_path): def main(): 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") parser.add_argument("--check-only", action="store_true", help="Validate config and dependencies without training") @@ -248,7 +159,7 @@ def check_setup(config_path): # Check 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" test_path = repo_root / "training" / "data" / "test.jsonl"