SLM-GPT (~144M)

Small Language Model decoder-only estilo GPT, treinado do zero. Arquitetura moderna com RoPE, atenção causal multi-head, MLP SwiGLU, RMSNorm (pré-normalização) e weight tying entre embedding e cabeça de saída. Otimizado com Muon (matrizes internas) + AdamW (embedding/norms).

• Parâmetros: 143.7M total · 111.5M sem embedding
• Tokenizador: BPE do GPT-2 (tiktoken, encoding gpt2)
• Contexto: 1024 tokens
• Idioma: inglês (treinado em FineWeb)

Arquitetura

Parâmetro	Valor
n_layer	16
n_head	10
n_embd	640
head_dim	64
vocab_size	50304
block_size	1024
rope_theta	10000
bias	False
dropout	0.0

• Positional encoding: RoPE (Rotary Position Embeddings)
• Atenção: Causal Multi-Head Self-Attention
• Ativação (MLP): SwiGLU
• Normalização: RMSNorm (pré-norm)

Detalhes de treino

Item	Valor
Dataset	FineWeb (~100B tokens de amostra)
Otimizador	Muon (lr 1e-3, momentum 0.95, match_rms_adamw) + AdamW
LR de pico (AdamW)	9e-4 · warmup 2500 · cosine até 50k
Weight decay	0.1 (AdamW) · 0.0 (Muon)
Grad clip	1.0
Batch efetivo	12 × 42 grad-accum × 1024 ≈ 500k tok/passo
Precisão	bfloat16
Iteração (checkpoint)	46.400
Melhor val loss	2.909

Curvas de treino em loss_curve.png / lr_schedule.png; métricas brutas em metrics.csv.

Conteúdo do checkpoint (ckpt.pt)

torch.load(..., weights_only=False) retorna um dict com:

Chave	Descrição
model	state dict dos pesos
optimizer	estado do(s) otimizador(es) — para retomar o treino
model_config	dict com a config da arquitetura (GPTConfig)
train_config	dict com a config de treino (TrainConfig)
iter_num	iteração em que foi salvo (46.400)
best_val_loss	melhor val loss até então (2.909)
loaders	posição dos data loaders (shard/offset) para resume

Como usar

O código do modelo (model/, data/tokenizer.py) está no repositório do projeto. Com ele no path:

import torch
from model import GPT, GPTConfig

# 1. carregar checkpoint
ckpt = torch.load("ckpt.pt", map_location="cpu", weights_only=False)

# 2. reconstruir o modelo a partir da config salva
config = GPTConfig(**ckpt["model_config"])
model = GPT(config)
model.load_state_dict(ckpt["model"])
model.eval()

# 3. gerar texto (o GPT já embute o tokenizer GPT-2)
prompt = "The history of artificial intelligence"
idx = torch.tensor([model.tokenizer.encode(prompt)], dtype=torch.long)
out = model.generate(idx, max_new_tokens=100, temperature=0.8, top_k=200)
print(model.tokenizer.decode(out[0].tolist()))

Retomar o treino

O optimizer e loaders no checkpoint permitem continuar exatamente de onde parou (ver training/checkpoint.py e training/train.py --resume).

Limitações

Modelo pequeno de pré-treino, sem fine-tuning de instrução nem alinhamento. Gera texto em inglês no estilo do FineWeb e pode produzir conteúdo incorreto, repetitivo ou enviesado. Uso educacional / de pesquisa.