transformer-toolkit 0.0.16

Minimal, modular transformer library for training your own LLM

transformer-toolkit

A modular, from-scratch transformer library for training and experimenting with modern LLM architectures. Swap attention types, positional encodings, FFN variants, and normalization — all from a single config object.

---

Installation

pip install transformer-toolkit

---

Quick Start

import torch
from transformer_toolkit.model import Transformer, TransformerConfig
from transformer_toolkit.c_tokenizers import RustBPETokenizer
from transformer_toolkit.dataloader import DataConfig, from_binary, save_binary
from transformer_toolkit.trainer import Trainer, TrainConfig

# tokenizer
tok = RustBPETokenizer()
tok.train(open("data.txt", encoding="utf-8").readlines(), vocab_size=8000)
tok.save("tokenizer.json")

# data
save_binary(tok.encode(open("data.txt", encoding="utf-8").read()), "data.bin")
train_dl, val_dl = from_binary("data.bin", DataConfig(seq_len=128, batch_size=32))

# model
model = Transformer(TransformerConfig(
    vocab_size = tok.vocab_size,
    dim        = 512,
    n_layers   = 8,
    n_heads    = 8,
    pos_enc    = "rope",
)).to("cuda")

# train
trainer = Trainer(model, train_dl, val_dl, tok.vocab_size, TrainConfig(max_steps=3000))
trainer.train()

---

Table of Contents

• Model
• Attention
• Feed-Forward Networks
• Positional Encodings
• Normalization
• Dataloader
• Tokenizers
• Trainer
• HuggingFace Hub
• Generation
• Full Examples

---

Model

TransformerConfig

All architecture decisions live in one dataclass. Pass it to Transformer().

from transformer_toolkit.model import TransformerConfig

cfg = TransformerConfig(
    # ── core ──────────────────────────────────────
    vocab_size = 32000,      # tokenizer vocabulary size
    dim        = 512,        # model embedding dimension
    n_layers   = 8,          # number of transformer blocks
    n_heads    = 8,          # attention heads
    max_seq    = 2048,       # maximum sequence length

    # ── attention ─────────────────────────────────
    attn       = "gqa",      # "mha" | "gqa" | "mqa" | "flash" | "mla"
    n_kv_heads = 4,          # gqa only — number of key/value heads
    latent_dim = 64,         # mla only — latent compression dim

    # ── feed-forward ──────────────────────────────
    ffn        = "swiglu",   # "ffn" | "swiglu" | "moe"
    hidden_dim = 2048,       # FFN inner dimension (default: dim * 4)
    n_experts  = 8,          # moe only — number of experts
    top_k      = 2,          # moe only — experts activated per token
    moe_aux_weight = 0.01,   # moe load-balancing loss coefficient

    # ── normalization ─────────────────────────────
    norm       = "rmsnorm",  # "rmsnorm" | "layernorm"
    eps        = 1e-6,

    # ── positional encoding ───────────────────────
    pos_enc    = "rope",     # "rope" | "sinusoidal" | "learned" | "alibi" | "none"

    # ── regularisation ────────────────────────────
    dropout    = 0.1,

    # ── output ────────────────────────────────────
    tie_weights = True,      # share embedding and output projection weights
)

Transformer

from transformer_toolkit.model import Transformer

model = Transformer(cfg).to("cuda")

print(model.n_params())   # "30.21M"

# forward pass — returns (logits, aux_loss)
# aux_loss is non-zero only for MoE; always add it to your training loss
logits, aux_loss = model(tokens)       # tokens: [B, T]  →  logits: [B, T, vocab_size]

# generation
output = model.generate(
    tokens      = prompt_tokens,   # [B, T]
    max_new     = 200,
    temperature = 0.8,
    top_k       = 40,
)

Weight tying — when tie_weights=True (default), the embedding and output projection share the same weight matrix. Use model.state_dict_for_save() instead of model.state_dict() when saving checkpoints, and model.load_state_dict_with_tie() when loading, to correctly preserve the tie across save/load cycles.

---

Attention

Five attention variants, all swappable via TransformerConfig.attn.

