from typing import Optional, Tuple, Union import torch from torch import nn import os from unimernet.common.config import Config import unimernet.tasks as tasks import argparse from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_causal_attention_mask class PatchedMBartLearnedPositionalEmbedding(nn.Module): def __init__(self, origin: nn.Module): super().__init__() self.offset = origin.offset self.embedding = nn.Embedding(origin.num_embeddings, origin.embedding_dim) self.embedding.weight.data = origin.weight.data def forward(self, input_ids: torch.Tensor, past_key_values_length: int = 0): """`input_ids' shape is expected to be [bsz x seqlen].""" bsz, seq_len = input_ids.shape[:2] positions = torch.arange(0, seq_len, dtype=torch.long, device=self.embedding.weight.device ) positions += past_key_values_length positions = positions.expand(bsz, -1) return self.embedding(positions + self.offset) class PatchedMBartDecoder(nn.Module): def __init__(self, origin: nn.Module, kvlen: torch.LongTensor): super().__init__() self.origin = origin self.kvlen = kvlen self.config = origin.config self.embed_tokens = origin.embed_tokens self.embed_scale = origin.embed_scale self._use_flash_attention_2 = origin._use_flash_attention_2 self.embed_positions = origin.embed_positions self.counting_context_weight = getattr(origin, 'counting_context_weight', None) self.layernorm_embedding = origin.layernorm_embedding self.layers = origin.layers self.layer_norm = origin.layer_norm self.patched_embed_positions = PatchedMBartLearnedPositionalEmbedding(self.embed_positions) def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, count_pred: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.LongTensor] = None, head_mask: Optional[torch.Tensor] = None, cross_attn_head_mask: Optional[torch.Tensor] = None, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]: run_origin = False if past_key_values is None: run_origin = True elif past_key_values[0][0].size(-2) < attention_mask.size(-1): run_origin = True if run_origin: return self.origin( input_ids=input_ids, attention_mask=attention_mask, count_pred=count_pred, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, head_mask=head_mask, cross_attn_head_mask=cross_attn_head_mask, 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, ) output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) 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 # retrieve input_ids and inputs_embeds if input_ids is not None and inputs_embeds is not None: raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") elif input_ids is not None: input = input_ids input_shape = input.size() input_ids = input_ids.view(-1, input_shape[-1]) elif inputs_embeds is not None: input_shape = inputs_embeds.size()[:-1] input = inputs_embeds[:, :, -1] else: raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") # past_key_values_length past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale if self._use_flash_attention_2: # 2d mask is passed through the layers attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None else: # 4d mask is passed through the layers attention_mask = _prepare_4d_causal_attention_mask( attention_mask, input_shape, inputs_embeds, past_key_values_length ) # expand encoder attention mask if encoder_hidden_states is not None and encoder_attention_mask is not None: if self._use_flash_attention_2: encoder_attention_mask = encoder_attention_mask if 0 in encoder_attention_mask else None else: # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] encoder_attention_mask = _prepare_4d_attention_mask( encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1] ) # embed positions positions = self.patched_embed_positions(input, self.kvlen) hidden_states = inputs_embeds + positions.to(inputs_embeds.device) # TODO: add counting context weight to hidden_states if count_pred is not None: count_context_weight = self.counting_context_weight(count_pred) hidden_states = hidden_states + 0.5 * count_context_weight.unsqueeze(1) hidden_states = self.layernorm_embedding(hidden_states) # decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None next_decoder_cache = () if use_cache else None # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]): if attn_mask is not None: if attn_mask.size()[0] != len(self.layers): raise ValueError( f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for" f" {attn_mask.size()[0]}." ) for idx, decoder_layer in enumerate(self.layers): # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) if output_hidden_states: all_hidden_states += (hidden_states,) past_key_value = past_key_values[idx] if past_key_values is not None else None layer_outputs = decoder_layer( hidden_states, attention_mask=attention_mask, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, layer_head_mask=(head_mask[idx] if head_mask is not None else None), cross_attn_layer_head_mask=( cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None ), past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, ) hidden_states = layer_outputs[0] if use_cache: next_decoder_cache += (layer_outputs[3 if output_attentions else 1],) if output_attentions: all_self_attns += (layer_outputs[1],) if encoder_hidden_states is not None: all_cross_attentions += (layer_outputs[2],) hidden_states = self.layer_norm(hidden_states) # add hidden states from the last decoder layer if output_hidden_states: all_hidden_states += (hidden_states,) next_cache = next_decoder_cache if use_cache else None if not return_dict: return tuple( v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attentions] if v is not None ) return BaseModelOutputWithPastAndCrossAttentions( last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=all_hidden_states, attentions=all_self_attns, cross_attentions=all_cross_attentions, ) class PatchedMBartAttention(nn.Module): def __init__(self, origin: nn.Module, kvlen: torch.LongTensor): super().__init__() self.embed_dim = origin.embed_dim self.num_heads = origin.num_heads self.dropout = origin.dropout self.head_dim = origin.head_dim self.config = origin.config self.scaling = origin.scaling self.is_decoder = origin.is_decoder self.is_causal = origin.is_causal self.k_proj = origin.k_proj self.v_proj = origin.v_proj self.q_proj = origin.q_proj self.out_proj = origin.out_proj self.kvlen = kvlen def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() def forward( self, hidden_states: torch.Tensor, key_value_states: Optional[torch.Tensor] = None, past_key_value: Optional[Tuple[torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, layer_head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: """Input shape: Batch x Time x Channel""" # if key_value_states are provided this layer is used as a cross-attention layer # for the decoder is_cross_attention = key_value_states is not None bsz, tgt_len, _ = hidden_states.size() # get query proj query_states = self.q_proj(hidden_states) * self.scaling # get key, value proj # `past_key_value[0].shape[2] == key_value_states.shape[1]` # is checking that the `sequence_length` of the `past_key_value` is the same as # the provided `key_value_states` to support prefix tuning if ( is_cross_attention and past_key_value is not None and past_key_value[0].shape[2] == key_value_states.shape[1] ): # reuse k,v, cross_attentions key_states = past_key_value[0] value_states = past_key_value[1] elif is_cross_attention: # cross_attentions key_states = self._shape(self.k_proj(key_value_states), -1, bsz) value_states = self._shape(self.v_proj(key_value_states), -1, bsz) elif past_key_value is not None: # reuse k, v, self_attention key_states = self._shape(self.k_proj(hidden_states), -1, bsz) value_states = self._shape(self.v_proj(hidden_states), -1, bsz) if past_key_value[0].size(-2) < attention_mask.size(-1): key_states = torch.cat([past_key_value[0], key_states], dim=2) value_states = torch.cat([past_key_value[1], value_states], dim=2) else: past_key_value[0][:, :, self.kvlen[None]] = key_states past_key_value[1][:, :, self.kvlen[None]] = value_states key_states = past_key_value[0] value_states = past_key_value[1] else: # self_attention key_states = self._shape(self.k_proj(hidden_states), -1, bsz) value_states = self._shape(self.v_proj(hidden_states), -1, bsz) if self.is_decoder: past_key_value = (key_states, value_states) proj_shape = (bsz * self.num_heads, -1, self.head_dim) query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape) key_states = key_states.reshape(*proj_shape) value_states = value_states.reshape(*proj_shape) src_len = key_states.size(1) attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len): raise ValueError( f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is" f" {attn_weights.size()}" ) if attention_mask is not None: if attention_mask.size() != (bsz, 1, tgt_len, src_len): raise ValueError( f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}" ) attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) attn_weights = nn.functional.softmax(attn_weights, dim=-1) if layer_head_mask is not None: if layer_head_mask.size() != (self.num_heads,): raise ValueError( f"Head mask for a single layer should be of size {(self.num_heads,)}, but is" f" {layer_head_mask.size()}" ) attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len) attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) if output_attentions: # this operation is a bit awkward, but it's required to # make sure that attn_weights keeps its gradient. # In order to do so, attn_weights have to be reshaped # twice and have to be reused in the following attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len) else: attn_weights_reshaped = None attn_probs = attn_weights attn_output = torch.bmm(attn_probs, value_states) if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim): raise ValueError( f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is" f" {attn_output.size()}" ) attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim) attn_output = attn_output.transpose(1, 2) # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be # partitioned across GPUs when using tensor-parallelism. attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim) # attn_output = self.out_proj(attn_output) attn_output = self.out_proj(attn_output) return attn_output, attn_weights_reshaped, past_key_value class PatchedMBartSqueezeAttention(nn.Module): def __init__(self, origin: nn.Module, kvlen: torch.LongTensor): super().__init__() self.embed_dim = origin.embed_dim self.num_heads = origin.num_heads self.dropout = origin.dropout self.head_dim = origin.head_dim self.squeeze_head_dim=origin.squeeze_head_dim self.config = origin.config self.scaling = origin.scaling self.is_decoder = origin.is_decoder self.scaling = origin.scaling self.q_proj = origin.q_proj self.k_proj = origin.k_proj self.v_proj = origin.v_proj self.out_proj = origin.out_proj self.kvlen = kvlen def _shape_qk(self, tensor: torch.Tensor, seq_len: int, bsz: int): return tensor.view(bsz, seq_len, self.num_heads, self.squeeze_head_dim).transpose(1, 2).contiguous() def _shape_v(self, tensor: torch.Tensor, seq_len: int, bsz: int): return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() def forward( self, hidden_states: torch.Tensor, key_value_states: Optional[torch.Tensor] = None, past_key_value: Optional[Tuple[torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, layer_head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: """Input shape: Batch x Time x Channel""" # if key_value_states are provided this layer is used as a cross-attention layer # for the decoder is_cross_attention = key_value_states is not None bsz, tgt_len, _ = hidden_states.size() # get query proj query_states = self.q_proj(hidden_states) * self.scaling # get key, value proj # `past_key_value[0].shape[2] == key_value_states.shape[1]` # is checking that the `sequence_length` of the `past_key_value` is the same as # the provided `key_value_states` to support prefix tuning if ( is_cross_attention and past_key_value is not None and past_key_value[0].shape[2] == key_value_states.shape[1] ): # reuse k,v, cross_attentions key_states = past_key_value[0] value_states = past_key_value[1] elif is_cross_attention: # cross_attentions key_states = self._shape_qk(self.k_proj(key_value_states), -1, bsz) value_states = self._shape_v(self.v_proj(key_value_states), -1, bsz) elif past_key_value is not None: # reuse k, v, self_attention key_states = self._shape_qk(self.k_proj(hidden_states), -1, bsz) value_states = self._shape_v(self.v_proj(hidden_states), -1, bsz) if past_key_value[0].size(-2) < attention_mask.size(-1): key_states = torch.cat([past_key_value[0], key_states], dim=2) value_states = torch.cat([past_key_value[1], value_states], dim=2) else: past_key_value[0][:, :, self.kvlen[None]] = key_states past_key_value[1][:, :, self.kvlen[None]] = value_states key_states = past_key_value[0] value_states = past_key_value[1] else: # self_attention key_states = self._shape_qk(self.k_proj(hidden_states), -1, bsz) value_states = self._shape_v(self.v_proj(hidden_states), -1, bsz) if self.is_decoder: # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. # Further calls to cross_attention layer can then reuse all cross-attention # key/value_states (first "if" case) # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of # all previous decoder key/value_states. Further calls to uni-directional self-attention # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) # if encoder bi-directional self-attention `past_key_value` is always `None` past_key_value = (key_states, value_states) proj_shape = (bsz * self.num_heads, -1, self.squeeze_head_dim) value_shape = (bsz * self.num_heads, -1, self.head_dim) query_states = self._shape_qk(query_states, tgt_len, bsz).view(*proj_shape) key_states = key_states.reshape(*proj_shape) value_states = value_states.reshape(*value_shape) src_len = key_states.size(1) attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len): raise ValueError( f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is" f" {attn_weights.size()}" ) if attention_mask is not None: if attention_mask.size() != (bsz, 1, tgt_len, src_len): raise ValueError( f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}" ) attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) attn_weights = nn.functional.softmax(attn_weights, dim=-1) if layer_head_mask is not None: if layer_head_mask.size() != (self.num_heads,): raise ValueError( f"Head mask for a single layer should be of size {(self.num_heads,)}, but is" f" {layer_head_mask.size()}" ) attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len) attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) if output_attentions: # this operation is a bit awkward, but it's required to # make sure that attn_weights keeps its gradient. # In order to do so, attn_weights have to be reshaped # twice and have to be reused in the following attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len) else: attn_weights_reshaped = None attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) attn_output = torch.bmm(attn_probs, value_states) if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim): raise ValueError( f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is" f" {attn_output.size()}" ) attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim) attn_output = attn_output.transpose(1, 2) # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be # partitioned across GPUs when using tensor-parallelism. attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim) attn_output = self.out_proj(attn_output) return attn_output, attn_weights_reshaped, past_key_value def patch_model(model: nn.Module, kvlen: torch.LongTensor): for name, child in model.named_children(): cls_name = type(child).__name__ if cls_name == 'MBartAttention': patched_child = PatchedMBartAttention(child, kvlen) model.register_module(name, patched_child) elif cls_name == 'MBartSqueezeAttention': patched_child = PatchedMBartSqueezeAttention(child, kvlen) model.register_module(name, patched_child) else: patch_model(child, kvlen) cls_name = type(model).__name__ if cls_name == 'CustomMBartDecoder': model = PatchedMBartDecoder(model, kvlen) return model def next_power_of_2(n: int): """Return the smallest power of 2 greater than or equal to n.""" n -= 1 n |= n >> 1 n |= n >> 2 n |= n >> 4 n |= n >> 8 n |= n >> 16 n |= n >> 32 n += 1 return n def get_graph_key(batch_size: int, kvlens: int): batch_size = next_power_of_2(batch_size) kvlens = next_power_of_2(kvlens) batch_size = max(8, batch_size) kvlens = max(32, kvlens) return batch_size, kvlens class GraphRunnerImpl: def __init__(self, model: nn.Module, graph: torch.cuda.CUDAGraph, input_buffers: dict, output_buffers: dict): self.model = model self.graph = graph self.input_buffers = input_buffers self.output_buffers = output_buffers @staticmethod def extract_input_buffers(input_buffers: dict, batch_size: int, kvlens: int): input_ids = input_buffers['input_ids'][:batch_size] attention_mask = input_buffers['attention_mask'][:batch_size, :kvlens] encoder_hidden_states = input_buffers['encoder_hidden_states'][:batch_size] kvlen=input_buffers['kvlen'] past_key_values = [] for past_key_value in input_buffers['past_key_values']: k0 = past_key_value[0][:batch_size, :, :kvlens] v0 = past_key_value[1][:batch_size, :, :kvlens] k1 = past_key_value[2][:batch_size] v1 = past_key_value[3][:batch_size] past_key_values.append((k0, v0, k1, v1)) input_buffers = dict( input_ids=input_ids, attention_mask=attention_mask, encoder_hidden_states=encoder_hidden_states, past_key_values=past_key_values, kvlen=kvlen, ) return input_buffers @staticmethod def fill_input_buffers( input_buffer: dict, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, ): batch_size = input_ids.size(0) kvlens = attention_mask.size(1) input_buffer['input_ids'][:batch_size] = input_ids if input_buffer['attention_mask'].data_ptr() != attention_mask.data_ptr(): input_buffer['attention_mask'].fill_(0) input_buffer['attention_mask'][:batch_size, :kvlens] = attention_mask input_buffer['encoder_hidden_states'][:batch_size] = encoder_hidden_states if past_key_values is not None: for buf_kv, kv in zip(input_buffer['past_key_values'], past_key_values): idx = 0 if buf_kv[idx].data_ptr() != kv[idx].data_ptr(): buf_kv[idx].fill_(0) buf_kv[idx][:batch_size, :, :kvlens-1] = kv[idx] idx = 1 if buf_kv[idx].data_ptr() != kv[idx].data_ptr(): buf_kv[idx].fill_(0) buf_kv[idx][:batch_size, :, :kvlens-1] = kv[idx] idx = 2 if buf_kv[idx].data_ptr() != kv[idx].data_ptr(): buf_kv[idx].fill_(0) buf_kv[idx][:batch_size] = kv[idx] idx = 3 if buf_kv[idx].data_ptr() != kv[idx].data_ptr(): buf_kv[idx].fill_(0) buf_kv[idx][:batch_size] = kv[idx] input_buffer['kvlen'].fill_(kvlens - 1) @classmethod @torch.inference_mode() def capture(cls, model: nn.Module, input_buffers: dict, pool, warmup: bool = False, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, count_pred: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.LongTensor] = None, head_mask: Optional[torch.Tensor] = None, cross_attn_head_mask: Optional[torch.Tensor] = None, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None,): batch_size = input_ids.size(0) kvlens = attention_mask.size(1) graph_key = get_graph_key(batch_size, kvlens) batch_size = graph_key[0] kvlens = graph_key[1] input_buffers = cls.extract_input_buffers(input_buffers, batch_size=batch_size, kvlens=kvlens) cls.fill_input_buffers(input_buffers, input_ids, attention_mask, encoder_hidden_states, past_key_values) input_ids = input_buffers['input_ids'] attention_mask = input_buffers['attention_mask'] encoder_hidden_states = input_buffers['encoder_hidden_states'] past_key_values = input_buffers['past_key_values'] if warmup: # warmup model( input_ids=input_ids, attention_mask=attention_mask, count_pred=count_pred, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, head_mask=head_mask, cross_attn_head_mask=cross_attn_head_mask, 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) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, pool=pool): outputs = model( input_ids=input_ids, attention_mask=attention_mask, count_pred=count_pred, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, head_mask=head_mask, cross_attn_head_mask=cross_attn_head_mask, 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) output_buffers = dict( last_hidden_state=outputs['last_hidden_state'], past_key_values=outputs['past_key_values'], ) return GraphRunnerImpl(model, graph, input_buffers, output_buffers) def __call__(self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, count_pred: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.LongTensor] = None, head_mask: Optional[torch.Tensor] = None, cross_attn_head_mask: Optional[torch.Tensor] = None, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ): batch_size = input_ids.size(0) kvlens = attention_mask.size(1) self.fill_input_buffers(self.input_buffers, input_ids, attention_mask, encoder_hidden_states, past_key_values) self.graph.replay() last_hidden_state = self.output_buffers['last_hidden_state'][:batch_size] past_key_values = [] for past_key_value in self.output_buffers['past_key_values']: k0 = past_key_value[0][:batch_size, :, :kvlens] v0 = past_key_value[1][:batch_size, :, :kvlens] k1 = past_key_value[2][:batch_size] v1 = past_key_value[3][:batch_size] past_key_values.append((k0, v0, k1, v1)) outputs = BaseModelOutputWithPastAndCrossAttentions( last_hidden_state=last_hidden_state, past_key_values=past_key_values, ) return outputs class GraphRunner(nn.Module): def __init__(self, model: nn.Module, max_batchs: int, max_kvlens: int, dtype:torch.dtype = torch.float16, device: torch.device = 'cuda'): super().__init__() self.kvlen = torch.tensor(0, dtype=torch.long, device=device) model = patch_model(model.to(dtype), self.kvlen) self.model = model self.max_batchs = max_batchs self.max_kvlens = max_kvlens self.device = device self.input_buffers = None self.impl_map = dict() self.graph_pool_handle = torch.cuda.graph_pool_handle() self.warmuped = False def create_buffers(self, encoder_kvlens: int, dtype: torch.dtype): max_batchs = self.max_batchs max_kvlens = self.max_kvlens device = self.device config = self.model.config d_model = config.d_model decoder_layers = config.decoder_layers num_heads = config.decoder_attention_heads head_dim = d_model // num_heads self_attn = self.model.layers[0].self_attn qk_head_dim = getattr(self_attn, 'squeeze_head_dim', head_dim) input_ids = torch.ones((max_batchs, 1), dtype=torch.int64, device=device) attention_mask = torch.zeros((max_batchs, max_kvlens), dtype=torch.int64, device=device) encoder_hidden_states = torch.zeros((max_batchs, encoder_kvlens, d_model), dtype=dtype, device=device) past_key_values = [] for _ in range(decoder_layers): k0 = torch.zeros((max_batchs, num_heads, max_kvlens, qk_head_dim), dtype=dtype, device=device) v0 = torch.zeros((max_batchs, num_heads, max_kvlens, head_dim), dtype=dtype, device=device) k1 = torch.zeros((max_batchs, num_heads, encoder_kvlens, qk_head_dim), dtype=dtype, device=device) v1 = torch.zeros((max_batchs, num_heads, encoder_kvlens, head_dim), dtype=dtype, device=device) past_key_values.append((k0, v0, k1, v1)) self.input_buffers = dict( input_ids=input_ids, attention_mask=attention_mask, encoder_hidden_states=encoder_hidden_states, past_key_values=past_key_values, kvlen=self.kvlen ) @torch.inference_mode() def forward(self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, count_pred: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.LongTensor] = None, head_mask: Optional[torch.Tensor] = None, cross_attn_head_mask: Optional[torch.Tensor] = None, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ): batch_size, qlens = input_ids.size() kvlens = attention_mask.size(1) eager_mode = False if qlens != 1: eager_mode = True if past_key_values is None: eager_mode = True else: for past_key_value in past_key_values: if past_key_value is None: eager_mode = True break if batch_size >= self.max_batchs or kvlens >= self.max_kvlens: eager_mode = True if eager_mode: return self.model( input_ids=input_ids, attention_mask=attention_mask, count_pred=count_pred, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, head_mask=head_mask, cross_attn_head_mask=cross_attn_head_mask, 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,) # create buffer if not exists. if self.input_buffers is None: encoder_kvlens = encoder_hidden_states.size(1) self.create_buffers(encoder_kvlens=encoder_kvlens, dtype=encoder_hidden_states.dtype) graph_key = get_graph_key(batch_size, kvlens) if graph_key not in self.impl_map: warmup = False if not self.warmuped: warmup = True self.warmuped = True impl = GraphRunnerImpl.capture( self.model, self.input_buffers, self.graph_pool_handle, warmup=warmup, input_ids=input_ids, attention_mask=attention_mask, count_pred=count_pred, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, head_mask=head_mask, cross_attn_head_mask=cross_attn_head_mask, 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, ) self.impl_map[graph_key] = impl impl = self.impl_map[graph_key] ret = impl( input_ids=input_ids, attention_mask=attention_mask, count_pred=count_pred, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, head_mask=head_mask, cross_attn_head_mask=cross_attn_head_mask, 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, ) return ret