feat: add CPU offloading for 70GB → 17GB conversion
This commit is contained in:
@@ -16,13 +16,13 @@ def quantize_model(model_path, output_path):
|
||||
)
|
||||
|
||||
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_path,
|
||||
quantization_config=bnb_config,
|
||||
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,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user