Part 1 — Build a Small Language Model

A language model becomes easier to understand when we follow the complete path from raw text to the probability distribution over the next token. That path contains several distinct ideas: a tokenizer turns text into discrete symbols, a decoder-only Transformer transforms those symbols into contextual representations, a causal objective teaches the model to predict the future from the past, and an autoregressive decoding loop turns the trained model back into text.

SmallLM-Forge implements each of these pieces at a scale where the equations, tensor shapes, and PyTorch operations can be inspected directly. This part walks through that stack from the bottom up.

From Raw Text to a Language Model

A causal language model learns a probability distribution over token sequences. For a sequence

the model factorizes the sequence probability as

Training therefore reduces to the same task at every position: predict the next token from the tokens that came before it.

The complete path can be summarized as

raw text
   ↓
tokenizer
   ↓
token ids
   ↓
token embeddings
   ↓
decoder-only Transformer
   ↓
vocabulary logits
   ↓
next-token probabilities

The same model is used during generation. We provide an initial sequence, read the probability distribution at its final position, choose one token, append it, and run the model again. Because each new token is fed back into the model as part of the context, this generation process is called autoregressive generation.

This gives us a useful way to organize the implementation:

representation
    ↓
architecture
    ↓
training objective
    ↓
optimization
    ↓
generation

We start with representation.

Build a Byte-Level BPE Tokenizer

Neural networks operate on numerical tensors, so text first has to be mapped to a finite vocabulary of integer token IDs.

A tokenizer has to balance two competing effects. Very small units such as bytes provide complete coverage of arbitrary input text, although they produce long sequences. Large units such as full words shorten sequences, although a fixed word vocabulary grows quickly and handles unseen forms poorly.

Byte Pair Encoding provides a practical middle ground. We begin with bytes and learn larger reusable units by repeatedly merging frequent adjacent pairs.

Start from Bytes

SmallLM-Forge uses a byte-level vocabulary as the base representation.

A UTF-8 string is first converted to bytes. Every possible byte value can then be represented by one of 256 initial symbols:

text
   ↓ UTF-8
bytes
   ↓
byte symbols

This gives the tokenizer complete coverage of the input space. New vocabulary items are created by composing these initial symbols.

The tokenizer training code uses a byte-to-Unicode mapping so the intermediate symbols can be represented as Python strings while preserving the original byte values.

Count and Merge Token Pairs

Suppose the corpus contains the sentence

the cat sat on the mat

Splitting each word into its symbols gives the token sequences

t h e
c a t
s a t
o n
t h e
m a t

The tokenizer counts every adjacent pair that appears, word by word:

(a, t)    3
(t, h)    2
(h, e)    2
(c, a)    1
(s, a)    1
(o, n)    1
(m, a)    1

The most frequent pair, (a, t), is merged into a new symbol at. The affected sequences become

t h e
c at
s at
o n
t h e
m at

The statistics are then recomputed and the process repeats, so other frequent fragments such as th and he would be merged on later rounds.

At a high level, training looks like

while len(vocab) < target_vocab_size:
    pair = most_frequent_pair(corpus)
    merges.append(pair)
    vocab.add(pair[0] + pair[1])
    corpus = apply_merge(corpus, pair)

The actual implementation keeps two main frequency structures:

words_by_tokens
    tokenized word → frequency

pair_frequencies
    adjacent token pair → frequency

and an ordered merges list that records the sequence of BPE operations.

The order matters. BPE encoding is determined by the learned merge priority, so the tokenizer needs both the final vocabulary and the ordered merge rules.

Add Special Tokens

Special symbols can be added after the learned vocabulary has been built. In the pretraining workflow, an end-of-sequence token gives the model an explicit marker for sequence termination.

Conceptually, the final tokenizer state contains

base byte vocabulary + learned BPE tokens + special tokens

Each token maps to one integer ID.

Encode Text with the Learned Merges

Training the tokenizer and using the tokenizer are separate operations.

During encoding, text is converted to its byte-level representation and the learned merge rules are applied in their learned order. The final symbols are mapped to token IDs.

Decoding reverses this process:

token ids
   ↓
token strings
   ↓
bytes
   ↓ UTF-8
text

