A minimal implementation of a Transformer-based language model, following nanoGPT very closely. This implementation emphasizes pedagogy, using JAX and Penzai's named array system to make the attention mechanism more intuitive and easier to understand. This project was developed for CSC2541: Large Models at the University of Toronto.
- Self-contained implementation of a Transformer model in JAX
- Uses named arrays to make tensor operations more semantically meaningful
- Trains on a subset of the TinyStories dataset for lightweight experimentation
The cleanest part of the code base is model.py, which contains the implementation of the Transformer model. We recommend you focus your efforts there.
We use Penzai's NamedArray to make Transformer's computations more clear. Although they require a bit of getting used to, named axes allow you to focus on the key semantics of the computation, while ignoring the structure of the data. The best place to start learning about Penzai's named axes is here.
For example, in this pseudo-code for a (simplified) attention mechanism, instead of keeping track of tensor axes, we explicitly name our dimensions. This is what it looks like:
# Every key interacts with every query via the dot of their embeddings
attn_logits = nmap(jnp.dot)(q.untag("embed"), k.untag("embed"))
# Compute a distribution over keys for each query
attn_dist = nmap(jax.nn.softmax)(attn_logits.untag("key))
attn_dist = attn_dist.tag("key")
# The value returned for each query is the average value
# indexed by the keys, under the attention distribution
out = nmap(jnp.dot)(attn_dist.untag("key"), v.untag("key"))
This makes it a bit more clear which axes are participating in which computation.
- Install dependencies:
jax,penzai,optax,tqdm,requests,numpy - Run the training script:
python train.py
The model will train on a subset of TinyStories and generate sample text.
This project is licensed under the GNU General Public License v3.0 - see the LICENSE file for details.