fix: use CompressedTensors for Ornith, add torch import, expand LoRA targets

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

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(