Upload model.py with huggingface_hub

4d2987c verified 7 months ago

10.8 kB

"""
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}")

	"""
	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}")