|
import math |
|
import torch |
|
import torch.nn as nn |
|
from transformers import PreTrainedModel, AutoModelForCausalLM |
|
from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions |
|
from transformers.file_utils import add_start_docstrings_to_model_forward |
|
from .configuration_spt import SPTConfig |
|
|
|
def repeat_kv(hidden_states, repeat_times): |
|
if repeat_times == 1: |
|
return hidden_states |
|
batch, n_kv_heads, seq_len, head_dim = hidden_states.shape |
|
hidden_states = hidden_states[:,:,None,:,:].expand(batch, n_kv_heads, repeat_times, seq_len, head_dim) |
|
return hidden_states.reshape(batch, n_kv_heads*repeat_times, seq_len, head_dim) |
|
|
|
class RMSNorm(nn.Module): |
|
def __init__(self, hidden_size, eps=1e-6): |
|
super().__init__() |
|
self.weight = nn.Parameter(torch.ones(hidden_size)) |
|
self.variance_epsilon = eps |
|
|
|
def forward(self, hidden_states): |
|
input_dtype = hidden_states.dtype |
|
hidden_states = hidden_states.to(torch.float32) |
|
variance = hidden_states.pow(2).mean(-1, keepdim=True) |
|
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) |
|
return self.weight * hidden_states.to(input_dtype) |
|
|
|
class Attention(nn.Module): |
|
def __init__(self, config): |
|
super().__init__() |
|
self.head_dim = config.hidden_size // config.n_attn_heads |
|
kv_size = config.n_kv_heads * self.head_dim |
|
self.hidden_size = config.hidden_size |
|
self.n_attn_heads = config.n_attn_heads |
|
self.n_kv_heads = config.n_kv_heads |
|
|
|
self.q = nn.Linear(config.hidden_size, config.hidden_size, bias=False) |
|
self.k = nn.Linear(config.hidden_size, kv_size, bias=False) |
|
self.v = nn.Linear(config.hidden_size, kv_size, bias=False) |
|
|
|
self.register_buffer('tril', torch.tril(torch.ones(config.max_len, config.max_len)).view(1, 1, config.max_len, config.max_len)) |
|
|
|
def forward(self, x): |
|
batch_size, seq_len, hidden_dim = x.shape |
|
|
|
q = self.q(x) |
|
k = self.k(x) |
|
v = self.v(x) |
|
|
|
q = q.view(batch_size, seq_len, self.n_attn_heads, self.head_dim).transpose(1, 2) |
|
k = k.view(batch_size, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) |
|
v = v.view(batch_size, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) |
|
|
|
k = repeat_kv(k, self.n_attn_heads//self.n_kv_heads) |
|
v = repeat_kv(v, self.n_attn_heads//self.n_kv_heads) |
|
|
|
attention = (q @ k.transpose(-2,-1)) * (1.0/math.sqrt(self.hidden_size)) |
|
attention = attention.masked_fill(self.tril[:,:,:seq_len,:seq_len]==0, float('-inf')) |
|
probs = nn.functional.softmax(attention,dim=-1) |
|
y = probs@v |
|
y = y.transpose(1,2).contiguous().reshape(batch_size, seq_len, -1) |
|
|
|
return y |
|
|
|
class MLP(nn.Module): |
|
def __init__(self, config): |
|
super().__init__() |
|
self.up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) |
|
self.gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) |
|
self.down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) |
|
self.act_fn = nn.GELU() |
|
|
|
def forward(self,x): |
|
up = self.up(x) |
|
gate = self.gate(x) |
|
return self.down(self.act_fn(up * gate)) |
|
|
|
class TransformerBlock(nn.Module): |
|
def __init__(self, config): |
|
super().__init__() |
|
self.attn = Attention(config) |
|
self.mlp = MLP(config) |
|
self.residual = config.residual |
|
self.norm = RMSNorm(config.hidden_size) if config.normalise else nn.Identity() |
|
|
|
def forward(self, x): |
|
if self.residual: |
|
x = x + self.attn(self.norm(x)) |
|
x = x + self.mlp(self.norm(x)) |
|
else: |
|
x = self.attn(self.norm(x)) |
|
x = self.mlp(self.norm(x)) |
|
return x |
|
|
|
class SPTModel(PreTrainedModel): |
|
config_class = SPTConfig |
|
|
|
def __init__(self, config): |
|
super().__init__(config) |
|
self.embedding = nn.Embedding(config.vocab_size, config.hidden_size) |
|
self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)]) |
|
self.norm = RMSNorm(config.hidden_size) if config.normalise else nn.Identity() |
|
|
|
def forward(self, input_ids): |
|
x = self.embedding(input_ids) |
|
for layer in self.layers: |
|
x = layer(x) |
|
x = self.norm(x) |
|
return x |
|
|
|
class SPTForCausalLM(PreTrainedModel): |
|
config_class = SPTConfig |
|
|
|
def __init__(self, config): |
|
super().__init__(config) |
|
self.model = SPTModel(config) |
|
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
|
|
|
def forward(self, input_ids, labels=None): |
|
x = self.model(input_ids) |
|
logits = self.lm_head(x) |
|
|
|
loss = None |
|
if labels is not None: |
|
loss_fct = nn.CrossEntropyLoss() |
|
loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1)) |
|
|
|
return CausalLMOutputWithCrossAttentions( |
|
loss=loss, |
|
logits=logits, |
|
hidden_states=x, |
|
) |
|
|
|
def prepare_inputs_for_generation(self, input_ids, **kwargs): |
|
return {"input_ids": input_ids} |
|
|
|
@staticmethod |
|
def _reorder_cache(past, beam_idx): |
|
return past |
|
|
|
|
|
AutoModelForCausalLM.register(SPTConfig, SPTForCausalLM) |
