Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add parallel sampling using vllm #409

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

add parallel sampling using vllm #409

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6694965
parallel sampling using vllm
daspartho Nov 29, 2023
3e6a87e
vllm_support default true
daspartho Dec 4, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion src/llm_vm/onsite_llm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
from peft import get_peft_model, LoraConfig, prepare_model_for_kbit_training
from trl import SFTTrainer
from sentence_transformers import SentenceTransformer
from vllm import LLM, SamplingParams



Expand Down Expand Up @@ -86,14 +87,15 @@ def __getitem__(self, idx):
return self.dataset[idx]

class BaseOnsiteLLM(ABC):
def __init__(self,model_uri=None, tokenizer_kw_args={}, model_kw_args={}):
def __init__(self,model_uri=None, vllm_support=True, tokenizer_kw_args={}, model_kw_args={}):
if model_uri != None :
self.model_uri= model_uri
if model_uri is None and self.model_uri is None:
raise ValueError('model_uri not found')
self.model_name : str = self.model_uri.split('/')[-1] # our default for deriving model name
self.model=self.model_loader(**model_kw_args)
self.tokenizer=self.tokenizer_loader(**tokenizer_kw_args)
self.vllm_support=vllm_support

# Move the model to the specified device(s)
if isinstance(device, list):
Expand Down Expand Up @@ -145,6 +147,12 @@ def generate(self,prompt,max_length=100, tokenizer_kwargs={}, generation_kwargs=
I think it takes about a week for the apple to grow.
"""

if generation_kwargs['num_return_sequences']>1 and self.vllm_support:
print("doing parallel sampling using vllm")
sampling_params = SamplingParams(n=generation_kwargs['num_return_sequences'], max_tokens=max_length)
llm = LLM(model=self.model_uri)
outputs = llm.generate(prompt, sampling_params)
return [outputs[0].outputs[i].text for i in range(generation_kwargs['num_return_sequences'])]

if isinstance(device, list):
# If multiple GPUs are available, use first one
Expand Down Expand Up @@ -442,6 +450,7 @@ class SmallLocalNeo(BaseOnsiteLLM):
generate: Generates a response from a given prompt with the loaded LLM and tokenizer
"""
model_uri="EleutherAI/gpt-neo-1.3B"
vllm_support = False

def model_loader(self):
return GPTNeoForCausalLM.from_pretrained(self.model_uri)
Expand Down Expand Up @@ -680,6 +689,7 @@ class SmallLocalFlanT5(BaseOnsiteLLM):
"""

model_uri="google/flan-t5-small"
vllm_support = False
def model_loader(self):
return AutoModelForSeq2SeqLM.from_pretrained(self.model_uri)
def tokenizer_loader(self):
Expand All @@ -704,6 +714,7 @@ class SmallLocalBERT(BaseOnsiteLLM):
"""

model_uri = "bert-base-cased"
vllm_support = False
def model_loader(self):
return AutoModelForMaskedLM.from_pretrained(self.model_uri)
def tokenizer_loader(self):
Expand Down