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}