feat: use BnB 4-bit model for QLoRA training

train.py

@@ -33,49 +33,26 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load bf16 model to CPU, then quantize with BnB
-    print(f"\n[INFO] Loading {config['base_model']} bf16 to CPU...")
-    model = AutoModelForCausalLM.from_pretrained(
-        config["base_model"],
-        device_map="cpu",
-        torch_dtype=torch.bfloat16,
-        trust_remote_code=True,
-        low_cpu_mem_usage=True,
-    )
-    print("✓ Model loaded to CPU (~70GB bf16)")
-    
-    # Apply BnB 4-bit quantization on CPU
-    print("   Applying BnB 4-bit quantization on CPU...")
+    # Load BnB 4-bit model (already quantized)
+    print(f"\n[INFO] Loading {config['base_model']} (BnB 4-bit)...")
     from transformers import BitsAndBytesConfig
     bnb_config = BitsAndBytesConfig(
         load_in_4bit=True,
         bnb_4bit_quant_type="nf4",
         bnb_4bit_compute_dtype=torch.float16,
-        bnb_4bit_use_double_quant=True,
     )
-    
-    # Reload with quantization config
-    del model
-    import gc
-    gc.collect()
-    torch.cuda.empty_cache()
-    
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        device_map="cpu",
         quantization_config=bnb_config,
+        device_map="auto",
         torch_dtype=torch.float16,
         trust_remote_code=True,
         low_cpu_mem_usage=True,
     )
-    print("✓ Model quantized to 4-bit on CPU (~17.5GB)")
-    
-    # Move to GPU
-    print("   Moving to GPU 0...")
-    model = model.to("cuda:0")
-    print("✓ Success: Model loaded to GPU 0 (4-bit)")
+    print("✓ Success: Model loaded (BnB 4-bit)")
     print(f"   GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
-    print(f"   Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB")
+    if torch.cuda.device_count() > 1:
+        print(f"   GPU 1: {torch.cuda.memory_allocated(1) / 1e9:.2f} GB")
     
     # Add LoRA
     lora_config = LoraConfig(

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

@@ -1,7 +1,7 @@
 # LoRA Training Configuration for Ornith-1.0-35B
 # Dataset: cyron_summary_lora_dataset (20k examples)
 
-base_model: /data/models/Ornith-1.0-35B
+base_model: /data/models/Ornith-1.0-35B-bnb-4bit
 model_type: Qwen3_5MoeForCausalLM
 tokenizer_type: AutoTokenizer