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

fix: preserve NVFP4 quantization, don't force bf16

Browse Source
This commit is contained in:
Christian Medina
2026-06-30 23:50:27 -04:00
parent 114e25fcfb
commit 905c2ba30d
1 changed files with 1 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
training/scripts/train.py
View File

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