diff --git a/training/scripts/train.py b/training/scripts/train.py index d1c1ab2..d66fc85 100755 --- a/training/scripts/train.py +++ b/training/scripts/train.py @@ -32,17 +32,17 @@ def train(config_path): print(f"Loading model: {config['base_model']}") - # Load model with BitsAndBytes INT4 (PEFT supports this) + # Load model and convert to bf16 (remove NVFP4 quantization) print(f"Loading model: {config['base_model']}") - from transformers import BitsAndBytesConfig - model = AutoModelForCausalLM.from_pretrained( config["base_model"], - quantization_config=BitsAndBytesConfig(load_in_4bit=True), + torch_dtype=torch.bfloat16, device_map="cpu", # Load to CPU first trust_remote_code=True, + # Override any quantization config (NVFP4 -> bf16) + _fast_init=False, ) - print("Model loaded with INT4 quantization.") + print("Model loaded and converted to bf16.") # Add LoRA lora_config = LoraConfig(