walsh-1-7b / modelling_walsh.py
dinalt's picture
Update modelling_walsh.py
35354b2 verified
# 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)