biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Add Model Revision Support (#1014)

Browse Source 
Co-authored-by: Jasmond Loh <Jasmond.Loh@hotmail.com>
Co-authored-by: Zhuohan Li <zhuohan123@gmail.com>
This commit is contained in:
Jasmond L
2023-09-14 06:20:02 +08:00
committed by GitHub
parent 9841d48a10
commit ab019eea75
20 changed files with 75 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
vllm/model_executor/models/gpt_bigcode.py
View File

@@ -259,14 +259,15 @@ class GPTBigCodeForCausalLM(nn.Module):
def load_weights(self,
model_name_or_path: str,
cache_dir: Optional[str] = None,
load_format: str = "auto"):
load_format: str = "auto",
revision: Optional[str] = None):
tensor_model_parallel_world_size = (
get_tensor_model_parallel_world_size())
tensor_model_parallel_rank = get_tensor_model_parallel_rank()
state_dict = self.state_dict()
for name, loaded_weight in hf_model_weights_iterator(
model_name_or_path, cache_dir, load_format):
model_name_or_path, cache_dir, load_format, revision):
if "lm_head.weight" in name:
# GPT-2 ties the weights of the embedding layer and the final
# linear layer.