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

fix(server): fix warpers on CPU

Browse Source
This commit is contained in:
OlivierDehaene 2023-06-19 17:44:53 +02:00
parent f59fb8b630
commit 82c9fadefe
2 changed files with 27 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
server/text_generation_server/models/__init__.py
View File

@ -230,20 +230,12 @@ def get_model(
) )
elif model_type == "t5": elif model_type == "t5":
if sharded:
return T5Sharded( return T5Sharded(
model_id, model_id,
revision, revision,
quantize=quantize, quantize=quantize,
trust_remote_code=trust_remote_code, trust_remote_code=trust_remote_code,
) )
else:
return Seq2SeqLM(
model_id,
revision,
quantize=quantize,
trust_remote_code=trust_remote_code,
)
if sharded: if sharded:
raise ValueError("sharded is not supported for AutoModel") raise ValueError("sharded is not supported for AutoModel")

6
server/text_generation_server/utils/logits_process.py
View File

@ -42,6 +42,7 @@ class StaticWarper:
self.static_next_logprob = None self.static_next_logprob = None
def __call__(self, scores): def __call__(self, scores):
if torch.cuda.is_available():
if self.cuda_graph is None: if self.cuda_graph is None:
self.static_scores = scores self.static_scores = scores
self.cuda_graph = torch.cuda.CUDAGraph() self.cuda_graph = torch.cuda.CUDAGraph()
@ -62,6 +63,11 @@ class StaticWarper:
return self.static_warped_scores, self.static_next_logprob return self.static_warped_scores, self.static_next_logprob
# CPU branch
for warper in self.warpers:
scores = warper(None, scores)
return scores, torch.log_softmax(scores, -1)
@lru_cache(10) @lru_cache(10)
def static_warper( def static_warper(