From 905c2ba30d72ae0092e6370cc9ab119754805f76 Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Tue, 30 Jun 2026 23:50:27 -0400
Subject: [PATCH] fix: preserve NVFP4 quantization, don't force bf16

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

diff --git a/training/scripts/train.py b/training/scripts/train.py
index 92c6aff..4e92eae 100755
--- a/training/scripts/train.py
+++ b/training/scripts/train.py
@@ -32,14 +32,11 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model in bf16 (fp8 not supported for training)
-    # Disable quantization config to train full precision
+    # Load model - preserve NVFP4 quantization
     # Use device_map="auto" to distribute across GPUs from the start
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
-        torch_dtype=torch.bfloat16,
         device_map="auto",  # Distribute across GPUs
-        quantization_config=None,  # Override any fp8 quantization
         trust_remote_code=True,
     )