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, FSDP shards later

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 08:26:23 -04:00
parent dbc97138e9
commit c607888d8f
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
train.py
View File

@@ -50,11 +50,11 @@ def train(config_path):
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
config["base_model"], config["base_model"],
quantization_config=bnb_config, quantization_config=bnb_config,
device_map="auto", device_map="cpu", # Load to CPU, FSDP shards later
trust_remote_code=True, trust_remote_code=True,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
) )
print("✓ Success: QLoRA 4-bit") print("✓ Success: QLoRA 4-bit loaded to CPU")
except Exception as e: except Exception as e:
errors.append(("QLoRA 4-bit", e)) errors.append(("QLoRA 4-bit", e))
print(f"✗ Failed: {e}") print(f"✗ Failed: {e}")