From de093712c902ce658dc85fa27aea338e57f1d45c Mon Sep 17 00:00:00 2001 From: LYMDLUT <70597027+LYMDLUT@users.noreply.github.com> Date: Thu, 18 Jan 2024 18:03:44 +0800 Subject: [PATCH] Update README.md --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 822ea3b..c1976a6 100644 --- a/README.md +++ b/README.md @@ -152,11 +152,11 @@ model_dir = snapshot_download('Shanghai_AI_Laboratory/internlm2-chat-7b') tokenizer = AutoTokenizer.from_pretrained(model_dir, device_map="auto", trust_remote_code=True) # Set `torch_dtype=torch.float16` to load model in float16, otherwise it will be loaded as float32 and might cause OOM Error. model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True, torch_dtype=torch.float16) -# (Optional) If on low resource devices, you can load model in 4bits or 8 bits to further save GPU memory. +# (Optional) If on low resource devices, you can load model in 4-bit or 8-bit to further save GPU memory via bitsandbytes. # InternLM 7B in 4bit will cost nearly 8GB GPU memory. # pip install -U bitsandbytes - # 8bit: model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True, load_in_8bit=True) - # 4bit: model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True, load_in_4bit=True) + # 8-bit: model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True, load_in_8bit=True) + # 4-bit: model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True, load_in_4bit=True) model = model.eval() response, history = model.chat(tokenizer, "hello", history=[]) print(response)