Learn how to implement LLaMA 3, a decoder-only transformer language model, using JAX in just 100 lines of code. This post covers various components such as tokenization with Byte Pair Encoding (BPE), embeddings, rotary positional encoding, group-query attention, and feed-forward layers. The guide aims to be educational with an emphasis on functional programming, initialization of weights, and training the model on a Shakespeare dataset using Stochastic Gradient Descent (SGD) for optimization.
Table of contents
Table of ContentsLLaMA3Model Weights InitializationTokenizationEmbeddingsRoot Mean Square Layer NormalizationRotary Positional EncodingGroup-Query AttentionFeed-forwardTransformer-blockForward-PassDatasetLoss FunctionUpdate functionTrainig-LoopSort: