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

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(