modeling_nexus_small.py

import math
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig, AutoConfig, AutoModel, AutoModelForCausalLM


class NexusSmAllConfig(PretrainedConfig):
    model_type = "nexus_small"

    def __init__(
        self,
        vocab_size=50304,
        max_seq_len=512,
        dim=768,
        num_layers=10,
        num_heads=12,
        num_kv_heads=4,
        multiple_of=256,
        ff_dim=2048,
        norm_eps=1e-6,
        rope_theta=500000.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.max_seq_len = max_seq_len
        self.dim = dim
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.multiple_of = multiple_of
        self.ff_dim = ff_dim
        self.norm_eps = norm_eps
        self.rope_theta = rope_theta


AutoConfig.register("nexus_small", NexusSmAllConfig)


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        rms = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
        return (x.float() * rms * self.weight.float()).type_as(x)


def precompute_freqs_cis(max_seq_len, dim, num_heads, rope_theta):
    head_dim = dim // num_heads
    freqs = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
    t = torch.arange(max_seq_len)
    freqs = torch.outer(t, freqs)
    return torch.polar(torch.ones_like(freqs), freqs)


class RotaryEmbedding(nn.Module):
    def __init__(self, max_seq_len, dim, num_heads, rope_theta):
        super().__init__()
        self.freqs_cis = precompute_freqs_cis(max_seq_len, dim, num_heads, rope_theta)

    def forward(self, x, start_pos=0):
        _, seq_len, _, head_dim = x.shape
        freqs_cis = self.freqs_cis[start_pos : start_pos + seq_len, : head_dim // 2].to(x.device)
        freqs_cis = freqs_cis.view(1, seq_len, 1, head_dim // 2)

        x_shaped = x.float().reshape(*x.shape[:-1], -1, 2)
        x_complex = torch.complex(x_shaped[..., 0], x_shaped[..., 1])
        x_rotated = x_complex * freqs_cis
        x_out = torch.stack([x_rotated.real, x_rotated.imag], dim=-1).reshape_as(x_shaped)
        return x_out.reshape_as(x).type_as(x)


class Attention(nn.Module):
    def __init__(self, dim, num_heads, num_kv_heads, max_seq_len, rope_theta):
        super().__init__()
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads or num_heads
        self.head_dim = dim // num_heads
        self.num_kv_groups = num_heads // self.num_kv_heads

        self.wq = nn.Linear(dim, dim, bias=False)
        self.wk = nn.Linear(dim, self.head_dim * self.num_kv_heads, bias=False)
        self.wv = nn.Linear(dim, self.head_dim * self.num_kv_heads, bias=False)
        self.wo = nn.Linear(dim, dim, bias=False)
        self.rotary = RotaryEmbedding(max_seq_len, dim, num_heads, rope_theta)

    def forward(self, x, start_pos=0, mask=None):
        bsz, seqlen, _ = x.shape

        q = self.wq(x).view(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.wk(x).view(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.wv(x).view(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2)

        q = self.rotary(q, start_pos)
        k = self.rotary(k, start_pos)

        if self.num_kv_groups > 1:
            k = k[:, :, None, :, :].expand(bsz, self.num_kv_heads, self.num_kv_groups, seqlen, self.head_dim)
            k = k.reshape(bsz, self.num_heads, seqlen, self.head_dim)
            v = v[:, :, None, :, :].expand(bsz, self.num_kv_heads, self.num_kv_groups, seqlen, self.head_dim)
            v = v.reshape(bsz, self.num_heads, seqlen, self.head_dim)

        scale = 1.0 / math.sqrt(self.head_dim)
        attn_weights = torch.matmul(q, k.transpose(-2, -1)) * scale

        if mask is not None:
            attn_weights = attn_weights + mask

        attn_weights = F.softmax(attn_weights.float(), dim=-1).type_as(q)
        attn_output = torch.matmul(attn_weights, v)
        attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
        return self.wo(attn_output)


class FeedForward(nn.Module):
    def __init__(self, dim, ff_dim, multiple_of):
        super().__init__()
        hidden_dim = int(2 * ff_dim / 3)
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
        self.w3 = nn.Linear(dim, hidden_dim, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class TransformerBlock(nn.Module):
    def __init__(self, dim, num_heads, num_kv_heads, ff_dim, multiple_of, norm_eps, max_seq_len, rope_theta):
        super().__init__()
        self.attention = Attention(dim, num_heads, num_kv_heads, max_seq_len, rope_theta)
        self.feed_forward = FeedForward(dim, ff_dim, multiple_of)
        self.attention_norm = RMSNorm(dim, norm_eps)
        self.ff_norm = RMSNorm(dim, norm_eps)

    def forward(self, x, start_pos=0, mask=None):
        h = x + self.attention(self.attention_norm(x), start_pos, mask)
        out = h + self.feed_forward(self.ff_norm(h))
        return out


class Nexus(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        self.token_embeddings = nn.Embedding(config.vocab_size, config.dim)
        self.layers = nn.ModuleList([
            TransformerBlock(
                config.dim, config.num_heads, config.num_kv_heads,
                config.ff_dim, config.multiple_of, config.norm_eps,
                config.max_seq_len, config.rope_theta,
            )
            for _ in range(config.num_layers)
        ])
        self.norm = RMSNorm(config.dim, config.norm_eps)
        self.output = nn.Linear(config.dim, config.vocab_size, bias=False)
        self.token_embeddings.weight = self.output.weight

    def forward(self, input_ids, start_pos=0):
        _, seqlen = input_ids.shape

        mask = torch.full(
            (1, 1, seqlen, start_pos + seqlen), float("-inf"),
            dtype=torch.float32, device=input_ids.device,
        )
        mask = torch.triu(mask, diagonal=start_pos + 1).type_as(input_ids)

        x = self.token_embeddings(input_ids)
        for layer in self.layers:
            x = layer(x, start_pos, mask)
        x = self.norm(x)
        return self.output(x)


class NexusForCausalLM(PreTrainedModel):
    config_class = NexusSmAllConfig
    base_model_prefix = "nexus_small"
    supports_gradient_checkpointing = False
    _no_split_modules = ["TransformerBlock"]

    def __init__(self, config):
        super().__init__(config)
        self.model = Nexus(config)
        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
        with torch.no_grad():
            self.lm_head.weight = self.model.output.weight

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, input_ids, attention_mask=None, **kwargs):
        return self.model(input_ids)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids}


AutoModel.register(NexusSmAllConfig, NexusForCausalLM)
AutoModelForCausalLM.register(NexusSmAllConfig, NexusForCausalLM)