mirror of https://github.com/InternLM/InternLM
Update README.md
parent
a8a2285561
commit
fefef768af
|
@ -152,7 +152,8 @@ model_dir = snapshot_download('Shanghai_AI_Laboratory/internlm2-chat-7b')
|
||||||
tokenizer = AutoTokenizer.from_pretrained(model_dir, device_map="auto", trust_remote_code=True)
|
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.
|
# 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)
|
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. InternLM 7B in 4bit will cost nearly 8GB GPU memory.
|
# (Optional) If on low resource devices, you can load model in 4bits or 8 bits to further save GPU memory.
|
||||||
|
# InternLM 7B in 4bit will cost nearly 8GB GPU memory.
|
||||||
# pip install -U bitsandbytes
|
# pip install -U bitsandbytes
|
||||||
# 8bit: model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True, load_in_8bit=True)
|
# 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)
|
# 4bit: model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True, load_in_4bit=True)
|
||||||
|
|
Loading…
Reference in New Issue