Delete models/medusa_heads.py with huggingface_hub

models/medusa_heads.py

@@ -1,409 +0,0 @@
-"""medusa_heads.py — Multi-Token Prediction (MTP) com Medusa Heads.
-
-Implementa Medusa (Cai et al., 2024): múltiplas cabeças de previsão que
-predizem tokens em posições futuras (t+1, t+2, ..., t+K), permitindo
-geração paralela de K tokens por forward pass.
-
-==============================================================================
-ANÁLISE MATEMÁTICA E LÓGICA
-==============================================================================
-
-Seja h_t ∈ R^d o hidden state produzido pelo modelo (por exemplo, saída do
-TransformerDecoderStack no tempo t). Em vez de prever apenas y_{t+1},
-Medusa adiciona K cabeças extras:
-
-    y_{t+k} = softmax(W_k h_t + b_k)    para k = 1, ..., K
-
-A cabeça principal (k=1) é o LM head original; as cabeças k=2..K são
-"cabeças Medusa" adicionais treinadas com a loss:
-
-    L_Medusa = sum_{k=1}^{K} lambda_k * CE(head_k(h_t), y_{t+k})
-
-onde lambda_k é um peso decrescente (tipicamente lambda_k = 1/k ou
-lambda_k = decay^k), pois previsões mais distantes são mais difíceis.
-
-ÁRVORE DE DECODIFICAÇÃO:
-    Em inferência, cada cabeça Medusa gera top-B candidatos, formando uma
-    árvore de filhos. A árvore é avaliada em paralelo (single forward) e
-    os candidatos são aceitos se coincidirem com a previsão do modelo
-    na posição correspondente (verificação via attention mask da árvore).
-
-    Speedup teórico: até K× (em prática 2-3× devido à taxa de aceitação <1).
-
-INTEGRAÇÃO COM CNN-BiGRU:
-    - As cabeças Medusa são aplicadas sobre o hidden state final do
-      TransformerDecoderStack (após self-attention + FFN).
-    - O hidden state h_t já contém contexto cooperativo (streams A+B fundidos).
-    - Em treinamento, adicionamos L_Medusa à loss total (com peso mu_medusa).
-    - Em inferência, ativamos tree-decoding quando use_medusa=True.
-
-VANTAGENS:
-    1. Acelera geração em ~2-3× sem perda de qualidade.
-    2. Não requer mudanças no backbone — apenas cabeças extras.
-    3. Treinável com poucas amostras (as cabeças convergem rápido).
-
-LIMITAÇÕES:
-    1. Aumenta o uso de memória (K * vocab_size params extras).
-    2. Taxa de aceitação depende da tarefa — previsões muito long-range
-       são menos precisas.
-
-Autor: CNN-BiGRU Project
-"""
-from __future__ import annotations
-
-import logging
-import math
-from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-logger = logging.getLogger(__name__)
-
-
-# ============================================================================
-# Configuração
-# ============================================================================
-
-@dataclass
-class MedusaConfig:
-    """Configuração das cabeças Medusa."""
-    vocab_size: int = 32000
-    embed_dim: int = 256
-    n_heads: int = 4                # K cabeças Medusa (além da principal)
-    head_hidden_mult: int = 2       # multiplicador da camada oculta de cada cabeça
-    dropout: float = 0.1
-    # Peso de cada cabeça na loss (lambda_k). Se None, usa 1/(k+1).
-    head_weights: Optional[List[float]] = None
-    # Top-B candidatos por cabeça para tree decoding
-    top_b_candidates: int = 4
-    # Se True, compartilha a primeira camada das cabeças (reduz params)
-    share_first_layer: bool = False
-    # Initialize biases to favor first token (BOS) — estabilidade
-    init_bias_to_bos: bool = False
-    bos_id: int = 0
-
-
-# ============================================================================
-# Cabeça Medusa individual
-# ============================================================================
-
-class MedusaHead(nn.Module):
-    """Uma única cabeça Medusa: prediz y_{t+k} dado h_t.
-
-    Estrutura: Linear(d, d*mult) -> SiLU -> Dropout -> Linear(d*mult, V)
-    """
-
-    def __init__(
-        self,
-        embed_dim: int,
-        vocab_size: int,
-        hidden_mult: int = 2,
-        dropout: float = 0.1,
-    ):
-        super().__init__()
-        hidden = embed_dim * hidden_mult
-        self.fc1 = nn.Linear(embed_dim, hidden)
-        self.fc2 = nn.Linear(hidden, vocab_size)
-        self.dropout = nn.Dropout(dropout)
-        # Inicialização
-        for layer in (self.fc1, self.fc2):
-            nn.init.xavier_uniform_(layer.weight)
-            nn.init.zeros_(layer.bias)
-
-    def forward(self, h: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            h: [B, T, D] ou [B, D] hidden states
-        Returns:
-            logits: [B, T, V] ou [B, V]
-        """
-        x = self.dropout(F.silu(self.fc1(h)))
-        return self.fc2(x)
-
-
-# ============================================================================
-# MTP com Medusa Heads
-# ============================================================================
-
-class MedusaMTP(nn.Module):
-    """Multi-Token Prediction com Medusa Heads.
-
-    Args:
-        config: configuração das cabeças
-    """
-
-    def __init__(self, config: MedusaConfig):
-        super().__init__()
-        self.config = config
-        self.K = config.n_heads
-        self.vocab_size = config.vocab_size
-        self.embed_dim = config.embed_dim
-
-        # Cabeças Medusa (k=1..K)
-        # k=1 é a cabeça "principal extra" (além do lm_head do modelo)
-        # k=2..K são cabeças adicionais
-        if config.share_first_layer:
-            # Compartilha fc1 entre todas as cabeças
-            shared_hidden = config.embed_dim * config.head_hidden_mult
-            self.shared_fc1 = nn.Linear(config.embed_dim, shared_hidden)
-            nn.init.xavier_uniform_(self.shared_fc1.weight)
-            nn.init.zeros_(self.shared_fc1.bias)
-            self.heads = nn.ModuleList([
-                nn.Linear(shared_hidden, config.vocab_size)
-                for _ in range(config.n_heads)
-            ])
-            for h in self.heads:
-                nn.init.xavier_uniform_(h.weight)
-                nn.init.zeros_(h.bias)
-            self._shared = True
-        else:
-            self.heads = nn.ModuleList([
-                MedusaHead(
-                    embed_dim=config.embed_dim,
-                    vocab_size=config.vocab_size,
-                    hidden_mult=config.head_hidden_mult,
-                    dropout=config.dropout,
-                )
-                for _ in range(config.n_heads)
-            ])
-            self._shared = False
-
-        # Pesos lambda_k (decrescentes)
-        if config.head_weights is not None:
-            assert len(config.head_weights) == config.n_heads, \
-                "head_weights deve ter n_heads elementos"
-            self.register_buffer(
-                "head_weights",
-                torch.tensor(config.head_weights, dtype=torch.float32),
-            )
-        else:
-            # default: 1/(k+1) para k=1..K
-            weights = [1.0 / (k + 1) for k in range(1, config.n_heads + 1)]
-            self.register_buffer(
-                "head_weights",
-                torch.tensor(weights, dtype=torch.float32),
-            )
-
-        # Inicialização bias para BOS (opcional)
-        if config.init_bias_to_bos:
-            self._init_biases_to_bos()
-
-    def _init_biases_to_bos(self):
-        """Inicializa biases das cabeças para favorecer BOS no início."""
-        bos_id = self.config.bos_id
-        with torch.no_grad():
-            if self._shared:
-                for h in self.heads:
-                    h.bias.zero_()
-                    h.bias[bos_id] = 1.0
-            else:
-                for h in self.heads:
-                    h.fc2.bias.zero_()
-                    h.fc2.bias[bos_id] = 1.0
-
-    def forward(self, h: torch.Tensor) -> List[torch.Tensor]:
-        """Computa logits de todas as cabeças.
-
-        Args:
-            h: [B, T, D] hidden states do backbone (saída do TransformerDecoder)
-
-        Returns:
-            list de K tensores, cada um [B, T, V] (logits da cabeça k)
-            heads[k-1] prediz y_{t+k} dado h_t
-        """
-        if h.dim() == 2:
-            h = h.unsqueeze(1)  # [B, D] -> [B, 1, D]
-        if h.size(-1) != self.embed_dim:
-            raise ValueError(
-                f"hidden dim mismatch: h={h.size(-1)}, config={self.embed_dim}"
-            )
-
-        all_logits = []
-        if self._shared:
-            shared = F.silu(self.shared_fc1(h))
-            for head in self.heads:
-                all_logits.append(head(shared))
-        else:
-            for head in self.heads:
-                all_logits.append(head(h))
-        return all_logits
-
-    def compute_loss(
-        self,
-        h: torch.Tensor,
-        target_ids: torch.Tensor,
-        ignore_index: int = -100,
-    ) -> Tuple[torch.Tensor, Dict[str, float]]:
-        """Computa a loss Medusa total.
-
-        Args:
-            h: [B, T, D] hidden states do backbone
-            target_ids: [B, T] tokens alvo (y_{t+1}, y_{t+2}, ...)
-                A cabeça k deve prever target_ids[:, t+k-1] dado h_t.
-            ignore_index: índice a ignorar na CE (default: -100, ou pad_id)
-
-        Returns:
-            loss: escalar (weighted sum of head losses)
-            stats: dict com loss por cabeça
-        """
-        if h.dim() == 2:
-            h = h.unsqueeze(1)
-        B, T, D = h.shape
-        if target_ids.dim() != 2:
-            raise ValueError(f"target_ids deve ser [B, T]; recebido {target_ids.shape}")
-
-        all_logits = self.forward(h)  # list de K tensores [B, T, V]
-
-        total_loss = torch.zeros(1, device=h.device, dtype=h.dtype)
-        stats = {}
-
-        for k in range(self.K):
-            # Cabeça k prediz y_{t+k+1} dado h_t (k=0 -> y_{t+1}, k=1 -> y_{t+2}, ...)
-            shift = k + 1
-            if shift >= T:
-                # Não há targets suficientes para esta cabeça
-                stats[f"head_{k+1}_loss"] = float("nan")
-                continue
-            # logits_k: [B, T, V] -> pegamos[:B, :T-shift, :] para alinhar com target[:, shift:]
-            logits_k = all_logits[k][:, :T - shift, :].contiguous()  # [B, T-shift, V]
-            target_k = target_ids[:, shift:T].contiguous()  # [B, T-shift]
-
-            loss_k = F.cross_entropy(
-                logits_k.view(-1, self.vocab_size),
-                target_k.view(-1),
-                ignore_index=ignore_index,
-                reduction="mean",
-            )
-            weight_k = self.head_weights[k].to(h.device)
-            total_loss = total_loss + weight_k * loss_k
-            stats[f"head_{k+1}_loss"] = float(loss_k.detach())
-
-        return total_loss.squeeze(), stats
-
-    def predict_top_b(
-        self,
-        h: torch.Tensor,
-        top_b: Optional[int] = None,
-    ) -> List[Tuple[torch.Tensor, torch.Tensor]]:
-        """Para cada cabeça, retorna top-B candidatos e suas probabilidades.
-
-        Usado em tree decoding (inferência).
-
-        Args:
-            h: [B, T, D] hidden states (ou [B, D] para último token)
-            top_b: número de candidatos (default: config.top_b_candidates)
-
-        Returns:
-            list de K tuplas (tokens, probs) onde:
-                tokens: [B, top_b] — IDs dos top-B tokens
-                probs: [B, top_b] — probabilidades (após softmax)
-        """
-        if h.dim() == 2:
-            h = h.unsqueeze(1)
-        # Usar apenas o último token para previsão
-        h_last = h[:, -1:, :]  # [B, 1, D]
-        all_logits = self.forward(h_last)  # list de K [B, 1, V]
-        b = top_b or self.config.top_b_candidates
-
-        results = []
-        for k, logits_k in enumerate(all_logits):
-            logits_k = logits_k[:, -1, :]  # [B, V]
-            probs = F.softmax(logits_k, dim=-1)
-            top_probs, top_ids = torch.topk(probs, k=b, dim=-1)
-            results.append((top_ids, top_probs))
-        return results
-
-
-# ============================================================================
-# Tree Decoding (esqueleto — implementação simplificada)
-# ============================================================================
-
-def medusa_tree_decode(
-    medusa: MedusaMTP,
-    base_logits: torch.Tensor,
-    h_last: torch.Tensor,
-    accept_threshold: float = 0.0,
-) -> Dict[str, torch.Tensor]:
-    """Decodificação por árvore Medusa (simplificada).
-
-    Em uma implementação completa, isto construiria uma árvore de candidatos
-    e faria um forward pass único para verificar aceitação. Aqui retornamos
-    apenas os candidatos de cada cabeça para uso posterior.
-
-    Args:
-        medusa: módulo MedusaMTP
-        base_logits: [B, V] logits da cabeça principal (já computados)
-        h_last: [B, D] hidden state do último token
-        accept_threshold: threshold de probabilidade para aceitar candidato
-
-    Returns:
-        dict com:
-            base_token: [B] argmax da cabeça principal
-            base_prob: [B] prob do argmax
-            medusa_candidates: list de K tuplas (tokens, probs)
-            accepted: list de K tensores booleanos [B] indicando aceitação
-    """
-    base_probs = F.softmax(base_logits, dim=-1)
-    base_token = base_probs.argmax(dim=-1)
-    base_prob = base_probs.max(dim=-1).values
-
-    candidates = medusa.predict_top_b(h_last)
-
-    accepted = []
-    for k, (tokens, probs) in enumerate(candidates):
-        # Aceita o primeiro candidato se prob >= threshold
-        first_prob = probs[:, 0]
-        accept = first_prob >= accept_threshold
-        accepted.append(accept)
-
-    return {
-        "base_token": base_token,
-        "base_prob": base_prob,
-        "medusa_candidates": candidates,
-        "accepted": accepted,
-    }
-
-
-# ============================================================================
-# Self-test
-# ============================================================================
-
-def _self_test():
-    torch.manual_seed(42)
-    config = MedusaConfig(
-        vocab_size=100,
-        embed_dim=32,
-        n_heads=4,
-        head_hidden_mult=2,
-        dropout=0.1,
-    )
-    medusa = MedusaMTP(config)
-    h = torch.randn(2, 8, 32)
-    target = torch.randint(0, 100, (2, 8))
-    loss, stats = medusa.compute_loss(h, target)
-    print(f"Loss: {loss.item():.4f}")
-    print(f"Stats: {stats}")
-
-    # Test top-B
-    cands = medusa.predict_top_b(h)
-    print(f"K={len(cands)}, top_b shape: {cands[0][0].shape}")
-
-    # Test tree decode
-    base_logits = torch.randn(2, 100)
-    result = medusa_tree_decode(medusa, base_logits, h[:, -1, :])
-    print(f"Base token: {result['base_token']}")
-
-
-if __name__ == "__main__":
-    _self_test()
-
-
-__all__ = [
-    "MedusaConfig",
-    "MedusaHead",
-    "MedusaMTP",
-    "medusa_tree_decode",
-]