slm-55m-llama-style

SLM-55M LLaMA-style

Small Language Model de ~55M parâmetros pré-treinado em 524M tokens do FineWeb-Edu. Projeto da disciplina Aprendizado Profundo II (PUCRS).

Detalhes

• Parâmetros: 55.257.600
• Arquitetura: decoder-only transformer (estilo LLaMA)
• Componentes: RoPE, GQA (8 Q heads / 4 KV heads), SwiGLU, RMSNorm
• Tokenizer: BPE GPT-2 (50.304 tokens com padding)
• Context length: 1024
• Treino: 4000 steps, batch efetivo 131K tokens, ~2h40 em RTX 5060 Ti

Resultados

• Val loss: 3.648
• Val perplexity: 38.4

Como usar

Código completo: https://github.com/cjfbr/slm-pretraining

Carregar o checkpoint requer a definição do modelo do repositório de código. Exemplo:

import torch
from slm.model import GPT
from slm.config import ModelConfig

cfg = ModelConfig()
model = GPT(cfg)
ckpt = torch.load("pretrain_final.pt", weights_only=False, map_location="cuda")
model.load_state_dict(ckpt["model"])
model.eval()

Curvas de treino

Disponíveis no Weights & Biases: https://wandb.ai/cjferrazzo-pucrs/slm-pretrain