From 05ace2d38fb3102ebf5086e90c6664259c2e34ca Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 07:27:13 -0400
Subject: [PATCH] fix: load NVFP4 as bf16 by overriding quantization config

---
 training/scripts/train.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

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(