Add CPU offloading and max_memory caps for FP8 model loading

2026-05-07 02:40:48 +00:00
parent aff12c6951
commit 717151b98c
1 changed files with 4 additions and 0 deletions
--- 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)