diff --git a/training/scripts/train.py b/training/scripts/train.py index e4874d8..d1c1ab2 100755 --- a/training/scripts/train.py +++ b/training/scripts/train.py @@ -32,15 +32,17 @@ def train(config_path): print(f"Loading model: {config['base_model']}") - # Load model - preserve NVFP4 quantization + # Load model with BitsAndBytes INT4 (PEFT supports this) print(f"Loading model: {config['base_model']}") + from transformers import BitsAndBytesConfig + model = AutoModelForCausalLM.from_pretrained( config["base_model"], - torch_dtype=torch.float16, # Keep NVFP4 quantized + quantization_config=BitsAndBytesConfig(load_in_4bit=True), device_map="cpu", # Load to CPU first trust_remote_code=True, ) - print("Model loaded to CPU.") + print("Model loaded with INT4 quantization.") # Add LoRA lora_config = LoraConfig(