Tiny Transformer

Tiny Causal Transformer — Built From Scratch

This PyTorch-only educational model demonstrates next-token training without a pretrained-model download. It is intentionally not a large language model.

flowchart TD
    Text["UTF-8 text"] --> Bytes["Byte token IDs: bytes + 2, EOS=1, PAD=0"]
    Bytes --> Emb["Token embedding: vocab 258, width 64"]
    Pos["Learned positions up to seq_len"] --> Sum["Sum embeddings"]
    Emb --> Sum
    Sum --> Layer["1 TransformerEncoderLayer: 4 attention heads; FFN width 128"]
    Mask["Upper-triangular causal attention mask"] --> Layer
    Layer --> LN["LayerNorm"] --> Head["LM head: 64 to 258 logits"]
    Head --> CE["Shifted next-token cross entropy"]
    Labels["Labels: pad positions ignored (-100)"] --> CE

Loading

Architecture

• Vocabulary size: 258 (PAD=0, EOS=1, bytes mapped to 2..257).
• Model width: 64; 4 attention heads; 1 Transformer encoder layer with a feed-forward dimension of 128.
• Learned positional embeddings; upper-triangular causal attention mask.
• LayerNorm and Linear head mapping hidden width 64 → 258 logits.
• Forward returns loss, not generated text. No text-generation loop or sampling API is implemented.

Training behavior

AdamW defaults to 3e-4 unless overridden. Training clips gradients to maximum norm 1.0. The single-process CPU smoke-test command is:

python llm_ddp_parallel_training.py --model tiny --cpu --max-steps 2

The tiny checkpoint includes model_state_dict, seq_len, width, and heads; this project does not provide a standalone loader/inference CLI for it.