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

Implement selective batching for vllm #9659

Merged
merged 26 commits into from
Dec 22, 2023
Merged

Implement selective batching for vllm #9659

Show file tree
Hide file tree
Changes from 25 commits
Commits
Show all changes
26 commits
Select commit Hold shift + click to select a range
b30eac7
add control to load hf model
gc-fu Dec 7, 2023
7dbe7ce
finish initial version of selective_batching
gc-fu Dec 12, 2023
dec985a
temp
gc-fu Dec 19, 2023
9132a5a
finish
gc-fu Dec 19, 2023
4baaf9e
Remove print statement
gc-fu Dec 19, 2023
59ad16a
fix error
gc-fu Dec 19, 2023
35a0ef8
Apply yang's optimization
gc-fu Dec 19, 2023
196e878
a version that works
gc-fu Dec 19, 2023
6fa0f68
We need to check kv_cache passed in, this could be an error. TODO: ad…
gc-fu Dec 20, 2023
6e883c4
format
gc-fu Dec 20, 2023
39ff898
temp solution: not batching prefill requests
gc-fu Dec 20, 2023
6ebc580
a version that works for prefill batching
gc-fu Dec 20, 2023
9f4435a
format
gc-fu Dec 21, 2023
16bd217
a solid version: works normally
gc-fu Dec 21, 2023
cc4cc9d
a temp version
gc-fu Dec 21, 2023
888eb70
Solid version: remove redundant functions
gc-fu Dec 21, 2023
39a8a11
fix format
gc-fu Dec 21, 2023
861c072
format
gc-fu Dec 21, 2023
29b2e60
solid: add option to enable selective_batching
gc-fu Dec 21, 2023
fa68542
remove logic for using transformer models
gc-fu Dec 21, 2023
b7d300a
format
gc-fu Dec 21, 2023
2f4abf1
format
gc-fu Dec 21, 2023
3e908f7
solid: enable argument VLLM_ENABLE_SELECTIVE_BATCHING
gc-fu Dec 21, 2023
f1216ae
format
gc-fu Dec 21, 2023
ec4a84e
finish
gc-fu Dec 22, 2023
a473fbc
format
gc-fu Dec 22, 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
18 changes: 18 additions & 0 deletions python/llm/src/bigdl/llm/transformers/convert.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
from .utils import logger
from typing import Union
import numpy as np
import os
from bigdl.llm.utils.common import invalidInputError


Expand Down Expand Up @@ -373,6 +374,8 @@ def convert_forward(m, target_m, new_forward):
def _optimize_post(model, lightweight_bmm=False):
from packaging import version
from bigdl.llm.transformers.models.llama import llama_attention_forward_4_31
from bigdl.llm.transformers.models.llama import llama_attention_selective_batching_forward_4_31
from bigdl.llm.transformers.models.llama import llama_model_selective_batching_forward_4_31
from bigdl.llm.transformers.models.llama import llama_rms_norm_forward
from bigdl.llm.transformers.models.llama import llama_mlp_forward
from transformers.modeling_utils import PreTrainedModel
Expand All @@ -383,6 +386,10 @@ def _optimize_post(model, lightweight_bmm=False):
"supported for further optimizations")
return model

vllm_selective_batching = os.getenv("VLLM_ENABLE_SELECTIVE_BATCHING")
enable_vllm_se_batching = vllm_selective_batching is not None
enable_vllm_se_batching = enable_vllm_se_batching and vllm_selective_batching.lower() == "true"

trans_version = transformers.__version__
if version.parse(trans_version) >= version.parse("4.31.0"):
convert_forward(
Expand All @@ -396,6 +403,17 @@ def _optimize_post(model, lightweight_bmm=False):
convert_forward(model,
transformers.models.llama.modeling_llama.LlamaMLP,
llama_mlp_forward)
if enable_vllm_se_batching:
convert_forward(
model,
transformers.models.llama.modeling_llama.LlamaModel,
llama_model_selective_batching_forward_4_31,
)
convert_forward(
model,
transformers.models.llama.modeling_llama.LlamaAttention,
llama_attention_selective_batching_forward_4_31,
)
else:
# todo implement 4.28.0 ~ 4.30.2
pass
Expand Down
Loading