# # Copyright 2016 The BigDL Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # Some parts of this file is adapted from # https://github.com/huggingface/transformers/blob/main/src/transformers/generation/utils.py # import torch from torch import nn from torch.nn import CrossEntropyLoss import torch.nn.functional as F import torch.distributed as dist import os import time import numpy as np from typing import Callable, List, Optional, Union, Tuple, Any from types import SimpleNamespace import transformers from transformers import GenerationConfig, LogitsProcessorList, StoppingCriteriaList from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ipex_llm.utils.common import invalidInputError from ipex_llm.ggml.quantize import ggml_tensor_qtype import logging logger = logging.getLogger(__name__) import asyncio import uuid import threading import pickle try: from pydantic import BaseModel except ImportError: from abc import ABCMeta BaseModel = ABCMeta # patch GenerationMixin.generate from transformers import GenerationMixin original_generate = GenerationMixin.generate class DummyLayer(nn.Module): def __init__(self, *args): super().__init__() # to avoid AttributeError in https://github.com/intel-analytics/ipex-llm/blob/main/ # python/llm/src/ipex_llm/transformers/models/llama.py#L2076 self.weight = nn.Parameter(torch.empty(0,), requires_grad=False) def forward(self, x): return x class Dummy_MLPLayer(nn.Module): def __init__(self, *args): super().__init__() # to avoid AttributeError in https://github.com/intel-analytics/ipex-llm/blob/main/ # python/llm/src/ipex_llm/transformers/models/llama.py#L119 self.up_proj = DummyLayer() self.down_proj = DummyLayer() self.shared_expert = SimpleNamespace() self.shared_expert.up_proj = DummyLayer() def forward(self, x): return x class Dummy_DecoderLayer(nn.Module): def __init__(self, *args): super().__init__() # to avoid AttributeError self.input_layernorm = DummyLayer() self.mlp = Dummy_MLPLayer() def forward(self, hidden_states, *args, **kwargs): past_key_value = kwargs.get('past_key_value', None) use_cache = kwargs.get('use_cache', False) outputs = (hidden_states,) if use_cache: outputs += (past_key_value,) return outputs class Dummy_GLMBlock(nn.Module): def __init__(self, *args): super().__init__() # to avoid AttributeError self.input_layernorm = DummyLayer() self.mlp = Dummy_MLPLayer() def forward( self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True, ): if kv_cache is None: return hidden_states, () return hidden_states, kv_cache def init_pipeline_parallel(): import oneccl_bindings_for_pytorch os.environ["MASTER_ADDR"] = os.environ.get("MASTER_ADDR", "127.0.0.1") os.environ["MASTER_PORT"] = os.environ.get("MASTER_PORT", "29500") dist.init_process_group('ccl') def low_mem_convert(model): from ipex_llm.transformers.convert import convert_forward import importlib if 'llama' in model.config.model_type: convert_forward( model, transformers.models.llama.modeling_llama.LlamaForCausalLM, llama_causallm_forward_4_37_lowmem) elif model.config.model_type == "chatglm" and not hasattr(model.config, "vision_config"): if model.config.num_layers == 40: # for glm4-9b modeling_module_name = model.__class__.__module__ module = importlib.import_module(modeling_module_name) convert_forward( model, module.ChatGLMForConditionalGeneration, glm4_conditional_generation_forward_lowmem) else: # for chatglm3-6b modeling_module_name = model.__class__.__module__ module = importlib.import_module(modeling_module_name) convert_forward( model, module.ChatGLMForConditionalGeneration, chatglm3_conditional_generation_forward_lowmem) return model def _check_quantize_kv_cache(model, idx, batch_size): # align use_quantize_kv_cache setting for different GPU in pipeline parallel pp_quantize_kv_cache = (os.environ.get("BIGDL_QUANTIZE_KV_CACHE", None) == "1") or \ (os.environ.get("IPEX_LLM_QUANTIZE_KV_CACHE", None) == "1") or \ (os.environ.get("IPEX_LLM_LOW_MEM", None) == "1") if model.config.model_type == "qwen" and hasattr(model.config, "visual"): # for Qwen-VL-Chat linear = model._modules['transformer'].h[idx].mlp.c_proj elif model.config.model_type == "chatglm": # for chatglm3-6b, glm-4-9b-chat linear = model._modules['transformer'].encoder.layers[idx].self_attention.query_key_value else: linear = model._modules['model'].layers[idx].mlp.up_proj pp_quantize_kv_cache = pp_quantize_kv_cache or (1 < batch_size and batch_size <= 8 and hasattr(linear, "qtype") and linear.qtype != ggml_tensor_qtype["fp16"] and linear.qtype != ggml_tensor_qtype["bf16"]) if pp_quantize_kv_cache: os.environ["IPEX_LLM_QUANTIZE_KV_CACHE"] = "1" else: os.environ["IPEX_LLM_QUANTIZE_KV_CACHE"] = "0" def pipeline_parallel(model, pipeline_parallel_stages, torch_dtype=torch.float32): global num_layers if hasattr(model.config, 'num_hidden_layers'): num_layers = model.config.num_hidden_layers elif hasattr(model.config, 'num_layers'): # for chatglm3-6b num_layers = model.config.num_layers slice_size = (num_layers + pipeline_parallel_stages - 1) // pipeline_parallel_stages local_rank = dist.get_rank() global layer_start global layer_end layer_start = slice_size * local_rank layer_end = layer_start + min(slice_size, num_layers - layer_start) if model.config.model_type == "qwen" and hasattr(model.config, "visual"): # for Qwen-VL-Chat for i in range(num_layers): if i < layer_start or i >= layer_end: model._modules['transformer'].h[i] = Dummy_DecoderLayer() if local_rank != 0: model._modules['transformer'].wte = DummyLayer() model._modules['transformer'].drop = DummyLayer() if local_rank != pipeline_parallel_stages - 1: model._modules['transformer'].ln_f = DummyLayer() model._modules['ln_f'] = DummyLayer() model._modules['lm_head'] = DummyLayer() elif model.config.model_type == "chatglm": # for chatglm3-6b, glm-4-9b-chat for i in range(num_layers): if i < layer_start or i >= layer_end: model._modules['transformer'].encoder.layers[i] = Dummy_GLMBlock() else: model._modules['transformer'].encoder.layers[i].self_attention.num_layers = \ i - layer_start if local_rank != 0: model._modules['transformer'].embedding = DummyLayer() if local_rank != pipeline_parallel_stages - 1: model._modules['transformer'].encoder.final_layernorm = DummyLayer() model._modules['transformer'].output_layer = DummyLayer() else: for i in range(num_layers): if i < layer_start or i >= layer_end: model._modules['model'].layers[i] = Dummy_DecoderLayer() else: model._modules['model'].layers[i].self_attn.layer_idx = i - layer_start if local_rank != 0: model._modules['model'].embed_tokens = DummyLayer() if local_rank != pipeline_parallel_stages - 1: model._modules['model'].norm = DummyLayer() model._modules['lm_head'] = DummyLayer() _enable_lowmem = os.getenv('IPEX_LLM_LOW_MEM') _enable_lowmem = (_enable_lowmem is not None) and (_enable_lowmem.lower() == "1") if _enable_lowmem: model = low_mem_convert(model) model.pipeline_parallel_stages = pipeline_parallel_stages model.layer_start = layer_start model.layer_end = layer_end model.num_layers = num_layers if torch_dtype == torch.float16: model = model.half() model = model.to(f'xpu:{local_rank}') return model @torch.no_grad() def generate( self, inputs: Optional[torch.Tensor] = None, generation_config: Optional[GenerationConfig] = None, logits_processor: Optional[LogitsProcessorList] = None, stopping_criteria: Optional[StoppingCriteriaList] = None, prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]]=None, synced_gpus: Optional[bool] = None, assistant_model: Optional["PreTrainedModel"] = None, streamer: Optional["BaseStreamer"] = None, negative_prompt_ids: Optional[torch.Tensor] = None, negative_prompt_attention_mask: Optional[torch.Tensor] = None, **kwargs, ): if hasattr(self, 'pipeline_parallel_stages') and self.pipeline_parallel_stages > 1: # priority: `generation_config` argument > `model.generation_config` if generation_config is None: if ( self.generation_config._from_model_config and self.generation_config._original_object_hash == hash(self.generation_config) and self.config._has_non_default_generation_parameters() ): new_generation_config = GenerationConfig.from_model_config(self.config) if new_generation_config != self.generation_config: self.generation_config = new_generation_config generation_config = self.generation_config if generation_config.pad_token_id is None and generation_config.eos_token_id is not None: eos_token_id = generation_config.eos_token_id if isinstance(eos_token_id, list): eos_token_id = eos_token_id[0] logger.warning("Setting `pad_token_id` to `eos_token_id`: " f"{eos_token_id} for open-end generation.") generation_config.pad_token_id = eos_token_id if generation_config is not None and generation_config.max_new_tokens is not None: max_new_tokens = generation_config.pop("max_new_tokens") else: max_new_tokens = kwargs.pop("max_new_tokens", None) return self.pipeline_parallel_generate(inputs=inputs, max_new_tokens=max_new_tokens, generation_config=generation_config, **kwargs) return original_generate(self, inputs=inputs, generation_config=generation_config, logits_processor=logits_processor, stopping_criteria=stopping_criteria, prefix_allowed_tokens_fn=prefix_allowed_tokens_fn, synced_gpus=synced_gpus, assistant_model=assistant_model, streamer=streamer, negative_prompt_ids=negative_prompt_ids, negative_prompt_attention_mask=negative_prompt_attention_mask, **kwargs) GenerationMixin.generate = generate @torch.no_grad() def pipeline_parallel_generate(self, inputs: Optional[torch.Tensor] = None, max_new_tokens: int = 32, generation_config: Optional[GenerationConfig] = None, **kwargs): model_kwargs = generation_config.update(**kwargs) inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs( inputs, generation_config.bos_token_id, model_kwargs ) bs = inputs_tensor.shape[0] if self.config.is_encoder_decoder: input_ids, model_kwargs = self._prepare_decoder_input_ids_for_generation( batch_size=bs, model_input_name=model_input_name, model_kwargs=model_kwargs, decoder_start_token_id=generation_config.decoder_start_token_id, bos_token_id=generation_config.bos_token_id, device=inputs_tensor.device, ) else: input_ids = inputs_tensor if model_input_name == "input_ids" \ else model_kwargs.pop("input_ids") local_rank = dist.get_rank() pre_rank = (local_rank - 1) % self.pipeline_parallel_stages next_rank = (local_rank + 1) % self.pipeline_parallel_stages global layer_start global layer_end global num_layers self.first_token_time = 0 self.next_token_time = [] pad_token_id = generation_config.pad_token_id eos_token_id = generation_config.eos_token_id if isinstance(eos_token_id, int): eos_token_id = [eos_token_id] eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) \ if eos_token_id is not None else None _input_ids = None _past_key_values = None bs = input_ids.shape[0] output_ids = input_ids.clone() _check_quantize_kv_cache(self, layer_start, bs) step = 0 # keep track of which sequences are already finished unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device) this_peer_finished = False while True: if step >= max_new_tokens: break if _input_ids is None: _input_ids = input_ids tic = time.time() if local_rank == 0: outputs = self(input_ids=_input_ids, inputs_embeds=None, past_key_values=_past_key_values, use_cache=True, **model_kwargs) else: _inputs_shape = _input_ids.shape + (self.config.hidden_size,) if step == 0 and self.config.model_type == "chatglm" \ and hasattr(self.config, "vision_config"): # for glm-4v, image features are mapped during 1st token # 1597 are computed according to computation process of conv _images_feature = 1597 + _input_ids.shape[0] * 2 + _input_ids.shape[1] _inputs_shape = (_input_ids.shape[0], _images_feature, self.config.hidden_size,) inputs_embeds = torch.empty(_inputs_shape, device=f'xpu:{local_rank}', dtype=self.dtype) dist.recv(inputs_embeds, src=pre_rank) outputs = self(input_ids=None, inputs_embeds=inputs_embeds, past_key_values=_past_key_values, use_cache=True, **model_kwargs) if local_rank == self.pipeline_parallel_stages - 1: logits = outputs.logits next_ids = torch.argmax(logits[:, -1:, :], dim=-1) dist.broadcast(next_ids, src=local_rank) else: dist.send(outputs[0].to(self.dtype), dst=next_rank) next_ids = torch.empty((bs, 1), device=f'xpu:{local_rank}', dtype=torch.int64) dist.broadcast(next_ids, src=self.pipeline_parallel_stages - 1) _input_ids = next_ids output_ids = torch.cat([output_ids, next_ids], dim=-1) # finished sentences should have their next token be a padding token next_ids = next_ids.squeeze() if eos_token_id is not None: if pad_token_id is None: invalidInputError(False, "If `eos_token_id` is defined, " "make sure that `pad_token_id` is defined.") next_ids = next_ids * unfinished_sequences + pad_token_id * (1 - unfinished_sequences) if self.config.model_type == "chatglm" and self.config.num_layers == 40 \ and not hasattr(self.config, "vision_config"): # for glm-4-9b-chat if step == 0: value_placeholder = torch.empty_like((outputs.past_key_values)[-1][0]) past_key_values_placeholder = tuple( (value_placeholder, value_placeholder) for _ in range(layer_start) ) + (outputs.past_key_values)[: layer_end - layer_start] + tuple( (value_placeholder, value_placeholder) for _ in range(layer_end, num_layers) ) _past_key_values = past_key_values_placeholder else: _past_key_values = outputs.past_key_values elif self.config.model_type in ["baichuan", "chatglm"] or \ (self.config.model_type == "qwen" and hasattr(self.config, "visual")): # for baichuan2, chatglm3, Qwen-VL-Chat, glm-4v-9b if local_rank != 0: value_placeholder = torch.empty_like((outputs.past_key_values)[-1][0]) past_key_values_placeholder = tuple( (value_placeholder, value_placeholder) for _ in range(layer_start) ) + (outputs.past_key_values)[layer_start:] _past_key_values = past_key_values_placeholder else: _past_key_values = outputs.past_key_values else: _past_key_values = outputs.past_key_values toc = time.time() if step == 0: self.first_token_time = toc - tic else: self.next_token_time.append(toc - tic) # if eos_token was found in one sentence, set sentence to finished if eos_token_id_tensor is not None: unfinished_sequences = unfinished_sequences.mul( next_ids.tile(eos_token_id_tensor.shape[0], 1) .ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0) ) # stop when each sentence is finished if unfinished_sequences.max() == 0: this_peer_finished = True if this_peer_finished: break step += 1 if self.device.type == 'xpu': torch.xpu.synchronize() self.rest_cost_mean = np.mean(self.next_token_time) return output_ids class PPConfig: """Configuration for ModelSlices during serving.""" def __init__(self, pp_rank: int, pp_world_size: int) -> None: self.pp_rank = pp_rank self.pp_world_size = pp_world_size self.is_head = self.pp_rank == 0 self.is_tail = self.pp_rank == self.pp_world_size - 1 class BatchTask(BaseModel): batch_id: str request_ids: List[str] max_tokens: int batch_size: int input_len: int prompt_lengths: List[int] stopped: bool prefilled_index: int partial_prefilling: int def make_attention_mask(prompt_lengths, device): max_length = max(prompt_lengths) batch_size = len(prompt_lengths) range_tensor = torch.arange(max_length, device=device).expand(batch_size, max_length) prompt_lengths_tensor = torch.tensor(prompt_lengths, device=device).unsqueeze(1) attention_mask = range_tensor >= max_length - prompt_lengths_tensor attention_mask = attention_mask.to(torch.int64) return attention_mask class PPModelWorker: """Implementation for pipeline parallel multi-stage serving.""" def __init__(self, checkpoint, rank, world_size, low_bit, max_num_seqs, max_prefilled_seqs, torch_dtype=torch.float16): self.pp_config = PPConfig(rank, world_size) self.dtype = torch_dtype start = time.perf_counter() model = self.load_model(checkpoint, world_size, low_bit) end = time.perf_counter() logger.info(f"Time to load weights: {end - start:.2f}s") self.model = model self.rank = rank self.world_size = world_size self.pre_rank = (self.rank - 1) % self.world_size self.next_rank = (self.rank + 1) % self.world_size self.hidden_size = self.model.config.hidden_size self.max_num_seqs = max_num_seqs self.on_going_batches = [None] * self.world_size self.input_ids_dict = {} self.past_key_values_dict = {} self.tokens = {} self.token_times = {} self.waiting_requests = asyncio.Queue() self.send_buff = None self.dict_lock = threading.Lock() self.streamer = {} self.token_cache = {} self.print_len = {} self.is_finish = {} self.model_name = checkpoint self.device = f"xpu:{self.rank}" # self.layer_start = 0 # self.layer_end = 0 self.max_prefilled_seqs = max_prefilled_seqs self.partial_output_dict = {} self.stream_tasks = {} def load_model(self, model_path, world_size, low_bit='sym_int4'): from ipex_llm.transformers import AutoModelForCausalLM, AutoModel try: model = AutoModelForCausalLM.from_pretrained(model_path, load_in_low_bit=low_bit, torch_dtype=self.dtype, cpu_embedding=True, optimize_model=True, trust_remote_code=True, use_cache=True, pipeline_parallel_stages=world_size) except: model = AutoModel.from_pretrained(model_path, load_in_low_bit=low_bit, torch_dtype=self.dtype, optimize_model=True, trust_remote_code=True, use_cache=True, pipeline_parallel_stages=world_size) model = model.eval() return model def prepare_batch(self, cur_batch): if self.rank == 0: cur_input_start = cur_batch.prefilled_index if self.max_prefilled_seqs > 0: if cur_input_start < cur_batch.batch_size: cur_input_end = cur_input_start + self.max_prefilled_seqs cur_input_end = min(cur_input_end, cur_batch.batch_size) cur_batch.partial_prefilling = cur_input_end - cur_input_start else: cur_batch.partial_prefilling = 0 return cur_batch def cat_kv_cache(self, model_type, kv_cache_1, kv_cache_2): if model_type in ["baichuan", "chatglm", "mixtral"]: result = [] for sub_tuple1, sub_tuple2 in zip(kv_cache_1, kv_cache_2): if sub_tuple1 is None: sub_result = [sub_tuple2] elif sub_tuple2 is None: sub_result = [sub_tuple1] else: sub_result = [] for t1, t2 in zip(sub_tuple1, sub_tuple2): if t1 is None: sub_result.append(t2) elif t2 is None: sub_result.append(t1) else: if model_type == "chatglm" and self.model.config.num_layers != 40: sub_result.append(torch.cat((t1, t2), dim=1)) else: sub_result.append(torch.cat((t1, t2), dim=0)) result.append(tuple(sub_result)) return tuple(result) else: # num_layers = self.model.layer_end - self.model.layer_start num_cache = min(len(kv_cache_1.key_cache), self.model.num_layers) for layer_idx in range(num_cache): kv_cache_1.key_cache[layer_idx] = \ torch.cat([kv_cache_1.key_cache[layer_idx], kv_cache_2.key_cache[layer_idx]], dim=0) kv_cache_1.value_cache[layer_idx] = \ torch.cat([kv_cache_1.value_cache[layer_idx], kv_cache_2.value_cache[layer_idx]], dim=0) return kv_cache_1 def update_kv_cache(self, kv_cache, prefill=False): layer_start = self.model.layer_start layer_end = self.model.layer_end num_layers = self.model.num_layers if self.model.config.model_type == "chatglm" and self.model.config.num_layers == 40: # for glm-4-9b-chat if prefill: value_placeholder = torch.empty_like((kv_cache)[-1][0]) past_key_values_placeholder = tuple( (value_placeholder, value_placeholder) for _ in range(layer_start) ) + (kv_cache)[:layer_end - layer_start] + tuple( (value_placeholder, value_placeholder) for _ in range(layer_end, num_layers) ) kv_cache = past_key_values_placeholder else: pass elif self.model.config.model_type in ["baichuan", "chatglm"] and self.rank > 0: value_placeholder = torch.empty_like((kv_cache)[-1][0]) past_key_values_placeholder = tuple( (value_placeholder, value_placeholder) for _ in range(layer_start) ) + (kv_cache)[layer_start:] kv_cache = past_key_values_placeholder else: pass return kv_cache @torch.no_grad() def model_step(self, input, cur_batch): if cur_batch is None or cur_batch.stopped or input is None: return None, cur_batch # logger.info(f"{self.rank} {cur_batch} {input.shape}") cur_id = cur_batch.batch_id _past_key_values = self.past_key_values_dict.get(cur_id, None) attention_mask = make_attention_mask(cur_batch.prompt_lengths, input.device) if self.rank == 0: input_ids = input inputs_embeds = None if cur_batch.partial_prefilling > 0: cur_input_start = cur_batch.prefilled_index cur_input_end = cur_input_start + cur_batch.partial_prefilling input_ids = input_ids[cur_input_start:cur_input_end] attention_mask = attention_mask[cur_input_start:cur_input_end] tmp_past_key_values = _past_key_values _past_key_values = None else: input_ids = None inputs_embeds = input if cur_batch.partial_prefilling > 0: cur_input_start = cur_batch.prefilled_index cur_input_end = cur_input_start + cur_batch.partial_prefilling attention_mask = attention_mask[cur_input_start:cur_input_end] tmp_past_key_values = _past_key_values _past_key_values = None torch.xpu.empty_cache() output = self.model(input_ids=input_ids, inputs_embeds=inputs_embeds, past_key_values=_past_key_values, attention_mask=attention_mask, use_cache=True,) if cur_batch.partial_prefilling > 0: cur_batch.prefilled_index = cur_input_end if tmp_past_key_values is None: tmp_past_key_values = output.past_key_values else: tmp_past_key_values = self.cat_kv_cache(self.model.config.model_type, tmp_past_key_values, output.past_key_values) # torch.xpu.empty_cache() if cur_batch.prefilled_index == cur_batch.batch_size: tmp_past_key_values = self.update_kv_cache(tmp_past_key_values, True) self.past_key_values_dict[cur_id] = tmp_past_key_values if self.pp_config.is_tail: _pre_output = self.partial_output_dict.get(cur_id, None) tmp_output = output.logits tmp_output = torch.argmax(tmp_output[:, -1:, :], dim=-1) if _pre_output is None: _pre_output = tmp_output else: _pre_output = torch.cat((_pre_output, tmp_output), dim=0) self.partial_output_dict[cur_id] = _pre_output else: _prefill = self.past_key_values_dict.get(cur_id, None) is None _past_key_values = self.update_kv_cache(output.past_key_values, prefill=_prefill) self.past_key_values_dict[cur_id] = _past_key_values torch.xpu.synchronize() if not self.pp_config.is_tail: _output = output[0] if _output.dtype != self.dtype: _output = _output.to(self.dtype) else: if cur_batch.partial_prefilling > 0 and \ cur_batch.prefilled_index == cur_batch.batch_size: _output = self.partial_output_dict.pop(cur_id, None) cur_batch.partial_prefilling = 0 else: _output = torch.argmax(output.logits[:, -1:, :], dim=-1) return _output, cur_batch def is_initialized(self): return True async def add_request(self, tokenizer): request_ids, prompt_requests = [], [] for _ in range(self.max_num_seqs): if self.waiting_requests.empty(): break tmp_result = await self.waiting_requests.get() request_id, prompt_request = tmp_result request_ids.append(request_id) prompt_requests.append(prompt_request) plain_texts = [req.inputs for req in prompt_requests] inputs = tokenizer(plain_texts, return_tensors="pt", padding=True) input_ids = inputs.input_ids.to(f'xpu:{self.rank}') attention_mask = inputs.attention_mask.to(f'xpu:{self.rank}') new_batch = BatchTask( batch_id="batch_" + str(uuid.uuid4()), request_ids=request_ids, max_tokens=max([req.parameters.max_new_tokens for req in prompt_requests]), batch_size=input_ids.size(0), input_len=input_ids.size(1), prompt_lengths=[sum(attention_mask[i, :]) for i in range(input_ids.size(0))], stopped=False, prefilled_index=0, partial_prefilling=0, ) self.input_ids_dict[new_batch.batch_id] = input_ids self.token_times[new_batch.batch_id] = [time.perf_counter()] return new_batch def clear_batch(self, cur_id): self.input_ids_dict.pop(cur_id, None) self.tokens.pop(cur_id, None) self.token_times.pop(cur_id, None) self.past_key_values_dict.pop(cur_id, None) self.is_finish.pop(cur_id, None) self.partial_output_dict.pop(cur_id, None) async def wait_stream_output(self, cur_id): cur_task = self.stream_tasks.pop(cur_id, None) if cur_task is not None: await cur_task def get_printable_text(self, cur_text, request_id): if cur_text.endswith("\n"): printable_text = cur_text[self.print_len[request_id]:] self.token_cache[request_id] = [] self.print_len[request_id] = 0 elif len(cur_text) > 0 and _is_chinese_char(ord(cur_text[-1])): printable_text = cur_text[self.print_len[request_id]:] self.print_len[request_id] += len(printable_text) self.token_cache[request_id] = [] self.print_len[request_id] = 0 else: r_index = cur_text.rfind(" ") + 1 if r_index > self.print_len[request_id]: printable_text = cur_text[self.print_len[request_id]: r_index] self.token_cache[request_id] = self.token_cache[request_id][-1:] self.print_len[request_id] = 0 else: printable_text = cur_text[self.print_len[request_id]: r_index] return printable_text async def stream_output(self, cur_batch, tokenizer, next_ids): cur_id = cur_batch.batch_id cur_cached_ids = [] _stream_tasks = [] for index, request_id in enumerate(cur_batch.request_ids): if not self.is_finish.get(request_id, False): if self.token_cache.get(request_id, None) is None: self.token_cache[request_id] = [] self.print_len[request_id] = 0 self.token_cache[request_id].extend(next_ids[index].tolist()) cur_cached_ids.append(self.token_cache[request_id]) for index, request_id in enumerate(cur_batch.request_ids): if not self.is_finish.get(request_id, False): remain = cur_batch.max_tokens - len(self.tokens[cur_id]) if self.streamer.get(request_id, None) is None: self.streamer[request_id] = asyncio.Queue() # Currently ignore eos for benchmark # if next_ids[index].int() == tokenizer.eos_token_id: # remain = 0 # self.is_finish[request_id] = True cur_text = tokenizer.decode(self.token_cache[request_id]) printable_text = self.get_printable_text(cur_text, request_id) if remain > 0: _stream_tasks.append(self.streamer[request_id].put((remain, printable_text))) else: printable_text = printable_text + cur_text[self.print_len[request_id]:] self.token_cache.pop(request_id, None) self.print_len.pop(request_id, None) _stream_tasks.append(self.streamer[request_id].put((remain, printable_text))) await asyncio.gather(*_stream_tasks) async def process_step(self, tokenizer, result_dict, processor=None): cur_batch = None torch.xpu.synchronize(self.device) if self.rank == 0: if self.on_going_batches[0] is not None: cur_batch = self.on_going_batches[0] cur_input = None if cur_batch is None: if not self.waiting_requests.empty(): # wait more requests to be put in self.waiting_requests await asyncio.sleep(0.01) cur_batch = await self.add_request(tokenizer) cur_input = self.input_ids_dict[cur_batch.batch_id] else: cur_batch = None cur_input = None if (cur_batch is not None) and (not cur_batch.stopped) and (cur_input is None): cur_id = cur_batch.batch_id if cur_batch.prefilled_index >= cur_batch.batch_size: cur_batch.partial_prefilling = 0 if cur_batch.partial_prefilling > 0: next_ids = torch.empty((cur_batch.partial_prefilling, 1,), device=f'xpu:{self.rank}', dtype=torch.int64) else: next_ids = torch.empty((cur_batch.batch_size, 1,), device=f'xpu:{self.rank}', dtype=torch.int64) # logger.info(f"recv {self.rank} {next_ids.shape}") dist.recv(next_ids, src=self.pre_rank) torch.xpu.synchronize(self.device) if cur_batch.partial_prefilling > 0: cur_input = self.input_ids_dict[cur_batch.batch_id] else: if self.tokens.get(cur_id, None) is None: self.tokens[cur_id] = [] if len(next_ids.shape) == 1: next_ids = next_ids.unsqueeze(0) self.tokens[cur_id].append(next_ids) self.token_times[cur_id].append(time.perf_counter()) cur_input = next_ids cur_batch.input_len = 1 cur_batch.prompt_lengths = [x + 1 for x in cur_batch.prompt_lengths] pre_task = self.stream_tasks.get(cur_id) if pre_task is not None: await pre_task del self.stream_tasks[cur_id] cur_task = asyncio.create_task( self.stream_output(cur_batch, tokenizer, next_ids) ) self.stream_tasks[cur_id] = cur_task if len(self.tokens[cur_id]) >= cur_batch.max_tokens: # Finish a batch outputs = torch.cat(self.tokens[cur_id], dim=1) outputs = outputs.cpu() output_strs = tokenizer.batch_decode(outputs, skip_special_tokens=False) for request_id, output_str in zip(cur_batch.request_ids, output_strs): with self.dict_lock: result_dict[request_id] = output_str cur_times = self.token_times[cur_id] first_token = cur_times[1] - cur_times[0] next_token = (cur_times[-1] - cur_times[1]) / (len(self.tokens[cur_id]) - 1) logger.info(f"First token latency: {first_token}, " f"next token latency: {next_token}") await self.wait_stream_output(cur_id) self.clear_batch(cur_id) cur_batch.stopped = True else: if (cur_batch is not None) and cur_batch.stopped: cur_batch = None if cur_batch is not None: cur_batch = self.prepare_batch(cur_batch) dist.broadcast_object_list([cur_batch], src=0) else: await asyncio.sleep(0) else: batch_list = [None] dist.broadcast_object_list(batch_list, src=0) cur_batch = batch_list[0] cur_input = None if cur_batch is not None: if cur_batch.stopped: self.clear_batch(cur_batch.batch_id) else: cur_batch = self.prepare_batch(cur_batch) cur_len = cur_batch.input_len if cur_batch.partial_prefilling: cur_input = torch.empty( (cur_batch.partial_prefilling, cur_len, self.hidden_size,), device=f'xpu:{self.rank}', dtype=self.dtype, ) else: cur_input = torch.empty( (cur_batch.batch_size, cur_len, self.hidden_size,), device=f'xpu:{self.rank}', dtype=self.dtype, ) # logger.info(f"recv {self.rank} {cur_input.shape}") dist.recv(cur_input, src=self.pre_rank) torch.xpu.synchronize(self.device) output, cur_batch = self.model_step(cur_input, cur_batch) torch.xpu.synchronize(self.device) if self.send_buff is not None: self.send_buff.wait() if output is not None: self.send_buff = dist.isend(output, dst=self.next_rank) if self.rank == 0: self.on_going_batches[:-1] = self.on_going_batches[1:] self.on_going_batches[self.world_size - 1] = cur_batch def _is_chinese_char(cp): """Checks whether CP is the codepoint of a CJK character.""" # This defines a "chinese character" as anything in the CJK Unicode block: # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block) # # Note that the CJK Unicode block is NOT all Japanese and Korean characters, # despite its name. The modern Korean Hangul alphabet is a different block, # as is Japanese Hiragana and Katakana. Those alphabets are used to write # space-separated words, so they are not treated specially and handled # like the all of the other languages. if ( (cp >= 0x4E00 and cp <= 0x9FFF) or (cp >= 0x3400 and cp <= 0x4DBF) # or (cp >= 0x20000 and cp <= 0x2A6DF) # or (cp >= 0x2A700 and cp <= 0x2B73F) # or (cp >= 0x2B740 and cp <= 0x2B81F) # or (cp >= 0x2B820 and cp <= 0x2CEAF) # or (cp >= 0xF900 and cp <= 0xFAFF) or (cp >= 0x2F800 and cp <= 0x2FA1F) # ): # return True return False def llama_causallm_forward_4_37_lowmem( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions # noqa output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states # noqa ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) hidden_states = outputs[0] # ipex-llm change starts if self.config.pretraining_tp > 1: lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0) # noqa logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)] # noqa logits = torch.cat(logits, dim=-1) else: # Only empty cache for first token if hidden_states.shape[1] > 1: torch.xpu.empty_cache() logits = self.lm_head(hidden_states) # Only empty cache for first token if hidden_states.shape[1] > 1: torch.xpu.empty_cache() # logits = logits.float() # ipex-llm change ends loss = None if labels is not None: # Shift so that tokens < n predict n shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() # Flatten the tokens loss_fct = CrossEntropyLoss() shift_logits = shift_logits.view(-1, self.config.vocab_size) shift_labels = shift_labels.view(-1) # Enable model parallelism shift_labels = shift_labels.to(shift_logits.device) loss = loss_fct(shift_logits, shift_labels) if not return_dict: output = (logits,) + outputs[1:] return (loss,) + output if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) def chatglm3_conditional_generation_forward_lowmem( self, input_ids: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[Tuple[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, return_last_logit: Optional[bool] = False, ): use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict transformer_outputs = self.transformer( input_ids=input_ids, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_hidden_states=output_hidden_states, return_dict=return_dict, ) hidden_states = transformer_outputs[0] if return_last_logit: hidden_states = hidden_states[-1:] # ipex-llm change starts torch.xpu.empty_cache() lm_logits = self.transformer.output_layer(hidden_states) torch.xpu.empty_cache() lm_logits = lm_logits.transpose(0, 1).contiguous() loss = None if labels is not None: # lm_logits = lm_logits.to(torch.float32) # Shift so that tokens < n predict n shift_logits = lm_logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() # Flatten the tokens loss_fct = CrossEntropyLoss(ignore_index=-100) loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) lm_logits = lm_logits.to(hidden_states.dtype) loss = loss.to(hidden_states.dtype) # ipex-llm change ends if not return_dict: output = (lm_logits,) + transformer_outputs[1:] return ((loss,) + output) if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=lm_logits, past_key_values=transformer_outputs.past_key_values, hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions, ) def glm4_conditional_generation_forward_lowmem( self, input_ids: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[Tuple[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, return_last_logit: Optional[bool] = False, ): use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict transformer_outputs = self.transformer( input_ids=input_ids, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_hidden_states=output_hidden_states, return_dict=return_dict, ) hidden_states = transformer_outputs[0] if return_last_logit: hidden_states = hidden_states[:, -1:] # ipex-llm change starts torch.xpu.empty_cache() lm_logits = self.transformer.output_layer(hidden_states) torch.xpu.empty_cache() loss = None if labels is not None: # lm_logits = lm_logits.to(torch.float32) # Shift so that tokens < n predict n shift_logits = lm_logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() # Flatten the tokens loss_fct = CrossEntropyLoss(ignore_index=-100) loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) lm_logits = lm_logits.to(hidden_states.dtype) loss = loss.to(hidden_states.dtype) # ipex-llm change ends if not return_dict: output = (lm_logits,) + transformer_outputs[1:] return ((loss,) + output) if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=lm_logits, past_key_values=transformer_outputs.past_key_values, hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions, )