Implementing LLaMA3 in 100 Lines of Pure Jax

In this post, we'll implement llama3 from scratch using pure jax in just 100 lines of code. Why jax? Because I think it has good aesthetics. Also jax looks like a NumPy wrapper but it has some cool features like xla; a linear algebra accelerator, jit, vmap, pmap etc., which makes your training go brr brr.

Jax is one of the first libraries which strongly focuses on the soul of pure functional programming which makes it more cool.¹

Note :

LLaMA3

So lets fucking go !! we're doing it, get your diet coke !! First, we will begin with setting up device⁵ and configuring the model.

So these are the hyperparameter we need to train approximately 2 million parameters model.

Model Weights Initialization

In pure JAX, we don't use classes like in PyTorch. We use only pure fucntions why ? cause it makes our code more predictable and easier to parallelize. A pure function always returns the same output for the same input and doesn’t cause any side effects.⁶ For example, if you call F(x), you'll always get the same y.

Since we aren’t using a framework like PyTorch’s nn.Module to automatically track parameters, we must initialize and update our weights manually.

Handling randomness is also different. Instead of relying on a single global seed as in NumPy or PyTorch, in jax we need to manage randomness with explicit pseudo-random number generator (PRNG) keys. Each random operation gets its own unique key, which is derived by splitting a parent key. This will help in reproducibility and parallelism.

For example, below you can see we are creating a key and splitting it into sub keys and then providing that key to the function which involves the randomness.

Now lets start with our Model Weights Initialization, first we create the random values for our parameters with normal ditribuition.

Next, we'll identify all the learnable parameters of our model(llama3), assign each a unique key to ensure reproducibility, and apply the initialization process to them.

Since weights are essentially numbers stored in arrays, we can use dictionaries to manage them as key-value pairs.

First we will start with attention module which has four trainable parameters.

Next we have our Feed-forward network which has 3 trainable parameters.

Then we combine our weights into transformer block, adding two additional parameters for two layers of RMSNorm.

Finally we assemble Model's Weights Initialization in one place.

Tokenization

Tokenization means dividing the text into words and subwords (tokens). We will be using Byte Pair Encoding (BPE) for training our model (BPE was used in training Llama 3).⁷ I will not build bpe from scratch we will use tiktoken library by openai for bpe.

Embeddings

We cannot provide tokens directly to a model because tokens are discrete, while neural networks operate on continuous numerical data this is important for performing mathematical operations. Therefore, we use an embedding layer to convert the discrete tokens into a continuous vector space. These embeddings also help encode the semantic and syntactic relationships between tokens.

There are two types of embeddings: static and dynamic. We use dynamic embeddings to train LLMs. Why? Because static embeddings work well for finding similarities between words and representing them in a similar vector space, as seen in the first image.

However, they suffer from semantic ambiguity, as shown in the second image. This is where Self-Attention comes in, it refines these embeddings to incorporate context. So, we start with random embeddings and update them according to the context.

Root Mean Square Layer Normalization

RMS normalization is an important layer in llama3 models. It helps keep the training stable by making sure that the numbers in the network don’t become too high or too low. This balance is very important, especially in deep networks.

Rotary Positional Encoding

Transformers don't naturally know the order of tokens, so we need to add some position info. In llama3 to solve this we have ROPE. It does this by “rotating” the query and key vectors based on their position.⁸

How It Works:

Precompute Rotation Factors: First we create a table of rotation factors using a range of frequencies. This means each token gets its own unique rotation angle.

Apply the Rotation:

Pair Up Features: we reshape the vectors so that every two numbers form a pair; imagine them as the real and imaginary parts of a complex number.

Rotate: We multiply these complex numbers by our precomputed rotation factors. This rotates each pair in the complex plane.

Convert Back: Finally, we split the rotated complex numbers back into their real and imaginary parts to restore the original shape.

