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

fix(server): fix warpers on CPU (#472)

Browse Source 
Closes #471
This commit is contained in:
OlivierDehaene 2023-06-20 11:06:10 +02:00 committed by GitHub
parent ece7ffa40a
commit 53aa9194c8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 27 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -237,20 +237,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")

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

@ -42,25 +42,31 @@ class StaticWarper:
self.static_next_logprob = None self.static_next_logprob = None
def __call__(self, scores): def __call__(self, scores):
if self.cuda_graph is None: if torch.cuda.is_available():
self.static_scores = scores if self.cuda_graph is None:
self.cuda_graph = torch.cuda.CUDAGraph() self.static_scores = scores
self.cuda_graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(self.cuda_graph, pool=mempool): with torch.cuda.graph(self.cuda_graph, pool=mempool):
local_scores = self.static_scores local_scores = self.static_scores
for warper in self.warpers: for warper in self.warpers:
local_scores = warper(None, local_scores) local_scores = warper(None, local_scores)
self.static_warped_scores = local_scores self.static_warped_scores = local_scores
# Compute logprobs # Compute logprobs
self.static_next_logprob = torch.log_softmax( self.static_next_logprob = torch.log_softmax(
self.static_warped_scores, -1 self.static_warped_scores, -1
) )
self.static_scores.copy_(scores) self.static_scores.copy_(scores)
self.cuda_graph.replay() self.cuda_graph.replay()
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)