From b8bdd00d1965b0472f0039ef8a382902efc7f62c Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 7 May 2026 02:49:24 +0000 Subject: [PATCH] Lower GPU max_memory to 100GiB, add CPU-only fallback for low_memory_mode --- quantize_modelopt.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/quantize_modelopt.py b/quantize_modelopt.py index ad9f801..1e01305 100644 --- a/quantize_modelopt.py +++ b/quantize_modelopt.py @@ -102,8 +102,11 @@ def main(): model_kwargs["device_map"] = "auto" model_kwargs["offload_folder"] = "offload" model_kwargs["offload_state_dict"] = True - model_kwargs["max_memory"] = {i: "160GiB" for i in range(8)} + model_kwargs["max_memory"] = {i: "100GiB" for i in range(8)} model_kwargs["max_memory"]["cpu"] = "2500GiB" + elif args.low_memory_mode: + # Load entirely on CPU, modelopt will handle placement + model_kwargs["device_map"] = {"": "cpu"} model = AutoModelForCausalLM.from_pretrained(args.model, **model_kwargs)