This round trip is important because the rest of the language-model stack only sees integer IDs. The tokenizer defines the boundary between raw text and the model’s numerical input space.

Prepare Sequences for Causal Language Modeling

Real text examples have different lengths, while a training batch is a dense tensor. The data pipeline therefore has to tokenize, truncate, pad, and record which positions contain real tokens.

SmallLM-Forge keeps these responsibilities separate.

TextDataset stores raw strings and tokenizes one sample when it is requested:

raw string
   ↓ tokenizer.encode(...)
list[int]

A collator then combines multiple tokenized examples into a batch.

Consider two sequences:

A = [12, 45, 81, 7]
B = [93, 20]

After padding:

input_ids

[12, 45, 81,   7]
[93, 20, PAD, PAD]

The corresponding attention mask is

attention_mask

[1, 1, 1, 1]
[1, 1, 0, 0]

The model can now distinguish real sequence positions from padding introduced only for batching.

Long examples are truncated to the configured context length. The resulting batch tensors have shape

where \(B\) is the batch size and \(T\) is the sequence length used for that batch.

These two tensors follow the sample through the rest of the training pipeline:

input_ids       → token identity
attention_mask  → valid sequence positions

Build the Decoder-Only Transformer

The language model maps token IDs to logits over the vocabulary.

At the highest level, the architecture is compact:

If

• \(B\) is the batch size,
• \(T\) is the sequence length,
• \(H\) is the hidden dimension,
• \(V\) is the vocabulary size,

then the main tensor shapes are

The first transformation is performed by the token embedding. The last is performed by a linear language-model head.

SmallLM-Forge provides three preset configurations:

The presets change the scale of the same architecture. The internal building blocks remain the same.

Before assembling a full Transformer block, we can examine the two components that define most of its behavior: positional encoding and causal attention.

Build Causal Self-Attention with GQA

The Transformer block receives hidden states

where \(B\) is the batch size, \(T\) is the sequence length, and \(H\) is the hidden dimension.

Self-attention first transforms each hidden state into three representations:

The query represents what a position is looking for, the key represents what a position can match against, and the value contains the information that will be aggregated once the attention weights are known.

For one attention head,

where \(d\) is the head dimension.

SmallLM-Forge uses Grouped-Query Attention (GQA) rather than giving every query head an independent key/value pair.

Share Key and Value Heads

In regular multi-head attention, the number of query, key, and value heads is the same. GQA keeps more query heads than key/value heads and lets several query heads share one KV head.

For example,

The model configuration therefore distinguishes

and

The size of one query group is

and the head dimension is

Build the Query Heads

The query path begins with

X
[B, T, H]

and projects it back into the same hidden width:

self.q_proj = nn.Linear(
    config.hidden_dim,
    config.hidden_dim,
    bias=False,
)

During the forward pass,

query = (
    self.q_proj(x)
    .view(
        bs,
        seq_len,
        self.config.n_head,
        self.head_dim,
    )
    .transpose(1, 2)
)

rearranges the tensor into

Q
[B, h_q, T, d]

so each query head can attend independently.

Build the Shared KV Heads

Keys and values are projected together:

self.kv_proj = nn.Linear(
    config.hidden_dim,
    2 * config.hidden_dim // self.q_per_kv,
    bias=False,
)

with

self.q_per_kv = (
    config.n_head
    // config.n_kv_head
)

The forward pass first produces

kv = (
    self.kv_proj(x)
    .view(
        bs,
        seq_len,
        self.config.n_kv_head,
        2 * self.head_dim,
    )
    .transpose(1, 2)
)

giving

KV
[B, h_kv, T, 2d]

The extra factor of two stores key and value features in the same projection.

To match each query group with its shared KV head, the implementation inserts a group dimension and expands it:

kv = (
    kv[:, :, None, :, :]
    .expand(
        bs,
        self.config.n_kv_head,
        self.q_per_kv,
        seq_len,
        2 * self.head_dim,
    )
    .reshape(
        bs,
        self.config.n_head,
        seq_len,
        2 * self.head_dim,
    )
)

The resulting tensor has shape

[B, h_q, T, 2d]

and can now be split into

key, value = torch.chunk(
    kv,
    2,
    dim=-1,
)

so that

