modeling_lowonmind.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

try:  # carregado como remote code (pacote)
    from .configuration_lowonmind import LowOnMindConfig
except ImportError:  # carregado como arquivo solto no sys.path
    from configuration_lowonmind import LowOnMindConfig


class LowOnMindRMSNorm(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)


class LowOnMindRotary(nn.Module):
    """RoPE com cos/sin pre-computados e cacheados.

    O DynamicMind-Mini recalculava inv_freq/cos/sin a cada forward. Aqui o cache
    e construido uma vez e reaproveitado; se aparecer uma sequencia mais longa
    (ou outro device) ele e reconstruido em vez de estourar num broadcast error.

    O cache NAO e um buffer registrado de proposito: buffers nao-persistentes
    criados no __init__ sao materializados com lixo/NaN pelo carregamento em
    meta-device do from_pretrained. Como atributo simples ele e sempre
    reconstruido no primeiro forward.
    """

    def __init__(self, head_dim, max_position_embeddings, base):
        super().__init__()
        self.head_dim = head_dim
        self.base = base
        self.max_position_embeddings = max_position_embeddings
        self._cos = None
        self._sin = None
        self._cached_len = 0

    def _build(self, seq_len, device):
        inv_freq = 1.0 / (
            self.base
            ** (torch.arange(0, self.head_dim, 2, device=device, dtype=torch.float32) / self.head_dim)
        )
        t = torch.arange(seq_len, device=device, dtype=torch.float32)
        freqs = torch.outer(t, inv_freq)
        self._cos = freqs.cos()[None, None]
        self._sin = freqs.sin()[None, None]
        self._cached_len = seq_len

    def forward(self, seq_len, dtype, device):
        if self._cos is None or seq_len > self._cached_len or self._cos.device != device:
            self._build(max(seq_len, self.max_position_embeddings, self._cached_len), device)
        return self._cos[:, :, :seq_len].to(dtype), self._sin[:, :, :seq_len].to(dtype)


def apply_rope(x, cos, sin):
    # x: [B, H, T, head_dim]; cos/sin: [1, 1, T, head_dim // 2]
    x_even, x_odd = x[..., 0::2], 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 LowOnMindAttention(nn.Module):
    def __init__(self, config, rotary):
        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.attention_dropout = config.attention_dropout
        self.rotary = rotary

        assert self.hidden_size % self.num_heads == 0
        assert self.num_heads % self.num_kv_heads == 0
        assert self.head_dim % 2 == 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(self.num_heads * self.head_dim, config.hidden_size, bias=False)
        self.o_proj._is_residual_proj = True

        if config.use_qk_norm:
            self.q_norm = LowOnMindRMSNorm(self.head_dim, config.rms_norm_eps)
            self.k_norm = LowOnMindRMSNorm(self.head_dim, config.rms_norm_eps)
        else:
            self.q_norm = None
            self.k_norm = None

    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)

        if self.q_norm is not None:
            q = self.q_norm(q)
            k = self.k_norm(k)

        cos, sin = self.rotary(seq_len, q.dtype, q.device)
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)

        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, -1)
        return self.o_proj(y)


class LowOnMindMLP(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)
        self.down_proj._is_residual_proj = True

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


class LowOnMindBlock(nn.Module):
    def __init__(self, config, rotary):
        super().__init__()
        self.input_layernorm = LowOnMindRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.self_attn = LowOnMindAttention(config, rotary)
        self.post_attention_layernorm = LowOnMindRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = LowOnMindMLP(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 LowOnMindPreTrainedModel(PreTrainedModel):
    config_class = LowOnMindConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    _no_split_modules = ["LowOnMindBlock"]

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            # projecoes que escrevem no residual: init escalado por 1/sqrt(2L)
            if getattr(module, "_is_residual_proj", False):
                std = std / math.sqrt(2 * self.config.num_hidden_layers)
            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)
        elif isinstance(module, LowOnMindRMSNorm):
            nn.init.ones_(module.weight)


class LowOnMindForCausalLM(LowOnMindPreTrainedModel, 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)

        head_dim = config.hidden_size // config.num_attention_heads
        self.rotary = LowOnMindRotary(head_dim, config.max_position_embeddings, config.rope_theta)

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(
            [LowOnMindBlock(config, self.rotary) for _ in range(config.num_hidden_layers)]
        )
        self.norm = LowOnMindRMSNorm(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),
                ignore_index=-100,
            )

        return CausalLMOutput(loss=loss, logits=logits)

    def state_dict(self, *args, **kwargs):
        sd = super().state_dict(*args, **kwargs)
        # lm_head.weight e tied com embed_tokens.weight; safetensors nao guarda
        # tensores compartilhados duplicados.
        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):
        # sem KV cache: a janela e truncada em max_position_embeddings
        input_ids = input_ids[:, -self.config.max_position_embeddings:]
        return {"input_ids": input_ids}