feat: implement QLoRA with 4-bit BitsAndBytes quantization

training/scripts/train.py

@@ -32,17 +32,24 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model and convert to bf16 (ignore FP8 quantization)
+    # Load model with QLoRA (4-bit quantization)
     print(f"Loading model: {config['base_model']}")
+    from transformers import BitsAndBytesConfig
+    
+    quantization_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+    )
+    
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        dtype=torch.bfloat16,  # Convert FP8 -> bf16
-        device_map="cpu",  # Load to CPU first
+        quantization_config=quantization_config,
+        device_map="auto",  # Distribute across GPUs
         trust_remote_code=True,
     )
-    # Remove quantization config to avoid SFTTrainer validation error
-    model.config.quantization_config = None
-    print("Model loaded and converted to bf16.")
+    print("Model loaded with QLoRA (4-bit).")
     
     # Add LoRA
     lora_config = LoraConfig(