Value	Class	Used in
"mha"	MultiHeadAttention	Original Transformer, BERT, GPT-2
"gqa"	GroupedQueryAttention	LLaMA 3, Mistral
"mqa"	MultiQueryAttention	Falcon, early Gemini
"flash"	FlashAttention	Any model on PyTorch ≥ 2.0
"mla"	MLAttention	DeepSeek-V2/V3

RoPE is applied inside attention to q and k after head-splitting — not to the residual stream. It is instantiated once and shared across all layers.

ALiBi bias is computed once per forward pass and passed as an additive mask to every block.

Example — switch to Flash Attention

cfg = TransformerConfig(
    dim     = 512,
    n_heads = 8,
    attn    = "flash",   # uses torch.nn.functional.scaled_dot_product_attention
)

Example — Grouped Query Attention (LLaMA-style)

cfg = TransformerConfig(
    dim        = 512,
    n_heads    = 8,
    attn       = "gqa",
    n_kv_heads = 2,   # 4 query heads share each kv head → 4x kv cache reduction
)

---

Feed-Forward Networks

Value	Class	Used in
"ffn"	FFN	Original Transformer, BERT
"swiglu"	SwiGLU	LLaMA, Mistral, PaLM
"moe"	MoE	Mixtral, GPT-4 (rumoured)

MoE — Mixture of Experts

When using ffn="moe", the model forward pass returns an auxiliary load-balancing loss. You must add it to your training loss or all tokens will collapse onto 1–2 experts within a few hundred steps.

cfg = TransformerConfig(
    ffn            = "moe",
    n_experts      = 8,
    top_k          = 2,
    moe_aux_weight = 0.01,   # weight of the load-balancing term
)

logits, aux_loss = model(tokens)
ce_loss  = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))
loss     = ce_loss + aux_loss   # aux_loss is 0.0 for non-MoE models

---

Positional Encodings

Value	Applied	Notes
"rope"	Inside attention, on q and k	LLaMA, Mistral, Qwen
"sinusoidal"	Residual stream before blocks	Original Transformer
"learned"	Residual stream before blocks	BERT, GPT-2
"alibi"	Additive bias on attention scores	ALiBi paper
"none"	Not applied	Bare model, no position info

Each encoding applies exactly once in exactly one place — there is no double-application.

---

Normalization

Value	Class	Notes
"rmsnorm"	RMSNorm	LLaMA, Mistral, Qwen — no mean subtraction, no bias
"layernorm"	LayerNorm	BERT, GPT-2 — classic with bias

---

Dataloader

DataConfig

from transformer_toolkit.dataloader import DataConfig

cfg = DataConfig(
    seq_len     = 128,    # sequence length fed to the model
    batch_size  = 32,     # samples per batch
    split       = 0.9,    # fraction of data used for training
    stride      = None,   # None = non-overlapping windows (recommended)
                          # set stride < seq_len for overlapping windows
    shuffle     = True,
    num_workers = 4,
    pin_memory  = True,
    debug       = False,  # print sample decoding before training
    debug_n     = 3,      # number of samples to show when debug=True
)

stride — the default stride=None (equivalent to stride=seq_len) produces non-overlapping windows. For a 1.86M token dataset with seq_len=128 this gives ~14,600 clean distinct samples. Setting stride=1 gives 1.86M heavily-overlapping samples and causes rapid overfitting on small datasets.

Loading from a binary file

from transformer_toolkit.dataloader import save_binary, from_binary

# tokenize once and save
save_binary(tok.encode(text), "data.bin")

# load — supports both raw uint16 binary and .npy
train_dl, val_dl = from_binary("data.bin", cfg, tokenizer=tok)

# optionally save the splits as .npy for memmap reuse on future runs
train_dl, val_dl = from_binary(
    "data.bin", cfg,
    train_path = "train.npy",
    val_path   = "val.npy",
    tokenizer  = tok,
)

Memmap — loading pre-split .npy files (zero RAM)

On the second run, load the pre-split files directly. The token file stays on disk — only the pages you access load into RAM. Scales to 100GB+ datasets.

from transformer_toolkit.dataloader import from_npy_split

