diff --git a/train.py b/train.py index 1660f36..cb3ceee 100644 --- a/train.py +++ b/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( diff --git a/training/configs/ornith-35b-lora.yaml b/training/configs/ornith-35b-lora.yaml index 2cf394e..32a05b4 100644 --- a/training/configs/ornith-35b-lora.yaml +++ b/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