feat: implement QLoRA with 4-bit BitsAndBytes quantization

training/configs/ornith-35b-lora.yaml

@@ -1,7 +1,7 @@
 # LoRA Training Configuration for Llama-2-7b
 # Dataset: cyron_summary_lora_dataset (20k examples)
 
-base_model: /data/models/Ornith-1.0-35B-FP8
+base_model: /data/models/Ornith-1.0-35B
 model_type: LlamaForCausalLM
 tokenizer_type: LlamaTokenizer
 
@@ -53,18 +53,9 @@ huggingface_hub:
 deepspeed_config:
   zero_optimization:
     stage: 3
-    offload_optimizer:
-      device: cpu
-      pin_memory: true
-    offload_param:
-      device: cpu
-      pin_memory: true
-      offload_params_device: cpu
   gradient_clipping: 1.0
   train_batch_size: auto
   train_micro_batch_size_per_gpu: auto
-  # Offload 32GB to RAM
-  zero_hierarchical_offload: true
 
 # Evaluation
 eval_strategy: steps

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
+        quantization_config=quantization_config,
-        device_map="cpu",  # Load to CPU first
+        device_map="auto",  # Distribute across GPUs
         trust_remote_code=True,
     )
-    # Remove quantization config to avoid SFTTrainer validation error
+    print("Model loaded with QLoRA (4-bit).")
-    model.config.quantization_config = None
-    print("Model loaded and converted to bf16.")
     
     # Add LoRA
     lora_config = LoraConfig(