train_dl, val_dl = from_npy_split("train.npy", "val.npy", cfg, tokenizer=tok)

Loading from text files

from transformer_toolkit.dataloader import from_files

train_dl, val_dl = from_files(
    paths      = ["data1.txt", "data2.txt"],
    tokenizer  = tok,
    cfg        = cfg,
    train_path = "train.npy",   # optional — save for future memmap reuse
    val_path   = "val.npy",
    bos_id     = tok.bos_id,    # optional — wrap documents with BOS/EOS
    eos_id     = tok.eos_id,
)

Loading from HuggingFace

from transformer_toolkit.dataloader import from_hf

# streaming — no full download, infinite dataset
cfg_stream = DataConfig(seq_len=512, batch_size=16, streaming=True)
train_dl, val_dl = from_hf("roneneldan/TinyStories", tok, cfg_stream)

# in-memory — downloads fully, then splits and optionally saves as .npy
train_dl, val_dl = from_hf(
    dataset_name = "roneneldan/TinyStories",
    tokenizer    = tok,
    cfg          = cfg,
    text_col     = "text",
    bos_id       = 1,
    eos_id        = 2,
    train_path   = "train.npy",
    val_path     = "val.npy",
)

Debug mode

cfg = DataConfig(seq_len=128, batch_size=32, debug=True, debug_n=3)
train_dl, val_dl = from_binary("data.bin", cfg, tokenizer=tok)

Prints before training starts:

──────────────────────────────────────────────────────────────
  🔍 Debug samples (train)
  seq_len=128  stride=128  batch_size=32
──────────────────────────────────────────────────────────────

  sample 1
  x ids : [23, 451, 12, 8, 1203, 44, 91 ...] ... +121
  y ids : [451, 12, 8, 1203, 44, 91, 7  ...] ... +121
  x text: 'ROMEO:\nBut soft, what light through yonder window...'
  y text: '\nBut soft, what light through yonder window breaks'
  ✓  x/y alignment correct (y = x shifted by 1)

---

Tokenizers

Three tokenizers available, all sharing the same interface.

from transformer_toolkit.c_tokenizers import (
    ByteLevelTokenizer,
    RustBPETokenizer,
    HFTokenizer,
)

ByteLevelTokenizer

Zero dependencies. Every byte is a token (vocab size = 256). Works on any language out of the box.

tok = ByteLevelTokenizer()
ids = tok.encode("Hello world")   # [72, 101, 108, 108, 111, 32, 119, 111, 114, 108, 100]
txt = tok.decode(ids)             # "Hello world"

RustBPETokenizer

BPE tokenizer backed by HuggingFace's Rust tokenizers library. Trains ~100x faster than pure Python BPE.

pip install tokenizers

tok = RustBPETokenizer()
tok.train(open("data.txt").readlines(), vocab_size=8000)
tok.save("tokenizer.json")

# later
tok.load("tokenizer.json")
ids = tok.encode("Hello world")
txt = tok.decode(ids)
print(tok.vocab_size)   # 8000

HFTokenizer

Thin wrapper around any HuggingFace tokenizer.

pip install transformers

tok = HFTokenizer("gpt2")
ids = tok.encode("Hello world")
txt = tok.decode(ids)
print(tok.vocab_size)   # 50257

---

Trainer

TrainConfig

from transformer_toolkit.trainer import TrainConfig

