kavin-ravi/slm-synthetic

SLM — 538M GPT (synthetic-corpus, from scratch)

A ~538M-parameter decoder-only GPT trained from scratch on a mix of synthetic mini-textbooks, Wikipedia, PG19, and FineWeb-Edu. Custom architecture (per-head QK-norm) with a custom 50k BPE tokenizer — this is not a 🤗 Transformers model, so load it with the self-contained snippet below.

Params	~538M
d_model / heads / layers	1280 / 20 / 24
d_ff / context	5120 / 2048
Tokenizer	custom 50k BPE (tokenizer.json)
Attention	QK-norm (per-head RMSNorm on Q and K)

Files

• best.pt — best-val checkpoint ({'model': state_dict, 'step', 'val_loss'})
• tokenizer.json — the 50k BPE tokenizer
• config.yaml — model + training config

Use it (copy-paste, runs on a free Colab CPU or GPU)

# pip install -q torch tokenizers huggingface_hub pyyaml
import yaml, torch, torch.nn as nn, torch.nn.functional as F
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download

REPO = "kavin-ravi/slm-synthetic"
device = 'cuda' if torch.cuda.is_available() else 'cpu'

config = yaml.safe_load(open(hf_hub_download(REPO, "config.yaml")))
mcfg = config['model']
tokenizer = Tokenizer.from_file(hf_hub_download(REPO, "tokenizer.json"))
vocab_size = tokenizer.get_vocab_size()


class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)
        self.q_norm = nn.RMSNorm(self.head_dim)
        self.k_norm = nn.RMSNorm(self.head_dim)

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=-1)
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        q, k = self.q_norm(q), self.k_norm(k)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.proj(out.transpose(1, 2).contiguous().view(B, T, C))


class FFN(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.up = nn.Linear(d_model, d_ff, bias=False)
        self.down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):
        return self.down(F.gelu(self.up(x)))


class Block(nn.Module):
    def __init__(self, d_model, n_heads, d_ff):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = FFN(d_model, d_ff)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x


class GPT(nn.Module):
    def __init__(self, vocab_size, d_model, n_heads, n_layers, d_ff, context_length):
        super().__init__()
        self.context_length = context_length
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(context_length, d_model)
        self.blocks = nn.ModuleList([Block(d_model, n_heads, d_ff) for _ in range(n_layers)])
        self.ln_f = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        self.tok_emb.weight = self.lm_head.weight

    def forward(self, idx):
        T = idx.shape[1]
        x = self.tok_emb(idx) + self.pos_emb(torch.arange(T, device=idx.device))
        for block in self.blocks:
            x = block(x)
        return self.lm_head(self.ln_f(x))


model = GPT(vocab_size, mcfg['d_model'], mcfg['n_heads'], mcfg['n_layers'],
            mcfg['d_ff'], mcfg['context_length']).to(device)
ckpt = torch.load(hf_hub_download(REPO, "best.pt"), weights_only=False, map_location=device)
model.load_state_dict(ckpt['model'])
model.eval()
print(f"Loaded step {ckpt.get('step', '?')}, val_loss {ckpt.get('val_loss', '?')}")


@torch.no_grad()
def generate(prompt, max_tokens=200, temperature=0.8, top_k=50):
    ids = torch.tensor(tokenizer.encode(prompt).ids, device=device).unsqueeze(0)
    for _ in range(max_tokens):
        logits = model(ids[:, -model.context_length:])[:, -1] / temperature
        if top_k:
            v, _ = torch.topk(logits, top_k)
            logits[logits < v[:, [-1]]] = float('-inf')
        ids = torch.cat([ids, torch.multinomial(F.softmax(logits, -1), 1)], dim=1)
    return tokenizer.decode(ids[0].tolist())


print(generate("Once upon a time"))

Trained on an RTX 5090 with QK-norm, bf16, torch.compile, gradient checkpointing. Source: https://github.com/kavinravi/SLM