Build a Large Language Model (From Scratch)

Tiny Aya is a 3.35B parameter multi-lingual decoder-style transformer model by Cohere. It is designed for local experimentation and research, though it is restricted to non-commercial use due to its licensing terms.

There are several specialized versions of the model optimized for different linguistic regions:

• tiny-aya-base: The base model.
• tiny-aya-global: Provides the best balance across all languages and regions (default in notebooks).
• tiny-aya-fire: Optimized for South Asian languages (e.g., Hindi, Bengali, Tamil, Telugu).
• tiny-aya-water: Optimized for Asia Pacific (e.g., Chinese, Japanese, Korean, Vietnamese) and European languages (e.g., English, French, German, Spanish).
• tiny-aya-earth: Optimized for West Asian (e.g., Arabic, Turkish, Hebrew) and African languages (e.g., Swahili, Yoruba, Zulu).

The code for Chapter 5 is organized into main chapter components and optional standalone scripts.

Main Chapter Code:

• ch05.ipynb: The primary Jupyter Notebook containing all code as presented in the chapter.
• previous_chapters.py: A required utility module containing the MultiHeadAttention module and GPTModel class. This must be imported into ch05.ipynb to facilitate GPT model pretraining.
• gpt_download.py: Contains utility functions for downloading pretrained GPT model weights.
• exercise-solutions.ipynb: Contains the solutions to the exercises provided in the chapter.

Optional Standalone Scripts:

• gpt_train.py: A standalone Python script summarizing the GPT model training implementation found in the notebook.
• gpt_generate.py: A standalone Python script for loading and using pretrained model weights from OpenAI.

A KV (Key-Value) cache is a mechanism used during LLM inference to store intermediate key (K) and value (V) computations from previous tokens.

Why use it? During text generation, LLMs produce one token at a time. Without a cache, the model must recompute the K and V vectors for all previous tokens in every single generation step, leading to quadratic $O(n^2)$ complexity. With a KV cache, you only compute K and V for the newly generated token and append it to the existing cache, reducing per-step complexity to linear $O(n)$.

Trade-offs:

• Pros: Significant increase in inference speed (e.g., ~5x speed-up in small models).
• Cons: Increased memory usage that grows linearly with sequence length, and added code complexity. It cannot be used during training.

Sliding Window Attention (SWA) is a local attention mechanism used as an alternative to global Multi-Head Attention (MHA). Instead of each token attending to every other token in a sequence, each token only attends to a fixed-size local window around its position.

This approach significantly reduces the size of the KV cache, leading to substantial memory and compute savings. Modern models like Google's Gemma 3 use a hybrid approach, combining SWA layers with full attention layers (e.g., a 5:1 ratio) to balance efficiency with the ability to model global context.

The CloudFormation template includes a lifecycle configuration script that automates the installation of a specialized Python environment. This ensures all deep learning dependencies are correctly configured with CUDA support for GPU acceleration.

Environment Details:

• Package Manager: A separate Miniconda installation in the user's home directory.
• Core Frameworks:
  • TensorFlow: version 2.15.0 (with CUDA support)
  • PyTorch: version 2.1.0 (with CUDA support)
• Additional Libraries: Jupyter Lab, Matplotlib, and other utility packages.
• Jupyter Integration: The custom environment is automatically registered as a Jupyter kernel, allowing you to select it directly from the Jupyter interface.

The Qwen3Model class is designed to be a drop-in replacement for the GPTModel class used in earlier chapters of this project. This allows you to use it for:

• Training: Following the patterns established in Chapter 5.
• Finetuning: Following the patterns established in Chapters 6 and 7.

Mixture of Experts (MoE) is a technique to increase a model's capacity without proportionally increasing inference compute. It replaces standard dense feed-forward (FFN) modules with multiple 'expert' layers (also FFNs).

Key concepts:

• Sparsity: Instead of activating all experts for every token, a router selects a small subset of experts (e.g., top_k out of num_experts) per token. This makes the module 'sparse'.
• Capacity vs. Efficiency: While the total parameter count increases (increasing knowledge capacity), the active parameters per token remain low, keeping inference efficient.
• Shared Experts: Some designs (like DeepSeek-V3) use a 'shared expert' that is always active for every token. This helps capture common patterns, allowing individual experts to specialize in more complex patterns.

Standard (Gated) Attention

• Complexity: Quadratic $O(n^2)$ relative to context length $n$.
• Mechanism: Computes an $n \times n$ attention matrix where every token attends to every other token.
• Output Gating: A sigmoid gate can be applied to the attention output to decide how much to keep.

Gated DeltaNet

• Complexity: Linear $O(n)$ relative to context length $n$.
• Mechanism: Recurrent state update. It maintains a running memory state $S$ that is updated token-by-token.
• Memory: Uses a fixed-size state $S$ to compress past context, acting as a bottleneck compared to the full pairwise modeling of standard attention.

The KVCache class manages the Key-Value cache across transformer layers to improve compute efficiency during autoregressive generation.

• get(layer_idx): Retrieves the cache for a specific layer.
• update(layer_idx, value): Updates the cache for a specific layer.
• reset(): Clears all cached values.

In the Qwen3Model.forward pass, if a cache object is provided, the model uses it to append new keys and values instead of recomputing the entire sequence history.

class KVCache:
    def __init__(self, n_layers):
        self.cache = [None] * n_layers

    def get(self, layer_idx):
        return self.cache[layer_idx]

    def update(self, layer_idx, value):
        self.cache[layer_idx] = value

    def get_all(self):
        return self.cache

    def reset(self):
        for i in range(len(self.cache)):
            self.cache[i] = None