Q  [B, h_q, T, d]
K  [B, h_q, T, d]
V  [B, h_q, T, d]

share the same execution layout.

The learned KV parameters still come from only \(h_{kv}\) heads. The expansion only prepares those shared heads for the batched attention computation.

Before computing \(QK^\top\), the model adds positional information to \(Q\) and \(K\).

Add Position with RoPE

The query and key tensors above describe token content, but the raw similarity

does not explicitly describe the positional relationship between token \(m\) and token \(n\).

The causal mask will later decide which positions are visible. RoPE adds information about where the visible positions are relative to each other.

SmallLM-Forge uses Rotary Positional Embeddings (RoPE) by rotating query and key vectors according to their token positions before their dot product is computed.

Why Use a Rotation?

Consider the two-dimensional rotation matrix

For a query at position \(m\) and a key at position \(n\), RoPE applies

Their inner product becomes

Because

we obtain

so the positional contribution to the attention score depends on the relative offset

This is the main reason rotations are useful here: each token receives an absolute position-dependent transformation, while the query-key comparison naturally exposes their relative displacement.

A rotation also preserves vector length,

so position changes the orientation of \(Q\) and \(K\) without directly rescaling their magnitude.

Apply Multiple Frequencies Across the Head

An attention head contains many dimensions. RoPE groups them into two-dimensional pairs:

[x0, x1, x2, x3, x4, x5, ...]
          ↓
[(x0, x1), (x2, x3), (x4, x5), ...]

Each pair \(i\) uses

and position \(m\) is represented by the rotation angle

The same token position is therefore represented at several angular frequencies across the head dimension.

For one pair,

the rotation is

Precompute the Rotation Angles

SmallLM-Forge builds the frequency table once:

def _compute_freqs(
    head_dim: int,
    max_seq_len: int,
) -> Tensor:
    theta = 10000 ** (
        -torch.arange(
            0,
            head_dim,
            2,
            dtype=torch.float,
        ) / head_dim
    )

    pos = torch.arange(
        max_seq_len,
        dtype=torch.float,
    )

    return torch.outer(pos, theta)

The result has shape

[max_seq_len, d / 2]

and stores

The attention module registers it as a buffer:

freqs = _compute_freqs(
    self.head_dim,
    config.max_seq_len,
)

self.register_buffer(
    "rope_freqs",
    freqs,
)

so it follows the model between devices without becoming a trainable parameter.

Rotate Q and K

During a forward pass, only the rows needed for the current sequence are used:

seq_len = q.shape[2]
freqs = freqs[:seq_len].to(q.device)

cos = freqs.cos()[None, None, :, :]
sin = freqs.sin()[None, None, :, :]

The sine and cosine tensors have shape

[1, 1, T, d / 2]

and broadcast over both the batch and head dimensions.

The final head dimension of \(Q\) and \(K\) is then grouped into pairs:

q_reshaped = q.float().reshape(
    *q.shape[:-1],
    -1,
    2,
)

k_reshaped = k.float().reshape(
    *k.shape[:-1],
    -1,
    2,
)

giving

[B, h_q, T, d / 2, 2]

and the two-dimensional rotation is applied directly:

q_rot = torch.stack(
    [
        q_reshaped[..., 0] * cos
        - q_reshaped[..., 1] * sin,

        q_reshaped[..., 1] * cos
        + q_reshaped[..., 0] * sin,
    ],
    dim=-1,
).reshape(q.shape)

The same operation is applied to \(K\), after which the original dtype is restored.

RoPE therefore changes the geometry of the attention tensors while preserving their original shape:

before: [B, h_q, T, d]
after:  [B, h_q, T, d]

Insert RoPE into the GQA Path

RoPE is applied after GQA has produced \(Q\), \(K\), and \(V\):

query, key = apply_rope(
    query,
    key,
    self.rope_freqs,
)

Only \(Q\) and \(K\) are rotated. \(V\) keeps the content that will eventually be aggregated.

Compute Causal Attention

With \(Q\), \(K\), and \(V\) prepared, the attention score is

attn_score = torch.matmul(
    query,
    key.transpose(2, 3),
) * self.scale

where

self.scale = self.head_dim ** -0.5

implements the factor

The score tensor has shape

[B, h_q, T, T]

