Writing an LLM from scratch, part 34b -- from bigrams to GPT-2, one component at a time (in JAX)

This blog post details the author's journey in building a large language model (LLM) from scratch using JAX, documenting the process through various implementations of the GPT-2 architecture, including insights on multi-head attention, layer normalization, and dropout. The author shares personal revelations derived from experimenting and modifying strategies to optimize the model's performance, ultimately concluding that JAX offers superior results compared to PyTorch, although results show variations in performance depending on factors like loss and architecture design.