Update README.md

README.md

@@ -217,6 +217,8 @@ We also provide OpenAI-Compatible API at DeepSeek Platform: [platform.deepseek.c
 
 ## 8. How to run locally
 **To utilize DeepSeek-V2 in BF16 format for inference, 80GB*8 GPUs are required.**
+
+**To utilize DeepSeek-V2-Lite in BF16 format for inference, 40GB*1 GPU is required.**
 ### Inference with Huggingface's Transformers
 You can directly employ [Huggingface's Transformers](https://github.com/huggingface/transformers) for model inference.
 
@@ -225,12 +227,9 @@ You can directly employ [Huggingface's Transformers](https://github.com/huggingf
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 
-model_name = "deepseek-ai/DeepSeek-V2"
+model_name = "deepseek-ai/DeepSeek-V2-Lite"
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-# `max_memory` should be set based on your devices
-max_memory = {i: "75GB" for i in range(8)}
-# `device_map` cannot be set to `auto`
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, device_map="sequential", torch_dtype=torch.bfloat16, max_memory=max_memory, attn_implementation="eager")
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.bfloat16).cuda()
 model.generation_config = GenerationConfig.from_pretrained(model_name)
 model.generation_config.pad_token_id = model.generation_config.eos_token_id
 
@@ -247,12 +246,9 @@ print(result)
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 
-model_name = "deepseek-ai/DeepSeek-V2-Chat"
+model_name = "deepseek-ai/DeepSeek-V2-Lite-Chat"
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-# `max_memory` should be set based on your devices
-max_memory = {i: "75GB" for i in range(8)}
-# `device_map` cannot be set to `auto`
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, device_map="sequential", torch_dtype=torch.bfloat16, max_memory=max_memory, attn_implementation="eager")
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.bfloat16).cuda()
 model.generation_config = GenerationConfig.from_pretrained(model_name)
 model.generation_config.pad_token_id = model.generation_config.eos_token_id
 
@@ -297,8 +293,8 @@ To utilize [vLLM](https://github.com/vllm-project/vllm) for model inference, ple
 from transformers import AutoTokenizer
 from vllm import LLM, SamplingParams
 
-max_model_len, tp_size = 8192, 8
-model_name = "deepseek-ai/DeepSeek-V2-Chat"
+max_model_len, tp_size = 8192, 1
+model_name = "deepseek-ai/DeepSeek-V2-Lite-Chat"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 llm = LLM(model=model_name, tensor_parallel_size=tp_size, max_model_len=max_model_len, trust_remote_code=True, enforce_eager=True)
 sampling_params = SamplingParams(temperature=0.3, max_tokens=256, stop_token_ids=[tokenizer.eos_token_id])