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

Set default dtype to half

Browse Source
This commit is contained in:
Woosuk Kwon
2023-02-23 21:31:39 +00:00
parent de0fabbc5c
commit 1ce1333573
3 changed files with 23 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
cacheflow/models/model_utils.py
View File

@@ -1,3 +1,6 @@
from typing import Union
import torch
import torch.nn as nn
from cacheflow.models.opt import OPTForCausalLM
@@ -6,9 +9,23 @@ MODEL_CLASSES = {
'opt': OPTForCausalLM,
}
STR_DTYPE_TO_TORCH_DTYPE = {
'half': torch.half,
'float': torch.float,
'float16': torch.float16,
'float32': torch.float32,
}
def get_model(model_name: str) -> nn.Module:
def get_model(
model_name: str,
dtype: Union[torch.dtype, str],
) -> nn.Module:
if isinstance(dtype, str):
torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype.lower()]
else:
torch_dtype = dtype
for model_class, model in MODEL_CLASSES.items():
if model_class in model_name:
return model.from_pretrained(model_name)
return model.from_pretrained(model_name, torch_dtype=torch_dtype)
raise ValueError(f'Invalid model name: {model_name}')