初始化项目，由ModelHub XC社区提供模型
Model: veyra-ai/Veyra-30M-Base Source: Original Platform
This commit is contained in:
ModelHub XC
245
modeling_veyra.py
Normal file
245
modeling_veyra.py
Normal file
@@ -0,0 +1,245 @@
|
||||
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}
|
||||
Reference in New Issue
Block a user