diff --git a/server/tests/models/test_bloom.py b/server/tests/models/test_bloom.py index 32ee6686..6e9e5205 100644 --- a/server/tests/models/test_bloom.py +++ b/server/tests/models/test_bloom.py @@ -5,7 +5,7 @@ from copy import copy from transformers import AutoTokenizer from text_generation_server.pb import generate_pb2 -from text_generation_server.models.causal_lm import CausalLMBatch +from text_generation_server.models.transformers_causal_lm import CausalLMBatch from text_generation_server.utils import weight_hub_files, download_weights from text_generation_server.models.bloom import BloomCausalLMBatch, BLOOMSharded diff --git a/server/tests/models/test_causal_lm.py b/server/tests/models/test_causal_lm.py index 6e6463bc..7d674947 100644 --- a/server/tests/models/test_causal_lm.py +++ b/server/tests/models/test_causal_lm.py @@ -5,7 +5,10 @@ from copy import copy from transformers import AutoTokenizer from text_generation_server.pb import generate_pb2 -from text_generation_server.models.causal_lm import CausalLM, CausalLMBatch +from text_generation_server.models.transformers_causal_lm import ( + TransformersCausalLM, + CausalLMBatch, +) @pytest.fixture(scope="session") diff --git a/server/tests/models/test_santacoder.py b/server/tests/models/test_santacoder.py index cb2622d9..19152659 100644 --- a/server/tests/models/test_santacoder.py +++ b/server/tests/models/test_santacoder.py @@ -1,7 +1,7 @@ import pytest from text_generation_server.pb import generate_pb2 -from text_generation_server.models.causal_lm import CausalLMBatch +from text_generation_server.models.transformers_causal_lm import CausalLMBatch from text_generation_server.models.santacoder import SantaCoder diff --git a/server/text_generation_server/models/__init__.py b/server/text_generation_server/models/__init__.py index c3be5d0d..5615de65 100644 --- a/server/text_generation_server/models/__init__.py +++ b/server/text_generation_server/models/__init__.py @@ -8,11 +8,13 @@ from transformers.models.auto import modeling_auto from huggingface_hub import hf_hub_download, HfApi from typing import Optional, List from pathlib import Path - +import transformers from text_generation_server.utils.speculate import get_speculate, set_speculate from text_generation_server.models.model import Model -from text_generation_server.models.causal_lm import CausalLM -from text_generation_server.models.causal_lm_ragged import CausalLMRagged +from text_generation_server.models.transformers_causal_lm import TransformersCausalLM +from text_generation_server.models.transformers_flash_causal_lm import ( + TransformersFlashCausalLM, +) from text_generation_server.models.flash_causal_lm import FlashCausalLM from text_generation_server.models.bloom import BLOOMSharded from text_generation_server.models.mpt import MPTSharded @@ -25,6 +27,8 @@ from text_generation_server.models.t5 import T5Sharded from text_generation_server.models.gpt_neox import GPTNeoxSharded from text_generation_server.models.phi import Phi +from text_generation_server.models.globals import USE_CUSTOM_MODELING + from text_generation_server.utils.import_utils import SYSTEM # The flag below controls whether to allow TF32 on matmul. This flag defaults to False @@ -289,6 +293,31 @@ def get_model( ) model_type = config_dict.get("model_type", None) + transformers_causal_lm_class = TransformersCausalLM + if ( + not USE_CUSTOM_MODELING + and model_type in modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES + ): + logger.info( + "TGI's flash enabled models could either not be loaded or are disabled, using Transformers fallback." + ) + transformers_model_class = getattr( + transformers, modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[model_type] + ) + + if ( + transformers_model_class._supports_flash_attn_2 + and transformers_model_class._supports_cache_class + ): + logger.info( + f"Transformers' {model_type} implementation supports custom cache and flash/paged attention. Using TransformersFlashCausalLM with ragged tensors (single dimension for batch and sequence length)." + ) + transformers_causal_lm_class = TransformersFlashCausalLM + else: + logger.info( + f"Transformers' {model_type} implementation supports custom cache and flash/paged attention. Using TransformersCausalLM with classic tensors with padding (two dimensions for batch size and sequence length)." + ) + speculator = None if "medusa_num_heads" in config_dict: medusa_model_id = model_id @@ -450,7 +479,7 @@ def get_model( or model_type == GPT2 and model_id.startswith("bigcode/") ): - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return FlashSantacoderSharded( model_id, revision, @@ -492,7 +521,7 @@ def get_model( trust_remote_code=trust_remote_code, ) elif model_type == GPT2: - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: try: return FlashGPT2( model_id, @@ -505,7 +534,8 @@ def get_model( except RuntimeError as e: # Lots of legacy models with various weight names. logger.warning(f"Couldn't load flash gpt2 variant: {e}") - return CausalLM( + + return transformers_causal_lm_class( model_id, revision, quantize=quantize, @@ -516,7 +546,7 @@ def get_model( elif sharded: raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded GPT-2")) else: - return CausalLM( + return transformers_causal_lm_class( model_id, revision, quantize=quantize, @@ -525,7 +555,7 @@ def get_model( trust_remote_code=trust_remote_code, ) elif model_type == GPT_NEOX: - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return FlashNeoXSharded( model_id, revision, @@ -544,7 +574,7 @@ def get_model( trust_remote_code=trust_remote_code, ) else: - return CausalLM( + return transformers_causal_lm_class( model_id, revision, quantize=quantize, @@ -554,7 +584,7 @@ def get_model( ) elif model_type == PHI: - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return FlashPhi( model_id, revision, @@ -564,7 +594,7 @@ def get_model( trust_remote_code=trust_remote_code, ) else: - return CausalLM( + return transformers_causal_lm_class( model_id, revision, quantize=quantize, @@ -574,7 +604,7 @@ def get_model( ) elif model_type == "phi-msft": - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: raise NotImplementedError( "Legacy phi-msft is not supported with Flash Attention" ) @@ -589,7 +619,7 @@ def get_model( ) elif model_type == LLAMA or model_type == BAICHUAN or model_type == PHI3: - if FLASH_ATTENTION and False: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return FlashLlama( model_id, revision, @@ -602,8 +632,7 @@ def get_model( elif sharded: raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Llama")) else: - logger.info("LOADING CAUSALLM!!!!!!!!!!!!!!!!!!") - return CausalLMRagged( + return transformers_causal_lm_class( model_id, revision, quantize=quantize, @@ -612,7 +641,7 @@ def get_model( trust_remote_code=trust_remote_code, ) if model_type == GEMMA: - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return FlashGemma( model_id, revision, @@ -624,7 +653,7 @@ def get_model( elif sharded: raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Gemma")) else: - return CausalLM( + return transformers_causal_lm_class( model_id, revision, quantize=quantize, @@ -634,7 +663,7 @@ def get_model( ) if model_type == COHERE: - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return FlashCohere( model_id, revision, @@ -646,7 +675,7 @@ def get_model( elif sharded: raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Cohere")) else: - return CausalLM( + return transformers_causal_lm_class( model_id, revision, quantize=quantize, @@ -656,7 +685,7 @@ def get_model( ) if model_type == DBRX: - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return FlashDbrx( model_id, revision, @@ -668,7 +697,7 @@ def get_model( elif sharded: raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded DBRX")) else: - return CausalLM( + return transformers_causal_lm_class( model_id, revision, quantize=quantize, @@ -679,7 +708,7 @@ def get_model( if model_type in ["RefinedWeb", "RefinedWebModel", FALCON]: if sharded: - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: if config_dict.get("alibi", False): raise NotImplementedError("sharded is not supported for this model") return FlashRWSharded( @@ -712,7 +741,7 @@ def get_model( ) if model_type == MISTRAL: - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return FlashMistral( model_id, revision, @@ -724,7 +753,7 @@ def get_model( elif sharded: raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Mistral")) else: - return CausalLM( + return transformers_causal_lm_class( model_id, revision, quantize=quantize, @@ -734,7 +763,7 @@ def get_model( ) if model_type == MIXTRAL: - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return FlashMixtral( model_id, revision, @@ -746,7 +775,7 @@ def get_model( elif sharded: raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Mixtral")) else: - return CausalLM( + return transformers_causal_lm_class( model_id, revision, quantize=quantize, @@ -756,7 +785,7 @@ def get_model( ) if model_type == STARCODER2: - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return FlashStarcoder2( model_id, revision, @@ -769,7 +798,7 @@ def get_model( FLASH_ATT_ERROR_MESSAGE.format("Sharded Starcoder2") ) else: - return CausalLM( + return transformers_causal_lm_class( model_id, revision, quantize=quantize, @@ -779,7 +808,7 @@ def get_model( ) if model_type == QWEN2: - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return FlashQwen2( model_id, revision, @@ -790,7 +819,7 @@ def get_model( elif sharded: raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Qwen2")) else: - return CausalLM( + return transformers_causal_lm_class( model_id, revision, quantize=quantize, @@ -819,7 +848,7 @@ def get_model( trust_remote_code=trust_remote_code, ) if model_type == IDEFICS: - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return IDEFICSSharded( model_id, revision, @@ -831,7 +860,7 @@ def get_model( else: raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Idefics")) if model_type == IDEFICS2: - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return Idefics2( model_id, revision, @@ -843,7 +872,7 @@ def get_model( else: raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Idefics")) if model_type == "paligemma": - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return PaliGemma( model_id, revision, @@ -856,7 +885,7 @@ def get_model( raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Idefics")) if model_type == LLAVA_NEXT: - if FLASH_ATTENTION: + if FLASH_ATTENTION and USE_CUSTOM_MODELING: return LlavaNext( model_id, revision, @@ -883,7 +912,7 @@ def get_model( elif quantize == "exl2": raise NotImplementedError("exl2 quantization is not supported for AutoModel") if model_type in modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES: - return CausalLM( + return transformers_causal_lm_class( model_id, revision, quantize=quantize, @@ -904,7 +933,7 @@ def get_model( auto_map = config_dict.get("auto_map", None) if trust_remote_code and auto_map is not None: if "AutoModelForCausalLM" in auto_map.keys(): - return CausalLM( + return transformers_causal_lm_class( model_id, revision, quantize=quantize, diff --git a/server/text_generation_server/models/bloom.py b/server/text_generation_server/models/bloom.py index 17aa12e8..88cb2bdf 100644 --- a/server/text_generation_server/models/bloom.py +++ b/server/text_generation_server/models/bloom.py @@ -12,8 +12,8 @@ from transformers import ( from text_generation_server.models.custom_modeling.bloom_modeling import ( BloomForCausalLM, ) -from text_generation_server.models import CausalLM -from text_generation_server.models.causal_lm import CausalLMBatch +from text_generation_server.models import TransformersCausalLM +from text_generation_server.models.transformers_causal_lm import CausalLMBatch from text_generation_server.pb import generate_pb2 from text_generation_server.utils import ( initialize_torch_distributed, @@ -36,7 +36,7 @@ class BloomCausalLMBatch(CausalLMBatch): return batch -class BLOOMSharded(CausalLM): +class BLOOMSharded(TransformersCausalLM): def __init__( self, model_id: str, @@ -89,7 +89,7 @@ class BLOOMSharded(CausalLM): model = BloomForCausalLM(config, weights) torch.distributed.barrier(group=self.process_group) - super(CausalLM, self).__init__( + super().__init__( model_id=model_id, model=model, tokenizer=tokenizer, diff --git a/server/text_generation_server/models/causal_lm.py b/server/text_generation_server/models/causal_lm.py deleted file mode 100644 index 10c64c66..00000000 --- a/server/text_generation_server/models/causal_lm.py +++ /dev/null @@ -1,787 +0,0 @@ -import torch -import time - -from dataclasses import dataclass -from opentelemetry import trace -from transformers import AutoTokenizer, AutoModelForCausalLM, PreTrainedTokenizerBase -from typing import Optional, Tuple, List, Type, Dict - -from text_generation_server.models import Model -from text_generation_server.utils.chunks import concat_text_chunks -from text_generation_server.utils.tokens import batch_top_tokens -from text_generation_server.models.types import ( - Batch, - Tokens, - Generation, - GeneratedText, -) -from text_generation_server.pb import generate_pb2 -from text_generation_server.utils import NextTokenChooser, StoppingCriteria, Sampling - -tracer = trace.get_tracer(__name__) - - -@dataclass -class CausalLMBatch(Batch): - batch_id: int - requests: List[generate_pb2.Request] - requests_idx_mapping: Dict[int, int] - - # Decoder values - input_ids: torch.Tensor - attention_mask: torch.Tensor - position_ids: torch.Tensor - past_key_values: Optional[List[Tuple]] - - # All tokens - all_input_ids: List[torch.Tensor] - - # Lengths of all generations present in the batch - input_lengths: List[int] - prefix_offsets: List[int] - read_offsets: List[int] - - # Generation helpers - next_token_choosers: List[NextTokenChooser] - stopping_criterias: List[StoppingCriteria] - top_n_tokens: List[int] - top_n_tokens_tensor: torch.Tensor - - # Metadata used for padding - max_input_length: int - padding_right_offset: int - - # Maximum number of tokens this batch will grow to - max_tokens: int - - # Past metadata - keys_head_dim_last: bool = True - - def to_pb(self) -> generate_pb2.CachedBatch: - return generate_pb2.CachedBatch( - id=self.batch_id, - request_ids=[r.id for r in self.requests], - size=len(self), - max_tokens=self.max_tokens, - ) - - @classmethod - def from_pb( - cls, - pb: generate_pb2.Batch, - tokenizer: PreTrainedTokenizerBase, - dtype: torch.dtype, - device: torch.device, - ) -> "CausalLMBatch": - inputs = [] - next_token_choosers = [] - stopping_criterias = [] - top_n_tokens = [] - prefix_offsets = [] - read_offsets = [] - requests_idx_mapping = {} - - # Parse batch - max_truncation = 0 - padding_right_offset = 0 - max_decode_tokens = 0 - for i, r in enumerate(pb.requests): - requests_idx_mapping[r.id] = i - inputs.append(concat_text_chunks(r.input_chunks.chunks)) - - next_token_choosers.append( - NextTokenChooser.from_pb(r.parameters, device, tokenizer) - ) - stopping_criteria = StoppingCriteria.from_pb( - r.stopping_parameters, tokenizer - ) - stopping_criterias.append(stopping_criteria) - top_n_tokens.append(r.top_n_tokens) - max_truncation = max(max_truncation, r.truncate) - max_decode_tokens += stopping_criteria.max_new_tokens - padding_right_offset = max( - padding_right_offset, stopping_criteria.max_new_tokens - ) - - tokenized_inputs = tokenizer( - inputs, - return_tensors="pt", - padding=True, - return_token_type_ids=False, - truncation=True, - max_length=max_truncation, - ).to(device) - for _ in pb.requests: - input_len = tokenized_inputs["input_ids"].shape[1] - prefix_offsets.append(input_len - 5) - read_offsets.append(input_len) - - input_lengths = tokenized_inputs["attention_mask"].sum(1) - max_input_length = input_lengths.max() - - input_ids = tokenized_inputs["input_ids"] - # Allocate maximum attention_mask - attention_mask = input_ids.new_zeros( - (pb.size, max_input_length + padding_right_offset) - ) - # Copy tokenizer attention_mask into fully allocated attention_mask - attention_mask[:, :max_input_length] = tokenized_inputs["attention_mask"] - - position_ids = tokenized_inputs["attention_mask"].long().cumsum(-1) - 1 - position_ids.masked_fill_(tokenized_inputs["attention_mask"] == 0, 1) - all_input_ids = tokenized_inputs["input_ids"].T.split(1, dim=1) - top_n_tokens_tensor = torch.tensor( - top_n_tokens, device=device, dtype=torch.int64 - ) - - max_tokens = len(inputs) * (max_input_length + max_decode_tokens) - - return cls( - batch_id=pb.id, - requests=pb.requests, - requests_idx_mapping=requests_idx_mapping, - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=None, - all_input_ids=list(all_input_ids), - input_lengths=input_lengths.tolist(), - prefix_offsets=prefix_offsets, - read_offsets=read_offsets, - next_token_choosers=next_token_choosers, - stopping_criterias=stopping_criterias, - top_n_tokens=top_n_tokens, - top_n_tokens_tensor=top_n_tokens_tensor, - max_input_length=max_input_length.item(), - padding_right_offset=padding_right_offset, - max_tokens=max_tokens, - ) - - @tracer.start_as_current_span("filter") - def filter(self, request_ids: List[int]) -> Optional["CausalLMBatch"]: - if len(request_ids) == 0: - raise ValueError("Batch must have at least one request") - if len(request_ids) == len(self): - return self - - keep_indices = [] - - # New values after filtering - requests_idx_mapping = {} - requests = [] - input_lengths = [] - prefix_offsets = [] - read_offsets = [] - all_input_ids = [] - max_input_length = 0 - - next_token_choosers = [] - stopping_criterias = [] - top_n_tokens = [] - - total_remaining_decode_tokens = 0 - new_padding_right_offset = 0 - - for i, request_id in enumerate(request_ids): - idx = self.requests_idx_mapping[request_id] - requests_idx_mapping[request_id] = i - keep_indices.append(idx) - - requests.append(self.requests[idx]) - prefix_offsets.append(self.prefix_offsets[idx]) - read_offsets.append(self.read_offsets[idx]) - all_input_ids.append(self.all_input_ids[idx]) - - request_input_length = self.input_lengths[idx] - input_lengths.append(request_input_length) - max_input_length = max(max_input_length, request_input_length) - - next_token_choosers.append(self.next_token_choosers[idx]) - stopping_criteria = self.stopping_criterias[idx] - stopping_criterias.append(stopping_criteria) - top_n_tokens.append(self.top_n_tokens[idx]) - remaining_decode_tokens = ( - stopping_criteria.max_new_tokens - stopping_criteria.current_tokens - ) - total_remaining_decode_tokens += remaining_decode_tokens - new_padding_right_offset = max( - new_padding_right_offset, remaining_decode_tokens - ) - - # Apply indices to input_ids, attention mask, past key values and other items that need to be cached - input_ids = self.input_ids[keep_indices] - position_ids = self.position_ids[keep_indices] - self.attention_mask = self.attention_mask[ - keep_indices, - -(self.padding_right_offset + max_input_length) : ( - self.attention_mask.shape[1] - self.padding_right_offset - ) - + new_padding_right_offset, - ] - - # Ensure that past_key_values tensors can be updated in-place - if type(self.past_key_values[0]) == tuple: - self.past_key_values = [list(layer) for layer in self.past_key_values] - - # Update tensors in-place to allow incremental garbage collection - past_kv_length = max_input_length - 1 - for layer in self.past_key_values: - past_keys, past_values = layer - if len(past_keys.shape) == 3: - # Force past to be of dim [self_size, num_heads, ...] for easy indexing - past_keys = past_keys.view(len(self), -1, *past_keys.shape[-2:]) - past_values = past_values.view(len(self), -1, *past_values.shape[-2:]) - if self.keys_head_dim_last: - layer[0] = past_keys[keep_indices, :, -past_kv_length:, :] - else: - layer[0] = past_keys[keep_indices, :, :, -past_kv_length:] - del past_keys - layer[1] = past_values[keep_indices, :, -past_kv_length:, :] - del past_values - - top_n_tokens_tensor = self.top_n_tokens_tensor[keep_indices] - max_tokens = len(request_ids) * max_input_length + total_remaining_decode_tokens - - self.requests = requests - self.requests_idx_mapping = requests_idx_mapping - self.input_ids = input_ids - self.position_ids = position_ids - self.all_input_ids = all_input_ids - self.input_lengths = input_lengths - self.prefix_offsets = prefix_offsets - self.read_offsets = read_offsets - self.next_token_choosers = next_token_choosers - self.stopping_criterias = stopping_criterias - self.top_n_tokens = top_n_tokens - self.top_n_tokens_tensor = top_n_tokens_tensor - self.max_input_length = max_input_length - self.padding_right_offset = new_padding_right_offset - self.max_tokens = max_tokens - - return self - - @classmethod - @tracer.start_as_current_span("concatenate") - def concatenate(cls, batches: List["CausalLMBatch"]) -> "CausalLMBatch": - # Used for padding - total_batch_size = 0 - max_input_length = 0 - padding_right_offset = 0 - for batch in batches: - total_batch_size += len(batch) - max_input_length = max(max_input_length, batch.max_input_length) - padding_right_offset = max(padding_right_offset, batch.padding_right_offset) - - # Batch attributes - requests = [] - requests_idx_mapping = {} - input_lengths = [] - prefix_offsets = [] - read_offsets = [] - all_input_ids = [] - next_token_choosers = [] - stopping_criterias = [] - top_n_tokens = [] - max_tokens = 0 - - # Batch tensors - input_ids = None - attention_mask = None - position_ids = None - past_key_values = [] - top_n_tokens_tensor = None - - # Used for slicing correctly inside the tensors - # Equivalent to a cumsum on batch sizes - start_index = 0 - for i, batch in enumerate(batches): - requests.extend(batch.requests) - input_lengths.extend(batch.input_lengths) - prefix_offsets.extend(batch.prefix_offsets) - read_offsets.extend(batch.read_offsets) - all_input_ids.extend(batch.all_input_ids) - next_token_choosers.extend(batch.next_token_choosers) - stopping_criterias.extend(batch.stopping_criterias) - top_n_tokens.extend(batch.top_n_tokens) - - if i == 0: - requests_idx_mapping = batch.requests_idx_mapping - else: - # We need to offset the mapping for each batch by the cumulative batch size - for k, v in batch.requests_idx_mapping.items(): - requests_idx_mapping[k] = v + start_index - - # Slicing end index for this batch - end_index = start_index + len(batch) - - # We only concatenate batches that did at least one step - if batch.past_key_values is None: - raise ValueError("only concatenate prefilled batches") - - # Create empty tensor - # input_ids is always of shape [batch_size, 1] - # We do not need to pad it - if input_ids is None: - input_ids = batch.input_ids.new_empty((total_batch_size, 1)) - # Copy to correct indices - input_ids[start_index:end_index] = batch.input_ids - - # Create padded tensor - if attention_mask is None: - attention_mask = batch.attention_mask.new_zeros( - (total_batch_size, max_input_length + padding_right_offset), - ) - - if top_n_tokens_tensor is None: - top_n_tokens_tensor = batches[0].top_n_tokens_tensor.new_zeros( - total_batch_size, - ) - top_n_tokens_tensor[start_index:end_index] = batch.top_n_tokens_tensor - - # We need to slice the attention mask to remove padding from previous steps - # and to remove unused allocated space - left_offset = max_input_length - batch.max_input_length - batch_left_offset = ( - batch.attention_mask.shape[1] - - batch.max_input_length - - batch.padding_right_offset - ) - attention_mask[ - start_index:end_index, - left_offset:-padding_right_offset, - ] = batch.attention_mask[ - :, - batch_left_offset : -batch.padding_right_offset, - ] - - # Create empty tensor - # position_ids is always of shape [batch_size, 1] - if position_ids is None: - position_ids = batch.position_ids.new_empty((total_batch_size, 1)) - position_ids[start_index:end_index] = batch.position_ids - - # Shenanigans to get dimensions because BLOOM outputs a past with a different shape - # BLOOM Keys: [batch_size * num_heads, head_dim, seq_length] - # BLOOM Values: [batch_size * num_heads, seq_length, head_dim] - # And ensure that we can update tensors in-place - if type(batch.past_key_values[0]) == tuple: - batch.past_key_values = [ - [t.view(len(batch), -1, *t.shape[-2:]) for t in layer] - for layer in batch.past_key_values - ] - elif len(batch.past_key_values[0][0].shape) == 3: - for layer in batch.past_key_values: - for k, t in enumerate(layer): - layer[k] = t.view(len(batch), -1, *t.shape[-2:]) - - # Add eventual padding tokens that were added while concatenating - max_tokens += batch.max_tokens + ( - max_input_length - batch.max_input_length - ) * len(batch) - - start_index = end_index - - first_past_kvs = batches[0].past_key_values - _, num_heads, padded_sequence_length, head_dim = first_past_kvs[0][1].shape - - padded_past_values_shape = ( - total_batch_size, - num_heads, - max_input_length - 1, - head_dim, - ) - - if batches[0].keys_head_dim_last: - padded_past_keys_shape = padded_past_values_shape - else: - # seq_length is last for BLOOM - padded_past_keys_shape = ( - total_batch_size, - num_heads, - head_dim, - max_input_length - 1, - ) - - # Iterate over attention layers - # Concatenate past key values layer by layer to allow incremental garbage collection - for j in range(len(first_past_kvs)): - padded_past_keys = first_past_kvs[j][0].new_zeros(padded_past_keys_shape) - start_index = 0 - for batch in batches: - past_keys = batch.past_key_values[j][0] - # Clear reference to the original tensor - batch.past_key_values[j][0] = None - - # Slicing end index for this batch - end_index = start_index + len(batch) - # We slice the keys to remove the padding from previous batches - past_seq_len = batch.max_input_length - 1 - if batch.keys_head_dim_last: - padded_past_keys[start_index:end_index, :, -past_seq_len:, :] = ( - past_keys[:, :, -past_seq_len:, :] - ) - else: - # BLOOM case - padded_past_keys[start_index:end_index, :, :, -past_seq_len:] = ( - past_keys[:, :, :, -past_seq_len:] - ) - del past_keys - - start_index = end_index - - padded_past_values = first_past_kvs[j][1].new_zeros( - padded_past_values_shape - ) - start_index = 0 - for batch in batches: - past_values = batch.past_key_values[j][1] - # Clear reference to the original tensor - batch.past_key_values[j][1] = None - - # Slicing end index for this batch - end_index = start_index + len(batch) - # We slice the past values to remove the padding from previous batches - past_seq_len = batch.max_input_length - 1 - padded_past_values[start_index:end_index, :, -past_seq_len:, :] = ( - past_values[:, :, -past_seq_len:, :] - ) - del past_values - - # Update values - start_index = end_index - - past_key_values.append([padded_past_keys, padded_past_values]) - - return cls( - batch_id=batches[0].batch_id, - requests=requests, - requests_idx_mapping=requests_idx_mapping, - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - all_input_ids=all_input_ids, - input_lengths=input_lengths, - prefix_offsets=prefix_offsets, - read_offsets=read_offsets, - next_token_choosers=next_token_choosers, - stopping_criterias=stopping_criterias, - top_n_tokens=top_n_tokens, - top_n_tokens_tensor=top_n_tokens_tensor, - max_input_length=max_input_length, - padding_right_offset=padding_right_offset, - keys_head_dim_last=batches[0].keys_head_dim_last, - max_tokens=max_tokens, - ) - - def __len__(self): - return len(self.requests) - - -class CausalLM(Model): - def __init__( - self, - model_id: str, - revision: Optional[str] = None, - quantize: Optional[str] = None, - speculator: Optional[str] = None, - dtype: Optional[torch.dtype] = None, - trust_remote_code: bool = False, - ): - if speculator: - raise RuntimeError("Speculator decoding is not enabled for AutoModel") - - if torch.cuda.is_available(): - device = torch.device("cuda") - dtype = torch.float16 if dtype is None else dtype - else: - if quantize: - raise ValueError("quantization is not available on CPU") - - device = torch.device("cpu") - dtype = torch.float32 if dtype is None else dtype - - tokenizer = AutoTokenizer.from_pretrained( - model_id, - revision=revision, - padding_side="left", - truncation_side="left", - trust_remote_code=trust_remote_code, - ) - model = AutoModelForCausalLM.from_pretrained( - model_id, - revision=revision, - torch_dtype=dtype, - device_map=( - "auto" - if torch.cuda.is_available() and torch.cuda.device_count() > 1 - else None - ), - load_in_8bit=quantize == "bitsandbytes", - trust_remote_code=trust_remote_code, - ) - if ( - torch.cuda.is_available() - and torch.cuda.device_count() == 1 - and quantize != "bitsandbytes" - ): - model = model.cuda() - - if tokenizer.pad_token_id is None: - if model.config.pad_token_id is not None: - tokenizer.pad_token_id = model.config.pad_token_id - elif model.config.eos_token_id is not None: - tokenizer.pad_token_id = model.config.eos_token_id - elif tokenizer.eos_token_id is not None: - tokenizer.pad_token_id = tokenizer.eos_token_id - else: - tokenizer.add_special_tokens({"pad_token": "[PAD]"}) - - super(CausalLM, self).__init__( - model_id=model_id, - model=model, - tokenizer=tokenizer, - requires_padding=True, - dtype=dtype, - device=device, - ) - - @property - def batch_type(self) -> Type[CausalLMBatch]: - return CausalLMBatch - - def decode(self, generated_ids: List[int]) -> str: - return self.tokenizer.decode( - generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False - ) - - def forward( - self, input_ids, attention_mask, position_ids, past_key_values: Optional = None - ) -> Tuple[ - torch.Tensor, Optional[torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]] - ]: - # Model Forward - kwargs = { - "input_ids": input_ids, - "attention_mask": attention_mask, - "past_key_values": past_key_values, - "use_cache": True, - "return_dict": True, - } - if self.has_position_ids: - kwargs["position_ids"] = position_ids - - outputs = self.model.forward(**kwargs) - if isinstance(outputs, tuple): - outputs, speculative_logits = outputs - else: - speculative_logits = None - return outputs.logits, speculative_logits, outputs.past_key_values - - @tracer.start_as_current_span("generate_token") - def generate_token( - self, batch: CausalLMBatch - ) -> Tuple[List[Generation], Optional[CausalLMBatch], Tuple[int, int]]: - start = time.time_ns() - # slice the attention mask to the correct shape - attention_mask = batch.attention_mask[:, : -batch.padding_right_offset] - - logits, speculative_logits, past = self.forward( - batch.input_ids, - attention_mask, - batch.position_ids, - batch.past_key_values, - ) - - # Results - generations: List[Generation] = [] - stopped = True - - # Speculation is not active for causal - accepted_ids = torch.ones_like(batch.input_ids)[:, 0] - batch_top_token_ids, batch_top_token_logprobs = batch_top_tokens( - batch.top_n_tokens, - batch.top_n_tokens_tensor, - torch.log_softmax(logits[:, -1], -1), - accepted_ids, - ) - - start_decode = time.time_ns() - - # Zipped iterator - iterator = zip( - batch.requests, - batch.input_lengths, - batch.prefix_offsets, - batch.read_offsets, - logits, - batch.next_token_choosers, - batch.stopping_criterias, - batch.all_input_ids, - batch.top_n_tokens, - batch_top_token_ids, - batch_top_token_logprobs, - ) - - # For each member of the batch - for i, ( - request, - input_length, - prefix_offset, - read_offset, - logits, - next_token_chooser, - stopping_criteria, - all_input_ids, - top_n_tokens, - top_token_ids, - top_token_logprobs, - ) in enumerate(iterator): - # Select next token - next_token_id, logprobs = next_token_chooser( - all_input_ids.view(1, -1), logits[-1:, :] - ) - - # Append next token to all tokens - all_input_ids = torch.cat([all_input_ids, next_token_id]) - new_input_length = input_length + 1 - - # Generated token - next_token_logprob = logprobs[-1, next_token_id] - next_token_id_squeezed = next_token_id.squeeze() - next_token_text, prefix_offset, read_offset = self.decode_token( - all_input_ids[:, 0], prefix_offset, read_offset - ) - - # Evaluate stopping criteria - stop, reason = stopping_criteria( - next_token_id_squeezed, - next_token_text, - ) - - if not stop: - stopped = False - - # Shard generations - # All generations will be appended in the rust sharded client - if i % self.world_size == self.rank: - if stop: - # Decode generated tokens - output_text, _, _ = self.decode_token( - all_input_ids[:, 0], - prefix_offset=len(all_input_ids) - - stopping_criteria.current_tokens - - 1, - read_offset=len(all_input_ids) - - stopping_criteria.current_tokens, - skip_special_tokens=True, - ) - # Get seed - if isinstance(next_token_chooser.choice, Sampling): - seed = next_token_chooser.choice.seed - else: - seed = None - - generated_text = GeneratedText( - output_text, stopping_criteria.current_tokens, reason, seed - ) - else: - generated_text = None - - # Prefill - if stopping_criteria.current_tokens == 1 and request.prefill_logprobs: - # Remove generated token to only have prefill and add nan for first prompt token - prefill_logprobs = [float("nan")] + torch.log_softmax( - logits, -1 - ).gather(1, all_input_ids[1:]).squeeze(1)[ - -new_input_length:-1 - ].tolist() - prefill_token_ids = all_input_ids[-new_input_length:-1] - prefill_texts = self.tokenizer.batch_decode( - prefill_token_ids, - clean_up_tokenization_spaces=False, - skip_special_tokens=False, - ) - prefill_tokens = Tokens( - prefill_token_ids, - prefill_logprobs, - prefill_texts, - is_special=[], - ) - else: - prefill_tokens = None - - if top_n_tokens > 0: - all_top_tokens = [] - for top_token_ids, top_token_logprobs in zip( - top_token_ids, top_token_logprobs - ): - toptoken_texts = self.tokenizer.batch_decode( - top_token_ids, - clean_up_tokenization_spaces=False, - skip_special_tokens=False, - ) - special_toptokens = [ - token_id in self.all_special_ids - for token_id in top_token_ids - ] - top_tokens = Tokens( - top_token_ids, - top_token_logprobs, - toptoken_texts, - special_toptokens, - ) - all_top_tokens.append(top_tokens) - top_tokens = all_top_tokens - else: - top_tokens = None - - generation = Generation( - request.id, - prefill_tokens, - Tokens( - [next_token_id_squeezed], - [next_token_logprob], - [next_token_text], - [next_token_id_squeezed.item() in self.all_special_ids], - ), - generated_text, - top_tokens, - ) - - generations.append(generation) - - # Update values - batch.next_token_choosers[i] = batch.next_token_choosers[i].advance_grammar( - next_token_id_squeezed.item() - ) - batch.input_ids[i, 0] = next_token_id - batch.all_input_ids[i] = all_input_ids - batch.input_lengths[i] = new_input_length - batch.prefix_offsets[i] = prefix_offset - batch.read_offsets[i] = read_offset - batch.max_input_length = max(batch.max_input_length, new_input_length) - - # We finished all generations in the batch; there is no next batch - if stopped: - forward_ns = start_decode - start - decode_ns = time.time_ns() - start_decode - return generations, None, (forward_ns, decode_ns) - - # Slice unused values from prefill - batch.input_ids = batch.input_ids[:, :1] - - # Update attention_mask as we added a new token to input_ids - batch.attention_mask[:, -batch.padding_right_offset] = 1 - # Decrease right offset - batch.padding_right_offset -= 1 - - # Update position_ids - batch.position_ids = batch.position_ids[:, -1:] + 1 - - # Update past key values - batch.past_key_values = past - - forward_ns = start_decode - start - decode_ns = time.time_ns() - start_decode - return generations, batch, (forward_ns, decode_ns) diff --git a/server/text_generation_server/models/causal_lm_ragged.py b/server/text_generation_server/models/causal_lm_ragged.py deleted file mode 100644 index 5ec16902..00000000 --- a/server/text_generation_server/models/causal_lm_ragged.py +++ /dev/null @@ -1,630 +0,0 @@ -import torch -import time - -from dataclasses import dataclass -from opentelemetry import trace -from transformers import AutoTokenizer, AutoModelForCausalLM, PreTrainedTokenizerBase -from typing import Optional, Tuple, List, Type, Dict -from text_generation_server.utils.import_utils import SYSTEM -from text_generation_server.models import Model -from text_generation_server.utils.chunks import concat_text_chunks -from text_generation_server.utils.tokens import batch_top_tokens -from text_generation_server.models.types import ( - Batch, - Tokens, - Generation, - GeneratedText, -) -from text_generation_server.pb import generate_pb2 -from text_generation_server.utils import NextTokenChooser, StoppingCriteria, Sampling -from text_generation_server.models.flash_causal_lm import FlashCausalLMBatch - -from text_generation_server.utils.import_utils import ( - empty_cache, - synchronize, - get_free_memory, -) -from text_generation_server.utils.speculate import get_speculate -from text_generation_server.utils.dist import MEMORY_FRACTION - -tracer = trace.get_tracer(__name__) - -from transformers.cache_utils import PagedCache - -from loguru import logger - -# Why define it here? -BLOCK_SIZE: int = 16 - - -class CausalLMRagged(Model): - def __init__( - self, - model_id: str, - revision: Optional[str] = None, - quantize: Optional[str] = None, - speculator: Optional[str] = None, - dtype: Optional[torch.dtype] = None, - trust_remote_code: bool = False, - ): - if speculator: - raise RuntimeError("Speculator decoding is not enabled for AutoModel") - - if torch.cuda.is_available(): - device = torch.device("cuda:0") # TODO felix: fix support for accelerate - dtype = torch.float16 if dtype is None else dtype - else: - if quantize: - raise ValueError("quantization is not available on CPU") - - device = torch.device("cpu") - dtype = torch.float32 if dtype is None else dtype - - tokenizer = AutoTokenizer.from_pretrained( - model_id, - revision=revision, - padding_side="left", - truncation_side="left", - trust_remote_code=trust_remote_code, - ) - model = AutoModelForCausalLM.from_pretrained( - model_id, - revision=revision, - torch_dtype=dtype, - device_map=None, - load_in_8bit=quantize == "bitsandbytes", - trust_remote_code=trust_remote_code, - attn_implementation="flash_attention_2", - ) - if ( - torch.cuda.is_available() - and torch.cuda.device_count() == 1 - and quantize != "bitsandbytes" - ): - model = model.cuda() - - self.kv_cache = [] - self.num_layers = len(model.model.layers) - self.num_kv_heads = model.config.num_key_value_heads - self.head_size = model.config.hidden_size // model.config.num_attention_heads - - if tokenizer.pad_token_id is None: - if model.config.pad_token_id is not None: - tokenizer.pad_token_id = model.config.pad_token_id - elif model.config.eos_token_id is not None: - tokenizer.pad_token_id = model.config.eos_token_id - elif tokenizer.eos_token_id is not None: - tokenizer.pad_token_id = tokenizer.eos_token_id - else: - tokenizer.add_special_tokens({"pad_token": "[PAD]"}) - - super().__init__( - model_id=model_id, - model=model, - tokenizer=tokenizer, - requires_padding=False, - dtype=dtype, - device=device, - ) - - def warmup(self, batch: FlashCausalLMBatch): - # The warmup batch is the biggest batch we could ever receive - empty_cache() - - try: - self.init_kv_cache( - batch.num_blocks, - self.num_layers, - self.num_kv_heads, - self.head_size, - self.dtype, - self.device, - ) - max_bt = batch.max_blocks - max_s = max_bt * BLOCK_SIZE - - _, batch, _ = self.generate_token(batch) - except torch.cuda.OutOfMemoryError as e: - raise RuntimeError( - f"Not enough memory to handle {len(batch.input_ids)} prefill tokens. " - f"You need to decrease `--max-batch-prefill-tokens`" - ) from e - - synchronize(self.device) - - # Inspired by the original implementation in [vllm](https://github.com/vllm-project/vllm) - # Calculate the number of blocks that can be allocated with the free memory - dtype_size = torch.tensor([], dtype=self.dtype).element_size() - cache_block_size = BLOCK_SIZE * self.num_kv_heads * self.head_size - total_cache_size = self.num_layers * cache_block_size * 2 * dtype_size - - free_memory = get_free_memory(self.device, MEMORY_FRACTION) - batch_num_blocks = batch.num_blocks if batch is not None else 0 - - num_blocks = ( - # Leave 5% for some wiggle room - int((free_memory * 0.95) // total_cache_size) - # Add batch.num_blocks as we allocated it above, so it is included in the peak memory. - + batch_num_blocks - ) - - del batch - - self.init_kv_cache( - num_blocks, - self.num_layers, - self.num_kv_heads, - self.head_size, - self.dtype, - self.device, - ) - - return int(num_blocks * BLOCK_SIZE) - - def init_kv_cache( - self, - num_blocks: int, - num_layers: int, - num_heads: int, - head_size: int, - dtype: torch.dtype, - device: torch.device, - ): - self.kv_cache = [] - empty_cache() - - element_size = torch.tensor([], dtype=dtype).element_size() - if SYSTEM == "ipex" and device.type == "xpu": - raise ValueError("Untested. Please open an issue") - else: - x = BLOCK_SIZE // element_size - - if SYSTEM == "ipex" and device == torch.device("cpu"): - raise ValueError("Untested. Please open an issue") - - self.kv_cache = [ - ( - torch.empty( - (num_blocks, num_heads, head_size // x, BLOCK_SIZE, x), - dtype=dtype, - device=device, - ), - torch.empty( - (num_blocks, num_heads, head_size, BLOCK_SIZE), - dtype=dtype, - device=device, - ), - ) - for _ in range(num_layers) - ] - - @property - def batch_type(self) -> Type[FlashCausalLMBatch]: - return FlashCausalLMBatch - - def decode(self, generated_ids: List[int]) -> str: - return self.tokenizer.decode( - generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False - ) - - def forward( - self, batch: FlashCausalLMBatch - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - # NOTE: adapter_data: not supported - - input_ids = batch.input_ids - position_ids = batch.position_ids - cu_seqlen_prefill = batch.cu_seqlen_prefill - kv_cache = self.kv_cache - block_tables = batch.block_tables_tensor - slots = batch.slots[batch.slot_indices] - input_lengths = batch.input_lengths_tensor - max_s = batch.max_seqlen - lm_head_indices = batch.prefill_head_indices - - # TODO felix: support window attention - # if cu_seqlen_prefill is None and self.max_past() is not None: - # # In decode, not prefill, we're actually overwriting the KV-cache - # # in a circular buffer mode. - # # This makes sure the max_s for the decode pass is correct. - # max_s = min(self.max_past(), max_s) - - bs = input_ids.shape[0] - - logits = self.model.forward( - input_ids=input_ids, - position_ids=position_ids, - past_key_values=PagedCache(), - cu_seqlen_prefill=cu_seqlen_prefill, - kv_cache=kv_cache, - block_tables=block_tables, - slots=slots, - input_lengths=input_lengths, - max_s=max_s, - prefill_cache_indices=batch.prefill_cache_indices, - lm_head_indices=lm_head_indices, - cache_position=False, - return_dict=False, - )[0] - - if lm_head_indices is not None: - logits = logits[lm_head_indices] - - if batch.prefill_cache_indices is not None: - batch.prefill_cache_indices = None - - speculative_logits = None - - return logits, speculative_logits - - @tracer.start_as_current_span("generate_token") - def generate_token( - self, batch: FlashCausalLMBatch - ) -> Tuple[List[Generation], Optional[FlashCausalLMBatch], Tuple[int, int]]: - start = time.time_ns() - prefill = batch.cu_seqlen_prefill is not None - prefill_logprobs = batch.prefill_next_token_indices is not None - - # Update adapter indices for speculative tokens (if present) - # adapter_meta = batch.adapter_meta - # if batch.speculative_ids is not None: - # B, speculative_length = batch.speculative_ids.shape - # new_length = speculative_length + 1 - # adapter_indices = ( - # adapter_meta.adapter_indices.unsqueeze(-1) - # .expand(B, new_length) - # .reshape(-1) - # ) - # adapter_segments = adapter_meta.adapter_segments * new_length - # adapter_meta = AdapterBatchMetadata( - # adapter_indices=adapter_indices, - # adapter_set=adapter_meta.adapter_set, - # adapter_segments=adapter_segments, - # segment_indices=adapter_meta.segment_indices, - # ) - - # Assign pointers to adapter weights - # TODO(travis): don't update this if indices haven't changed - # adapter_data = AdapterBatchData.from_meta( - # adapter_meta, - # self.layer_to_adapter_weights, - # prefill, - # batch.prefill_head_indices, - # ) - - logger.info(f"batch.input_ids {batch.input_ids}") - out, speculative_logits = self.forward(batch) - - logger.info(f"out {out.shape}") - logger.info(f"speculative_logits {speculative_logits}") - - if prefill: - next_token_logits = ( - out[batch.prefill_next_token_indices] if prefill_logprobs else out - ) - if speculative_logits is not None: - speculative_logits = ( - speculative_logits[batch.prefill_next_token_indices] - if prefill_logprobs - else speculative_logits - ) - # next_adapter_indices = batch.adapter_meta.adapter_indices.new_empty( - # len(batch) - # ) - - else: - next_token_logits = out - # next_adapter_indices = batch.adapter_meta.adapter_indices - - speculate = get_speculate() - ( - next_input_ids, - next_token_logprobs, - logprobs, - accepted_ids, - speculative_ids, - ) = batch.next_token_chooser( - batch.all_input_ids_tensor[:, : batch.max_seqlen], - next_token_logits, - speculate, - batch.speculative_ids, - speculative_logits, - ) - - batch_top_token_ids, batch_top_token_logprobs = batch_top_tokens( - batch.top_n_tokens, batch.top_n_tokens_tensor, logprobs, accepted_ids - ) - - if prefill: - if len(batch) > 1 and prefill_logprobs: - # We create the prefill_tokens_indices tensor that will be used to gather prefill logprobs - # When batch == 1, we will just use the batch.input_ids values directly - prefill_tokens_indices = batch.input_ids.new_zeros(len(out)) - - next_position_ids = batch.position_ids.new_empty(len(batch)) - batch.slot_indices = batch.slot_indices[batch.cu_seqlen_prefill[1:] - 1] - # We do not need cu_seqlen_prefill anymore - batch.cu_seqlen_prefill = None - else: - prefill_logprobs = None - next_position_ids = batch.position_ids - - # Cumulative length - cumulative_length = 0 - - # Results - generations: List[Generation] = [] - stopped = True - - # Zipped iterator - iterator = zip(batch.input_lengths, batch.all_input_ids, accepted_ids) - - # We do two for loops as the first one can run completely asynchronously from the GPU while for the second - # one, we need to first do a GPU <-> CPU sync - # It is faster if we delay this sync for the maximum amount of time - - # For each member of the batch - index = 0 - for i, (input_length, all_input_ids, n_accepted_ids) in enumerate(iterator): - # Indexing metadata - start_index = cumulative_length - end_index = cumulative_length + input_length - - if prefill: - # Indexing metadata - out_start_index = batch.prefill_cu_outlens[i] - out_end_index = batch.prefill_cu_outlens[i + 1] - out_length = out_end_index - out_start_index - - # Initialize position_ids - # In decode, we do not need this as we can just increment position ids - next_position_ids[i] = batch.position_ids[end_index - 1] - - # Initialize adapter indices - # In decode, we only have one token per row in the batch, so grab last index - # next_adapter_indices[i] = batch.adapter_meta.adapter_indices[ - # end_index - 1 - # ] - - # Used to gather prefill logprobs - # Copy batch.input_ids to prefill_token_indices - if prefill_logprobs: - if len(batch) > 1: - prefill_tokens_indices[out_start_index : out_end_index - 1] = ( - batch.input_ids[start_index + 1 : start_index + out_length] - ) - else: - # Set prefill_tokens_indices to the correct slice - prefill_tokens_indices = batch.input_ids[ - start_index + 1 : start_index + out_length - ] - - for j in range(n_accepted_ids): - batch.all_input_ids_tensor[i, input_length + j] = next_input_ids[index] - index += 1 - - cumulative_length += input_length - - logger.info(f"batch.input_lengths_tensor {batch.input_lengths_tensor}") - logger.info(f"accepted_ids {accepted_ids}") - logger.info(f"batch.all_input_ids {batch.all_input_ids}") - - # Update values - batch.input_ids = next_input_ids[accepted_ids.cumsum(dim=-1) - 1] - batch.speculative_ids = speculative_ids - batch.position_ids = next_position_ids + accepted_ids - batch.input_lengths_tensor += accepted_ids - batch.slot_indices += accepted_ids - # batch.adapter_meta.adapter_indices = None - - # if prefill: - # # adjust segment lengths to account for all request lengths being 1 during decoding - # adapter_segments, _ = find_segments(batch.adapter_meta.adapter_indices) - # batch.adapter_meta.adapter_segments = torch.tensor( - # adapter_segments, - # dtype=torch.int32, - # device=batch.adapter_meta.adapter_segments.device, - # ) - - if prefill and prefill_logprobs: - # Get prefill logprobs - prefill_logprobs_tensor = torch.log_softmax(out, -1) - prefill_logprobs = torch.gather( - prefill_logprobs_tensor, 1, prefill_tokens_indices.view(-1, 1) - ) - # GPU <-> CPU sync - prefill_logprobs = prefill_logprobs.view(-1).tolist() - - # GPU <-> CPU sync - next_token_logprobs = next_token_logprobs.tolist() - next_token_ids = next_input_ids.tolist() - accepted_ids = accepted_ids.tolist() - start_decode = time.time_ns() - - # Zipped iterator - iterator = zip( - batch.requests, - batch.input_lengths, - batch.prefix_offsets, - batch.read_offsets, - batch.stopping_criterias, - batch.all_input_ids, - batch.next_token_chooser.do_sample, - batch.next_token_chooser.seeds, - batch.top_n_tokens, - accepted_ids, - batch_top_token_ids, - batch_top_token_logprobs, - ) - - # For each member of the batch - index = 0 - for i, ( - request, - input_length, - prefix_offset, - read_offset, - stopping_criteria, - all_input_ids, - do_sample, - seed, - top_n_tokens, - n_accepted_ids, - top_token_ids, - top_token_logprobs, - ) in enumerate(iterator): - # Append next token to all tokens - next_token_texts = [] - left = 0 - - if n_accepted_ids > 1: - if RANK == 0: - logger.debug(f"Speculated ids {n_accepted_ids - 1}") - - current_stopped = False - for j in range(index, index + n_accepted_ids): - # Generated token - next_token_id = next_token_ids[j] - all_input_ids.append(next_token_id) - next_token_text, prefix_offset, read_offset = self.decode_token( - all_input_ids, - prefix_offset, - read_offset, - ) - next_token_texts.append(next_token_text) - - stop, reason = stopping_criteria( - next_token_id, - next_token_text, - ) - - if stop: - left = index + n_accepted_ids - j - 1 - current_stopped = True - break - else: - current_stopped = False - stopped = stopped and current_stopped - - _next_token_ids = next_token_ids[index : index + n_accepted_ids - left] - _next_token_logprobs = next_token_logprobs[ - index : index + n_accepted_ids - left - ] - index += n_accepted_ids - - # Shard generations - # All generations will be appended in the rust sharded client - if i % self.world_size == self.rank: - if stop: - # Decode generated tokens - output_text, _, _ = self.decode_token( - all_input_ids, - prefix_offset=len(all_input_ids) - - stopping_criteria.current_tokens - - 1, - read_offset=len(all_input_ids) - - stopping_criteria.current_tokens, - skip_special_tokens=True, - ) - generated_text = GeneratedText( - output_text, - stopping_criteria.current_tokens, - reason, - seed if do_sample else None, - ) - else: - generated_text = None - - # Prefill - if prefill and request.prefill_logprobs: - out_start_index = batch.prefill_cu_outlens[i] - out_end_index = batch.prefill_cu_outlens[i + 1] - - # Remove generated token to only have prefill and add nan for first prompt token - request_prefill_logprobs = [float("nan")] + prefill_logprobs[ - out_start_index : out_end_index - 1 - ] - prefill_token_ids = all_input_ids[:-1] - prefill_texts = self.tokenizer.batch_decode( - prefill_token_ids, - clean_up_tokenization_spaces=False, - skip_special_tokens=False, - ) - - prefill_tokens = Tokens( - prefill_token_ids, - request_prefill_logprobs, - prefill_texts, - is_special=[], - ) - else: - prefill_tokens = None - - if top_n_tokens > 0: - all_top_tokens = [] - for top_token_ids, top_token_logprobs in zip( - top_token_ids, top_token_logprobs - ): - toptoken_texts = self.tokenizer.batch_decode( - top_token_ids, - clean_up_tokenization_spaces=False, - skip_special_tokens=False, - ) - special_toptokens = [ - token_id in self.all_special_ids - for token_id in top_token_ids - ] - top_tokens = Tokens( - top_token_ids, - top_token_logprobs, - toptoken_texts, - special_toptokens, - ) - all_top_tokens.append(top_tokens) - top_tokens = all_top_tokens - else: - top_tokens = None - - generation = Generation( - request.id, - prefill_tokens, - Tokens( - _next_token_ids, - _next_token_logprobs, - next_token_texts, - [nid in self.all_special_ids for nid in _next_token_ids], - ), - generated_text, - top_tokens, - ) - - generations.append(generation) - - # accept each new token for this specific request since we may - # have more than one new token per request with speculative decoding - for next_token_id in _next_token_ids: - batch.next_token_chooser = ( - batch.next_token_chooser.advance_grammar_single(i, next_token_id) - ) - - # Update values - batch.input_lengths[i] = input_length + n_accepted_ids - if batch.input_lengths[i] > batch.max_seqlen: - batch.max_seqlen = batch.input_lengths[i] - batch.prefix_offsets[i] = prefix_offset - batch.read_offsets[i] = read_offset - batch.all_input_ids[i] = all_input_ids - - if stopped: - # No need to return a batch if we know that all requests stopped - forward_ns = start_decode - start - decode_ns = time.time_ns() - start_decode - return generations, None, (forward_ns, decode_ns) - - batch.prefill_cu_outlens = None - batch.prefill_head_indices = None - batch.prefill_next_token_indices = None - - forward_ns = start_decode - start - decode_ns = time.time_ns() - start_decode - return generations, batch, (forward_ns, decode_ns) diff --git a/server/text_generation_server/models/galactica.py b/server/text_generation_server/models/galactica.py index 30c92d90..0f9ffd3b 100644 --- a/server/text_generation_server/models/galactica.py +++ b/server/text_generation_server/models/galactica.py @@ -9,8 +9,8 @@ from transformers import ( AutoConfig, PreTrainedTokenizerBase, ) -from text_generation_server.models import CausalLM -from text_generation_server.models.causal_lm import CausalLMBatch +from text_generation_server.models import TransformersCausalLM +from text_generation_server.models.transformers_causal_lm import CausalLMBatch from text_generation_server.pb import generate_pb2 from text_generation_server.models.custom_modeling.opt_modeling import OPTForCausalLM from text_generation_server.utils import ( @@ -164,7 +164,7 @@ class GalacticaCausalLMBatch(CausalLMBatch): ) -class GalacticaSharded(CausalLM): +class GalacticaSharded(TransformersCausalLM): def __init__( self, model_id: str, @@ -211,7 +211,7 @@ class GalacticaSharded(CausalLM): model = OPTForCausalLM(config, weights) torch.distributed.barrier(group=self.process_group) - super(CausalLM, self).__init__( + super().__init__( model_id=model_id, model=model, tokenizer=tokenizer, diff --git a/server/text_generation_server/models/globals.py b/server/text_generation_server/models/globals.py index cc2f172a..157c88b9 100644 --- a/server/text_generation_server/models/globals.py +++ b/server/text_generation_server/models/globals.py @@ -44,3 +44,7 @@ ADAPTER_TO_INDEX: Dict[str, int] = None def set_adapter_to_index(adapter_to_index: Dict[str, int]): global ADAPTER_TO_INDEX ADAPTER_TO_INDEX = adapter_to_index + + +USE_CUSTOM_MODELING = os.getenv("USE_CUSTOM_MODELING", "true") +USE_CUSTOM_MODELING = USE_CUSTOM_MODELING == "true" or USE_CUSTOM_MODELING == "1" diff --git a/server/text_generation_server/models/gpt_neox.py b/server/text_generation_server/models/gpt_neox.py index c37cfb7d..a707c833 100644 --- a/server/text_generation_server/models/gpt_neox.py +++ b/server/text_generation_server/models/gpt_neox.py @@ -7,7 +7,7 @@ from transformers import ( AutoTokenizer, AutoConfig, ) -from text_generation_server.models import CausalLM +from text_generation_server.models import TransformersCausalLM from text_generation_server.models.custom_modeling.neox_modeling import ( GPTNeoxForCausalLM, ) @@ -18,7 +18,7 @@ from text_generation_server.utils import ( ) -class GPTNeoxSharded(CausalLM): +class GPTNeoxSharded(TransformersCausalLM): def __init__( self, model_id: str, @@ -64,7 +64,7 @@ class GPTNeoxSharded(CausalLM): model = GPTNeoxForCausalLM(config, weights) torch.distributed.barrier(group=self.process_group) - super(CausalLM, self).__init__( + super().__init__( model_id=model_id, model=model, tokenizer=tokenizer, diff --git a/server/text_generation_server/models/mpt.py b/server/text_generation_server/models/mpt.py index 1e79b25f..355c257f 100644 --- a/server/text_generation_server/models/mpt.py +++ b/server/text_generation_server/models/mpt.py @@ -8,8 +8,8 @@ from transformers import AutoTokenizer, PretrainedConfig, PreTrainedTokenizerBas from huggingface_hub import hf_hub_download import json -from text_generation_server.models import CausalLM -from text_generation_server.models.causal_lm import CausalLMBatch +from text_generation_server.models import TransformersCausalLM +from text_generation_server.models.transformers_causal_lm import CausalLMBatch from text_generation_server.pb import generate_pb2 from text_generation_server.models.custom_modeling.mpt_modeling import ( MPTForCausalLM, @@ -37,7 +37,7 @@ class MPTCausalLMBatch(CausalLMBatch): return batch -class MPTSharded(CausalLM): +class MPTSharded(TransformersCausalLM): def __init__( self, model_id: str, @@ -89,7 +89,7 @@ class MPTSharded(CausalLM): model = MPTForCausalLM(config, weights) torch.distributed.barrier(group=self.process_group) - super(CausalLM, self).__init__( + super().__init__( model_id=model_id, model=model, tokenizer=tokenizer, diff --git a/server/text_generation_server/models/opt.py b/server/text_generation_server/models/opt.py index 6d7d07f5..4f53faaf 100644 --- a/server/text_generation_server/models/opt.py +++ b/server/text_generation_server/models/opt.py @@ -8,7 +8,7 @@ from transformers import ( AutoConfig, ) from text_generation_server.models.custom_modeling.opt_modeling import OPTForCausalLM -from text_generation_server.models import CausalLM +from text_generation_server.models import TransformersCausalLM from text_generation_server.utils import ( initialize_torch_distributed, weight_files, @@ -16,7 +16,7 @@ from text_generation_server.utils import ( ) -class OPTSharded(CausalLM): +class OPTSharded(TransformersCausalLM): def __init__( self, model_id: str, @@ -62,7 +62,7 @@ class OPTSharded(CausalLM): model = OPTForCausalLM(config, weights) torch.distributed.barrier(group=self.process_group) - super(CausalLM, self).__init__( + super().__init__( model_id=model_id, model=model, tokenizer=tokenizer, diff --git a/server/text_generation_server/models/phi.py b/server/text_generation_server/models/phi.py index 93d42b2b..92aab9fb 100644 --- a/server/text_generation_server/models/phi.py +++ b/server/text_generation_server/models/phi.py @@ -4,7 +4,7 @@ import torch.distributed from transformers import AutoConfig, AutoTokenizer from typing import Optional, List, Tuple -from text_generation_server.models import CausalLM +from text_generation_server.models import TransformersCausalLM from text_generation_server.models.custom_modeling.phi_modeling import ( PhiConfig, PhiForCausalLM, @@ -16,7 +16,7 @@ from text_generation_server.utils import ( ) -class Phi(CausalLM): +class Phi(TransformersCausalLM): def __init__( self, model_id: str, @@ -59,7 +59,7 @@ class Phi(CausalLM): weights = Weights(filenames, device, dtype, process_group=self.process_group) model = PhiForCausalLM(config, weights) torch.distributed.barrier(group=self.process_group) - super(CausalLM, self).__init__( + super().__init__( model_id=model_id, model=model, tokenizer=tokenizer, diff --git a/server/text_generation_server/models/rw.py b/server/text_generation_server/models/rw.py index 37ca277b..78513760 100644 --- a/server/text_generation_server/models/rw.py +++ b/server/text_generation_server/models/rw.py @@ -3,10 +3,10 @@ import torch from transformers import AutoTokenizer, AutoModelForCausalLM from typing import List, Optional, Tuple -from text_generation_server.models import CausalLM +from text_generation_server.models import TransformersCausalLM -class RW(CausalLM): +class RW(TransformersCausalLM): def __init__( self, model_id: str, @@ -61,7 +61,7 @@ class RW(CausalLM): else: tokenizer.add_special_tokens({"pad_token": "[PAD]"}) - super(CausalLM, self).__init__( + super().__init__( model_id=model_id, model=model, tokenizer=tokenizer, diff --git a/server/text_generation_server/models/santacoder.py b/server/text_generation_server/models/santacoder.py index caddbe19..b595718d 100644 --- a/server/text_generation_server/models/santacoder.py +++ b/server/text_generation_server/models/santacoder.py @@ -4,7 +4,7 @@ import torch.distributed from typing import Optional, List from transformers import AutoTokenizer, AutoModelForCausalLM -from text_generation_server.models import CausalLM +from text_generation_server.models import TransformersCausalLM FIM_PREFIX = "" FIM_MIDDLE = "" @@ -13,7 +13,7 @@ FIM_PAD = "" EOD = "<|endoftext|>" -class SantaCoder(CausalLM): +class SantaCoder(TransformersCausalLM): def __init__( self, model_id: str, @@ -61,7 +61,7 @@ class SantaCoder(CausalLM): trust_remote_code=trust_remote_code, ) - super(CausalLM, self).__init__( + super().__init__( model_id=model_id, model=model, tokenizer=tokenizer,