| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| @dataclass |
| class SovynConfig: |
| name: str = "SOVYN-120M-Cortex" |
| vocab_size: int = 32000 |
| max_seq_len: int = 1024 |
| n_layers: int = 12 |
| hidden_size: int = 768 |
| n_heads: int = 12 |
| n_kv_heads: int = 4 |
| ffn_size: int = 2688 |
| dropout: float = 0.0 |
| rope_theta: float = 10000.0 |
| tie_embeddings: bool = True |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float = 1e-6): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dim)) |
| self.eps = eps |
|
|
| def forward(self, x): |
| normed = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) |
| return normed * self.weight |
|
|
|
|
| def precompute_rope(head_dim: int, max_seq_len: int, theta: float): |
| inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim)) |
| t = torch.arange(max_seq_len).float() |
| freqs = torch.outer(t, inv_freq) |
| return torch.cos(freqs), torch.sin(freqs) |
|
|
|
|
| def apply_rope(x, cos, sin): |
| cos = cos[None, :, None, :] |
| sin = sin[None, :, None, :] |
| x_even = x[..., 0::2] |
| x_odd = x[..., 1::2] |
| out = torch.empty_like(x) |
| out[..., 0::2] = x_even * cos - x_odd * sin |
| out[..., 1::2] = x_even * sin + x_odd * cos |
| return out |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, cfg: SovynConfig): |
| super().__init__() |
| if cfg.n_heads % cfg.n_kv_heads != 0: |
| raise ValueError("n_heads must be divisible by n_kv_heads") |
| self.n_heads = cfg.n_heads |
| self.n_kv_heads = cfg.n_kv_heads |
| self.head_dim = cfg.hidden_size // cfg.n_heads |
| self.repeat = cfg.n_heads // cfg.n_kv_heads |
|
|
| kv_dim = cfg.n_kv_heads * self.head_dim |
| self.q_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False) |
| self.k_proj = nn.Linear(cfg.hidden_size, kv_dim, bias=False) |
| self.v_proj = nn.Linear(cfg.hidden_size, kv_dim, bias=False) |
| self.o_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False) |
| self.dropout = cfg.dropout |
|
|
| def forward(self, x, cos, sin): |
| bsz, seq_len, hidden = x.shape |
| q = self.q_proj(x).view(bsz, seq_len, self.n_heads, self.head_dim) |
| k = self.k_proj(x).view(bsz, seq_len, self.n_kv_heads, self.head_dim) |
| v = self.v_proj(x).view(bsz, seq_len, self.n_kv_heads, self.head_dim) |
|
|
| q = apply_rope(q, cos[:seq_len], sin[:seq_len]) |
| k = apply_rope(k, cos[:seq_len], sin[:seq_len]) |
|
|
| k = k.repeat_interleave(self.repeat, dim=2) |
| v = v.repeat_interleave(self.repeat, dim=2) |
|
|
| q = q.transpose(1, 2) |
| k = k.transpose(1, 2) |
| v = v.transpose(1, 2) |
|
|
| y = F.scaled_dot_product_attention( |
| q, |
| k, |
| v, |
| attn_mask=None, |
| dropout_p=self.dropout if self.training else 0.0, |
| is_causal=True, |
| ) |
| y = y.transpose(1, 2).contiguous().view(bsz, seq_len, hidden) |
| return self.o_proj(y) |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, cfg: SovynConfig): |
| super().__init__() |
| self.gate = nn.Linear(cfg.hidden_size, cfg.ffn_size, bias=False) |
| self.up = nn.Linear(cfg.hidden_size, cfg.ffn_size, bias=False) |
| self.down = nn.Linear(cfg.ffn_size, cfg.hidden_size, bias=False) |
|
|
| def forward(self, x): |
| return self.down(F.silu(self.gate(x)) * self.up(x)) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, cfg: SovynConfig): |
| super().__init__() |
| self.attn_norm = RMSNorm(cfg.hidden_size) |
| self.attn = Attention(cfg) |
| self.ffn_norm = RMSNorm(cfg.hidden_size) |
| self.ffn = SwiGLU(cfg) |
|
|
| def forward(self, x, cos, sin): |
| x = x + self.attn(self.attn_norm(x), cos, sin) |
| x = x + self.ffn(self.ffn_norm(x)) |
| return x |
|
|
|
|
| class SovynForCausalLM(nn.Module): |
| def __init__(self, cfg: SovynConfig): |
| super().__init__() |
| self.cfg = cfg |
| self.embed = nn.Embedding(cfg.vocab_size, cfg.hidden_size) |
| self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)]) |
| self.norm = RMSNorm(cfg.hidden_size) |
| self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False) |
| if cfg.tie_embeddings: |
| self.lm_head.weight = self.embed.weight |
|
|
| cos, sin = precompute_rope( |
| cfg.hidden_size // cfg.n_heads, |
| cfg.max_seq_len, |
| cfg.rope_theta, |
| ) |
| self.register_buffer("rope_cos", cos, persistent=False) |
| self.register_buffer("rope_sin", sin, persistent=False) |
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, module): |
| if isinstance(module, nn.Linear): |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) |
| elif isinstance(module, nn.Embedding): |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) |
|
|
| def forward(self, input_ids, labels=None): |
| if input_ids.size(1) > self.cfg.max_seq_len: |
| raise ValueError("Sequence length exceeds max_seq_len") |
|
|
| x = self.embed(input_ids) |
| for block in self.blocks: |
| x = block(x, self.rope_cos, self.rope_sin) |
| x = self.norm(x) |
| logits = self.lm_head(x) |
|
|
| loss = None |
| if labels is not None: |
| loss = F.cross_entropy( |
| logits.view(-1, logits.size(-1)), |
| labels.view(-1), |
| ignore_index=-100, |
| ) |
| return {"loss": loss, "logits": logits} |
|
|
| @torch.no_grad() |
| def generate( |
| self, |
| input_ids, |
| max_new_tokens=96, |
| temperature=0.8, |
| top_k=50, |
| eos_id=None, |
| stop_ids=None, |
| suppress_ids=None, |
| ): |
| self.eval() |
| stop_ids = set(stop_ids or []) |
| suppress_ids = list(suppress_ids or []) |
| for _ in range(max_new_tokens): |
| x = input_ids[:, -self.cfg.max_seq_len :] |
| logits = self(x)["logits"][:, -1, :] |
| if suppress_ids: |
| logits[:, suppress_ids] = -float("inf") |
| if temperature <= 0: |
| next_id = torch.argmax(logits, dim=-1, keepdim=True) |
| else: |
| logits = logits / temperature |
| if top_k > 0: |
| values, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| logits[logits < values[:, [-1]]] = -float("inf") |
| probs = F.softmax(logits, dim=-1) |
| next_id = torch.multinomial(probs, num_samples=1) |
| input_ids = torch.cat([input_ids, next_id], dim=1) |
| token_id = next_id.item() |
| if eos_id is not None and token_id == eos_id: |
| break |
| if token_id in stop_ids: |
| break |
| return input_ids |
|
|
