diff --git a/train.py b/train.py index 289752e..a547fb8 100644 --- a/train.py +++ b/train.py @@ -33,16 +33,46 @@ def train(config_path): print(f"Loading model: {config['base_model']}") - # Load BnB 4-bit model to single GPU - print(f"\n[INFO] Loading {config['base_model']} (BnB 4-bit) to GPU 0...") + # Load bf16 to CPU first + print(f"\n[INFO] Loading {config['base_model']} bf16 to CPU...") model = AutoModelForCausalLM.from_pretrained( config["base_model"], - device_map="cuda:0", + device_map="cpu", + torch_dtype=torch.bfloat16, + trust_remote_code=True, + low_cpu_mem_usage=True, + ) + print("✓ Model loaded to CPU (~70GB bf16)") + + # Quantize with BnB on CPU + print(" Quantizing with BnB 4-bit on CPU...") + from transformers import BitsAndBytesConfig + bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.float16, + ) + + # Reload with quantization + del model + import gc + gc.collect() + torch.cuda.empty_cache() + + model = AutoModelForCausalLM.from_pretrained( + config["base_model"], + quantization_config=bnb_config, + device_map="cpu", torch_dtype=torch.float16, trust_remote_code=True, low_cpu_mem_usage=True, ) - print("✓ Success: Model loaded to GPU 0 (BnB 4-bit)") + print("✓ Model quantized to 4-bit on CPU") + + # Move to GPU + print(" Moving to GPU 0...") + model = model.to("cuda:0") + print("✓ Success: Model loaded to GPU 0 (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")