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

fix: load QLoRA model to CPU first, DeepSpeed distributes

Browse Source
This commit is contained in:
Christian Medina
2026-07-01 09:37:20 -04:00
parent 2fa5173c56
commit 2fa52cd9d4
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
training/scripts/train.py
View File

@@ -47,7 +47,7 @@ def train(config_path):
config["base_model"], config["base_model"],
quantization_config=bnb_config, quantization_config=bnb_config,
dtype=torch.bfloat16, dtype=torch.bfloat16,
device_map="auto", device_map="cpu", # Load to CPU first
trust_remote_code=True, trust_remote_code=True,
) )
print("Model loaded with QLoRA (4-bit).") print("Model loaded with QLoRA (4-bit).")