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

fix: disable quantization config for fp8 model training

Browse Source
This commit is contained in:
Christian Medina
2026-06-30 22:41:28 -04:00
parent a8c9cd682e
commit 3bb7328e2a
1 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
training/scripts/train.py
View File

@@ -32,13 +32,14 @@ def train(config_path):
print(f"Loading model: {config['base_model']}") print(f"Loading model: {config['base_model']}")
# Load model - let the model's own quantization config handle it # Load model in bf16 (fp8 not supported for training)
# (Ornith uses CompressedTensors, not BitsAndBytes) # Disable quantization config to train full precision
# Load on CPU first, then DeepSpeed will distribute
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
dtype=torch.bfloat16, torch_dtype=torch.bfloat16,
device_map="cpu", # Load on CPU, DeepSpeed distributes device_map="cpu",
quantization_config=None, # Override any fp8 quantization
trust_remote_code=True,
) )
# Add LoRA # Add LoRA