Implementing a GPT Model

Everything hinges on two ideas that make deep stacks trainable. The first keeps activations well-behaved: layer normalization.

The second is the block's actual workhorse — a feed-forward network with a smooth activation, made trainable at depth by shortcut connections.

Stack twelve of those blocks, add the embeddings and output head, and the model is complete. Untrained, it already runs — it just has nothing to say yet.

The lab

The parameter count is pure arithmetic — no weights needed. Build any GPT and see where its 163 million (or 1.5 billion) weights actually live.

And the forward pass itself, running live: a real — if tiny and untrained — GPT, tokenizing your text as bytes and computing genuine logits in your browser.

The real code

The demos run a TypeScript port; here is the Python they mirror. A transformer block is the two-shortcut sandwich — normalize, attend, add; normalize, feed-forward, add:

class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.att = MultiHeadAttention(
            d_in=cfg["emb_dim"], d_out=cfg["emb_dim"],
            context_length=cfg["context_length"], num_heads=cfg["n_heads"],
            dropout=cfg["drop_rate"], qkv_bias=cfg["qkv_bias"])
        self.ff = FeedForward(cfg)
        self.norm1 = LayerNorm(cfg["emb_dim"])
        self.norm2 = LayerNorm(cfg["emb_dim"])
        self.drop_shortcut = nn.Dropout(cfg["drop_rate"])
 
    def forward(self, x):
        x = x + self.drop_shortcut(self.att(self.norm1(x)))   # attention + shortcut
        x = x + self.drop_shortcut(self.ff(self.norm2(x)))    # feed-forward + shortcut
        return x

The full model is embeddings → a Sequential of blocks → final norm → output head, and generation is a greedy loop over its logits:

class GPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
        self.trf_blocks = nn.Sequential(*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])
        self.final_norm = LayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)
 
    def forward(self, in_idx):
        b, seq_len = in_idx.shape
        x = self.tok_emb(in_idx) + self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = self.trf_blocks(self.drop_emb(x))
        return self.out_head(self.final_norm(x))
 
 
def generate_text_simple(model, idx, max_new_tokens, context_size):
    for _ in range(max_new_tokens):
        logits = model(idx[:, -context_size:])[:, -1, :]   # last position only
        idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # greedy
        idx = torch.cat((idx, idx_next), dim=1)
    return idx

class LayerNorm(nn.Module):
    def __init__(self, emb_dim):
        super().__init__()
        self.eps = 1e-5
        self.scale = nn.Parameter(torch.ones(emb_dim))
        self.shift = nn.Parameter(torch.zeros(emb_dim))
 
    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)  # biased, like GPT-2
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        return self.scale * norm_x + self.shift
 
 
class GELU(nn.Module):
    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(
            torch.sqrt(torch.tensor(2.0 / torch.pi)) *
            (x + 0.044715 * torch.pow(x, 3))))
 
 
class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),   # expand 4×
            GELU(),
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),   # project back
        )
 
    def forward(self, x):
        return self.layers(x)

Four things to remember

1. A GPT is one block, repeated. Twelve identical transformer blocks do the real work; scaling up mostly means stacking more of them.
2. Normalize, transform, add. Every sub-layer follows the same recipe — LayerNorm in front, a shortcut connection around — so gradients survive the depth.
3. The feed-forward net is where the parameters are. Its 4× expansion holds roughly twice the weights of attention in each block.
4. The architecture works before it's smart. The untrained model generates fluent-shaped gibberish. Giving it something to say is pretraining — the next chapter.

Adapted from Build a Large Language Model (From Scratch) by Sebastian Raschka — original notebook (Apache 2.0). The parameter counts, LayerNorm and GELU values, gradient means, and the model's generated output on this page are torch's real seed-123 output, verified number-for-number against the Python.