cfg = TrainConfig(
    # ── steps ─────────────────────────────────────
    max_steps        = 10000,   # total optimizer steps
    eval_every       = 500,     # run validation every N steps
    save_every       = 1000,    # save step_N.pt every N steps
    log_every        = 50,      # print loss to console every N steps
    interruptible    = True,    # Ctrl+C saves a clean checkpoint

    # ── optimiser ─────────────────────────────────
    lr               = 3e-4,   # peak learning rate
    min_lr           = 3e-5,   # floor lr after cosine decay
    weight_decay     = 0.1,
    beta1            = 0.9,
    beta2            = 0.95,
    grad_clip        = 1.0,    # max gradient norm

    # ── lr schedule ───────────────────────────────
    warmup_steps     = 200,    # linear warmup from 0 to lr

    # ── efficiency ────────────────────────────────
    grad_accum_steps = 4,      # effective batch = batch_size × grad_accum_steps
    mixed_precision  = True,   # bf16/fp16 on CUDA, fp32 on CPU
    grad_checkpoint  = False,  # recompute activations to save VRAM (~20% slower)

    # ── checkpoints ───────────────────────────────
    ckpt_dir         = "checkpoints",
    save_best        = True,   # save best.pt when val loss improves
    save_step_ckpts  = True,   # save step_N.pt every save_every steps

    # ── huggingface hub ───────────────────────────
    hf_repo          = "username/my-model",
    hf_private       = True,
    hf_push_best     = True,
    hf_push_every_n  = False,
    hf_push_end      = True,
    hf_push_on_pause = True,
)

Trainer

from transformer_toolkit.trainer import Trainer

trainer = Trainer(
    model      = model,
    train_dl   = train_dl,
    val_dl     = val_dl,
    vocab_size = tok.vocab_size,
    cfg        = cfg_train,
    tokenizer  = tok,          # optional — used for HF hub uploads
)

trainer.train()

# resume from a checkpoint
trainer.train(resume_from="checkpoints/step_2000.pt")

Training output:

──────────────────────────────────────────────────────────────
  ⚡ Transformer Toolkit Trainer
  steps=3000  lr=0.0003  warmup=200  accum=4
  mixed_precision=True  grad_clip=1.0
──────────────────────────────────────────────────────────────

  step    100/3000  ████████░░░░░░░░░░░░░░░░░░░░  loss 3.1423  lr 1.5e-04  eta 4m
  step    200/3000  ████████████░░░░░░░░░░░░░░░░  loss 2.8901  lr 3.0e-04  eta 3m

  ● eval  step 300  val_loss 2.7130  ppl 15.07  ▼0.1823  ★ best

---

HuggingFace Hub

Login

from transformer_toolkit.hf_hub import login

login(token="hf_your_token_here")
# or
login(username="you", password="your_password")

Push to Hub

from transformer_toolkit.hf_hub import push_to_hub

push_to_hub(
    repo_id  = "username/my-model",
    model    = model,
    cfg      = cfg_model,
    tokenizer = tok,
    metrics  = {"val_loss": 1.83, "perplexity": 6.23},
    step     = 3000,
    private  = True,
)

Pull from Hub

from transformer_toolkit.hf_hub import pull_from_hub

pull_from_hub("username/my-model", save_dir="checkpoints")
# downloads model.pt, tokenizer.json, config.json, metrics.json

---

Generation

from transformer_toolkit.model import Transformer, TransformerConfig
from transformer_toolkit.c_tokenizers import RustBPETokenizer
from transformer_toolkit.trainer import load_ckpt
import torch

DEVICE = torch.device("cuda")

tok = RustBPETokenizer()
tok.load("tokenizer.json")

cfg = TransformerConfig(
    vocab_size = tok.vocab_size,
    dim        = 512,
    n_layers   = 8,
    n_heads    = 8,
    attn       = "gqa",
    n_kv_heads = 4,
    ffn        = "swiglu",
    hidden_dim = 2048,
    norm       = "rmsnorm",
    pos_enc    = "rope",
    dropout    = 0.0,   # always 0 at inference
)
model = Transformer(cfg).to(DEVICE)
load_ckpt("checkpoints/best.pt", model)
model.eval()

def generate(prompt, max_new=200, temperature=0.8, top_k=40):
    ids    = tok.encode(prompt)
    tokens = torch.tensor([ids], dtype=torch.long, device=DEVICE)
    out    = model.generate(tokens, max_new=max_new,
                             temperature=temperature, top_k=top_k)
    return tok.decode(out[0].tolist())

print(generate("ROMEO:"))

Generation parameters:

Parameter	Effect	Recommended
temperature	Higher = more random, lower = more repetitive	0.7 – 1.0
top_k	Only sample from top-k tokens	20 – 50
max_new	Number of new tokens to generate	100 – 500

---

Full Examples

Small model — Shakespeare (< 2M tokens, any GPU)

