EvaByte-SFT / eva_pt_ref.py
linzheng's picture
Upload EvaByteForCausalLM
474addc verified
raw
history blame
18.1 kB
from typing import Optional, Tuple, Union
import torch
from torch import nn
MASK_MIN_VALUE = -10e10
def rotate_half(x: torch.Tensor) -> torch.Tensor:
"""
Rotates half the hidden dims (last dim) of the input.
Args:
x: Rotary embedded tensor
Return:
Tensor with half of last dim negated and rotated to the front.
"""
x1, x2 = x.split(x.shape[-1] // 2, dim=-1)
return torch.cat((-x2, x1), dim=-1)
def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
position_ids: torch.Tensor) -> torch.Tensor:
"""
Apply rotary embedding (cos, sin) to the query and key tensor on the sequence dimension.
The legends for dimensions are defined as:
num_heads: number of attention heads
current_seq_len: the current batch's sequence length, should be either 1 or max_seq_len
max_seq_len: the static sequence length, different from current_seq_len in cached inference case where it is always
maximum lenghth, e.g. the length of static sequence length of KV cache
Args:
q: Query tensor, of size (batch_size, num_heads, current_seq_len, head_dim)
k: Key tensor, of size (batch_size, num_key_value_heads, current_seq_len, head_dim)
cos: Cosine base of rotary embedding, of size (max_seq_len, head_dim)
sin: Sine base of rotary embedding, of size (max_seq_len, head_dim)
position_ids: The position indices of the tokens corresponding to the query and key tensors. It has a size of
(batch_size, current_seq_len).
Returns:
Embedded query and key tensor of same size as input.
"""
bs, nheads, cur_seq_len, head_dim = q.shape
assert len(
k.shape) == 4, f"k should be of shape (batch_size, num_heads, current_seq_len, head_dim), got {k.shape} instead"
assert k.shape[0] == bs, f"k has a different batch_size {k.shape[0]} compared to q {bs}"
assert list(k.shape[2:]) == [cur_seq_len,
head_dim], f"k has different current_seq_len and/or head_dim compared to q"
assert cos.shape[3] == head_dim, f"cos should have dim of head dim {head_dim}, got {cos.shape[3]} instead"
assert list(position_ids.shape) in [[bs, cur_seq_len], [1, cur_seq_len]],\
f"position_ids should be of shape {[bs, cur_seq_len]} or {[1, cur_seq_len]}, got {position_ids.shape} instead"
q_embed = (q * cos) + (rotate_half(q) * sin)
k_embed = (k * cos) + (rotate_half(k) * sin)
return q_embed, k_embed
def attention_op(
q,
k,
v,
attn_mask,
mixedp_attn,
head_dim_scaling
):
attn = torch.matmul(q, k.transpose(-2, -1))
if mixedp_attn:
attn = attn.to(torch.float)
attn = attn * head_dim_scaling
if attn_mask is not None:
attn = attn.masked_fill(attn_mask, MASK_MIN_VALUE)
attn_weights = torch.softmax(attn, dim=-1).to(q.dtype)
attn_output = torch.matmul(attn_weights, v)
return attn_output
def prm_projection(
x: torch.Tensor,
projection_matrix: torch.Tensor,
mixedp_attn: bool = False
):
"""
Constructs nonnegative kernel features for fast softmax attention.
Args:
x: input for which features are computed
projection_matrix: random matrix used to compute features
Returns:
Random features for fast attention.
"""
# x : [..., m, d]
# proj : [..., r, d]
scaling_factor = (x.shape[-1] ** -0.5)
proj_x = torch.matmul(projection_matrix, x.transpose(-1, -2)) # [..., r, m]
norm = torch.sum(x ** 2, dim=-1).unsqueeze(-2) * 0.5 # [..., 1]
if mixedp_attn:
proj_x = proj_x.to(torch.float)
norm = norm.to(torch.float)
phi_x = scaling_factor * (proj_x - norm)
return phi_x
class EvaAttention(nn.Module):
def __init__(self, config, layer_idx: Optional[int] = None):
super().__init__()
self.config = config
self.layer_idx = layer_idx
self.hidden_size = config.hidden_size
self.num_heads = config.num_attention_heads
self.head_dim = self.hidden_size // self.num_heads
self.head_dim_scaling = self.head_dim ** -0.5
self.max_position_embeddings = config.max_position_embeddings
if (self.head_dim * self.num_heads) != self.hidden_size:
raise ValueError(
f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
f" and `num_heads`: {self.num_heads})."
)
self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
self.window_size = config.window_size
self.num_chunks = config.num_chunks
self.chunk_size = config.chunk_size
if self.chunk_size is not None:
assert self.window_size >= self.chunk_size and self.window_size % self.chunk_size == 0
# chunk_size overrides the number of landmarks
self.num_chunks = None
self.chunks_per_window = int(self.window_size // self.chunk_size)
self.random_feature_dim = 1
self.adaptive_phi = nn.Parameter(
torch.randn(
1,
self.num_heads,
1,
1,
self.head_dim
).clamp(-1., 1.) * self.head_dim_scaling
)
self.adaptive_mu_k = nn.Parameter(
torch.randn(
1,
self.num_heads,
1,
1,
self.head_dim
).clamp(-1., 1.) * self.head_dim_scaling
)
def _generate_feature_map(self, rf_q, rf_k, rf_v):
rf_k_logits = torch.sum(self.adaptive_mu_k.to(rf_k.dtype) * rf_k, dim=-1, keepdim=True) # b h c m 1
if self.config.mixedp_attn:
rf_k_logits = rf_k_logits.to(torch.float)
rf_k_weights = torch.softmax(rf_k_logits, dim=-2).to(rf_k.dtype)
rf_k_bar = torch.sum(rf_k_weights * rf_k, dim=-2)
weights = self.adaptive_phi.to(rf_k.dtype)
return weights, rf_k_bar
def _calculate_chunk_rfa_cache(self, rf_q, rf_k, rf_v, weights, rf_mask=None):
proj_x = torch.sum(weights * rf_k, dim=-1, keepdim=True)
norm = torch.sum(rf_k ** 2, dim=-1, keepdim=True) * 0.5 # [..., 1]
if self.config.mixedp_attn:
proj_x = proj_x.to(torch.float)
norm = norm.to(torch.float)
log_phi_k = self.head_dim_scaling * (proj_x - norm)
if rf_mask is not None:
log_phi_k = log_phi_k.masked_fill(rf_mask, MASK_MIN_VALUE)
# [b, h, c, m, r]
softmax_phi_k = torch.softmax(log_phi_k, dim=-2).to(rf_k.dtype)
softmax_phi_k_v = torch.sum(softmax_phi_k * rf_v, dim=-2)
# [b, h, c, r, m] [b, h, c, m, d] -> [b, h, c, r, d]
# softmax_phi_k_v = torch.matmul(softmax_phi_k.transpose(-1, -2), rf_v).squeeze(-2)
log_sum_phi_k = None
return softmax_phi_k_v, log_sum_phi_k
def _calculate_chunk_rfa(self, q, softmax_phi_k_v, log_sum_phi_k, weights):
if self.random_feature_dim == 1:
# when r = 1, the snis weights becomes 1, so this takes no effect
# [b, h, c, r, d] -> [b, h, c, d]
return softmax_phi_k_v
else:
# [b, h, c, r, d] [b, h, 1, s, d] -> [b, h, c, r, s]
log_phi_q = prm_projection(q.unsqueeze(-3), weights, self.config.mixedp_attn)
# [b, h, c, r, s] [b, h, c, r, 1] -> [b, h, c, r, s]
sniw = torch.softmax(log_phi_q + log_sum_phi_k, dim=-1).to(q.dtype)
# [b, h, c, r, s] [b, h, c, r, d] -> [b, h, c, s, d] -> [b, h, s, c, d]
rfa_per_chunk = torch.matmul(sniw.transpose(-1, -2), softmax_phi_k_v).transpose(-3, -2)
return rfa_per_chunk
def window_partition(self, x, window_size=None):
window_size = window_size if window_size is not None else self.window_size
gw, d = x.shape[-2:]
leading_dims = x.shape[:-2]
n_groups = gw // window_size
return x.reshape(*leading_dims, n_groups, window_size, d)
def window_merge(self, x, window_size=None):
g, w, d = x.shape[-3:]
leading_dims = x.shape[:-3]
return x.reshape(*leading_dims, g * w, d)
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_value: Optional[Tuple[torch.Tensor]] = None,
output_attentions: bool = False,
use_cache: bool = False,
cos: Optional[torch.Tensor] = None,
sin: Optional[torch.Tensor] = None,
multibyte_decoding: Optional[bool] = False,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
assert not output_attentions
bsz, q_len, _ = hidden_states.size()
############################################
# initialize past states if not provided
############################################
if use_cache and past_key_value is None:
raise ValueError
if use_cache and multibyte_decoding:
raise NotImplementedError("Multibyte decoding is not supported for PyTorch native implementation")
# assert isinstance(attention_mask, tuple)
if len(attention_mask) == 4:
assert use_cache
prev_causal_mask, cur_causal_mask, chunk_causal_mask, intra_chunk_mask = attention_mask
elif len(attention_mask) == 3:
assert not use_cache
window_causal_mask, chunk_causal_mask, intra_chunk_mask = attention_mask
else:
raise NotImplementedError("Only attention-mask tuple with length 2 or 3 is supported")
############################################
# compute q, k, v from hidden states
############################################
# [b, h, q_len, d]
q = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
# [b, h, kv_len, d]
k = self.k_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
# [b, h, kv_len, d]
v = self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
if use_cache:
past_key_value.update_past_len(q.shape[-2], self.layer_idx)
############################################
# apply rotary positional embeddings to q, k
############################################
q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
############################################
# compute q, k, v stats for the local window
############################################
if use_cache:
(prev_w_q, prev_w_k, prev_w_v), (cur_w_q, cur_w_k, cur_w_v) = past_key_value.update_singletons(
q,
k,
v,
self.layer_idx,
self.window_size,
self.singleton_update
)
else:
prev_w_q = self.window_partition(q) # [b, h, w, i, d]
prev_w_k = self.window_partition(k) # [b, h, w, j, d]
prev_w_v = self.window_partition(v) # [b, h, w, j, d]
# during training, we assume window_size divides seq_len so no remainders
cur_w_q = cur_w_k = cur_w_v = None
############################################
# compute q, k, v stats for chunk-level RFAs
############################################
if use_cache:
dump_q, dump_k, dump_v = past_key_value.update_chunks(q, k, v, self.layer_idx, self.chunk_size)
else:
dump_q, dump_k, dump_v = q, k, v
if use_cache:
prev_s_mask, cur_s_mask, prev_chunk_mask, cur_chunk_mask, dump_rf_mask = past_key_value.update_mask(
prev_s_mask=prev_causal_mask,
cur_s_mask=cur_causal_mask,
chunk_mask=chunk_causal_mask,
rf_mask=intra_chunk_mask,
layer_idx=self.layer_idx,
window_size=self.window_size,
chunk_size=self.chunk_size,
singleton_update=self.singleton_update
)
else:
prev_s_mask = window_causal_mask # [1, 1, w, i, j]
cur_s_mask = None
prev_chunk_mask = self.window_partition(chunk_causal_mask)
cur_chunk_mask = None
dump_rf_mask = intra_chunk_mask
if prev_s_mask.shape[-3] == 1:
# need to expand
prev_s_mask = prev_s_mask.expand(-1, -1, prev_chunk_mask.shape[-3], -1, -1)
if (
dump_q is not None and
dump_k is not None and
dump_v is not None
):
# [b, h, c, j, d]
rf_q = self.window_partition(dump_q, window_size=self.chunk_size)
# [b, h, c, j, d]
rf_k = self.window_partition(dump_k, window_size=self.chunk_size)
# [b, h, c, j, d]
rf_v = self.window_partition(dump_v, window_size=self.chunk_size)
if dump_rf_mask is not None:
rf_mask = self.window_partition(dump_rf_mask, window_size=self.chunk_size)
rf_q = rf_q.masked_fill(rf_mask, 0.)
rf_k = rf_k.masked_fill(rf_mask, 0.)
rf_v = rf_v.masked_fill(rf_mask, 0.)
else:
rf_mask = None
else:
rf_q = None
rf_k = None
rf_v = None
rf_mask = None
if rf_q is not None:
# import pdb; pdb.set_trace()
weights, rf_k_bar = self._generate_feature_map(rf_q, rf_k, rf_v)
softmax_phi_k_v, log_sum_phi_k = self._calculate_chunk_rfa_cache(rf_q, rf_k, rf_v, weights, rf_mask=rf_mask)
if use_cache:
softmax_phi_k_v, log_sum_phi_k, rf_k_bar = past_key_value.update_chunk_rfas(
softmax_phi_k_v, log_sum_phi_k, rf_k_bar, self.layer_idx, 1
)
elif use_cache:
weights = None
softmax_phi_k_v, log_sum_phi_k, rf_k_bar = past_key_value.get_chunk_rfas(self.layer_idx)
else:
weights = None
softmax_phi_k_v = None
log_sum_phi_k = None
rf_k_bar = None
if rf_k_bar is not None:
rfa_per_chunk = self._calculate_chunk_rfa(q, softmax_phi_k_v, log_sum_phi_k, weights)
############################################
# compute meta-attention weights for
# - group-wise RFAs and
# - singletons (equivalent to exact local attention)
############################################
if prev_w_k is not None:
if rf_k_bar is not None:
num_windows = prev_w_k.shape[-3]
# rf_k_bar and rfa_per_chunk take the shape [b, h, c, d]
# -> [b, h, 1, c, d] -> [b, h, w, c, d]
prev_rf_k_bar = rf_k_bar.unsqueeze(-3).expand(-1, -1, num_windows, -1, -1)
prev_rfa_per_chunk = rfa_per_chunk.unsqueeze(-3).expand(-1, -1, num_windows, -1, -1)
prev_agg_k = torch.cat([prev_w_k, prev_rf_k_bar], dim=-2)
prev_agg_v = torch.cat([prev_w_v, prev_rfa_per_chunk], dim=-2)
prev_attn_mask = torch.cat([prev_s_mask, prev_chunk_mask], dim=-1)
else:
prev_agg_k = prev_w_k
prev_agg_v = prev_w_v
prev_attn_mask = prev_s_mask
prev_attn_output = attention_op(
q=prev_w_q,
k=prev_agg_k,
v=prev_agg_v,
attn_mask=prev_attn_mask,
mixedp_attn=self.config.mixedp_attn,
head_dim_scaling=self.head_dim_scaling
)
prev_attn_output = self.window_merge(prev_attn_output)
if cur_w_k is not None:
if rf_k_bar is not None:
# rf_k_bar and rfa_per_chunk take the shape [b, h, c, d]
# cur_w_k and cur_w_v also has shape [b, h, m, d]
cur_agg_k = torch.cat([cur_w_k, rf_k_bar], dim=-2)
cur_agg_v = torch.cat([cur_w_v, rfa_per_chunk], dim=-2)
cur_attn_mask = torch.cat([cur_s_mask, cur_chunk_mask], dim=-1)
else:
cur_agg_k = cur_w_k
cur_agg_v = cur_w_v
cur_attn_mask = cur_s_mask
cur_attn_output = attention_op(
q=cur_w_q,
k=cur_agg_k,
v=cur_agg_v,
attn_mask=cur_attn_mask,
mixedp_attn=self.config.mixedp_attn,
head_dim_scaling=self.head_dim_scaling
)
if prev_w_k is not None and cur_w_k is not None:
attn_output = torch.cat([prev_attn_output, cur_attn_output], dim=-2)
elif prev_w_k is not None:
attn_output = prev_attn_output
elif cur_w_k is not None:
attn_output = cur_attn_output
else:
raise ValueError("There must be some bug")
if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
raise ValueError(
f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
f" {attn_output.size()}"
)
attn_output = attn_output.transpose(1, 2).reshape(bsz, q_len, self.hidden_size)
attn_output = self.o_proj(attn_output)
attn_weights = None
return attn_output, attn_weights, past_key_value