and contains one query-key score for every visible head and token pair.

Apply the Causal Mask

A decoder cannot use future tokens when predicting the next token.

SmallLM-Forge creates an upper-triangular mask:

causal = torch.triu(
    torch.full(
        (seq_len, seq_len),
        float("-inf"),
        device=x.device,
    ),
    diagonal=1,
)

attn_score = attn_score + causal

giving the visibility pattern

          key position
          1  2  3  4

query 1   ✓  ×  ×  ×
      2   ✓  ✓  ×  ×
      3   ✓  ✓  ✓  ×
      4   ✓  ✓  ✓  ✓

If a padding mask is provided, padded key positions are also assigned \(-\infty\).

The two mechanisms have separate roles:

RoPE
    tells the score how positions relate

causal mask
    tells the score which positions are allowed

Turn Scores into an Attention Output

Softmax converts the remaining scores into probabilities:

attn_probs = F.softmax(
    attn_score,
    dim=-1,
)

attn_probs = self.attn_dropout(
    attn_probs,
)

The values are then aggregated:

out = torch.matmul(
    attn_probs,
    value,
)

producing

[B, h_q, T, d]

The heads are merged back into the hidden dimension:

out = (
    out.transpose(1, 2)
    .reshape(
        bs,
        seq_len,
        self.config.hidden_dim,
    )
)

and the final output projection returns

return self.out_proj(out)

with shape

[B, T, H].

Complete the Transformer Block

The attention layer is one of two main transformations inside each decoder block. SmallLM-Forge uses a pre-normalization structure with RMSNorm, residual connections, causal GQA, and a SwiGLU feed-forward network.

The implementation mirrors this structure directly:

def forward(
    self,
    x,
    attention_mask=None,
):
    x = x + self.res_dropout_1(
        self.attn(
            self.ln_1(x),
            attention_mask,
        )
    )

    x = x + self.res_dropout_2(
        self.mlp(
            self.ln_2(x)
        )
    )

    return x

Normalize with RMSNorm

RMSNorm rescales a hidden vector using its root-mean-square magnitude:

The output is

where \(g_i\) is a learned scale parameter.

SmallLM-Forge implements this as

def forward(self, x):
    input_dtype = x.dtype

    x = x.to(torch.float32)

    rms = torch.sqrt(
        torch.mean(
            x**2,
            dim=-1,
            keepdim=True,
        )
        + self.eps
    )

    return (
        self.scale * x / rms
    ).to(input_dtype)

The normalization is computed in float32 and then returned to the original dtype.

Because the block is pre-normalized, RMSNorm is applied before both attention and the feed-forward network.

Transform Features with SwiGLU

The second sub-layer is a gated feed-forward network.

SwiGLU computes

Conceptually,

SmallLM-Forge combines the first two projections into one linear layer:

self.fc1 = nn.Linear(
    config.hidden_dim,
    2 * config.intermediate_dim,
    bias=False,
)

self.fc2 = nn.Linear(
    config.intermediate_dim,
    config.hidden_dim,
    bias=False,
)

and splits the result:

x1, x3 = torch.chunk(
    self.fc1(x),
    2,
    dim=-1,
)

return self.fc2(
    F.silu(x1) * x3
)

The feed-forward path therefore expands the hidden representation into the intermediate dimension, applies a learned gate, and projects it back to \(H\).

Stack the Blocks

The complete decoder repeats this block \(N\) times:

token embedding
      ↓
Transformer block
      ↓
Transformer block
      ↓
...
      ↓
Transformer block
      ↓
final RMSNorm
      ↓
language-model head

Every block preserves the hidden-state shape

which makes the stack compositional.

After the final block, the model applies one more RMSNorm and projects each hidden state into vocabulary space:

These logits are the input to the causal next-token training objective.

Train with Next-Token Prediction

Once the model produces vocabulary logits, pretraining needs a target.

For a token sequence

The  cat  sat  on  the  mat

the model can be trained with a one-position shift:

input context position:
The  cat  sat  on  the

target:
cat  sat  on   the  mat

In tensor form:

targets = input_ids[:, 1:]
predicted_logits = logits[:, :-1, :]

Each position is therefore trained to predict the next token.

The causal language-model objective is the cross-entropy

