slm-pt100m — Small Language Model em Português-BR (~100M params)

SLM decoder-only estilo LLaMA, treinado do zero em Português-BR para o trabalho de Aprendizado Profundo II (PUCRS — Escola Politécnica).

Pipeline completo das 4 etapas: Pré-treino → Mid-training → SFT → Avaliação.

Arquitetura

Componente	Escolha
Camadas	12
Dimensão (n_embd)	768
Heads de Query	12
Heads de Key/Value	4 (GQA, ratio 3:1)
Head dim	64
FFN	SwiGLU, hidden 2048 (~8/3 × n_embd)
Normalização	RMSNorm (pre-norm)
Positional encoding	RoPE (cos/sin real-valued)
Atenção	F.scaled_dot_product_attention (Flash Attention 2 nativo)
Tokenizer	BPE byte-level treinado em PT-BR (vocab 32.000)
Embeddings	Compartilhadas com LM head (tied)
Block size	1024
Mixed precision	bfloat16 (autocast)

Total: 100.073.472 parâmetros (~62.9M sem embeddings).

Etapas de treino

Etapa	Dataset	Tokens	Otimizador	LR	Tempo	Métrica final
Pré-treino	FineWeb-2 (por_Latn)	~1.97B	AdamW fused	6e-4 → 6e-5 cosine	10h30	val_ppl 22.48
Mid-training	Canarim + alpaca-pt-br	~131M	AdamW	2e-4 → 2e-5 cosine	44.6 min	val_ppl 4.04
SFT (loss mask)	Mesmos do mid-train	~52M	AdamW	2e-5 → 2e-6 cosine	18 min	val_ppl 4.94 (best, step 200)

Hardware: 1 GPU RTX 5060 Ti (16 GB) + torch.compile + bfloat16 autocast.

Avaliação (Etapa 4)

Estágio	val_ppl (web PT)	ENEM completion acc_norm	ENEM chat acc_mean
pretrain	22.48	27.67%	19.29%
midtrain	54.13	23.27%	20.41%
sft (best)	57.96	23.13%	20.55%
random	—	20.00%	20.00%

ENEM Challenge (eduagarcia/enem_challenge, 1431 questões, 5 alternativas). Sobre o "alignment tax" e por que pretrain vence em completion enquanto SFT vence em chat, ver README do repositório.

Arquivos neste repo

• midtrain/final.pt — após mid-training em instruções PT-BR (val_ppl 4.04)
• midtrain/best.pt — checkpoint com menor val_loss durante o mid-training
• sft/best.pt — recomendado para uso — best val_loss durante o SFT (val_ppl 4.94 no step 200)
• sft/final.pt — último step do SFT (val_loss um pouco maior; overfit leve)
• tokenizer/tokenizer_ptbr.json — BPE byte-level treinado em PT-BR (~700 KB)

Como carregar

Os checkpoints não estão no formato transformers — são state_dict do PyTorch + o ModelConfig salvo no próprio dict. Clone o código fonte para carregar:

git clone https://github.com/cjfbr/slm-pretraining.git
cd slm-pretraining
pip install -r requirements.txt

import torch
from huggingface_hub import hf_hub_download
from slm.model import GPT
from slm.config import ModelConfig
from slm.tokenizer import Tokenizer

# 1. Baixar checkpoint + tokenizer do Hub
ckpt_path = hf_hub_download(
    repo_id="cjfb75/slm-pt100m", filename="sft/best.pt",
)
tok_path = hf_hub_download(
    repo_id="cjfb75/slm-pt100m", filename="tokenizer/tokenizer_ptbr.json",
)

# 2. Reconstruir modelo a partir do ModelConfig salvo no checkpoint
ckpt = torch.load(ckpt_path, weights_only=False, map_location="cuda")
cfg = ModelConfig(**{{k: v for k, v in ckpt["model_config"].items()
                     if k in ModelConfig.__dataclass_fields__}})
model = GPT(cfg).to("cuda")
model.load_state_dict(ckpt["model"])
model.eval()

tokenizer = Tokenizer(tok_path)

Como gerar (formato chat)

O modelo foi treinado com o template Usuário: ... \nAssistente: .... Para conversar:

prompt = "Usuário: Qual a capital do Brasil?\nAssistente:"
ids = tokenizer.encode(prompt, add_eot=False)
input_ids = torch.tensor([ids], dtype=torch.long, device="cuda")

out = model.generate(
    input_ids,
    max_new_tokens=120,
    temperature=0.8,
    top_k=50,
    repetition_penalty=1.3,
    eot_token=tokenizer.eot_token,
)
gen_ids = out[0].cpu().tolist()[len(ids):]
if tokenizer.eot_token in gen_ids:
    gen_ids = gen_ids[:gen_ids.index(tokenizer.eot_token)]
print(tokenizer.decode(gen_ids))
# -> " A capital do Brasil é Brasília."

Demo interativa

O repositório de código traz um app Streamlit:

streamlit run app.py

Limitações honestas

• 100M params é **70× menor que LLaMA-7B**. Respostas factuais simples saem bem, mas explicações longas têm verbosidade e ocasionais alucinações.
• Aderência à quantidade pedida em listas ("liste exatamente 3") ainda é fraca — o modelo tende a listar 4-5 itens.
• Datasets do SFT são single-turn; coerência multi-turn cai após 2-3 turnos.

Citação / créditos

Projeto da disciplina Aprendizado Profundo II (PUCRS). Implementação inspirada em nanoGPT (Karpathy), com receitas do SmolLM (HuggingFace) e Muon (Keller Jordan). Datasets de instrução: Canarim (dominguesm/Canarim-Instruct-PTBR-Dataset) e Alpaca-PT-BR (dominguesm/alpaca-data-pt-br).