cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: load NVFP4 as bf16 by overriding quantization config

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 07:27:13 -04:00
parent be7b589c6b
commit 05ace2d38f
1 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
training/scripts/train.py
View File

@@ -32,17 +32,17 @@ def train(config_path):
print(f"Loading model: {config['base_model']}") 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']}") print(f"Loading model: {config['base_model']}")
from transformers import BitsAndBytesConfig
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
quantization_config=BitsAndBytesConfig(load_in_4bit=True), torch_dtype=torch.bfloat16,
device_map="cpu", # Load to CPU first device_map="cpu", # Load to CPU first
trust_remote_code=True, 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 # Add LoRA
lora_config = LoraConfig( lora_config = LoraConfig(