Create src/veronica/modeling_veronica.py

src/veronica/modeling_veronica.py

@@ -0,0 +1,425 @@
+from typing import Optional, Tuple, List
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PreTrainedModel
+from transformers.generation.utils import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_veronica import VeronicaConfig
+from .modeling_components import PolymorphicMLP, router_aux_loss, Fp32LayerNorm, apply_rotary_pos_emb
+
+
+class MultiHeadSelfAttention(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, dropout: float = 0.0, max_position_embeddings: int = 4096, rope_theta: float = 10000.0):
+        super().__init__()
+        assert hidden_size % num_heads == 0, "hidden_size must be divisible by n_head"
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.scale = 1.0 / math.sqrt(self.head_dim)
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+
+        self.qkv = nn.Linear(hidden_size, 3 * hidden_size)
+        self.out_proj = nn.Linear(hidden_size, hidden_size)
+        self.attn_drop = nn.Dropout(dropout)
+        self.resid_drop = nn.Dropout(dropout)
+        
+        # Precomputa RoPE frequencies
+        self._rope_cached_seq_len = 0
+        self._rope_cos_cached = None
+        self._rope_sin_cached = None
+
+    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+        x = x.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)  # (B, nh, T, hd)
+        return x
+
+    def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:
+        B, nh, T, hd = x.shape
+        return x.transpose(1, 2).contiguous().view(B, T, nh * hd)
+    
+    def _get_rope_cos_sin(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Genera o recupera dalla cache cos/sin per RoPE."""
+        if seq_len <= self._rope_cached_seq_len and self._rope_cos_cached is not None:
+            return self._rope_cos_cached[:, :, :seq_len, :].to(device=device, dtype=dtype), \
+                   self._rope_sin_cached[:, :, :seq_len, :].to(device=device, dtype=dtype)
+        
+        # Genera nuove frequenze
+        self._rope_cached_seq_len = max(seq_len, self.max_position_embeddings)
+        
+        # inv_freq: (hd/2,)
+        dim = self.head_dim
+        inv_freq = 1.0 / (self.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
+        
+        # t: (seq_len,)
+        t = torch.arange(self._rope_cached_seq_len, dtype=torch.float32, device=device)
+        
+        # freqs: (seq_len, hd/2)
+        freqs = torch.outer(t, inv_freq)
+        
+        # Duplica per avere shape (seq_len, hd)
+        emb = torch.cat([freqs, freqs], dim=-1)  # (seq_len, hd)
+        
+        # cos, sin: (1, 1, seq_len, hd)
+        cos = emb.cos().unsqueeze(0).unsqueeze(0)
+        sin = emb.sin().unsqueeze(0).unsqueeze(0)
+        
+        self._rope_cos_cached = cos
+        self._rope_sin_cached = sin
+        
+        return cos[:, :, :seq_len, :].to(dtype=dtype), sin[:, :, :seq_len, :].to(dtype=dtype)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,  # additive mask [B,1,T,S] in float32
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+        position_offset: int = 0,  # offset per posizione (per KV cache)
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        B, T, C = x.shape
+        qkv = self.qkv(x)
+        q, k, v = qkv.split(C, dim=-1)
+        q = self._split_heads(q)  # (B, nh, T, hd)
+        k = self._split_heads(k)
+        v = self._split_heads(v)
+        
+        # Applica RoPE a q e k
+        cos, sin = self._get_rope_cos_sin(position_offset + T, q.device, q.dtype)
+        # Prendi solo le posizioni rilevanti [position_offset : position_offset+T]
+        cos = cos[:, :, position_offset:position_offset+T, :]
+        sin = sin[:, :, position_offset:position_offset+T, :]
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        present = None
+        if past_key_value is not None:
+            pk, pv = past_key_value  # (B, nh, Tp, hd)
+            k = torch.cat([pk, k], dim=-2)
+            v = torch.cat([pv, v], dim=-2)
+        if use_cache:
+            present = (k, v)
+
+        att = (q @ k.transpose(-2, -1)) * self.scale  # (B, nh, T, S)
+        att = att.float()
+        if attn_mask is not None:
+            att = att + attn_mask  # additive bias: -inf on masked pos
+        att = F.softmax(att, dim=-1)
+        att = self.attn_drop(att)
+        att = att.to(v.dtype)  # Cast back to match v's dtype (BF16/FP16/FP32)
+        y = att @ v  # (B, nh, T, hd)
+        y = self._merge_heads(y)
+        y = self.out_proj(y)
+        y = self.resid_drop(y)
+        return y, present
+
+
+class VeronicaBlock(nn.Module):
+    def __init__(self, config: VeronicaConfig):
+        super().__init__()
+        self.ln_1 = Fp32LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.attn = MultiHeadSelfAttention(
+            config.n_embd, 
+            config.n_head, 
+            dropout=config.dropout,
+            max_position_embeddings=config.max_position_embeddings,
+            rope_theta=getattr(config, 'rope_theta', 10000.0)
+        )
+        self.ln_2 = Fp32LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.mlp = PolymorphicMLP(
+            hidden_size=config.n_embd,
+            mlp_mult=config.mlp_mult,
+            num_funcs=config.num_funcs,
+            router_dim=config.router_dim,
+            dropout=config.dropout,
+            use_channel_attention=config.use_channel_attention,
+            router_tau=config.router_tau,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+        position_offset: int = 0,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        h = self.ln_1(x)
+        attn_out, present = self.attn(h, attn_mask, past_key_value=past_key_value, use_cache=use_cache, position_offset=position_offset)
+        x = x + attn_out
+        h = self.ln_2(x)
+        y, alpha = self.mlp(h)
+        x = x + y
+        return x, alpha, present
+
+
+class VeronicaModel(PreTrainedModel):
+    config_class = VeronicaConfig
+
+    def __init__(self, config: VeronicaConfig):
+        super().__init__(config)
+        self.embed_dim = config.n_embd
+        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
+        # RoPE sostituisce positional embeddings assoluti (wpe rimosso)
+        self.drop = nn.Dropout(config.dropout)
+        self.blocks = nn.ModuleList([VeronicaBlock(config) for _ in range(config.n_layer)])
+        self.ln_f = Fp32LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+
+        self.register_buffer(
+            "causal_mask",
+            torch.tril(
+                torch.ones(
+                    config.max_position_embeddings,
+                    config.max_position_embeddings,
+                    dtype=torch.uint8,
+                )
+            ).view(1, 1, config.max_position_embeddings, config.max_position_embeddings),
+            persistent=False,
+        )
+
+        # Monitoring
+        self.router_alpha_entropy: Optional[torch.Tensor] = None
+        self.router_alpha_mean: Optional[torch.Tensor] = None
+
+        self._use_gradient_checkpointing: bool = getattr(config, "gradient_checkpointing", False)
+
+    def get_input_embeddings(self):
+        return self.wte
+
+    def set_input_embeddings(self, value):
+        self.wte = value
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        self._use_gradient_checkpointing = True
+
+    def gradient_checkpointing_disable(self):
+        self._use_gradient_checkpointing = False
+
+    def _build_attn_mask(
+        self,
+        attention_mask: Optional[torch.Tensor],
+        seq_len: int,
+        past_kv_len: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        # Causal mask additiva in float32
+        T, P = seq_len, past_kv_len
+        causal = torch.full((T, T + P), float("-inf"), device=device, dtype=dtype)
+        causal = torch.triu(causal, diagonal=1 + P)  # -inf per future, 0 altrove
+
+        if attention_mask is None:
+            return causal.unsqueeze(0).unsqueeze(1)  # [1,1,T,T+P]
+
+        # attention_mask shape: [B, T+P] (0 pad, 1 valid)
+        attn_full = attention_mask.to(dtype)
+        pad_add = (1.0 - attn_full) * torch.finfo(dtype).min  # [B, T+P]
+        pad_add = pad_add.unsqueeze(1).unsqueeze(2)  # [B,1,1,T+P]
+        causal = causal.unsqueeze(0).unsqueeze(1)    # [1,1,T,T+P]
+        return causal + pad_add
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_router_stats: bool = True,
+        past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
+        use_cache: Optional[bool] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[Tuple[torch.Tensor, torch.Tensor]]]]:
+        device = input_ids.device
+        B, T = input_ids.shape
+
+        if use_cache is None:
+            use_cache = False if self.training else True
+
+        pkv_list: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = [] if use_cache else None
+
+        P = 0
+        if (
+            past_key_values is not None
+            and len(past_key_values) > 0
+            and past_key_values[0] is not None
+            and isinstance(past_key_values[0], (tuple, list))
+            and past_key_values[0][0] is not None
+        ):
+            P = past_key_values[0][0].size(-2)
+
+        # Solo token embeddings (RoPE gestisce le posizioni)
+        x = self.wte(input_ids)
+        x = self.drop(x)
+
+        # attention_mask full [B, T+P]
+        attn_full = None
+        if attention_mask is not None:
+            if attention_mask.size(-1) == T + P:
+                attn_full = attention_mask
+            elif attention_mask.size(-1) == T:
+                if P > 0:
+                    ones = torch.ones((B, P), dtype=attention_mask.dtype, device=attention_mask.device)
+                    attn_full = torch.cat([ones, attention_mask], dim=-1)
+                else:
+                    attn_full = attention_mask
+            else:
+                attn_full = None
+
+        attn_bias = self._build_attn_mask(attn_full, T, P, device, torch.float32)
+
+        alpha_list: List[torch.Tensor] = []
+        if self.training:
+            self._acc_aux_sum = 0.0
+            self._acc_aux_count = 0
+
+        if getattr(self, "_use_gradient_checkpointing", False) and self.training:
+            def create_custom_forward(module, pkv):
+                def custom_forward(x):
+                    out_x, out_alpha, _ = module(x, attn_bias, past_key_value=pkv, use_cache=False, position_offset=P)
+                    return out_x, out_alpha
+
+                return custom_forward
+
+            if past_key_values is not None:
+                curr_past = [
+                    pkv
+                    if (pkv is not None and isinstance(pkv, (tuple, list)) and pkv[0] is not None and pkv[1] is not None)
+                    else None
+                    for pkv in past_key_values
+                ]
+            else:
+                curr_past = [None] * len(self.blocks)
+            for layer_idx, block in enumerate(self.blocks):
+                x, alpha = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block, curr_past[layer_idx]), x, use_reentrant=False
+                )
+                alpha_list.append(alpha)
+                if self.training and getattr(block.mlp, "last_aux", None) is not None:
+                    self._acc_aux_sum = self._acc_aux_sum + block.mlp.last_aux
+                    self._acc_aux_count += 1
+        else:
+            if past_key_values is not None:
+                curr_past = [
+                    pkv
+                    if (pkv is not None and isinstance(pkv, (tuple, list)) and pkv[0] is not None and pkv[1] is not None)
+                    else None
+                    for pkv in past_key_values
+                ]
+            else:
+                curr_past = [None] * len(self.blocks)
+            for layer_idx, block in enumerate(self.blocks):
+                x, alpha, present = block(x, attn_bias, past_key_value=curr_past[layer_idx], use_cache=use_cache, position_offset=P)
+                alpha_list.append(alpha)
+                if self.training and getattr(block.mlp, "last_aux", None) is not None:
+                    self._acc_aux_sum = self._acc_aux_sum + block.mlp.last_aux
+                    self._acc_aux_count += 1
+                if use_cache and pkv_list is not None:
+                    pkv_list.append(present)
+
+        x = self.ln_f(x)
+
+        # Router stats
+        if output_router_stats and len(alpha_list) > 0:
+            alpha_stack = torch.stack(alpha_list, dim=0)  # (L, B, T, K)
+            alpha_mean = alpha_stack.mean(dim=(0, 1, 2))  # (K,)
+            self.router_alpha_mean = alpha_mean.detach()
+            self.router_alpha_entropy = router_aux_loss(alpha_stack.mean(dim=0))
+
+        # Aux-loss medio su profondità
+        if hasattr(self, "_acc_aux_sum"):
+            if self._acc_aux_count > 0:
+                self._last_router_aux = self._acc_aux_sum / self._acc_aux_count
+            else:
+                self._last_router_aux = None
+            delattr(self, "_acc_aux_sum")
+            delattr(self, "_acc_aux_count")
+
+        return x, pkv_list
+
+
+class VeronicaForCausalLM(VeronicaModel, GenerationMixin):
+    def __init__(self, config: VeronicaConfig):
+        super().__init__(config)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def tie_weights(self):
+        self._tie_or_clone_weights(self.lm_head, self.get_input_embeddings())
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        if past_key_values is not None and len(past_key_values) > 0:
+            input_ids = input_ids[:, -1:]
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "attention_mask": attention_mask,
+            "use_cache": True,
+        }
+
+    def _reorder_cache(self, past_key_values, beam_idx: torch.LongTensor):
+        if past_key_values is None:
+            return past_key_values
+        reordered = []
+        for (k, v) in past_key_values:
+            reordered.append((k.index_select(0, beam_idx), v.index_select(0, beam_idx)))
+        return reordered
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        hidden_states, present = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=None,
+            use_cache=use_cache,
+            past_key_values=past_key_values,
+            **kwargs,
+        )  # (B, T, H)
+        logits = self.lm_head(hidden_states)
+
+        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,
+            )
+
+            aux = getattr(self, "_last_router_aux", None)
+            if aux is not None and getattr(self.config, "router_aux_weight", 0.0) > 0:
+                if not torch.is_tensor(aux):
+                    aux = torch.as_tensor(aux, device=logits.device, dtype=logits.dtype)
+                else:
+                    aux = aux.to(device=logits.device, dtype=logits.dtype)
+                aux = aux.clamp_min(0.0)
+                loss = loss + float(self.config.router_aux_weight) * aux
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=present if use_cache else None,
+            hidden_states=None,
+            attentions=None,
+        )