huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 04:14:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

post process exllama model

Browse Source
This commit is contained in:
IlyasMoutawwakil 2024-02-01 12:48:17 +01:00
parent 12c1f54525
commit 5766c55b7a
1 changed files with 20 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
server/text_generation_server/server.py
View File

@ -63,20 +63,27 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
return generate_pb2.FilterBatchResponse(batch=filtered_batch.to_pb())
async def Warmup(self, request, context):
if self.quantize == "gptq":
try:
# When using GPTQ, Exllama kernels need some global kernels
# For which we have the finale shapes only after the model has loaded
# This will allocate those buffers.
from text_generation_server.utils.layers import (
create_exllama_buffers,
set_device,
)
if self.quantize in ["gptq", "awq"]:
has_exllama_layers = False
for _, module in self.model.model.named_modules():
if hasattr(module, "QUANT_TYPE"):
has_exllama_layers = True
break
set_device(self.model.device)
create_exllama_buffers(request.max_prefill_tokens)
except ImportError:
pass
if has_exllama_layers:
try:
# When using GPTQ or AWQ, Exllama kernels need some global kernels
# For which we have the finale shapes only after the model has loaded
# This will allocate those buffers.
from text_generation_server.utils.layers import (
create_exllama_buffers,
set_device,
)
set_device(self.model.device)
create_exllama_buffers(request.max_prefill_tokens)
except ImportError:
pass
if (
self.model.batch_type == IdeficsCausalLMBatch