Minimal GPT-style language model class · CS336 Language Modeling From Scratch

Deliverableartifact

A GPTModel class that wires embeddings, N Transformer blocks, and an output head, with a script that prints parameter count and verifies a dummy forward pass.

Goals

Instantiate your GPT model and run a forward pass on a batch of token ids, confirming output shapes (batch, seq_len, vocab_size).
Count the total number of parameters for given hyperparameters (layers, heads, d_model, vocab size).
Show that shared token embedding and output projection weights can be tied if you choose to do so.
Explain how sequence length and number of layers affect both memory usage and FLOPs.

Resources

📄stanford-cs336/assignment1-basics

Overview: what you’re building

You’ll build a minimal GPTModel class in PyTorch that does four things:

Turns token ids and positions into embeddings
Runs them through N stacked Transformer blocks
Maps the hidden states to vocabulary-sized logits with a final linear layer
Supports an autoregressive forward pass and a small script to check parameter count and a dummy forward

Here’s the rough structure we’re aiming for:

Rendering diagram…

You’ll end up with a file you can actually run, something like:

bash

python gpt_model.py

that prints the number of parameters and runs a dummy forward pass without errors.

1 / 6