# See: https://huggingface.co/docs/transformers/custom_models from typing import Optional, Tuple, Union, List import math import copy import sys from importlib import import_module import torch from torch import nn, Tensor import torch.nn.init as init from torch.nn import functional as F from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutput, CausalLMOutputWithPast from transformers import ( PreTrainedModel, PretrainedConfig, AutoConfig, AutoModel, AutoModelForCausalLM, ) from transformers.utils import logging from transformers.cache_utils import Cache, DynamicCache from transformers.utils import ( is_flash_attn_2_available, is_flash_attn_greater_or_equal_2_10, ) if is_flash_attn_2_available(): from flash_attn import flash_attn_qkvpacked_func, flash_attn_func logger = logging.get_logger(__name__) # The model type string to bind. model_type = "walsh-causal-v1" class Config(PretrainedConfig): model_type = model_type attribute_map = { "hidden_size": "d_embed", } def __init__( # All of these MUST have defaults, even if unused. self, vocab_size=16000, pad_index=None, hidden_size=1024, num_attention_heads=8, num_hidden_layers=6, max_sequence_length=2048, dim_feedforward = 4096, dropout=0.1, loss_function = "causal_loss", # Default class to use for each of these components. positional_encoder_cls='.PositionalEncoder', attention_cls='.CausalSelfAttention', activation_cls='torch.nn.ReLU', feedforward_cls='.FeedforwardLayer', layer_stack_cls='.TransformerLayerStack', layer_cls='.PostLayerNorm', transformer_cls='.Transformer', norm_cls='torch.nn.LayerNorm', embdding_cls='torch.nn.Embedding', output_proj_cls='torch.nn.Linear', positional_encoder_args={ 'd_model': 1024, 'max_seq_len': 2048, }, # Arg groups, passed to factory classes above. transformer_args=dict(), attention_args=dict(), feedforward_args=dict(), activation_args=dict(), norm_args={ 'normalized_shape': 1024, }, layer_stack_args=dict(), layer_args=dict(), embedding_args=dict(), output_proj_args=dict(), output_attentions=False, output_hidden_states=False, use_cache=True, **kwargs, ): self.vocab_size = vocab_size self.pad_index = pad_index self.hidden_size = hidden_size self.num_attention_heads = num_attention_heads self.num_hidden_layers = num_hidden_layers self.max_sequence_length = max_sequence_length self.loss_function = loss_function self.dim_feedforward = dim_feedforward self.dropout = dropout self.positional_encoder_cls = positional_encoder_cls self.attention_cls = attention_cls self.activation_cls = activation_cls self.feedforward_cls = feedforward_cls self.layer_stack_cls = layer_stack_cls self.layer_cls = layer_cls self.transformer_cls = transformer_cls self.norm_cls = norm_cls self.embdding_cls = embdding_cls self.output_proj_cls = output_proj_cls self.positional_encoder_args = positional_encoder_args self.transformer_args = transformer_args self.attention_args = attention_args self.feedforward_args = feedforward_args self.activation_args = activation_args self.norm_args = norm_args self.layer_stack_args = layer_stack_args self.layer_args = layer_args self.embedding_args = embedding_args self.output_proj_args = output_proj_args self.output_attentions = output_attentions self.output_hidden_states = output_hidden_states self.use_cache = use_cache super().__init__(**kwargs) def causal_loss(logits: Tensor, labels: Tensor, input_ids: Tensor, ignore_index=-100) -> Tensor: """ Compute and return the loss using logits and labels. """ # Shift so that tokens < n predict n shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss = torch.nn.functional.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=ignore_index, reduction='mean', ) return loss.nan_to_num() # Learning to Break the Loop: Analyzing and Mitigating Repetitions for Neural Text Generation # https://arxiv.org/abs/2206.02369 def ditto_loss(logits: Tensor, labels: Tensor, input_ids: Tensor) -> Tensor: batch_size, seq_len, vocab_size = logits.shape rep_reduce_gamma = 0.5 ditto_weight = 1.0e5 probs = torch.softmax(logits, dim=-1) total_loss = None for i in range(batch_size): context_len = labels[i, 0].item() sentence_len = labels[i, 1].item() n_repeats = labels[i, 2].item() # For readability context_end = context_len sentence_start = context_len sentence_end = sentence_start + sentence_len target_start = sentence_end # Get causal loss for context tokens causal_ids = input_ids[i:i+1, :context_end] c_loss = causal_loss( logits=logits[i:i+1, :context_end], labels=causal_ids, input_ids=causal_ids ) # Slice out target probabilities target_probs = probs[i , target_start:, :] # Slice out first instance of repeated sentence, detach is (prevents back-prop), repeat in N times, # and trim to length of target_probs. baseline_probs = probs[i, sentence_start:sentence_end, :].detach().repeat(n_repeats, 1)[:target_probs.size(0), :] # Compute DITTO loss. one_minus_probs = torch.clamp((1.0 - torch.abs((target_probs - baseline_probs * rep_reduce_gamma))), min=1e-20) r_loss = -torch.log(one_minus_probs).mean() * ditto_weight # Combine repitition and causal loss loss = c_loss + r_loss # Add this to the total if total_loss is None: total_loss = loss else: total_loss += loss return total_loss / batch_size # Dynamically lookup class name and return factory for class. def get_dynamic_class(name): try: module_path, class_name = name.rsplit('.', 1) if module_path == "": return getattr(sys.modules[__name__], class_name) module = import_module(module_path) return getattr(module, class_name) except (ImportError, AttributeError) as e: raise ImportError(name) # An easily extensible dynamic transformer class # Many variations can be specified entirely in the configuration, without touching this code. class HFCausalModel(PreTrainedModel): config_class = Config model_type = 'Transformer' supports_gradient_checkpointing = True # Presently needs to be manually set to match transformer layer class... _no_split_modules = ["DeepNetLayer"] _supports_flash_attn_2 = True _supports_sdpa = True _supports_cache_class = True _skip_keys_device_placement = "past_key_values" def __init__(self, config): super().__init__(config) self.d_model = config.hidden_size self.transformer_head = self._make_transformer(config) self.loss_function = get_dynamic_class(config.loss_function) self.gradient_checkpointing = False self.post_init() def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, labels: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, **kwargs, ) -> (Tensor, dict[str, Tensor]): batch_size, seq_len = input_ids.shape 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 if use_cache: # If legacy cache, convert to DynamicCache use_legacy_cache = not isinstance(past_key_values, Cache) if use_legacy_cache: past_key_values = DynamicCache.from_legacy_cache(past_key_values) if self.gradient_checkpointing and self.training: if use_cache: logger.warning_once( "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." ) use_cache = False gradient_checkpointing_func = self._gradient_checkpointing_func else: gradient_checkpointing_func = None outputs = self.transformer_head( input_ids=input_ids, position_ids=position_ids, output_attentions=output_attentions, gradient_checkpointing_func=gradient_checkpointing_func, past_key_values=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states, ) logits = outputs["logits"].float() attentions = outputs["attentions"] # Compute loss. if labels is not None: loss = self.loss_function(logits=logits, labels=labels, input_ids=input_ids) else: loss = None # Convert back to legacy cache, if that's what we received new_cache = outputs["past_key_values"] if use_cache and new_cache is not None and use_legacy_cache: new_cache = new_cache.to_legacy_cache() return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=new_cache, hidden_states=outputs["hidden_states"], attentions=outputs["attentions"], ) # Implementation from Huggingface Transformers, # https://github.com/huggingface/transformers/blob/main/src/transformers/models/mistral/modeling_mistral.py # Note: We do not implement attention mask at present, so some of this code is not applicable # TODO: Reenable attention mask support for batch inference.. def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs ): # Omit tokens covered by past_key_values if past_key_values is not None: if isinstance(past_key_values, Cache): cache_length = past_key_values.get_seq_length() past_length = past_key_values.seen_tokens max_cache_length = past_key_values.get_max_length() else: cache_length = past_length = past_key_values[0][0].shape[2] max_cache_length = None # Keep only the unprocessed tokens: # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as # input) if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard # input_ids based on the past_length. elif past_length < input_ids.shape[1]: input_ids = input_ids[:, past_length:] # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. # If we are about to go beyond the maximum cache length, we need to crop the input attention mask. if ( max_cache_length is not None and attention_mask is not None and cache_length + input_ids.shape[1] > max_cache_length ): attention_mask = attention_mask[:, -max_cache_length:] position_ids = kwargs.get("position_ids", None) if attention_mask is not None and position_ids is None: # create position_ids on the fly for batch generation position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) if past_key_values: position_ids = position_ids[:, -input_ids.shape[1] :] # if `inputs_embeds` are passed, we only want to use them in the 1st generation step # NOTE: Injecting positional embeddings is not yet supported. if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: model_inputs = {"input_ids": input_ids} model_inputs.update( { "position_ids": position_ids, "past_key_values": past_key_values, "use_cache": kwargs.get("use_cache"), "attention_mask": attention_mask, } ) return model_inputs @staticmethod def _reorder_cache(past_key_values, beam_idx): reordered_past = () for layer_past in past_key_values: reordered_past += ( tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past), ) return reordered_past def _make_embedding(self, config): embedding_cls = get_dynamic_class(config.embdding_cls) return embedding_cls(config.vocab_size, self.d_model, config.pad_index, **config.embedding_args) def _make_pos_encoder(self, config): pos_enc_cls = get_dynamic_class(config.positional_encoder_cls) return pos_enc_cls(**config.positional_encoder_args) def _make_output_projection(self, config): output_proj_cls = get_dynamic_class(config.output_proj_cls) return output_proj_cls(self.d_model, config.vocab_size, **config.output_proj_args) def _make_dropout(self, config): return nn.Dropout(config.dropout) def _make_activation(self, config): activation_cls = get_dynamic_class(config.activation_cls) return activation_cls(**config.activation_args) def _make_norm(self, config): norm_cls = get_dynamic_class(config.norm_cls) return norm_cls(self.d_model) def _make_self_attention(self, layer_idx, config): attention_cls = get_dynamic_class(config.attention_cls) # Map HF _attn_implementation to attn_type match config._attn_implementation: case "flash_attention_2": if is_flash_attn_2_available(): if not is_flash_attn_greater_or_equal_2_10(): raise Exception("flash_attn_2 >= 2.10 is required") attn_type = "flash2" else: attn_type = "torch" case "sdpa": attn_type = "torch" case "eager": attn_type = "native" case _: raise Exception(f"Unimplemented attention type '{config._attn_implementation}'") return attention_cls( d_model=self.d_model, num_heads=config.num_attention_heads, attn_type=attn_type, layer_idx=layer_idx, config=config, **config.attention_args, ) def _make_feedforward(self, layer_idx, config): feedforward_cls = get_dynamic_class(config.feedforward_cls) return feedforward_cls( d_model=self.d_model, feedforward_dim=config.dim_feedforward, dropout=config.dropout, activation=self._make_activation(config), layer_idx=layer_idx, **config.feedforward_args, ) def _make_layer(self, layer_idx, config): layer_cls = get_dynamic_class(config.layer_cls) return layer_cls( d_model=self.d_model, dropout=self._make_dropout(config), attention=self._make_self_attention(layer_idx, config), feedforward=self._make_feedforward(layer_idx, config), norm1=self._make_norm(config), norm2=self._make_norm(config), layer_idx=layer_idx, **config.layer_args, ) def _make_layer_stack(self, config): layer_stack_cls = get_dynamic_class(config.layer_stack_cls) return layer_stack_cls( layers=nn.ModuleList([ self._make_layer(layer_idx, config) for layer_idx in range(config.num_hidden_layers) ]), **config.layer_stack_args, ) def _make_transformer(self, config): transformer_cls = get_dynamic_class(config.transformer_cls) return transformer_cls( d_model=self.d_model, embedding=self._make_embedding(config), positional_encoder=self._make_pos_encoder(config), layer_stack=self._make_layer_stack(config), output_projection=self._make_output_projection(config), **config.transformer_args, ) @torch.no_grad() def _init_weights(self, module): pass # Register model type and configuration AutoConfig.register(model_type, Config) AutoModelForCausalLM.register(Config, HFCausalModel) # A generic container class for standard transformer components. class Transformer(nn.Module): def __init__(self, d_model, embedding, positional_encoder, layer_stack, output_projection, **kwargs): super().__init__() self.embedding = embedding self.positional_encoder = positional_encoder self.layer_stack = layer_stack self.output_projection = output_projection self.d_model = d_model self.sqrt_d_model = d_model**0.5 self.reset_parameters() def forward( self, input_ids, position_ids, output_attentions, gradient_checkpointing_func, past_key_values, use_cache, output_hidden_states, ): outputs = self.layer_stack( self.positional_encoder(self.embedding(input_ids) * self.sqrt_d_model, position_ids), output_attentions=output_attentions, gradient_checkpointing_func=gradient_checkpointing_func, past_key_values=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states, ) # Translate output states to logits. outputs["logits"] = self.output_projection(outputs["last_hidden_state"]) del outputs["last_hidden_state"] return outputs def reset_parameters(self): init.xavier_uniform_(self.output_projection.weight) init.constant_(self.output_projection.bias, 0.) init.normal_(self.embedding.weight, std=self.d_model**-0.5) # Converts a torch array of integers into their equivalent binary codes. def binary_tensor(x, bits): mask = 2**torch.arange(bits).to(x.device, x.dtype) return x.unsqueeze(-1).bitwise_and(mask).ne(0).byte() def hadamard_walsh_matrix(k: int): # k: The dimension of the matrix is 2^k assert k > 0 # Start with Hadamard H2^1 matrix. h1 = torch.tensor([[1, 1], [1, -1]], dtype=torch.float) # The series of matrices can be computed by recurisvely applying the Kronecker product, # starting with h1. # # This will produce the series of Hadamard-Wlash matrices in natural order. w = h1 for _ in range(k-1): w = torch.kron(h1, w) return w # This positional encoder adds absolute binary positions to the embedding, encoded via # Hadamard-Walsh matrix. # See: https://en.wikipedia.org/wiki/Hadamard_code # Each bit in the binary code word is encoded via a row the Hadamard-Walsh matrix, with a # 1 being encoded by the presense of the row and a 0 by its absence. While training, the base # sequence offset is randomly selected, which appears to allow the model to generalize to # sequences longer than it was trained on. This is similar to what is described here: # https://arxiv.org/pdf/2305.16843.pdf # I have tried this approach and found that my approach works better for generalization. # # Note: Without random shifting, the early performance of this encoder is exceptionally good. # The drawback is that the model can't generalize to longer sequences than it was trained on # and can't easily accomidate additonal bits later in the training process. class RSWalshPositionalEncoder(nn.Module): def __init__(self, d_embed, max_seq, gain=0.333): super().__init__() self.max_seq = max_seq self.d_embed = d_embed # Hadamard-Walsh k, where the dimension of the matrix is 2^k k = math.ceil(math.log2(d_embed)) # The number of bits required to encode max_seq bits = math.ceil(math.log2(max_seq)) # Gain controls the weight given to the encodings. # When a trainable parameter, the value appears to settle at around 0.333. self.gain = gain assert bits <= d_embed, "max_seq exceeds n-bits available for d_embed" # Generate sequential binary codes for absolute positionals. # The implementation originally used Grey codes, which where successive symbols # differ by by only one bit. See: https://en.wikipedia.org/wiki/Gray_code # This, along with a few other coding schemes were tested, with a simple # binary code having the best performance. binary_code = binary_tensor(torch.arange(0, max_seq, 1), bits) self.register_buffer('binary_code', binary_code, persistent=False) # Each bit is encoded via a row of a Hadamard-Walsh matrix. # We slice off the unused rows and columns -- ideally, d_embed should be # the same dimension as the matrix. walsh = hadamard_walsh_matrix(k)[:bits,:d_embed] * self.gain # This alternative appears superior to the original. # If starting from scratch, this use this. # walsh = (hadamard_walsh_matrix(k)[:bits,:d_embed] -0.5) * self.gain self.register_buffer('walsh', walsh, persistent=False) def forward(self, x, position_ids=None): seq_len = x.size(-2) # Get sequence of binary codes... # We use a random base offset when training. # This results in slower initial gains, but appears to allow the model to generalize to # the value of max_seq, even if never trained with sequences of this length. I also have # a suspicion that this has a regularizing effect on training, similar to dropout. Models with # random base offset shifting, despite slower initial improvement, appear to perform better in the long-run. # TODO: Setup a controlled experiment to test this hypothesis. if self.training: shift = torch.randint(self.max_seq - seq_len + 1, (1,)).item() seq = self.binary_code[shift:seq_len + shift,:] # When the cache is used for generation, after the first call, we are only passed a single token at a time, # with the remaining tokens being in the cache. We need to make sure that the newly injected tokens have the # correct relative position by indexing the codes with the position_ids. elif position_ids != None: seq = self.binary_code[position_ids, :] # Disable shifting when not training. This does not appear to change the evaluation loss, but # it does makes predictions easier to analyse when the attention weights are not shifting with each step. else: seq = self.binary_code[:seq_len,:] # For reasons I have yet to identify, when the model is running in Textgenwebui, the matrix appears # to evade conversion to bfloat16, despite everything else having been converted. # This is a work-around for this. self.walsh = self.walsh.to(dtype=x.dtype) # Encode binary sequence with Hadamard-Walsh codes and apply to embeddings. # If nothing else, the Walsh encodings make the positional information exceptionally # robust with respect to dropout and other adversities. They can still be easily detected # at the final layer. return x + (seq.to(dtype=x.dtype) @ self.walsh) # A generic stack of transformer layers. class TransformerLayerStack(nn.Module): def __init__(self, layers): super().__init__() self.layers = layers def forward( self, hidden_states, output_attentions, past_key_values, use_cache, output_hidden_states, gradient_checkpointing_func=None, ): present_key_value = None all_attentions = [] if output_attentions else None all_hidden_states = [hidden_states] if output_hidden_states else None for layer in self.layers: if gradient_checkpointing_func is not None: layer_outputs = gradient_checkpointing_func( layer.__call__, hidden_states, output_attentions, past_key_values, use_cache, use_reentrant=False, ) else: layer_outputs = layer( hidden_states, output_attentions, past_key_values, use_cache, ) hidden_states = layer_outputs["hidden_states"] if output_hidden_states: all_hidden_states.append(hidden_states) if use_cache: present_key_value = layer_outputs["past_key_values"] if output_attentions: all_attentions.append(layer_outputs["attentions"]) return dict( last_hidden_state=hidden_states, past_key_values=present_key_value, hidden_states=hidden_states, attentions=all_attentions, ) # DeepNet: Scaling Transformers to 1,000 Layers # https://arxiv.org/abs/2203.00555 # Note: This is a type of Pre-Layer-Norm Transformer layer. class DeepnetLayer(nn.Module): def __init__( self, d_model, attention, feedforward, norm1, norm2, dropout, layer_idx, alpha=1.0, ): super().__init__() self.d_model = d_model self.attention = attention self.feedforward = feedforward self.norm1 = norm1 self.norm2 = norm2 self.dropout = dropout # Deepnet alpha self.alpha = alpha self.layer_idx = layer_idx def forward( self, hidden_states, output_attentions, past_key_values, use_cache, ): # Keep input as residual residual = hidden_states * self.alpha # Compute attention attn_outputs = self.attention( hidden_states, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions ) hidden_states = attn_outputs["hidden_states"] # Add attention with residual and normalize. hidden_states = self.norm1(residual + self.dropout(hidden_states)) # Keep output as next residual. residual = hidden_states * self.alpha # Pass through feedforward network. hidden_states = self.feedforward(hidden_states) # Combine residual and ff output, then normalize again. hidden_states = self.norm2(residual + self.dropout(hidden_states)) return dict( hidden_states=hidden_states, attentions=attn_outputs["attentions"], past_key_values=attn_outputs["past_key_values"] ) # A vanilla MLP transfomer layer. class FeedforwardLayer(nn.Module): def __init__( self, d_model: int, feedforward_dim: int, dropout, layer_idx, activation=nn.ReLU(), beta=1.0, bias=True, ): super().__init__() self.d_model = d_model self.beta = beta self.activation = activation self.linear1 = nn.Linear(d_model, feedforward_dim, bias=bias) self.linear2 = nn.Linear(feedforward_dim, d_model, bias=bias) self.dropout = nn.Dropout(dropout) self.reset_parameters() def forward(self, x): return self.linear2(self.dropout(self.activation(self.linear1(x)))) def reset_parameters(self): init.xavier_uniform_(self.linear1.weight, gain=self.beta) init.xavier_uniform_(self.linear2.weight, gain=self.beta) init.constant_(self.linear1.bias, 0.) init.constant_(self.linear2.bias, 0.) class CausalSelfAttention(nn.Module): def __init__( self, d_model, num_heads, # values: # native: Use local impementation; slowest option; good for debugging; useful when experimenting with non-standard stuff. # torch: Use pytorch "scaled_dot_product_attention()"; faster; generally good compatibility; does not support returning attn weights. # flash2: Use Flash-Attention2 implementation; fastest; limited to int16 and bfloat16 types; least memory usage. attn_type, layer_idx, config, beta=1.0, dropout=0.1, ): super().__init__() self.d_model = d_model self.num_heads = num_heads self.beta = beta self.attn_type = attn_type self.layer_idx = layer_idx self.config = config assert d_model % num_heads == 0, "d_model must be evenly divisible by num_heads" # The dimension of each head. self.d_head = d_model // num_heads # We scale the attention scores by the inverse-square-root of the head dimension # this shifts the temerature of softmax. self.dot_product_scale = 1.0 / math.sqrt(self.d_head) self.in_proj = nn.Linear(self.d_model, 3 * self.d_model, bias=True) self.output_linear = nn.Linear(self.d_model, self.d_model, bias=True) self.dropout = nn.Dropout(dropout) self.reset_parameters() def extra_repr(self) -> str: return f'd_model={self.d_model}, num_heads={self.num_heads}, beta={self.beta}, attn_type={self.attn_type}, dropout={self.dropout}' def reset_parameters(self): # Deepnet initialization # https://arxiv.org/pdf/2203.00555.pdf q, k, v = self.in_proj.weight.chunk(3) init.xavier_uniform_(q, gain=1.0) init.xavier_uniform_(k, gain=1.0) init.xavier_uniform_(v, gain=self.beta) init.xavier_uniform_(self.output_linear.weight, gain=self.beta) init.constant_(self.in_proj.bias, 0.) init.constant_(self.output_linear.bias, 0.) # Project QKV input through input matrices, reshape to (batch_size, n_heads, seq_len, d_model), and apply cache. def project_input(self, qkv, past_key_values): batch_size, seq_len, d_embed = qkv.shape proj = self.in_proj(qkv) query, key, value = proj.chunk(chunks=3, dim=-1) # Split projections into multiple heads and swap position of sequence / heads dimension query = query.view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2) key = key.view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2) value = value.view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2) # Update the cache values. if past_key_values is not None: key, value = past_key_values.update(key, value, self.layer_idx) return query, key, value def forward( self, qkv, output_attentions, past_key_values, use_cache, ): attn_type = self.attn_type if output_attentions and attn_type != "native": logger.warning_once( "CausalSelfAttention(output_attentions=True) and attn_type is not 'native': " "Forcing native attention." ) attn_type = "native" if attn_type == "flash2": if use_cache is None or use_cache == False: return self.flash2_forward(qkv) else: return self.flash2_forward_cached(qkv, past_key_values) # qkv: (batch_size, seq_len, d_embed) batch_size, seq_len, d_embed = qkv.shape # Feed the inputs through the K, Q, V matrices. query, key, value = self.project_input(qkv, past_key_values) kv_seq_len = key.shape[-2] # Default to returning empty attention weights. attentions = None # https://github.com/pytorch/pytorch/issues/112577 if attn_type == "torch": # This context manager can be used to force which implementation to use. #with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False): attended_values = F.scaled_dot_product_attention( query, key, value, attn_mask=None, dropout_p=self.dropout.p if self.training else 0.0, is_causal=(seq_len > 1), scale=self.dot_product_scale ) # "native" scaled-dot-product attention implementation. else: # Compute attention scores scores = torch.matmul(query, key.transpose(-2, -1)) * self.dot_product_scale # Mask future positions from the past if seq_len > 1: scores.masked_fill_( torch.tril( torch.ones(seq_len, kv_seq_len, dtype=torch.bool, device=qkv.device), diagonal=0, ).logical_not(), float('-inf'), ) # Calculate the attention weights; avoid NANs that might emerge from zeros in softmax's denominator attentions = self.dropout(torch.softmax(scores, dim=-1).clamp(min=1e-10)) del scores # Use the attention weights to get a weighted combination of value vectors attended_values = torch.matmul(attentions, value) if not output_attentions: del attentions attentions = None # Concatenate attention heads and project to original embedding size using the output linear layer attended_values = attended_values.transpose(1, 2).contiguous().view(batch_size, seq_len, d_embed) # Project the concatenated output through the output matrix. attended_values = self.output_linear(attended_values) return dict( hidden_states=attended_values, attentions=attentions, past_key_values=past_key_values ) # No cache support, but faster def flash2_forward( self, qkv, ): batch_size, seq_len, d_embed = qkv.shape # Feed the inputs through the K, Q, V matrices. # query : (batch_size, seq_len, d_model) # qkv : (batch_size, seq_len, 3, num_heads, d_kq) # Feed the inputs through the K, Q, V matrices. # query : (batch_size, seq_len, d_model) # qkv : (batch_size, seq_len, 3, num_heads, d_kq) qkv = self.in_proj(qkv).unflatten( -1, (3, self.num_heads, self.d_head) ) attended_values = flash_attn_qkvpacked_func( self._downcast_to_float16(qkv)[0], dropout_p=self.dropout.p if self.training else 0.0, softmax_scale=self.dot_product_scale, causal=True, ) # attended_values: (batch_size, seqlen, nheads, headdim) # Concatentate heads back into d_embed attended_values = attended_values.view(batch_size, seq_len, d_embed) # Project the concatenated output through the output matrix. attended_values = self.output_linear(attended_values) return dict( hidden_states=attended_values, attentions=None, past_key_values=None ) # See https://github.com/huggingface/transformers/blob/main/src/transformers/cache_utils.py #https://huggingface.co/docs/transformers/internal/generation_utils def flash2_forward_cached( self, qkv, past_key_values, ): batch_size, seq_len, d_embed = qkv.shape # Feed the inputs through the K, Q, V matrices. query, key, value = self.project_input(qkv, past_key_values) query, key, value = self._downcast_to_float16(query, key, value) # Expected inputs to flash2: # q: (batch_size, seqlen, nheads, headdim) # k: (batch_size, seqlen, nheads_k, headdim) # v: (batch_size, seqlen, nheads_k, headdim) query = query.transpose(1, 2) key = key.transpose(1, 2) value = value.transpose(1, 2) attended_values = flash_attn_func( q=query, k=key, v=value, dropout_p=self.dropout.p if self.training else 0.0, softmax_scale=self.dot_product_scale, causal=True, ) # attended_values: (batch_size, seqlen, nheads, headdim) # Concatentate heads back into d_embed attended_values = attended_values.view(batch_size, seq_len, d_embed) # Project the concatenated output through the output matrix. attended_values = self.output_linear(attended_values) return dict( hidden_states=attended_values, attentions=None, past_key_values=past_key_values ) def _downcast_to_float16(self, *args): if args[0].dtype != torch.float32: return args if torch.is_autocast_enabled(): target_dtype = torch.get_autocast_gpu_dtype() # Handle the case where the model is quantized elif hasattr(self.config, "_pre_quantization_dtype"): target_dtype = self.config._pre_quantization_dtype else: target_dtype = self.output_linear.weight.dtype logger.warning_once( f"The input hidden states seems to be silently casted in float32, this might be related to" f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" f" {target_dtype}." ) return (arg.to(target_dtype) for arg in args)