import copy import logging import os import sys import time import traceback from bisect import bisect_left from enum import Enum from typing import Dict, List, Optional, Tuple import torch import torch.multiprocessing as mp import torch_xla.core.xla_model as xm import torch_xla.distributed.xla_multiprocessing as xmp from transformers import AutoTokenizer, PreTrainedTokenizerBase from transformers.generation import GenerationConfig import optimum.tpu.xla_logger as logger from optimum.tpu.generation import TokenSelector from optimum.tpu.modeling import AutoModelForCausalLM from optimum.tpu.static_cache_xla import StaticCacheXla from optimum.tpu.xla_mp_comm import AgentMailbox, RootMailbox from .generator_base import Generator from .pb.generate_pb2 import ( Batch, CachedBatch, FinishReason, GeneratedText, Generation, InfoResponse, NextTokenChooserParameters, Request, StoppingCriteriaParameters, Tokens, ) # Disable optimum-tpu warnings as it seems to block the server after a while optimum_logger = logging.getLogger("optimum.tpu") optimum_logger.setLevel("CRITICAL") # These will do some bucketing on prefill lengths to avoid too many different sizes PREFILL_LENGTHS = list(range(6, 16)) + [ 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, ] def take_nearest_length(length: int) -> int: """Gets the nearest length to the right in a set of lengths.""" pos = bisect_left(PREFILL_LENGTHS, length) if pos == len(PREFILL_LENGTHS): return PREFILL_LENGTHS[-1] return PREFILL_LENGTHS[pos] class Slot: """Represents a slot in a static batch""" class State(Enum): EMPTY = 0 PAUSE = 1 READY = 2 def __init__(self, id: int, tokenizer: PreTrainedTokenizerBase, device: [str, torch.device]): self._id = id self._tokenizer = tokenizer self.clear() self._device = device def clear(self): """Clear the slot and mark it as available.""" self._state = Slot.State.EMPTY self._batch_id = None self._request_id = None self._inputs = "" self._generation_config = None self._tokens = [] self._mask = None self._selector = None self._generated_tokens = 0 self._next_text_token_start = 0 self._next_text_token_end = 0 self._generated_text = "" self._next_text = "" self._kv_cache = None self._truncate = 0 self._position_id = 0 @property def id(self) -> int: return self._id @property def state(self) -> "Slot.State": return self._state @property def batch_id(self) -> int: return self._batch_id @property def request_id(self) -> int: return self._request_id @property def cached_text(self) -> str: return self._inputs + self._generated_text @property def generation_config(self) -> GenerationConfig: return self._generation_config @property def generated_tokens(self) -> int: return self._generated_tokens @property def truncate(self) -> int: return self._truncate @property def position_id(self) -> int: return self._position_id @position_id.setter def position_id(self, cur_pos: int): self._position_id = cur_pos @property def cache_position(self) -> int: # This corresponds to the cache position for this slot return self._next_text_token_start def assign(self, batch_id: int, request: Request, generation_config: GenerationConfig): """Assign a request to a slot. Args: batch_id (`int`): The id of the batch containing the request. request (`Request`): The request to be assigned. Contains the inputs and tokens selection parameters. generation_config (`transformers.GenerationConfig`): The base generation config (might be modified by the request generation parameters). """ self._state = Slot.State.READY self._batch_id = batch_id self._request_id = request.id self._inputs = request.inputs self._generation_config = copy.deepcopy(generation_config) # Update generation config with token chooser parameters self._generation_config.temperature = request.parameters.temperature self._generation_config.top_k = request.parameters.top_k self._generation_config.top_p = request.parameters.top_p self._generation_config.typical_p = request.parameters.typical_p self._generation_config.do_sample = request.parameters.do_sample self._generation_config.repetition_penalty = request.parameters.repetition_penalty self._truncate = request.truncate self.seed = request.parameters.seed # TODO: watermark self._generation_config.max_new_tokens = request.stopping_parameters.max_new_tokens self._max_new_tokens = self._generation_config.max_new_tokens # TODO: stop_sequences, ignore_eos_token def reset(self, input_ids: torch.LongTensor, attention_mask: torch.LongTensor = None, selector: TokenSelector = None): """Reset the slot for the next generation. Args: input_ids: (`torch.LongTensor`): The new input_ids to use to generate the next token. attention_mask: (`torch.LongTensor`): The new attention_mask to use to generate the next token. selector: (`TokenSelector`): An object implementing the updated token selection logic. """ self._tokens = input_ids.cpu() self._next_text_token_start = 0 self._next_text_token_end = torch.numel(self._tokens) self._next_text = "" if attention_mask is not None: self._mask = attention_mask.clone() else: self._mask = None self._selector = selector def pause(self): """Mark the current slot as paused for generation. Note that the KV cache for this slot will still be filled. """ # Drop the last token as it will be added back when resuming the slot self._generated_tokens -= 1 # Since generated tokens are now part of the prefill, we need to reevaluate # max_new_tokens for the next generation self._generation_config.max_new_tokens = self._max_new_tokens - self._generated_tokens self._state = Slot.State.PAUSE def resume(self): """Mark the slot as ready for generation.""" self._state = Slot.State.READY def _decode_next_tokens( self, ) -> str: """Hack to hopefully support generate_stream for the maximum number of tokenizers""" # Copy the tokens to CPU to avoid recompilation on TPU. Post-processing is quite fast anyway. tokens = self._tokens.cpu() # We need to include the tokens that produced the last text to defeat cleanup algorithms in the decode # which decide to add a space or not depending on the surrounding ids. new_text = self._tokenizer.decode(tokens[self._next_text_token_start :], skip_special_tokens=False) if new_text.endswith("�"): # utf-8 char at the end means it's a potential unfinished byte sequence # from byte fallback tokenization. return "" # Compare the generated text with the one using only the tokens producing the last one last_text = self._tokenizer.decode( tokens[self._next_text_token_start : self._next_text_token_end], skip_special_tokens=False, ) if len(new_text) == len(last_text): # Nothing new was actually generated return "" # Return the decoded text and store its token offsets self._next_text_token_start = self._next_text_token_end self._next_text_token_end = torch.numel(tokens) return new_text[len(last_text) :] def append(self, next_token: int) -> str: """Append a new generated token to this slot The new token is added to the list of generated tokens, which impacts directly the generated_text and stopped property. The new token is however not added immediately to the slot inputs: it will be added later on when it has effectively been used to produce the next token. Args: next_token (`int`): The newly generated token. Return: The corresponding decoded text (if any). """ self._tokens = torch.cat([self._tokens, torch.tensor([next_token], dtype=self._tokens.dtype)]) # Update mask only if it was set previously if self._mask is not None: self._mask = torch.cat([self._mask, torch.tensor([1], dtype=self._mask.dtype)]) self._generated_tokens += 1 next_text = self._decode_next_tokens() # Now that a new token has been generated, we can append the previous one to the generated text self._generated_text += self._next_text self._next_text = next_text return next_text def select(self, input_ids: torch.LongTensor, logits: torch.Tensor) -> torch.LongTensor: """Select the next token from the candidate logits. Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): The sequence used as a prompt for the generation (not used in all generation modes). logits (`torch.Tensor` of shape `(batch_size, sequence_length)`): The logits corresponding to the generated tokens. Return: `torch.LongTensor`: A scalar torch.LongTensor` containing the selected token. """ return self._selector.select(input_ids, logits)[0] @property def stopped(self) -> bool: # unsqueeze tokens to avoid problems with stopping criteria tokens = self._tokens.unsqueeze(0) return bool(torch.all(self._selector.stopping_criteria(tokens, None))) @property def generated_text(self) -> str: return self._generated_text + self._next_text @property def next_token(self) -> int: return None if len(self._tokens) == 0 else self._tokens[-1] @property def attention_mask(self) -> torch.LongTensor: return self._mask @property def max_token(self) -> int: return self._generation_config.max_length @property def max_new_tokens(self) -> int: # The current value of max_new_tokens: might be different of the target max_new_tokens # if the slot has been paused and resumed. return self._generation_config.max_new_tokens class TpuGeneratorSingleThread(Generator): """A Generator for models running on TPU, single threaded.""" def __init__( self, model, tokenizer: PreTrainedTokenizerBase, ): self.model = model # Specify padding options for decoder-only architecture tokenizer.pad_token_id = tokenizer.eos_token_id tokenizer.padding_side = "left" tokenizer.truncation_side = "left" self.tokenizer = tokenizer self.special_tokens = self.tokenizer.all_special_ids # The token selector will use the model's generation mixin internal variables to select the next token, and it # expects special tokens to be initialized in the model. model._prepare_special_tokens(generation_config=model.generation_config, device=model.device) # Slots are empty to begin with, they will be populated as new batches arrive self.slots = [] self.batch_id = 0 # Note: this index will _never_ be decremented, and that's fine. self.slot_index = 0 self.past_key_values = None # _supports_static_cache is specific to some models (e.g.: Gemma and Llama). self._supports_static_cache = True if getattr(self.model, "_supports_static_cache", False) is False: logger.warning( f"Static cache not available for {self.model.__class__.__name__}. Performance will be affected" ) self._supports_static_cache = False # compile model when possible to accelerate decoding if model.device.type == "xla" and ("DBG_COMPILE" in os.environ): self.model_one_token = torch.compile(model, backend="openxla") logger.debug("Model compiled for decoding") else: self.model_one_token = model @property def info(self) -> InfoResponse: """Returns the expected InfoResponse.""" dtype = getattr(self.model.config, "torch_dtype", "float32") return InfoResponse( requires_padding=True, dtype=str(dtype), device_type="xla", ) def _create_dummy_request(self, max_tokens: int) -> Batch: """Create a dummy request for warmup.""" # Generate a random input with slightly more tokens than requested, because special tokens are going to be # skipped. MARGIN = 10 input_tokens = torch.randint(self.model.config.vocab_size, (1, max_tokens + MARGIN), dtype=torch.int64) text = self.tokenizer.decode(input_tokens[0], skip_special_tokens=True) # These are just dummy params to allow Request creation parameters = NextTokenChooserParameters( temperature=1.0, top_k=None, top_p=None, do_sample=False, seed=None, repetition_penalty=1.0, typical_p=1.0, ) stopping_parameters = StoppingCriteriaParameters(max_new_tokens=20, ignore_eos_token=True) dummy_request = Request( id=0, inputs=text, truncate=max_tokens, parameters=parameters, stopping_parameters=stopping_parameters, ) return dummy_request def warmup(self, batch: Batch) -> int: """Verify if the hardware can support the target load. Args: batch (`Batch`): A batch corresponding to the maximum number of concurrent requests. Return: The maximum number of tokens the model supports. """ logger.debug("Warming up the model") start = time.time() # Just check that the warmup request parameters match the model capacity # NOTE: later self.model.config.batch_size might become self.model.config.max_batch_size. if self.model.config.batch_size is not None: batch_size = self.model.config.batch_size else: # batch size is not set, just assume it's unlimited and accept all requests batch_size = len(batch.requests) if len(batch.requests) > batch_size: raise ValueError( f"Inconsistent server configuration: please make sure max-prefill-tokens does not exceed {batch_size} x max-input-length." ) # Counter-intuitively, now we ignore the input batch. Instead, we create dummy batches to cover all possible # batch sizes and sequence lengths. seq_len = self.model.config.sequence_length if os.environ.get("SKIP_WARMUP", "0") == "1": logger.debug("Skipping warmup") return batch_size * seq_len bucket_seq_len = take_nearest_length(seq_len) requests = [self._create_dummy_request(seq_len) for _ in range(batch_size)] for _ in reversed(range(batch_size)): # Prefill with different truncate sizes to test all prefill lengths. List is reversed so first longest # sequences are tested and, if there is a memory failure, that will appear sooner. for l in reversed(PREFILL_LENGTHS): # Skip all the unsupported lengths if l > bucket_seq_len: continue # Set all truncate values for all requests for r in requests: r.truncate = l r.stopping_parameters.max_new_tokens = 10 warmup_batch = Batch(id=0, requests=requests, size=len(requests), max_tokens=batch.max_tokens) logger.debug(f"Warmup for {len(requests)} requests, truncate value {l} seq_len {seq_len}") _generations, next_batch = self.prefill(warmup_batch) if next_batch is not None: self.decode([next_batch]) else: logger.debug(f"No decode on warmup for {len(requests)}x{l}") self.clear() # remove the last requests to decrease the batch size requests.pop() elapsed = time.time() - start logger.debug(f"Warmup done, took {elapsed:.2f}s") return batch_size * seq_len @torch.no_grad def prefill(self, batch: Batch) -> Tuple[List[Generation], CachedBatch]: """Prefill new requests. Args: batch (`Batch`): A batch containing the new requests. Return: A list of `Generation` for each request and a `CachedBatch` containing all pending requests. """ slots = {state: [] for state in Slot.State} for slot in self.slots: slots[slot.state].append(slot) active_slots = slots[Slot.State.READY] # Delete all empty slots, no need to have them anymore empty_slots = slots[Slot.State.EMPTY] model_batch_size = self.model.config.batch_size if model_batch_size is not None and model_batch_size < len(active_slots) + len(batch.requests): # If raising an error here wouldn't crash the server, we could raise a ValueError error = ValueError( f"Cannot prefill {len(batch.requests)} new request(s)." f" Maximum batch size supported is: {model_batch_size}." ) # but since it's not possible, we just log the error and return an empty generation logger.error(error) return [], None for slot in empty_slots: self.slots.remove(slot) # Assign each request to an empty slot logger.debug(f"Prefilling {len(batch.requests)} new request(s) adding to {len(active_slots)} active slot(s)") for request in batch.requests: # Dynamically create a new slot for each request slot = Slot(self.slot_index, self.tokenizer, self.model.device) self.slot_index += 1 slot.assign(self.batch_id, request, self.model.generation_config) self.slots.append(slot) logger.debug(f"Request {slot.request_id} assigned to slot {slot.id}") logger.debug( f"Request {slot.request_id} assigned to slot {slot.id} with and max_new_tokens {slot.max_new_tokens}" ) # Reconstruct the full inputs (without padding) as seen by the model. # This comprises: # - the inputs for new requests, # - the inputs and the generated text that has already been cached (i.e. excluding the last generated token) # for unfinished requests. # Prepare inputs. They need to be tokenized and truncated afterwards. max_len = 0 batch_inputs = [] for slot in self.slots: batch_inputs.append(slot.cached_text) max_len = max(max_len, slot.truncate) if max_len == 0: max_len = self.model.config.sequence_length tokenized_inputs = self.tokenizer(batch_inputs, return_tensors="pt", padding=True, truncation=True, max_length=max_len) seq_length = tokenized_inputs.input_ids.size(-1) seq_length = min(seq_length, self.model.config.sequence_length) batch_size = len(self.slots) # Initialize input_ids and attention_mask with padding (to make them all the same size) input_ids = torch.full((batch_size, seq_length), self.tokenizer.pad_token_id, dtype=torch.int64) attention_mask = torch.full((batch_size, seq_length), 0, dtype=torch.int64) # Pause previously active slots during generation and store their last token. next_tokens = [] for slot in active_slots: next_tokens.append(slot.next_token) slot.pause() # Each slot must be reset with the padded inputs and masks for i, slot in enumerate(self.slots): assert slot.state != slot.state.EMPTY truncation = min(tokenized_inputs.input_ids.size(-1), input_ids.size(-1)) if slot.truncate > 0: truncation = min(truncation, slot.truncate) input_ids[i, -truncation:] = tokenized_inputs.input_ids[i, -truncation:] slot_input_ids = input_ids[i : i + 1, :] # Padded input ids are also required to set logits processors and stopping criterias try: selector = TokenSelector.create( slot_input_ids, slot.generation_config, self.model, self.model.config.sequence_length, seed=slot.seed, ) except ValueError as e: # This is very unlikely, but it seems it could be possible if router does not check values beforehand. # In that case, we just skip the slot, and mark it as empty. This should prevent returning this to the # router. logger.error(f"Invalid generation parameters for slot {slot.id}. Skipping it. Error: {e}") slot.clear() continue slot_input_ids = slot_input_ids.squeeze(dim=0).type(torch.int64) attention_mask[i, -truncation:] = tokenized_inputs.attention_mask[i, -truncation:] if self._supports_static_cache: # Attention mask does not need to be tracked when using static cache slot_attention_mask = None else: slot_attention_mask = attention_mask[i] slot.reset(slot_input_ids, slot_attention_mask, selector) # Clear KV cache self.past_key_values = None # Obtain position ids using attention mask. position_ids = (attention_mask.cumsum(-1) - 1).masked_fill(attention_mask == 0, 0) # Save position id for every slot for slot, position_id in zip(self.slots, position_ids): slot.position_id = position_id.max().item() + 1 extra_args = {} if self._supports_static_cache: self.past_key_values = StaticCacheXla( config=self.model.config, max_batch_size=len(self.slots), max_cache_len=self.model.config.sequence_length, device=self.model.device, dtype=self.model.dtype, ) extra_args["cache_position"] = torch.arange(seq_length, device=self.model.device) extra_args["past_key_values"] = self.past_key_values else: # Reset/clear KV cache self.past_key_values = None generation, next_batch = self._generate_token( self.batch_id, input_ids.to(self.model.device), self.model, attention_mask=attention_mask.to(self.model.device), position_ids=position_ids.to(self.model.device), **extra_args, ) self.batch_id += 1 # Reactivate previously active slots for the next decode, and append # back their next token. for slot, next_token in zip(active_slots, next_tokens): slot.append(next_token) slot.resume() logger.debug("Model ready for decoding") if next_batch is not None: logger.debug(f"Next batch is {next_batch.id} with requests: {next_batch.request_ids}") return generation, next_batch @torch.no_grad def decode(self, batches: List[CachedBatch]) -> Tuple[List[Generation], CachedBatch]: """Decode the specified prefilled requests. Args: batches (`List[CachedBatch]`): A list of previous batches containing the prefilled requests. Return: A list of `Generation` for each request and a `CachedBatch` containing all pending requests. """ # batches contains a list composed of: # - the batch id returned by the last decode, # - the batch id(s) returned by the last prefill(s) # Batches are always concatenated during prefill, so we can # just carry on with decoding. We adopt the id of the first # batch in the list as our next batch id. next_batch_id = batches[0].id request_ids = [] for batch in batches: request_ids += batch.request_ids cleared_request_ids = [] for slot in self.slots: if slot.state == slot.State.READY and slot.request_id not in request_ids: cleared_request_ids.append(slot.request_id) slot.clear() if len(cleared_request_ids) > 0: logger.info(f"Clearing slot for requests {cleared_request_ids} as they are not requested.") active_slots = [slot for slot in self.slots if slot.state == slot.State.READY] if len(active_slots) < len(request_ids): logger.error("Unable to decode tokens for non-prefilled batches (probably due to a previous failure)") # Reconstruct input_ids and attention_mask from slots input_ids = None attention_mask = None batch_size = len(self.slots) position_ids = torch.zeros( [batch_size, 1], dtype=torch.int64, ) # init pad_token_id and input_ids pad_token_id = self.tokenizer.pad_token_id if pad_token_id is None: if isinstance(self.tokenizer.eos_token_id, list): pad_token_id = self.tokenizer.eos_token_id[0] else: pad_token_id = self.tokenizer.eos_token_id # Create blank inputs covering all slots (even empty ones) input_ids = torch.full( [batch_size, 1], fill_value=pad_token_id, dtype=torch.int64, ) cache_position = torch.zeros([1], dtype=torch.int64) for i, slot in enumerate(self.slots): if slot.state != Slot.State.EMPTY: # input_ids are simply the tokens generated by the last decode or prefill requests (other tokens are cached) input_ids.index_put_([torch.tensor([i])], slot.next_token) if not self._supports_static_cache: # When using dynamic cache, the whole attention mask needs to be passed over to the model at each iteration. if attention_mask is None: # Create default mask covering all slots (even empty ones) attention_mask = torch.zeros( [batch_size, slot.attention_mask.size(-1)], dtype=torch.int64, ) attention_mask.index_put_([torch.tensor([i])], slot.attention_mask) position_ids.index_put_([torch.tensor([i])], torch.tensor(slot.position_id)) cache_position = torch.maximum(cache_position, torch.tensor([slot.cache_position])) if input_ids is None: raise ValueError("Unable to decode tokens for non-prefilled batches (probably due to a previous failure)") extra_args = {} if self._supports_static_cache: extra_args["cache_position"] = position_ids.max().unsqueeze(0).to(self.model.device) else: extra_args["attention_mask"] = attention_mask.to(self.model.device) extra_args["past_key_values"] = self.past_key_values generations, next_batch = self._generate_token( next_batch_id, input_ids.to(self.model.device), self.model_one_token, position_ids=position_ids.to(self.model.device), **extra_args, ) for slot, gen in zip(self.slots, generations): slot.position_id += len(gen.tokens.ids) return generations, next_batch def _generate_token( self, next_batch_id: int, input_ids: torch.LongTensor, model: torch.nn.Module, **forward_extra_params ) -> Tuple[List[Generation], CachedBatch]: # Add barrier to allow next graph step to always be the same xm.mark_step() # Forward outputs = model( input_ids, return_dict=True, use_cache=True, **forward_extra_params, ) if not self._supports_static_cache: # Save KV cache self.past_key_values = outputs.past_key_values # Barrier for XLA model xm.mark_step() ret = self._post_generate(outputs, next_batch_id, input_ids) return ret def _post_generate( self, outputs: Dict, next_batch_id: int, input_ids: torch.LongTensor ) -> Tuple[List[Generation], CachedBatch]: generations = [] active_slots = False for i, slot in enumerate(self.slots): if slot.state != Slot.State.READY: continue request_id = slot.request_id next_token_logits = outputs.logits[i : i + 1, -1, :] slot_input_ids = input_ids[i : i + 1, :] next_token = slot.select(slot_input_ids, next_token_logits) next_token_text = slot.append(next_token) generated_text = None finish_reason = None if next_token == self.tokenizer.eos_token_id: finish_reason = FinishReason.FINISH_REASON_EOS_TOKEN elif slot.stopped: if slot.generated_tokens == slot.max_new_tokens: finish_reason = FinishReason.FINISH_REASON_LENGTH else: finish_reason = FinishReason.FINISH_REASON_STOP_SEQUENCE if finish_reason is not None: # We must include the generated text for each finished sequence in the response generated_text = GeneratedText( text=slot.generated_text, generated_tokens=slot.generated_tokens, finish_reason=finish_reason ) logger.debug(f"Decode complete for request {request_id} with {slot.generated_tokens} tokens") # This slot is now empty, it will be removed from the list of # active slots once a new prefill is requested slot.clear() else: active_slots = True generations.append( Generation( request_id=request_id, prefill_tokens=None, tokens=Tokens( ids=[next_token], logprobs=[0], texts=[next_token_text], is_special=[next_token in self.special_tokens], ), generated_text=generated_text, ) ) batch = None if active_slots: # Whatever initial batch these requests came from, we always return all pending requests in a single batch request_ids = [slot.request_id for slot in self.slots if slot.state == Slot.State.READY] batch = self._cached_batch(next_batch_id, request_ids) else: logger.debug("No more pending requests") return generations, batch def _cached_batch(self, batch_id: int, request_ids: List): size = len(request_ids) max_tokens = size * self.model.config.sequence_length return CachedBatch(id=batch_id, request_ids=request_ids, size=size, max_tokens=max_tokens) def filter(self, batch_id: int, keep_request_ids: List[int]) -> CachedBatch: """Remove requests that are not listed from the specified batch Args: batch_id (`int`): The id of a cached batch. request_ids(`List[int]`): The list of requests that must be kept. Return: A `CachedBatch` containing the pending requests. """ keep_slot_ids = [slot.id for slot in self.slots if slot.request_id in keep_request_ids] self._clear(keep_slot_ids) return self._cached_batch(batch_id, keep_request_ids) def clear(self, batch_id: Optional[int] = None): """Remove a subset or all requests from the generator""" keep_ids = [] if batch_id is not None: keep_ids = [slot.id for slot in self.slots if slot.batch_id != batch_id] return self._clear(keep_ids) def _clear(self, keep_slot_ids: List): for slot in self.slots: if slot.state != Slot.State.EMPTY and slot.id not in keep_slot_ids: logger.debug(f"Removing slot {slot.id} with request {slot.request_id}") slot.clear() @classmethod def from_pretrained(cls, model_path: str, revision: str, max_batch_size: int, max_sequence_length: int): """Instantiate a TpuGenerator. Args: model_path (`str`): The path to a local model. This path must also contain a Tokenizer. revision (`str`): The revision of the model. max_batch_size (`int`): The maximum batch size. max_sequence_length (`int`): The maximum sequence length. Returns: A TpuGenerator. """ logger.info("Loading model (this can take a few minutes).") start = time.time() model = AutoModelForCausalLM.from_pretrained( model_path, revision=revision, batch_size=max_batch_size, sequence_length=max_sequence_length ) end = time.time() logger.info(f"Model successfully loaded in {end - start:.2f} s.") tokenizer = AutoTokenizer.from_pretrained(model_path) return cls(model, tokenizer) class GeneratorCommand(Enum): INFO = 0 WARMUP = 1 PREFILL = 2 DECODE = 3 FILTER = 4 CLEAR = 5 DELETE = -1 def _mp_fn( rank, model_path: str, revision: str, max_batch_size: int, max_sequence_length: int, root_mailbox: RootMailbox ): device = xm.xla_device() world_size = xm.xrt_world_size() # create agent mailbox out of root's one mailbox = AgentMailbox(root_mailbox) # re-init logger for each child process logger_level = os.environ.get("LOGGER_LEVEL_GENERATOR", "DEBUG") logger.logger.remove() logger.logger.add( sys.stdout, filter="text_generation_server", level=logger_level, backtrace=True, diagnose=False, ) logger.debug( f"Rank {rank} on {device} real device {xm.xla_real_devices([device])} ordinal {xm.get_ordinal()} " + f"world size {world_size}" ) generator = TpuGeneratorSingleThread.from_pretrained(model_path, revision, max_batch_size, max_sequence_length) # TODO: maybe model_config can be removed from mailbox def return_to_caller(*data): # consider adding a rendezvous here if rank == 0: xm.mark_step() mailbox.send(*data) while True: xm.rendezvous("start") if rank == 0: mailbox.agent_ready.set() mailbox.receive() # Wait for rank 0 to receive command xm.rendezvous("wait_command") command, data = mailbox.command_data logger.debug(f"Generator@{rank} {command.name}") try: if command == GeneratorCommand.INFO: info = generator.info return_to_caller(info.SerializeToString()) if command == GeneratorCommand.WARMUP: batch = Batch.FromString(data[0]) return_to_caller(generator.warmup(batch=batch)) if command == GeneratorCommand.PREFILL: batch = Batch.FromString(data[0]) generations, cached_batch = generator.prefill(batch=batch) s_cached_batch = cached_batch.SerializeToString() if cached_batch is not None else None return_to_caller([g.SerializeToString() for g in generations], s_cached_batch) if command == GeneratorCommand.DECODE: batches = [CachedBatch.FromString(b) for b in data[0]] generations, cached_batch = generator.decode(batches=batches) s_cached_batch = cached_batch.SerializeToString() if cached_batch is not None else None return_to_caller([g.SerializeToString() for g in generations], s_cached_batch) if command == GeneratorCommand.FILTER: batch_id, request_ids = data cached_batch = generator.filter(batch_id, request_ids) return_to_caller(cached_batch.SerializeToString()) if command == GeneratorCommand.CLEAR: batch_id = data[0] generator.clear(batch_id) if command == GeneratorCommand.DELETE: if rank == 0: # Set agent to ready mailbox.agent_ready.set() break except Exception as e: logger.error(f"Error in command {command.name}") mailbox.agent_error.set() mailbox.agent_ready.set() exc_info = sys.exc_info() logger.error(''.join(traceback.format_exception(*exc_info))) raise e # If error was only happening on one of the threads, all of them should exit if mailbox.agent_error.is_set(): return def model_loop_fn(*args): """Spawn processes in the TPUs forwarding arguments""" xmp.spawn(_mp_fn, args=(args), join=True, daemon=False) class TpuGenerator(Generator): """A Generator for models running on TPU. This generator actually spawns several processes to handle the requests in sharded models whenever possible. """ def __init__(self, model_path: str, revision: str, max_batch_size: int, max_sequence_length: int): manager = mp.Manager() self.mailbox = RootMailbox(manager) # Disable parallelism on tokenizers to avoid deadlocks on TPU threads os.environ["TOKENIZERS_PARALLELISM"] = "false" self.model_loop = mp.Process( target=model_loop_fn, args=(model_path, revision, max_batch_size, max_sequence_length, self.mailbox) ) self.model_loop.start() @property def info(self) -> InfoResponse: s_info = self.mailbox.send(GeneratorCommand.INFO, None)[0] return InfoResponse.FromString(s_info) def warmup(self, batch: Batch) -> int: return self.mailbox.send(GeneratorCommand.WARMUP, batch.SerializeToString())[0] def prefill(self, batch: Batch) -> Tuple[List[Generation], CachedBatch]: s_generations, s_cached_batch = self.mailbox.send(GeneratorCommand.PREFILL, batch.SerializeToString()) generations = [Generation.FromString(g) for g in s_generations] cached_batch = CachedBatch.FromString(s_cached_batch) if s_cached_batch is not None else None return generations, cached_batch def decode(self, batches: List[CachedBatch]) -> Tuple[List[Generation], CachedBatch]: s_batches = [b.SerializeToString() for b in batches] s_generations, s_cached_batch = self.mailbox.send(GeneratorCommand.DECODE, s_batches) generations = [Generation.FromString(g) for g in s_generations] cached_batch = CachedBatch.FromString(s_cached_batch) if s_cached_batch is not None else None return generations, cached_batch def filter(self, batch_id: int, request_ids: List[int]) -> CachedBatch: s_cached_batch = self.mailbox.send(GeneratorCommand.FILTER, batch_id, request_ids)[0] return CachedBatch.FromString(s_cached_batch) def clear(self, batch_id: Optional[int] = None): self.mailbox.send(GeneratorCommand.CLEAR, batch_id) def leave(self): if self.mailbox is None: return self.mailbox.send(GeneratorCommand.DELETE) # Use Loguru's logger directly, to avoid errors whyle TPU is shutting down logger.logger.debug("Joining...") self.model_loop.join() logger.logger.debug("Generator loop finished") self.mailbox = None @property def config(self): return self.mailbox.config def __del__(self): self.leave() @classmethod def from_pretrained(cls, model_path: str, revision: str, max_batch_size: int, max_sequence_length: int): """Instantiate a Generator distributed on as many cores as possible. Args: model_path (`str`): The path to a local model. This path must also contain a Tokenizer. revision (`str`): The revision of the model. max_batch_size (`int`): The maximum batch size. max_sequence_length (`int`): The maximum sequence length. Returns: A TpuGenerator. """ return cls(model_path, revision, max_batch_size, max_sequence_length)