modeling_dynamicmind.py

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers.modeling_utils import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutput
from .configuration_dynamicmind import DynamicMindConfig


class DynamicMindRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

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


def apply_rope(q, k, rope_theta):
    # q: [B, H, T, D], k: [B, KVH, T, D]
    device = q.device
    dtype = q.dtype
    seq_len = q.size(-2)
    head_dim = q.size(-1)

    inv_freq = 1.0 / (
        rope_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)

    cos = freqs.cos()[None, None, :, :].to(dtype)
    sin = freqs.sin()[None, None, :, :].to(dtype)

    def rotate(x):
        x_even = x[..., 0::2]
        x_odd = x[..., 1::2]
        x_rot_even = x_even * cos - x_odd * sin
        x_rot_odd = x_even * sin + x_odd * cos
        return torch.stack((x_rot_even, x_rot_odd), dim=-1).flatten(-2)

    return rotate(q), rotate(k)


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

        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.rope_theta = config.rope_theta
        self.attention_dropout = config.attention_dropout

        assert self.hidden_size % self.num_heads == 0
        assert self.num_heads % self.num_kv_heads == 0

        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)

    def forward(self, x):
        bsz, seq_len, _ = x.shape

        q = self.q_proj(x).view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

        q, k = apply_rope(q, k, self.rope_theta)

        if self.num_kv_heads != self.num_heads:
            repeats = self.num_heads // self.num_kv_heads
            k = k.repeat_interleave(repeats, dim=1)
            v = v.repeat_interleave(repeats, dim=1)

        y = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=None,
            dropout_p=self.attention_dropout if self.training else 0.0,
            is_causal=True,
        )

        y = y.transpose(1, 2).contiguous().view(bsz, seq_len, self.hidden_size)
        return self.o_proj(y)


class DynamicMindMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class DynamicMindBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.input_layernorm = DynamicMindRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.self_attn = DynamicMindAttention(config)
        self.post_attention_layernorm = DynamicMindRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = DynamicMindMLP(config)

    def forward(self, x):
        x = x + self.self_attn(self.input_layernorm(x))
        x = x + self.mlp(self.post_attention_layernorm(x))
        return x


class DynamicMindPreTrainedModel(PreTrainedModel):
    config_class = DynamicMindConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    _no_split_modules = ["DynamicMindBlock"]

    def _init_weights(self, module):
        std = 0.02
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=std)


class DynamicMindForCausalLM(DynamicMindPreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "embed_tokens.weight"}
    _keys_to_ignore_on_load_missing = [r"lm_head.weight"]
    def __init__(self, config):
        super().__init__(config)

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([DynamicMindBlock(config) for _ in range(config.num_hidden_layers)])
        self.norm = DynamicMindRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        if config.tie_word_embeddings:
            self.lm_head.weight = self.embed_tokens.weight

        self.post_init()

    def tie_weights(self, *args, **kwargs):
        if getattr(self.config, "tie_word_embeddings", True):
            self.lm_head.weight = self.embed_tokens.weight

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, value):
        self.lm_head = value

    def forward(self, input_ids=None, labels=None, **kwargs):
        x = self.embed_tokens(input_ids)

        for layer in self.layers:
            x = layer(x)

        x = self.norm(x)
        logits = self.lm_head(x)

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
            )

        return CausalLMOutput(loss=loss, logits=logits)

    def state_dict(self, *args, **kwargs):
        sd = super().state_dict(*args, **kwargs)
        # lm_head.weight is tied to embed_tokens.weight. Safetensors cannot store
        # duplicate shared tensors unless one key is removed.
        if getattr(self.config, "tie_word_embeddings", True):
            for k in list(sd.keys()):
                if k == "lm_head.weight" or k.endswith(".lm_head.weight"):
                    del sd[k]
        return sd

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