where \(m_t\) is the attention mask for valid target positions.

Padding introduced during batching contributes no loss.

Run the Optimization Loop

The training loop follows a standard sequence:

batch
  ↓
model forward
  ↓
masked next-token cross-entropy
  ↓
backward
  ↓
gradient clipping
  ↓
AdamW update
  ↓
learning-rate scheduler

In simplified form:

for input_ids, attention_mask in train_loader:
    logits = model(input_ids, attention_mask)
    loss = causal_lm_loss(input_ids, attention_mask, logits)

    optimizer.zero_grad()
    loss.backward()
    clip_grad_norm_(model.parameters(), max_norm)
    optimizer.step()
    scheduler.step()

SmallLM-Forge uses AdamW, gradient clipping, and a linear schedule with warmup and decay. Validation is run periodically with gradients disabled.

That gives the training loop two measurements:

training loss
validation loss

The goal of this stage is straightforward: improve next-token prediction on the training distribution while checking that the same objective also improves on held-out text.

Generate Autoregressively

Training evaluates every sequence position in parallel. Generation works one token at a time.

The decoding loop is

prompt
  ↓
model
  ↓
last-position logits
  ↓
sampling rule
  ↓
next token
  ↓
append
  ↓
repeat

Only the logits from the final sequence position are needed because that position represents

The generated token is appended to the context and becomes part of the next forward pass.

SmallLM-Forge supports several decoding controls.

Greedy Decoding

Greedy decoding selects

This is deterministic for a fixed model and prompt.

Temperature

Temperature rescales logits before softmax:

Lower temperatures concentrate probability on the most likely tokens. Higher temperatures flatten the distribution and increase diversity.

Top-k Sampling

Top-k keeps only the \(k\) highest-logit candidates:

full vocabulary
      ↓
keep k best tokens
      ↓
renormalize
      ↓
sample

All other logits are set to negative infinity before softmax.

Top-p Sampling

Top-p, or nucleus sampling, sorts candidates by probability and keeps the smallest set whose cumulative probability reaches a threshold \(p\).

The size of the candidate set therefore adapts to the shape of the model’s distribution.

A confident prediction may need only a few tokens. A flatter distribution may retain many more.

The implementation can apply top-k and top-p in the same decoding step.

Stop at EOS

Generation ends when

• the maximum number of new tokens has been generated, or
• the tokenizer’s EOS token is produced.

If the context grows beyond the configured context window, the generator keeps the most recent tokens that fit inside the model’s maximum sequence length.

Put Everything Together

We can now assemble the complete pretraining stack. The training path ends at the loss, and the inference path continues from the same logits:

Vocabulary Logits
   ↓
Temperature
   ↓
Top-k / Top-p
   ↓
Next Token
   ↓
Append and Repeat

Each component has a narrow responsibility:

This modular view is useful because each design choice can be changed without having to reinterpret the whole model.

Example: Pretraining on WikiText-103

The repository includes a small-scale pretraining workflow on WikiText-103 to exercise the complete stack.

The experiment is an application of the components described above. It uses the same byte-level BPE implementation, decoder-only Transformer, causal objective, and generation code.

Train the Tokenizer

A byte-level BPE tokenizer is trained on the text corpus and then reused by the language model.

The resulting tokenizer has been released on Hugging Face as:

liuhailin0123/wiki103-bpe-tokenizer

This artifact contains the learned vocabulary and merge rules needed to encode and decode text consistently with the pretraining run.

Train the Mini Model

The example uses the mini Transformer configuration:

6 Transformer blocks
6 query heads
3 key/value heads
hidden size 384
SwiGLU intermediate size 1024

The model is trained with causal next-token prediction using the tokenizer produced by the previous stage.

The resulting model has been released as:

liuhailin0123/wiki103-llama-mini

The accompanying wiki103_pretrain.ipynb notebook shows how the tokenizer, dataset, model, trainer, and generator are connected for this experiment.

The important outcome of the example is the integration path:

corpus
  ↓
train tokenizer
  ↓
tokenize text
  ↓
pretrain model
  ↓
validate
  ↓
generate
  ↓
save reusable artifacts

The same components can be reused with other corpora, vocabulary sizes, model configurations, and training schedules.