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

feat: add CPU offloading for 70GB → 17GB conversion

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 22:33:36 -04:00
parent bee6683396
commit 410718bc73
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
quantize_model.py
View File

@@ -16,13 +16,13 @@ def quantize_model(model_path, output_path):
) )
print("Loading model with 4-bit quantization...") print("Loading model with 4-bit quantization...")
print(" This will distribute across both GPUs automatically\n") print(" Using CPU offloading to handle 70GB bf16 → 17GB 4-bit conversion\n")
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
model_path, model_path,
quantization_config=bnb_config, quantization_config=bnb_config,
device_map="auto", device_map="auto",
max_memory={0: "28GiB", 1: "28GiB"}, # Leave room for training max_memory={0: "28GiB", 1: "28GiB", "cpu": "120GiB"},
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
trust_remote_code=True, trust_remote_code=True,
) )