Math Behind It: For each pair $(x_{2i},x_{2i+1})$, the rotation is given by:
$$\left(\begin{array}{c}{x}_{2i}^{\prime }\\ x_{2i+1}^{\prime }\end{array}\right)=\left(\begin{array}{cc}\cos \left({\theta }_i\right)& -\sin \left({\theta }_i\right)\\ \sin \left({\theta }_i\right)& \cos \left({\theta }_i\right)\end{array}\right)\left(\begin{array}{c}{x}_{2i}\\ x_{2i+1}\end{array}\right)$$ where ${\theta }_i$ is the rotation angle for that token. In short, ROPE embeds positional information directly into the token features by rotating them. This way attention module gets the info about token order without extra position vectors.

Group-Query Attention

Now it's time for attention. Grouped Query Attention (GQA) is an optimized version of Multi-Head Attention that improves efficiency by sharing key and value representations among multiple query heads. This reduces computational overhead and memory usage, enabling faster inference and better scaling for transformer models. At it's core, it's just self-attention but with some modification.

Scaled Dot-Product Attention:

$$A=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_h}}\right)V$$

# Attention mechanism with grouped-query attention
def attention(params, x, mask, freqs_cis, n_heads, n_kv_heads, cache=None, position=0):
    # Get input dimensions
    B, T, C = x.shape
    head_dim = C // n_heads
    
    # Project inputs to queries, keys, and values
    q = jnp.dot(x, params['wq']).reshape(B, T, n_heads, head_dim)
    k = jnp.dot(x, params['wk']).reshape(B, T, n_kv_heads, head_dim)
    v = jnp.dot(x, params['wv']).reshape(B, T, n_kv_heads, head_dim)
    
    # Apply rotary embeddings
    q, k = apply_rotary_emb(q, k, freqs_cis[position:position + T])
    
    # Handle cache for inference
    if cache is not None:
        k = jnp.concatenate([cache[0], k], axis=-1])
        v = jnp.concatenate([cache[1], v], axis=-1])
    new_cache = (k, v)
    
    # Repeat k/v heads for grouped-query attention
    k = repeat_kv(k, n_heads // n_kv_heads)
    v = repeat_kv(v, n_heads // n_kv_heads)
    
    # Compute attention scores and apply attention
    q, k, v = map(lambda x: x.transpose(0, 2, 1, 3), (q, k, v))
    scores = jnp.matmul(q, k.transpose(0, 1, 3, 2)) / math.sqrt(head_dim)
    
    # Apply attention mask if provided
    if mask is not None:
        scores = scores + mask[:, :, :T, :T]
    
    # Compute attention weights and final output
    scores = jax.nn.softmax(scores, axis=-1)
    output = jnp.matmul(scores, v)
    output = output.transpose(0, 2, 1, 3).reshape(B, T, -1)
    
    return jnp.dot(output, params['wo']), new_cache

KV-cache : It stores previously computed key (K) and value (V) tensors from past tokens. We can cache this kv-cache during inference.

if cache is not None:
    k = jnp.concatenate([cache[0], k], axis=-1)  # Concatenate cached keys with new keys
    v = jnp.concatenate([cache[1], v], axis=-1)  # Concatenate cached values with new values
new_cache = (k, v)  # Create new cache with updated k/v pairs

Feed-forward

This is simple feed-forward network with SiLU activation function.

def feed_forward(params, x):
    
    w3_ = jnp.dot(x, params['w3'])

    # SwiGLU(a,b)=SiLU(a)⊙b 
    activated = jax.nn.silu(w3_)
    
    
    w1_ = jnp.dot(x, params['w1'])
    
    
    combined = activated * w1_
    
    # Final output projection with w2
    output = jnp.dot(combined, params['w2'])
    
    return output

Transformer-block

This is where all the important components come together in the transformer block. We unpack the pre-initialized weights and distribute them to their respective layers. The transformer blocks include attention, normalization, feed-forward processing layers and residual connections.

Forward-Pass

The forward pass takes your data through the entire model from converting input tokens into embeddings, through a series of transformer blocks, and finally to the output layer. In other words, it connects all the layers (embedding, transformers, and output) to produce the final predictions.

Dataset

Now the model part is complete so its time to train our model on shakespeare dataset. First we will read our data from .txt file then we will encode our data with bpe and then convert it into jax array.

Get Batches

The get_batch function creates training batches from our Shakespeare dataset. We need to feed our model with chunks of data. For each batch, we randomly select starting positions in the text, this way the model sees a variety of contexts.

Now, here's where JAX's cool vmap feature comes into play. Instead of writing a loop to extract each chunk, we use vmap to automate.

How does it work ?

vmap is like a vectorized loop; it takes a function that processes a single index (using lax.dynamic_slice to get a sequence of tokens) and applies it to every element in our array of indices. This means our input sequences (x) and corresponding target sequences (y, which are shifted by one token for next-word prediction) are created in one go.

def get_batch(key, data, batch_size, seq_len):
    # Generate random starting indices
    ix = random.randint(key, (batch_size,), 0, len(data) - seq_len)
    
    # Vectorized operation to get input and target sequences
    x = vmap(lambda i: lax.dynamic_slice(data, (i,), (seq_len,)))(ix)
    y = vmap(lambda i: lax.dynamic_slice(data, (i + 1,), (seq_len,)))(ix)
    
    return x, y

Loss Function

This function computes the cross-entropy loss for a batch during training. It first performs a forward pass using the model to generate logits for the input data. Then, it reshapes both the logits and targets to merge the batch and sequence dimensions. After applying the log softmax to the logits, it extracts the log probabilities corresponding to the correct target tokens and computes their negative mean as the final loss value.

The cross-entropy loss is defined as:

$$L=-\frac{1}{N}\sum _{i=1}^N\log P\left(y_i\right)$$

Where:

$$P\left(y_i\right)=\frac{e^{z_i}}{{\sum }_j{e}^{z_j}}$$

# Compute cross-entropy loss
def compute_loss(params, batch):
    # Split batch into inputs and targets
    inputs, targets = batch
    # Forward pass to get logits
    logits, = model_forward(params, inputs, config)
    # Reshape for loss computation
    logits = logits.reshape(-1, config.vocab_size)
    targets = targets.reshape(-1)
    # Calculate negative log likelihood
    loss = -jnp.mean(jnp.take_along_axis(jax.nn.log_softmax(logits),
    targets[:, None], axis=1))
    return loss

def sqr(x): 
    print("HI jiited") # side effect 
    return x * x

print(sqr(2)) 
print(sqr(3)) 
print(sqr(4))

@jax.jit
def sqr(x): 
    print("HI jiited") # side effect  
    return x * x

print(sqr(2)) 
print(sqr(3)) 
print(sqr(4))

@jax.jit
def update_step(params, batch):
    # Compute both loss and gradients in a single pass using value_and_grad
    # This is more efficient than computing them separately
    loss, grads = jax.value_and_grad(compute_loss)(params, batch)

    # Update parameters using gradient descent
    # jax.tree.map applies the update rule to each parameter in the model
    # The lambda function implements: p_new = p_old - learning_rate * gradient
    params = jax.tree.map(
        lambda p, g: p - config.learning_rate * g,
        params,
        grads
    )

    # Return updated parameters and the loss value for monitoring training
    return params, loss

for epoch in range(num_epochs):
 
   
   epoch_loss = 0.0

   for step in range(steps_per_epoch):
   
      # Generate new random keys for reproducibility
      key, batch_key = random.split(key)
      
      # Sample random batch of sequences
      batch = get_batch(batch_key, data, config.batch_size, config.max_seq_len)
      
      # Forward pass, compute loss and update parameters
      params_state, loss = update_step(params_state, batch)
     
      # loss for epoch average
      epoch_loss += loss
      
   
      if step % 100 == 0:
            print(f"epoch {epoch + 1}, step {step}/{steps_per_epoch}: loss = {loss:.4f}")
      

  avg_epoch_loss = epoch_loss / steps_per_epoch
     
 
  epoch_losses.append(avg_epoch_loss)
      
  
  print(f"\nepoch {epoch + 1} | average loss: {avg_epoch_loss:.4f}")