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

fix: load 4-bit model AS-IS without additional quantization

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 13:53:43 -04:00
parent e967bd74f1
commit b300235be0
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
training/scripts/train.py
View File

@@ -36,15 +36,15 @@ def train(config_path):
print(f"Loading model: {config['base_model']}") print(f"Loading model: {config['base_model']}")
from transformers import BitsAndBytesConfig, AutoConfig from transformers import BitsAndBytesConfig, AutoConfig
# Skip quantization - use DeepSpeed ZeRO-3 with CPU offload # Load 4-bit model AS-IS (don't apply additional quantization)
print(f"Loading model: {config['base_model']}") print(f"Loading model: {config['base_model']}")
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
torch_dtype=torch.bfloat16, torch_dtype=torch.float16,
device_map="cpu", # Load to CPU first device_map="auto",
trust_remote_code=True, trust_remote_code=True,
) )
print("Model loaded to CPU. DeepSpeed will distribute.") print("Model loaded (4-bit). DeepSpeed will distribute.")
# Prepare model for k-bit training # Prepare model for k-bit training
from peft import prepare_model_for_kbit_training from peft import prepare_model_for_kbit_training