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

feat(server): optimize flash causal lm

Browse Source
This commit is contained in:
OlivierDehaene 2023-05-03 20:52:27 +02:00
parent 411b0d4e1f
commit 812de7ee50
1 changed files with 80 additions and 59 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

139
server/text_generation_server/models/flash_causal_lm.py
View File

@ -431,19 +431,15 @@ class FlashCausalLM(Model):
# Shortcut when batch_size == 1 # Shortcut when batch_size == 1
if len(batch) == 1: if len(batch) == 1:
input_ids = batch.input_ids[0].view(-1) input_ids = batch.input_ids[0].view(-1)
# No need to slice as flash attention will take care of it with cu_seqlens
past_key_values = batch.past_key_values
else: else:
# Concatenate tensors # Concatenate tensors
input_ids = torch.cat(batch.input_ids).view(-1) if not isinstance(batch.input_ids, torch.Tensor):
past_key_values = ( input_ids = torch.cat(batch.input_ids).view(-1)
torch.cat(batch.past_key_values, dim=1) else:
if batch.past_key_values is not None input_ids = batch.input_ids.view(-1)
else None
)
# if prefill and bs == 1 # if prefill and bs == 1
if past_key_values is None and len(batch) == 1: if batch.past_key_values is None and len(batch) == 1:
# Ask to pre-allocate kv to its max size # Ask to pre-allocate kv to its max size
# == number of tokens + max_new_tokens # == number of tokens + max_new_tokens
pre_allocate_past_size = ( pre_allocate_past_size = (
@ -453,11 +449,11 @@ class FlashCausalLM(Model):
pre_allocate_past_size = None pre_allocate_past_size = None
# Concatenate when prefill, torch.tensor when decode # Concatenate when prefill, torch.tensor when decode
position_ids = ( if batch.past_key_values is None:
torch.tensor(batch.position_ids, device=self.device) position_ids = torch.cat(batch.position_ids)
if batch.past_key_values is not None else:
else torch.cat(batch.position_ids) position_ids = batch.position_ids
)
cu_seqlens = torch.tensor( cu_seqlens = torch.tensor(
batch.cu_seqlens, device=self.device, dtype=torch.int32 batch.cu_seqlens, device=self.device, dtype=torch.int32
) )
@ -467,28 +463,10 @@ class FlashCausalLM(Model):
position_ids, position_ids,
cu_seqlens, cu_seqlens,
batch.max_seqlen, batch.max_seqlen,
past_key_values, batch.past_key_values,
pre_allocate_past_size, pre_allocate_past_size,
) )
# Initialize past_key_values in prefill
if batch.past_key_values is None:
# Initialize past padding tensor
if self.past_pad is None:
self.past_pad = present.new_zeros(
present.shape[0], 1, *present.shape[2:]
)
# Set in batch in case it needs to be used later in concatenate()
batch.past_pad = self.past_pad
if len(batch) == 1:
# present is already pre-padded
batch.past_key_values = present
else:
# Add padding after each sequence
# This will have the correct shape after the final past_key_values concatenation before the model
# forward
batch.past_key_values = [None, self.past_pad] * len(batch)
# Cumulative length # Cumulative length
cumulative_length = 0 cumulative_length = 0
@ -508,6 +486,9 @@ class FlashCausalLM(Model):
batch.all_input_ids_tensor, batch.all_input_ids_tensor,
) )
next_input_ids = input_ids.new_empty(len(batch.requests))
past_indices = []
# For each member of the batch # For each member of the batch
for i, ( for i, (
request, request,
@ -538,14 +519,56 @@ class FlashCausalLM(Model):
all_input_ids_tensor[None, :input_length], logits all_input_ids_tensor[None, :input_length], logits
) )
next_token_id_squeezed = next_token_id.squeeze() next_token_id_squeezed = next_token_id.squeeze()
next_token_id_item = next_token_id_squeezed.item() all_input_ids_tensor[input_length] = next_token_id_squeezed
next_input_ids[i] = next_token_id_squeezed
past_indices.extend([j for j in range(start_index + i, end_index + i)])
# Initialize past_key_values in prefill
if batch.past_key_values is None and len(batch) == 1:
# present is already pre-padded
batch.past_key_values = present
if len(batch) > 1:
batch.past_key_values = present.new_empty((present.shape[0], present.shape[1] + len(batch.requests), *present.shape[2:]))
batch.past_key_values[:, past_indices] = present
if prefill:
batch.position_ids = torch.tensor(batch.input_lengths, device=self.device)
else:
batch.position_ids = batch.position_ids + 1
next_token_ids = next_input_ids.tolist()
# Zipped iterator
iterator = zip(
batch.requests,
batch.input_lengths,
batch.offsets,
batch.token_offsets,
batch.next_token_choosers,
batch.stopping_criterias,
batch.all_input_ids,
batch.all_input_ids_tensor,
)
# For each member of the batch
for i, (
request,
input_length,
offset,
token_offset,
next_token_chooser,
stopping_criteria,
all_input_ids,
all_input_ids_tensor,
) in enumerate(iterator):
next_token_id_item = next_token_ids[i]
# Append next token to all tokens # Append next token to all tokens
all_input_ids.append(next_token_id_item) all_input_ids.append(next_token_id_item)
all_input_ids_tensor[input_length] = next_token_id_item
# Generated token # Generated token
next_token_logprob = logprobs[-1, next_token_id_item] next_token_logprob = 0.0
next_token_text, offset, token_offset = self.decode_token( next_token_text, offset, token_offset = self.decode_token(
all_input_ids, all_input_ids,
offset, offset,
@ -576,23 +599,24 @@ class FlashCausalLM(Model):
stopped = False stopped = False
generated_text = None generated_text = None
# Prefill prefill = stopping_criteria.current_tokens == 0
if prefill: # # Prefill
# Remove generated token to only have prefill and add nan for first prompt token # if prefill:
prefill_logprobs = [float("nan")] + logprobs.gather( # # Remove generated token to only have prefill and add nan for first prompt token
1, all_input_ids_tensor[1:input_length].unsqueeze(1) # prefill_logprobs = [float("nan")] + logprobs.gather(
).squeeze(1)[:-1].tolist() # 1, all_input_ids_tensor[1:input_length].unsqueeze(1)
prefill_token_ids = all_input_ids[:-1] # ).squeeze(1)[:-1].tolist()
prefill_texts = self.tokenizer.batch_decode( # prefill_token_ids = all_input_ids[:-1]
prefill_token_ids, # prefill_texts = self.tokenizer.batch_decode(
clean_up_tokenization_spaces=False, # prefill_token_ids,
skip_special_tokens=False, # clean_up_tokenization_spaces=False,
) # skip_special_tokens=False,
prefill_tokens = PrefillTokens( # )
prefill_token_ids, prefill_logprobs, prefill_texts # prefill_tokens = PrefillTokens(
) # prefill_token_ids, prefill_logprobs, prefill_texts
else: # )
prefill_tokens = None # else:
prefill_tokens = None
generation = Generation( generation = Generation(
request.id, request.id,
@ -609,19 +633,16 @@ class FlashCausalLM(Model):
new_input_length = input_length + 1 new_input_length = input_length + 1
# Update values # Update values
batch.input_ids[i] = next_token_id
batch.position_ids[i] = input_length
batch.input_lengths[i] = new_input_length batch.input_lengths[i] = new_input_length
batch.offsets[i] = offset batch.offsets[i] = offset
batch.token_offsets[i] = token_offset batch.token_offsets[i] = token_offset
batch.all_input_ids[i] = all_input_ids batch.all_input_ids[i] = all_input_ids
batch.all_input_ids_tensor[i] = all_input_ids_tensor
batch.max_seqlen = max(batch.max_seqlen, new_input_length) batch.max_seqlen = max(batch.max_seqlen, new_input_length)
if len(batch) != 1:
# Add each sequence before its padding
batch.past_key_values[i * 2] = present[:, start_index:end_index]
# Cumulative sum # Cumulative sum
batch.cu_seqlens[(i + 1)] = batch.cu_seqlens[i] + new_input_length batch.cu_seqlens[(i + 1)] = batch.cu_seqlens[i] + new_input_length
batch.input_ids = next_input_ids
# No need to return a batch if we know that all requests stopped # No need to return a batch if we know that all requests stopped
return generations, batch if not stopped else None return generations, batch if not stopped else None