初始化项目，由ModelHub XC社区提供模型
Model: bench-labs/cagliostro-v1 Source: Original Platform
This commit is contained in:
ModelHub XC
170
model.py
Normal file
170
model.py
Normal file
@@ -0,0 +1,170 @@
|
||||
from __future__ import annotations
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
|
||||
def __init__(self, dim, eps=1e-06):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def forward(self, x):
|
||||
dtype = x.dtype
|
||||
x = x.float()
|
||||
x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||||
return x.to(dtype) * self.weight
|
||||
|
||||
def rope_freqs(seq_len, head_dim, theta=100000.0, device=None):
|
||||
inv_freq = 1.0 / theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)
|
||||
t = torch.arange(seq_len, device=device).float()
|
||||
freqs = torch.outer(t, inv_freq)
|
||||
return torch.cat([freqs, freqs], dim=-1)
|
||||
|
||||
def rotate_half(x):
|
||||
x1, x2 = x.chunk(2, dim=-1)
|
||||
return torch.cat([-x2, x1], dim=-1)
|
||||
|
||||
def apply_rope(x, cos, sin):
|
||||
return x * cos + rotate_half(x) * sin
|
||||
|
||||
class Attention(nn.Module):
|
||||
|
||||
def __init__(self, dim, n_heads, n_kv_heads, head_dim, qk_norm_eps=1e-06):
|
||||
super().__init__()
|
||||
if n_heads % n_kv_heads != 0:
|
||||
raise ValueError(f'n_heads ({n_heads}) must be divisible by n_kv_heads ({n_kv_heads}); valid choices for {n_heads} heads: {[k for k in range(1, n_heads + 1) if n_heads % k == 0]}')
|
||||
self.n_heads = n_heads
|
||||
self.n_kv_heads = n_kv_heads
|
||||
self.head_dim = head_dim
|
||||
self.n_rep = n_heads // n_kv_heads
|
||||
self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
|
||||
self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
|
||||
self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
|
||||
self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
|
||||
self.q_norm = RMSNorm(head_dim, eps=qk_norm_eps)
|
||||
self.k_norm = RMSNorm(head_dim, eps=qk_norm_eps)
|
||||
|
||||
def forward(self, x, cos, sin):
|
||||
B, T, _ = x.shape
|
||||
q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim)
|
||||
k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim)
|
||||
v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim)
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
q = apply_rope(q, cos, sin)
|
||||
k = apply_rope(k, cos, sin)
|
||||
q = q.transpose(1, 2)
|
||||
k = k.transpose(1, 2).repeat_interleave(self.n_rep, dim=1)
|
||||
v = v.transpose(1, 2).repeat_interleave(self.n_rep, dim=1)
|
||||
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
|
||||
out = out.transpose(1, 2).contiguous().view(B, T, -1)
|
||||
return self.o_proj(out)
|
||||
|
||||
class SwiGLU(nn.Module):
|
||||
|
||||
def __init__(self, dim, hidden):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(dim, hidden, bias=False)
|
||||
self.up_proj = nn.Linear(dim, hidden, bias=False)
|
||||
self.down_proj = nn.Linear(hidden, dim, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
|
||||
def __init__(self, dim, n_heads, n_kv_heads, head_dim, mlp_hidden, norm_eps=1e-06):
|
||||
super().__init__()
|
||||
self.attn_norm = RMSNorm(dim, eps=norm_eps)
|
||||
self.attn = Attention(dim, n_heads, n_kv_heads, head_dim)
|
||||
self.mlp_norm = RMSNorm(dim, eps=norm_eps)
|
||||
self.mlp = SwiGLU(dim, mlp_hidden)
|
||||
|
||||
def forward(self, x, cos, sin):
|
||||
x = x + self.attn(self.attn_norm(x), cos, sin)
|
||||
x = x + self.mlp(self.mlp_norm(x))
|
||||
return x
|
||||
|
||||
class LogosModel(nn.Module):
|
||||
|
||||
def __init__(self, vocab_size=32768, dim=640, n_layers=30, n_heads=10, n_kv_heads=5, mlp_hidden=1728, max_seq_len=4096, rope_theta=100000.0, norm_eps=1e-06, tie_embeddings=True):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.n_layers = n_layers
|
||||
self.head_dim = dim // n_heads
|
||||
self.max_seq_len = max_seq_len
|
||||
self.rope_theta = rope_theta
|
||||
self.embed_tokens = nn.Embedding(vocab_size, dim)
|
||||
self.layers = nn.ModuleList([DecoderLayer(dim, n_heads, n_kv_heads, self.head_dim, mlp_hidden, norm_eps) for _ in range(n_layers)])
|
||||
self.norm_out = RMSNorm(dim, eps=norm_eps)
|
||||
self.lm_head = nn.Linear(dim, vocab_size, bias=False)
|
||||
if tie_embeddings:
|
||||
self.lm_head.weight = self.embed_tokens.weight
|
||||
self.apply(self._init_weights)
|
||||
for layer in self.layers:
|
||||
nn.init.normal_(layer.mlp.down_proj.weight, mean=0.0, std=0.02 / math.sqrt(2 * n_layers))
|
||||
nn.init.normal_(layer.attn.o_proj.weight, mean=0.0, std=0.02 / math.sqrt(2 * n_layers))
|
||||
cos, sin = self._build_rope_cache(max_seq_len)
|
||||
self.register_buffer('rope_cos', cos, persistent=False)
|
||||
self.register_buffer('rope_sin', sin, persistent=False)
|
||||
|
||||
def _init_weights(self, m):
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, mean=0.0, std=0.02)
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
elif isinstance(m, nn.Embedding):
|
||||
nn.init.normal_(m.weight, mean=0.0, std=0.02)
|
||||
|
||||
def _build_rope_cache(self, seq_len):
|
||||
freqs = rope_freqs(seq_len, self.head_dim, self.rope_theta)
|
||||
return (freqs.cos()[None, :, None, :], freqs.sin()[None, :, None, :])
|
||||
|
||||
def forward(self, input_ids, labels=None, loss_chunk_size=2048):
|
||||
B, T = input_ids.shape
|
||||
x = self.embed_tokens(input_ids)
|
||||
cos = self.rope_cos[:, :T].to(x.dtype)
|
||||
sin = self.rope_sin[:, :T].to(x.dtype)
|
||||
for layer in self.layers:
|
||||
x = layer(x, cos, sin)
|
||||
x = self.norm_out(x)
|
||||
if labels is None:
|
||||
return (self.lm_head(x), None)
|
||||
shift_x = x[:, :-1].reshape(-1, x.size(-1))
|
||||
shift_labels = labels[:, 1:].reshape(-1)
|
||||
n = shift_x.size(0)
|
||||
total_loss = x.new_zeros((), dtype=torch.float32)
|
||||
total_count = x.new_zeros((), dtype=torch.float32)
|
||||
for start in range(0, n, loss_chunk_size):
|
||||
end = min(start + loss_chunk_size, n)
|
||||
chunk_labels = shift_labels[start:end]
|
||||
valid = chunk_labels != -100
|
||||
count = valid.sum()
|
||||
if count == 0:
|
||||
continue
|
||||
chunk_logits = self.lm_head(shift_x[start:end]).float()
|
||||
chunk_loss = F.cross_entropy(chunk_logits, chunk_labels, ignore_index=-100, reduction='sum')
|
||||
total_loss = total_loss + chunk_loss
|
||||
total_count = total_count + count
|
||||
loss = total_loss / total_count.clamp(min=1)
|
||||
return (None, loss)
|
||||
|
||||
def num_params(self, exclude_embeddings=False):
|
||||
n = sum((p.numel() for p in self.parameters()))
|
||||
if exclude_embeddings:
|
||||
n -= self.embed_tokens.weight.numel()
|
||||
if self.lm_head.weight is not self.embed_tokens.weight:
|
||||
n -= self.lm_head.weight.numel()
|
||||
return n
|
||||
if __name__ == '__main__':
|
||||
m = LogosModel()
|
||||
print(f'total params: {m.num_params():,}')
|
||||
print(f'non-embedding params: {m.num_params(exclude_embeddings=True):,}')
|
||||
x = torch.randint(0, 32768, (2, 128))
|
||||
logits, loss = m(x, labels=x)
|
||||
print('loss:', loss.item() if loss is not None else None)
|
||||
logits, _ = m(x)
|
||||
print('logits shape (no labels):', logits.shape)
|
||||
Reference in New Issue
Block a user