import torch, os
from transformer_toolkit.model import Transformer, TransformerConfig
from transformer_toolkit.c_tokenizers import RustBPETokenizer
from transformer_toolkit.dataloader import DataConfig, from_binary, save_binary
from transformer_toolkit.trainer import Trainer, TrainConfig

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tok = RustBPETokenizer()
tok.train(open("shakespeare.txt", encoding="utf-8").readlines(), vocab_size=4000)
tok.save("tokenizer.json")

if not os.path.exists("data.bin"):
    save_binary(tok.encode(open("shakespeare.txt", encoding="utf-8").read()), "data.bin")

train_dl, val_dl = from_binary("data.bin", DataConfig(
    seq_len=128, batch_size=32, split=0.9, stride=None
), tokenizer=tok)

model = Transformer(TransformerConfig(
    vocab_size = tok.vocab_size,
    dim        = 384,
    n_layers   = 6,
    n_heads    = 6,
    n_kv_heads = 3,
    attn       = "gqa",
    ffn        = "swiglu",
    hidden_dim = 1536,
    norm       = "rmsnorm",
    pos_enc    = "rope",
    dropout    = 0.1,
)).to(DEVICE)
print(model.n_params())   # ~15M

trainer = Trainer(model, train_dl, val_dl, tok.vocab_size, TrainConfig(
    max_steps        = 3000,
    warmup_steps     = 200,
    eval_every       = 300,
    lr               = 3e-4,
    grad_accum_steps = 4,
    mixed_precision  = True,
    save_best        = True,
    save_step_ckpts  = True,
))
trainer.train()

Large dataset — HuggingFace streaming

from transformer_toolkit.dataloader import DataConfig, from_hf, from_npy_split
from transformer_toolkit.c_tokenizers import HFTokenizer

tok = HFTokenizer("HuggingFaceTB/SmolLM-135M")

cfg = DataConfig(seq_len=512, batch_size=16, stride=None, num_workers=4)

# first run — tokenizes and saves splits as memmap .npy
train_dl, val_dl = from_hf(
    dataset_name = "roneneldan/TinyStories",
    tokenizer    = tok,
    cfg          = cfg,
    bos_id       = tok._tok.bos_token_id,
    eos_id       = tok._tok.eos_token_id,
    train_path   = "train.npy",
    val_path     = "val.npy",
)

# second+ runs — zero RAM, instant load
train_dl, val_dl = from_npy_split("train.npy", "val.npy", cfg, tokenizer=tok)

MoE model

model = Transformer(TransformerConfig(
    vocab_size     = tok.vocab_size,
    dim            = 512,
    n_layers       = 8,
    n_heads        = 8,
    attn           = "flash",
    ffn            = "moe",
    n_experts      = 8,
    top_k          = 2,
    moe_aux_weight = 0.01,
    pos_enc        = "rope",
    dropout        = 0.1,
)).to("cuda")

# model.forward() returns (logits, aux_loss)
# TrainConfig handles this automatically — no changes needed in training code
trainer = Trainer(model, train_dl, val_dl, tok.vocab_size, TrainConfig(
    max_steps = 5000,
    lr        = 3e-4,
))
trainer.train()

---

Architecture Reference

Input tokens [B, T]
      │
      ▼
Embedding [B, T, dim]  +  positional encoding (sinusoidal / learned)
      │
      ▼  × n_layers
┌─────────────────────────────────┐
│  Norm → Attention (+ RoPE/ALiBi)│
│  Residual                       │
│  Norm → FFN / SwiGLU / MoE      │
│  Residual                       │
└─────────────────────────────────┘
      │
      ▼
Final Norm → Linear head [B, T, vocab_size]

---

Requirements

• Python ≥ 3.10
• PyTorch ≥ 2.0
• tokenizers — for RustBPETokenizer (pip install tokenizers)
• transformers — for HFTokenizer (pip install transformers)
• datasets — for from_hf() (pip install datasets)
• huggingface_hub — for hub push/pull (pip install huggingface_hub)
• pydantic — for TrainConfig validation (pip install pydantic)

---

License

MIT