From 1b0a7b30c6a2b33802f298cc9b8ce5dfb5532c3f Mon Sep 17 00:00:00 2001 From: root Date: Mon, 23 Oct 2023 03:12:54 +0000 Subject: [PATCH] added low_cpu_mem_usage Signed-off-by: root --- flagai/auto_model/auto_loader.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/flagai/auto_model/auto_loader.py b/flagai/auto_model/auto_loader.py index 77831231..51d14c60 100755 --- 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