Add torch_dtype to README example

README.md

@@ -75,18 +75,20 @@ Overall, InternLM-20B comprehensively outperforms open-source models in the 13B
 ## Import from Transformers
 To load the InternLM 20B model using Transformers, use the following code:
 ```python
->>> from transformers import AutoTokenizer, AutoModelForCausalLM
->>> tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-20b", trust_remote_code=True)
->>> model = AutoModelForCausalLM.from_pretrained("internlm/internlm-20b", trust_remote_code=True).cuda()
->>> model = model.eval()
->>> inputs = tokenizer(["Coming to the beautiful nature, we found"], return_tensors="pt")
->>> for k,v in inputs.items():
-        inputs[k] = v.cuda()
->>> gen_kwargs = {"max_length": 128, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.05}
->>> output = model.generate(**inputs, **gen_kwargs)
->>> output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
->>> print(output)
-Coming to the beautiful nature, we found not only various mountains, rivers, trees, and flowers but also many birds and beasts. Birds are the ones we are most familiar with; some are soaring in the sky, some are hopping on the ground, while others perch on trees...
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-20b", trust_remote_code=True)
+# Set `torch_dtype=torch.bfloat16` to load model in bfloat16, otherwise it will be loaded as float32 and cause OOM Error.
+model = AutoModelForCausalLM.from_pretrained("internlm/internlm-20b", torch_dtype=torch.bfloat16, trust_remote_code=True).cuda()
+model = model.eval()
+inputs = tokenizer(["Coming to the beautiful nature, we found"], return_tensors="pt")
+for k,v in inputs.items():
+    inputs[k] = v.cuda()
+gen_kwargs = {"max_length": 128, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.05}
+output = model.generate(**inputs, **gen_kwargs)
+output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
+print(output)
+# Coming to the beautiful nature, we found not only various mountains, rivers, trees, and flowers but also many birds and beasts. Birds are the ones we are most familiar with; some are soaring in the sky, some are hopping on the ground, while others perch on trees...
 ```
 
 **Limitations:** Although we have made efforts to ensure the safety of the model during the training process and to encourage the model to generate text that complies with ethical and legal requirements, the model may still produce unexpected outputs due to its size and probabilistic generation paradigm. For example, the generated responses may contain biases, discrimination, or other harmful content. Please do not propagate such content. We are not responsible for any consequences resulting from the dissemination of harmful information.
@@ -145,18 +147,20 @@ InternLM 20B 在模型结构上选择了深结构，层数设定为60层，超
 ## 通过 Transformers 加载
 通过以下的代码加载 InternLM 20B 模型
 ```python
->>> from transformers import AutoTokenizer, AutoModelForCausalLM
->>> tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-20b", trust_remote_code=True)
->>> model = AutoModelForCausalLM.from_pretrained("internlm/internlm-20b", trust_remote_code=True).cuda()
->>> model = model.eval()
->>> inputs = tokenizer(["来到美丽的大自然，我们发现"], return_tensors="pt")
->>> for k,v in inputs.items():
-        inputs[k] = v.cuda()
->>> gen_kwargs = {"max_length": 128, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.05}
->>> output = model.generate(**inputs, **gen_kwargs)
->>> output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
->>> print(output)
-来到美丽的大自然，我们发现，这里不仅有大大小小的山川河流和树木花草，而且还有很多飞鸟走兽。我们最熟悉的就是鸟类了，它们有的在天上飞翔，有的在地上跳跃，还有的在树上栖息……
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-20b", trust_remote_code=True)
+# `torch_dtype=torch.bfloat16` 可以令模型以 bfloat16 精度加载，否则 transformers 会将模型加载为 float32，导致显存不足
+model = AutoModelForCausalLM.from_pretrained("internlm/internlm-20b", torch_dtype=torch.bfloat16, trust_remote_code=True).cuda()
+model = model.eval()
+inputs = tokenizer(["来到美丽的大自然，我们发现"], return_tensors="pt")
+for k,v in inputs.items():
+    inputs[k] = v.cuda()
+gen_kwargs = {"max_length": 128, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.05}
+output = model.generate(**inputs, **gen_kwargs)
+output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
+print(output)
+# 来到美丽的大自然，我们发现，这里不仅有大大小小的山川河流和树木花草，而且还有很多飞鸟走兽。我们最熟悉的就是鸟类了，它们有的在天上飞翔，有的在地上跳跃，还有的在树���栖息……
 ```
 
 **局限性：** 尽管在训练过程中我们非常注重模型的安全性，尽力促使模型输出符合伦理和法律要求的文本，但受限于模型大小以及概率生成范式，模型可能会产生各种不符合预期的输出，例如回复内容包含偏见、歧视等有害内容，请勿传播这些内容。由于传播不良信息导致的任何后果，本项目不承担责任。