""" GPT pequeno com duas variantes de residual: - "baseline" : soma residual fixa (padrão Transformer) - "attnres" : Block Attention Residuals (atenção sobre estados anteriores no bloco) Arquitetura: d_model = 256 n_layers = 8 n_heads = 8 d_ff = 1024 max_seq = 512 vocab_size = 16384 ~25M parâmetros Pequeno o suficiente para treinar rápido na RTX 3080 e mostrar diferenças de convergência entre as variantes. """ import math import torch import torch.nn as nn import torch.nn.functional as F from dataclasses import dataclass, field @dataclass class GPTConfig: vocab_size : int = 16_384 max_seq : int = 512 d_model : int = 512 # 256→512: ~4× mais parâmetros por camada n_layers : int = 12 # 8→12: mais profundidade, mais estados para AttnRes explorar n_heads : int = 8 d_ff : int = 2_048 # 1024→2048 dropout : float = 0.1 variant : str = "baseline" block_size : int = 4 # Tamanhos de referência: # small : d=256, L=8, ff=1024 → ~10M params (experimento rápido) # medium : d=512, L=12, ff=2048 → ~50M params (treino sério, ~6h na 3080) # large : d=768, L=16, ff=3072 → ~120M params (cabível na 3080 com seq=256) # ── Blocos base ──────────────────────────────────────────────────────────────── class CausalSelfAttention(nn.Module): def __init__(self, cfg: GPTConfig): super().__init__() assert cfg.d_model % cfg.n_heads == 0 self.n_heads = cfg.n_heads self.d_head = cfg.d_model // cfg.n_heads self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False) self.out = nn.Linear(cfg.d_model, cfg.d_model, bias=False) self.drop = nn.Dropout(cfg.dropout) def forward(self, x: torch.Tensor) -> torch.Tensor: B, T, D = x.shape q, k, v = self.qkv(x).split(D, dim=-1) def reshape(t): return t.view(B, T, self.n_heads, self.d_head).transpose(1, 2) q, k, v = reshape(q), reshape(k), reshape(v) y = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=self.drop.p if self.training else 0.0) y = y.transpose(1, 2).contiguous().view(B, T, D) return self.out(y) class FeedForward(nn.Module): def __init__(self, cfg: GPTConfig): super().__init__() self.net = nn.Sequential( nn.Linear(cfg.d_model, cfg.d_ff, bias=False), nn.GELU(), nn.Linear(cfg.d_ff, cfg.d_model, bias=False), nn.Dropout(cfg.dropout), ) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.net(x) # ── Residual: variante clássica ──────────────────────────────────────────────── class FixedResidual(nn.Module): """x + sublayer(x) — residual padrão.""" def forward(self, x: torch.Tensor, states: list, out: torch.Tensor) -> torch.Tensor: return x + out # ── Residual: variante AttnRes ───────────────────────────────────────────────── class AttnResidual(nn.Module): """ Substitui a soma residual fixa por uma agregação aprendida sobre hidden states anteriores no bloco. Design "zero-init gate" (inspirado em ReZero/Zero-Init): - alpha inicializado em 0 → começa idêntico ao residual clássico - modelo abre o canal AttnRes gradualmente conforme aprende - gradientes fluem normalmente desde o primeiro step Fórmula: h = x + alpha * (attn_aggregate(states) - x) + out = x + out quando alpha=0 (residual clássico) = attn_agg + out quando alpha=1 (AttnRes completo) """ def __init__(self, d_model: int): super().__init__() d_proj = max(d_model // 8, 16) self.q_proj = nn.Linear(d_model, d_proj, bias=False) self.k_proj = nn.Linear(d_model, d_proj, bias=False) self.scale = d_proj ** -0.5 # Gate: sigmoid(-4) ≈ 0.018 → começa quase como residual clássico # O modelo abre o canal AttnRes gradualmente conforme aprende self.alpha = nn.Parameter(torch.full((1,), -4.0)) nn.init.normal_(self.q_proj.weight, std=0.02) nn.init.normal_(self.k_proj.weight, std=0.02) def forward(self, x: torch.Tensor, states: list[torch.Tensor], out: torch.Tensor) -> torch.Tensor: if not states: return x + out # Empilha estados anteriores: (B, T, N, D) stacked = torch.stack(states, dim=2) B, T, N, D = stacked.shape q = self.q_proj(x) # (B, T, d_proj) k = self.k_proj(stacked.reshape(B * T, N, D)) # (B*T, N, d_proj) k = k.reshape(B, T, N, -1) scores = torch.einsum("btd,btnd->btn", q, k) * self.scale # (B, T, N) weights = torch.softmax(scores, dim=-1) # (B, T, N) attn_res = torch.einsum("btn,btnd->btd", weights, stacked) # (B, T, D) # Interpola entre residual clássico (x) e residual por atenção (attn_res) # alpha=0 → x + out (clássico); alpha → 1 → attn_res + out (AttnRes) alpha = torch.sigmoid(self.alpha) residual = x + alpha * (attn_res - x) return residual + out # ── Bloco Transformer ───────────────────────────────────────────────────────── class TransformerBlock(nn.Module): def __init__(self, cfg: GPTConfig): super().__init__() self.ln1 = nn.LayerNorm(cfg.d_model) self.attn = CausalSelfAttention(cfg) self.ln2 = nn.LayerNorm(cfg.d_model) self.ff = FeedForward(cfg) if cfg.variant == "attnres": self.res_a = AttnResidual(cfg.d_model) self.res_m = AttnResidual(cfg.d_model) else: self.res_a = FixedResidual() self.res_m = FixedResidual() # Preenchido pelo GPTModel durante o forward self.block_states: list[torch.Tensor] = [] def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.res_a(x, self.block_states, self.attn(self.ln1(x))) x = self.res_m(x, self.block_states, self.ff(self.ln2(x))) return x # ── Modelo GPT ──────────────────────────────────────────────────────────────── class GPT(nn.Module): def __init__(self, cfg: GPTConfig): super().__init__() self.cfg = cfg self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model) self.pos_emb = nn.Embedding(cfg.max_seq, cfg.d_model) self.drop = nn.Dropout(cfg.dropout) self.blocks = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.n_layers)]) self.ln_f = nn.LayerNorm(cfg.d_model) self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) # Weight tying: embedding e head compartilham pesos self.head.weight = self.embed.weight self._init_weights() def _init_weights(self): for m in self.modules(): if isinstance(m, nn.Linear): nn.init.normal_(m.weight, std=0.02) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, std=0.02) def _reset_block_states(self): """Limpa buffers de estado antes de cada forward.""" for i, block in enumerate(self.blocks): block.block_states = [] def _update_block_states(self, layer_idx: int, state: torch.Tensor): """Adiciona o estado atual ao buffer do bloco correto. Sem detach: mantém o grafo de computação para que q_proj e k_proj recebam gradientes e possam aprender a selecionar estados relevantes. """ if self.cfg.variant != "attnres": return block_idx = layer_idx // self.cfg.block_size next_idx = layer_idx + 1 if (next_idx < self.cfg.n_layers and next_idx // self.cfg.block_size == block_idx): self.blocks[next_idx].block_states = ( self.blocks[layer_idx].block_states + [state] ) def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None): B, T = idx.shape assert T <= self.cfg.max_seq, f"Sequência {T} > max_seq {self.cfg.max_seq}" pos = torch.arange(T, device=idx.device) x = self.drop(self.embed(idx) + self.pos_emb(pos)) self._reset_block_states() for i, block in enumerate(self.blocks): x = block(x) self._update_block_states(i, x) x = self.ln_f(x) logits = self.head(x) # (B, T, V) loss = None if targets is not None: loss = F.cross_entropy( logits.reshape(-1, self.cfg.vocab_size), targets.reshape(-1), ignore_index=-1, ) return logits, loss @torch.no_grad() def generate(self, idx: torch.Tensor, max_new: int, temperature: float = 0.8, top_k: int = 40) -> torch.Tensor: for _ in range(max_new): idx_cond = idx[:, -self.cfg.max_seq:] logits, _ = self(idx_cond) logits = logits[:, -1, :] / temperature if top_k: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = float("-inf") probs = F.softmax(logits, dim=-1) next_t = torch.multinomial(probs, num_samples=1) idx = torch.cat([idx, next_t], dim=1) return idx def n_params(self) -> int: return sum(p.numel() for p in self.parameters()) # ── Utilitários ──────────────────────────────────────────────────────────────── def make_model(variant: str = "baseline", **overrides) -> GPT: cfg = GPTConfig(variant=variant, **overrides) return GPT(cfg) if __name__ == "__main__": for v in ("baseline", "attnres"): m = make_model(v) x = torch.randint(0, 16384, (2, 64)) logits, loss = m(x[:, :-1], x[:, 1:]) print(f"{v:10s} params: {m.n_params():,} logits: {logits.shape} loss: {loss.item():.4f}")