added low_cpu_mem_usage

Signed-off-by: root <[email protected]>
FlagAI-Open · Oct 23, 2023 · 1b0a7b3 · 1b0a7b3
1 parent 78671a2
commit 1b0a7b3
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/flagai/auto_model/auto_loader.py b/flagai/auto_model/auto_loader.py
@@ -305,9 +305,11 @@ def __init__(self,
                         math.ceil(model_max_length / orig_ctx_len))
                     config.rope_scaling = {"type": "linear", "factor": scaling_factor}
                 config.use_cache = False
-                model = AquilaForCausalLM.from_pretrained(download_path,
+                model = AquilaForCausalLM.from_pretrained(download_path,low_cpu_mem_usage=low_cpu_mem_usage,
                                                         **kwargs)
-
+                # from accelerate import load_checkpoint_and_dispatch
+                # model = load_checkpoint_and_dispatch(
+                #     model, download_path, device_map="auto", no_split_module_classes=["AquilaDecoderLayer"])
             tokenizer = AutoTokenizer.from_pretrained(download_path)
             self.model = model 
             self.tokenizer = tokenizer