From 717151b98c8554b0a5134a984de0554a65d46544 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 7 May 2026 02:40:48 +0000 Subject: [PATCH] Add CPU offloading and max_memory caps for FP8 model loading --- quantize_modelopt.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/quantize_modelopt.py b/quantize_modelopt.py index d3c5de1..ad9f801 100644 --- a/quantize_modelopt.py +++ b/quantize_modelopt.py @@ -100,6 +100,10 @@ def main(): } if args.use_seq_device_map: 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"]["cpu"] = "2500GiB" model = AutoModelForCausalLM.from_pretrained(args.model, **model_kwargs)