fix: use float16 dtype to preserve NVFP4 quantization

training/scripts/train.py

@@ -32,14 +32,15 @@ def train(config_path):
     
     print(f"Loading model: {config['base_model']}")
     
-    # Load model to CPU first, DeepSpeed will distribute
-    print("Loading model to CPU...")
+    # Load model - preserve NVFP4 quantization
+    print(f"Loading model: {config['base_model']}")
     model = AutoModelForCausalLM.from_pretrained(
         config["base_model"],
+        torch_dtype=torch.float16,  # Keep NVFP4 quantized
         device_map="cpu",  # Load to CPU first
         trust_remote_code=True,
     )
-    print(f"Model loaded to CPU. Moving to GPU with DeepSpeed...")
+    print("Model loaded to CPU.")
     
     # Add LoRA
     lora_config = LoraConfig(