diff --git a/training/configs/llama2-7b-lora.yaml b/training/configs/llama2-7b-lora.yaml index 3cd82b8..d951891 100644 --- a/training/configs/llama2-7b-lora.yaml +++ b/training/configs/llama2-7b-lora.yaml @@ -5,11 +5,8 @@ base_model: deepreinforce-ai/Ornith-1.0-35B-FP8 model_type: LlamaForCausalLM tokenizer_type: LlamaTokenizer -# Quantization (QLoRA) -load_in_4bit: true -bnb_4bit_compute_dtype: bfloat16 -bnb_4bit_quant_type: nf4 -use_nested_quant: false +# Model is already quantized (Ornith uses CompressedTensors) +# No need for BitsAndBytes configuration # LoRA Configuration lora_r: 16 @@ -20,6 +17,9 @@ target_modules: - v_proj - k_proj - o_proj + - gate_proj + - up_proj + - down_proj lora_task_type: CAUSAL_LM # Dataset diff --git a/training/scripts/train.py b/training/scripts/train.py index 76d27e9..83424a6 100755 --- a/training/scripts/train.py +++ b/training/scripts/train.py @@ -2,13 +2,16 @@ """ Train LoRA adapter on Cyron summary dataset. -Uses Hugging Face TRL for SFT training with QLoRA. +Uses Hugging Face TRL for SFT training. """ import argparse +import os import yaml from pathlib import Path +import torch + def train(config_path): """Train LoRA adapter using TRL.""" @@ -31,23 +34,13 @@ def train(config_path): print(f"Loading model: {config['base_model']}") - # Load model with quantization - bnb_config = BitsAndBytesConfig( - load_in_4bit=config.get("load_in_4bit", True), - bnb_4bit_compute_dtype=config.get("bnb_4bit_compute_dtype", "bfloat16"), - bnb_4bit_quant_type=config.get("bnb_4bit_quant_type", "nf4"), - use_nested_quant=config.get("use_nested_quant", False), - ) - - # Use all available GPUs - device_map = "auto" if torch.cuda.device_count() == 1 else "balanced" - + # Load model - let the model's own quantization config handle it + # (Ornith uses CompressedTensors, not BitsAndBytes) model = AutoModelForCausalLM.from_pretrained( config["base_model"], - quantization_config=bnb_config, - device_map=device_map, + device_map="auto", + torch_dtype=torch.bfloat16, ) - model = prepare_model_for_kbit_training(model) # Add LoRA lora_config = LoraConfig(