246 lines
7.6 KiB
Python
246 lines
7.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from transformers import PreTrainedModel, GenerationMixin
|
|
from transformers.modeling_outputs import CausalLMOutputWithPast
|
|
|
|
from .configuration_veyra import VeyraConfig
|
|
|
|
|
|
class RMSNorm(nn.Module):
|
|
def __init__(self, dim, eps=1e-6):
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.ones(dim))
|
|
self.eps = eps
|
|
|
|
def forward(self, x):
|
|
return self.weight * x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
|
|
|
|
|
|
class RotaryEmbedding(nn.Module):
|
|
def __init__(self, dim, max_seq_len, base=10000.0):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.max_seq_len = max_seq_len
|
|
self.base = base
|
|
|
|
def forward(self, x, positions):
|
|
# No buffers. Compute everything fresh to avoid HF remote-code buffer corruption.
|
|
half_dim = self.dim // 2
|
|
|
|
inv_freq = 1.0 / (
|
|
self.base ** (
|
|
torch.arange(0, self.dim, 2, device=x.device, dtype=torch.float32) / self.dim
|
|
)
|
|
)
|
|
|
|
positions = positions.to(device=x.device, dtype=torch.float32)
|
|
freqs = torch.outer(positions, inv_freq)
|
|
|
|
cos = torch.cos(freqs).to(dtype=x.dtype)[None, None, :, :]
|
|
sin = torch.sin(freqs).to(dtype=x.dtype)[None, None, :, :]
|
|
|
|
x_even = x[..., 0::2]
|
|
x_odd = x[..., 1::2]
|
|
|
|
out_even = x_even * cos - x_odd * sin
|
|
out_odd = x_even * sin + x_odd * cos
|
|
|
|
return torch.stack((out_even, out_odd), dim=-1).flatten(-2)
|
|
|
|
|
|
class SwiGLU(nn.Module):
|
|
def __init__(self, d_model, intermediate_size):
|
|
super().__init__()
|
|
self.gate_proj = nn.Linear(d_model, intermediate_size, bias=False)
|
|
self.up_proj = nn.Linear(d_model, intermediate_size, bias=False)
|
|
self.down_proj = nn.Linear(intermediate_size, d_model, bias=False)
|
|
|
|
def forward(self, x):
|
|
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
|
|
|
|
|
|
class GQAAttention(nn.Module):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
|
|
assert config.d_model % config.n_q_heads == 0
|
|
assert config.n_q_heads % config.n_kv_heads == 0
|
|
|
|
self.d_model = config.d_model
|
|
self.n_q_heads = config.n_q_heads
|
|
self.n_kv_heads = config.n_kv_heads
|
|
self.head_dim = config.d_model // config.n_q_heads
|
|
self.n_rep = config.n_q_heads // config.n_kv_heads
|
|
|
|
self.q_proj = nn.Linear(config.d_model, config.n_q_heads * self.head_dim, bias=False)
|
|
self.k_proj = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=False)
|
|
self.v_proj = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=False)
|
|
self.o_proj = nn.Linear(config.n_q_heads * self.head_dim, config.d_model, bias=False)
|
|
|
|
self.rope = RotaryEmbedding(self.head_dim, config.max_seq_len, config.rope_theta)
|
|
|
|
def _shape_q(self, x):
|
|
b, t, _ = x.shape
|
|
return x.view(b, t, self.n_q_heads, self.head_dim).transpose(1, 2)
|
|
|
|
def _shape_kv(self, x):
|
|
b, t, _ = x.shape
|
|
return x.view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)
|
|
|
|
def _repeat_kv(self, x):
|
|
if self.n_rep == 1:
|
|
return x
|
|
return x.repeat_interleave(self.n_rep, dim=1)
|
|
|
|
def forward(self, x):
|
|
b, t, _ = x.shape
|
|
positions = torch.arange(t, device=x.device)
|
|
|
|
q = self._shape_q(self.q_proj(x))
|
|
k = self._shape_kv(self.k_proj(x))
|
|
v = self._shape_kv(self.v_proj(x))
|
|
|
|
q = self.rope(q, positions)
|
|
k = self.rope(k, positions)
|
|
|
|
k = self._repeat_kv(k)
|
|
v = self._repeat_kv(v)
|
|
|
|
y = F.scaled_dot_product_attention(
|
|
q,
|
|
k,
|
|
v,
|
|
attn_mask=None,
|
|
dropout_p=0.0,
|
|
is_causal=True,
|
|
)
|
|
|
|
y = y.transpose(1, 2).contiguous().view(b, t, self.d_model)
|
|
|
|
return self.o_proj(y)
|
|
|
|
|
|
class AttentionBlock(nn.Module):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.input_norm = RMSNorm(config.d_model, config.rms_norm_eps)
|
|
self.attn = GQAAttention(config)
|
|
self.post_attn_norm = RMSNorm(config.d_model, config.rms_norm_eps)
|
|
self.mlp = SwiGLU(config.d_model, config.intermediate_size)
|
|
|
|
def forward(self, x):
|
|
x = x + self.attn(self.input_norm(x))
|
|
x = x + self.mlp(self.post_attn_norm(x))
|
|
return x
|
|
|
|
|
|
class MLPOnlyBlock(nn.Module):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.norm = RMSNorm(config.d_model, config.rms_norm_eps)
|
|
self.mlp = SwiGLU(config.d_model, config.intermediate_size)
|
|
|
|
def forward(self, x):
|
|
return x + self.mlp(self.norm(x))
|
|
|
|
|
|
class VeyraPreTrainedModel(PreTrainedModel):
|
|
config_class = VeyraConfig
|
|
base_model_prefix = ""
|
|
supports_gradient_checkpointing = False
|
|
_no_split_modules = ["AttentionBlock", "MLPOnlyBlock"]
|
|
|
|
|
|
class VeyraForCausalLM(VeyraPreTrainedModel, GenerationMixin):
|
|
config_class = VeyraConfig
|
|
_tied_weights_keys = {"lm_head.weight": "token_emb.weight"}
|
|
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
self.config = config
|
|
|
|
self.token_emb = nn.Embedding(config.vocab_size, config.d_model)
|
|
|
|
blocks = []
|
|
for kind in config.layer_pattern:
|
|
if kind == "A":
|
|
blocks.append(AttentionBlock(config))
|
|
elif kind == "M":
|
|
blocks.append(MLPOnlyBlock(config))
|
|
else:
|
|
raise ValueError(f"Unknown layer kind: {kind}")
|
|
|
|
self.blocks = nn.ModuleList(blocks)
|
|
self.norm = RMSNorm(config.d_model, config.rms_norm_eps)
|
|
self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
|
|
|
|
if config.tie_word_embeddings:
|
|
self.lm_head.weight = self.token_emb.weight
|
|
|
|
# Needed by newer Transformers internals, but we intentionally avoid post_init().
|
|
self.all_tied_weights_keys = {"lm_head.weight": "token_emb.weight"}
|
|
|
|
def tie_weights(self, *args, **kwargs):
|
|
if self.config.tie_word_embeddings:
|
|
self.lm_head.weight = self.token_emb.weight
|
|
|
|
def get_input_embeddings(self):
|
|
return self.token_emb
|
|
|
|
def set_input_embeddings(self, value):
|
|
self.token_emb = value
|
|
if self.config.tie_word_embeddings:
|
|
self.lm_head.weight = self.token_emb.weight
|
|
|
|
def get_output_embeddings(self):
|
|
return self.lm_head
|
|
|
|
def set_output_embeddings(self, new_embeddings):
|
|
self.lm_head = new_embeddings
|
|
|
|
def forward(
|
|
self,
|
|
input_ids=None,
|
|
attention_mask=None,
|
|
labels=None,
|
|
return_dict=True,
|
|
**kwargs,
|
|
):
|
|
if input_ids is None:
|
|
raise ValueError("input_ids must be provided")
|
|
|
|
x = self.token_emb(input_ids)
|
|
|
|
for block in self.blocks:
|
|
x = block(x)
|
|
|
|
logits = self.lm_head(self.norm(x))
|
|
|
|
loss = None
|
|
if labels is not None:
|
|
loss = F.cross_entropy(
|
|
logits[:, :-1, :].contiguous().view(-1, logits.size(-1)),
|
|
labels[:, 1:].contiguous().view(-1),
|
|
ignore_index=-100,
|
|
)
|
|
|
|
if not return_dict:
|
|
output = (logits,)
|
|
return ((loss,) + output) if loss is not None else output
|
|
|
|
return CausalLMOutputWithPast(
|
|
loss=loss,
|
|
logits=logits,
|
|
past_key_values=None,
|
|
hidden_states=None,
|
|
attentions=None,
|
|
)
|
|
|
|
def prepare_inputs_for_generation(self, input_ids, **kwargs):
|
|
max_len = self.config.max_seq_len
|
|
if input_ids.shape[1] > max_len:
|
|
input_ids = input_ids[:, -max_len:]
|
|
return {"input_ids": input_ids}
|