Amit Bahree's (useless?) insight!

🎉Announcing My New Book: Generative AI in Action📚

Mon, 16 Sep 2024 00:00:00 +0000

In today’s rapidly evolving tech world, mastering Generative AI isn’t just an advantage—it’s a necessity. Are you ready to harness its power to transform your business and solve real-world challenges? I’m excited to announce that my new book, Generative AI in Action, is now available in print and ebook formats from Manning Publications . 📖

Special Launch Offer 🌟

As a thank-you to my early supporters, I’m offering an exclusive discount. Use the code pbbahree at checkout to receive 45% off your purchase of Generative AI in Action in all formats (valid through Sept. 30, 2024)!

Get your discounted copy 👉 » here « .

Cover of Generative AI in Action

Why You Should Read Generative AI in Action ✨

Comprehensive Coverage

This book offers an in-depth introduction to Generative AI, covering foundation models, large language models (#LLMs), small language models (#SLMs), and practical applications. From the basics to advanced topics like prompt engineering, Generative AI in Action provides everything you need to start building and scaling AI solutions. Whether you’re a beginner or a seasoned professional, you’ll find valuable insights to accelerate your AI journey.

Real-World Examples

Discover how enterprises across industries are leveraging Generative AI to innovate and solve complex problems. Whether it’s improving customer engagement or optimizing operations, the practical examples provided can be directly applied to your projects for immediate impact.

Hands-On Techniques

Dive into step-by-step guides and hands-on examples for integrating AI models into your workflows. Learn techniques such as:

Prompt Engineering: Craft effective prompts to unlock the full potential of AI models like GPT-4.
Retrieval-Augmented Generation (RAG): Enhance your AI models with real-time data for improved accuracy.
Model Adaptation: Fine-tune AI models to meet your organization’s specific needs.

Ethical AI and Best Practices

As AI becomes more critical in decision-making, understanding its ethical implications is crucial. Generative AI in Action covers topics like privacy, security, and bias mitigation—ensuring your AI deployments are fair, transparent, and aligned with your organizational values.

Expert Insights

Drawing from my experience helping build the Azure AI platform, I share insider knowledge on leveraging the latest advancements in AI for your projects. This book provides you with the tools to make the most of cutting-edge technologies like large language models (LLMs) and small language models (SLMs).

Advanced Techniques Covered in the Book 🔥

Beyond the basics, Generative AI in Action delves into advanced techniques essential for mastering Generative AI in modern enterprise environments:

Prompt Engineering: Strategies like zero-shot, few-shot, and many-shot learning, along with chain-of-thought reasoning, to optimize AI outputs.
Retrieval-Augmented Generation (RAG): Combine retrieval-based methods with generative models for real-time, relevant data integration.
Model Adaptation and Fine-Tuning: Customize generative models to specific tasks using techniques such as low-rank adaptation and reinforcement learning from human feedback (RLHF).
Chatting with Your Data: Build AI-powered chat systems that interact with enterprise data using vector databases and retrieval techniques.
Scaling and Production Deployment: Strategies for scaling AI solutions while ensuring performance, reliability, and compliance with enterprise standards.
Evaluations and Benchmarks: Learn to evaluate and benchmark AI models using traditional metrics and cutting-edge frameworks.

Explore the GitHub Repository 💾

For those eager to dive into the code, the book has a companion GitHub repository filled with examples and projects to get you started. Check it out at bit.ly/GenAIBook . Explore the code, experiment, and start building your AI-powered solutions today.

Get Your Copy Today! 🛒

Don’t miss this opportunity to lead the AI revolution within your organization. Order your copy of Generative AI in Action and use the code pbbahree to receive 45% off (valid through Sept. 30, 2024). Transform your organization’s AI capabilities today!

With gratitude 💚

Amit Bahree.

In the world of AI, there’s a thrill,
With “Generative AI in Action,” you’ll skill.
From prompts to fine-tuning,
Your projects are blooming,
Grab your copy, and master the drill!

Building a microkernel in Rust (Part 4): Memory management and beyond

Thu, 09 Apr 2026 00:00:00 +0000

5-Part Series:

Part 0: Why build an OS from scratch?
Part 1: Foundations
Part 2: Communication
Part 3: Concurrency
Part 4 (this): Memory and beyond

GitHub Repository: bahree/rust-microkernel - full source code and build scripts

Docker Image: amitbahree/rust-microkernel - prebuilt dev environment with Rust, QEMU, and source code

Recap from Part 3 : we added timer interrupts and preemptive multitasking. The ARM Generic Timer fires every 100ms, the GIC routes interrupts to our handler, and a full context switch (31 registers + SP + ELR + SPSR) lets the OS forcibly alternate between two threads that never yield. The kernel is now in charge of scheduling, not the tasks.

Every program you’ve ever written has been lying to you about memory addresses. When your code reads at address 0x1000, it’s not actually reading the physical byte at 0x1000 in RAM. There’s a hardware translator sitting between your code and memory, silently remapping every address. We’re about to build that translator.

This is the final part of the series. Virtual memory is the foundation that makes modern operating systems possible: process isolation, memory protection, shared libraries, and even swap space. And the core mechanism is surprisingly elegant. You build a lookup table, point the hardware at it, flip a bit, and suddenly every memory access in the system goes through your table.

Let’s build it. 😃

TL;DR

We implement virtual memory on AArch64:

Frame allocator: A bump allocator that hands out 4KB physical memory pages (in this post, we’ll use a simple bump allocator, but in a real OS, you’ll need a more sophisticated allocator)
Page tables: 4-level translation structures that map virtual addresses to physical ones (in this post, we’ll use a 4-level page table, but in a real OS, you’ll need a more sophisticated page table)
MMU enablement: Configure the hardware translation unit and flip it on
Verification: Write through a virtual address, read it back, and prove the translation worked
TLB: A translation lookaside buffer that caches recent VA-to-PA translations (in this post, we’ll use a simple TLB, but in a real OS, you’ll need a more sophisticated TLB)

1. Why virtual memory?

Virtual memory is the foundation that makes modern operating systems possible: process isolation, memory protection, shared libraries, and even swap space. And the core mechanism is surprisingly elegant. You build a lookup table, point the hardware at it, flip a bit, and suddenly every memory access in the system goes through your table. It’s an awesome example of how hardware and software can work together to create a powerful abstraction.

1.1 The MMU is hardware, not software

Before we dive in, let’s ensure we’re on the same page. The Memory Management Unit (MMU) is a physical circuit inside the CPU, sitting between the processor core and the memory bus. When the CPU reads or writes any address, the MMU intercepts that access and translates it using a lookup table (the page table) before the request reaches RAM.

This is not a software layer; it’s dedicated hardware that runs on every single memory access, at full speed, with no CPU involvement once configured. All we need to do is build the lookup table and tell the MMU where to find it.

1.2 The problem with physical addressing

Right now, our kernel uses physical addresses directly. Address 0x4000_0000 in our code refers to physical byte 0x4000_0000 in RAM. This works fine for a single kernel, but it falls apart when you want multiple tasks to run simultaneously and independently. Let’s say we want to run two tasks simultaneously: Task A and Task B. Task A wants to use address 0x1000, and Task B wants to use the same address. If we use physical addressing, we’ll have a problem. Task A will write to address 0x1000, and Task B will write to address 0x1000. This is a problem because Task A and Task B are running simultaneously and independently, and they shouldn’t be able to write to each other’s memory.

This is where virtual memory comes in. Virtual memory is a system that allows each task to have its own view of the address space. Task A thinks it’s using address 0x1000, but the MMU translates that to physical address 0x4000_1000. Task B also thinks it’s using address 0x1000, but the MMU translates that to 0x4010_1000. Neither task is aware of the other’s memory.

Task A sees:                    Physical memory:
0x0000 -> 0x4000_0000 (RAM)      0x4000_0000: Task A code
0x1000 -> 0x4001_0000 (RAM)      0x4010_0000: Task B code
                                  0x5000_0000: Kernel
Task B sees:
0x0000 -> 0x4010_0000 (RAM)      ^ MMU translates
0x1000 -> 0x4011_0000 (RAM)      | every access

Isolation: Tasks can’t access each other. They have different translation tables. Relocation: Every task sees the same virtual addresses. Protection: You can mark pages read-only, non-executable, or kernel-only. Flexibility: Sparse virtual address spaces waste almost no physical memory. The flexibility of virtual memory is one of its most powerful features.

2. Pages and frames

The MMU doesn’t translate individual bytes. That would require billions of table entries (one per byte of address space). Instead, it groups addresses into fixed-size chunks called pages. We use 4KB pages (4096 bytes = 2^12), which has been the standard since the 1980s.

A page is a 4KB block in the virtual address space. A frame is a 4KB block in physical memory. The page table maps pages to frames. Think of it like a library catalog: a page is a shelf label (where you look), and a frame is the actual physical shelf (where the books are). The catalog maps labels to shelves, and you can rearrange which label points to which shelf without moving any books.

Think of the relationship between a page and a frame as a one-to-one mapping - one page maps to one frame. However, a page and frame can be different sizes. For example, a page can be 4KB, but a frame can be 2MB; this is because the page table is a tree, and the levels of the tree can be different sizes.

The bottom 12 bits of an address (the offset within the page) pass through untranslated, since both the virtual page and the physical frame share the same internal layout. For example, if virtual address 0x4000_1ABC maps to physical frame 0x7000_1000, then the offset 0xABC (the low 12 bits) is the same on both sides. The MMU only translates the upper bits, replacing 0x4000_1 with 0x7000_1, and the offset passes straight through.

Why 4KB specifically? It’s a trade-off. Smaller pages (e.g., 512 bytes) give finer-grained control over permissions and sharing, but need 8x more table entries to cover the same address range. On the other hand, larger pages (say 64KB) waste space when a program only needs a few hundred bytes - the rest of the page sits allocated but unused (this is called internal fragmentation). 4KB balances these concerns well for general-purpose systems. ARM also supports 16KB and 64KB granules, but Linux uses 4 KB on ARM, and we’ll use the same.

3. The page table walk

Let us walk through how the hardware translates a 48-bit virtual address (VA). First yt splits the address into five fields, and each field indexes into a different level of the page table. Think of it like a postal address: country, city, street, house number. Each part narrows down the search. At the end of the day, it’s a lot like a tree, with the top level being the most general, and the bottom level being the most specific.

The hardware starts at the root (L0), uses the first 9 bits to locate the L1 table, then uses the next 9 bits to locate the L2 table, then the next 9 bits to locate the L3 table, and finally uses the last 9 bits to locate the page within that table. The final 12 bits are the offset within the page.

48-bit Virtual Address (VA):
// VA = virtual address; PA = physical address. Low 12 bits are the page offset and pass through unchanged during translation.
+------+------+------+------+--------------+
|  L0  |  L1  |  L2  |  L3  | Page Offset  |
|  9b  |  9b  |  9b  |  9b  |    12b       |
+------+------+------+------+--------------+
 47:39  38:30  29:21  20:12      11:0

Each 9-bit index can hold values 0 through 511, so each page table has exactly 512 entries. At 8 bytes per entry, that’s 512 x 8 = 4096 bytes. One page table fits exactly in one 4KB page. Not a coincidence - the hardware designers intentionally made it this way to simplify memory management.

The hardware walks the tree on every memory access, starting from the root (L0) and following the pointers down to the leaf (L3) that contains the physical frame address. If any entry along the way is invalid, the MMU raises a page fault exception, which the OS can handle (e.g., by loading a page from disk or killing the offending process).

flowchart TD
    A["Virtual Address: 0x8000_1234"] --> B["L0 index = bits 47:39 = 0"]
    B --> C["L0 Table entry 0"]
    C --> D["Points to L1 Table"]
    D --> E["L1 index = bits 38:30 = 2"]
    E --> F["L1 Table entry 2"]
    F --> G["Points to L2 Table"]
    G --> H["L2 index = bits 29:21 = 0"]
    H --> I["L2 Table entry 0"]
    I --> J["Points to L3 Table"]
    J --> K["L3 index = bits 20:12 = 0"]
    K --> L["L3 Table entry 0"]
    L --> M["Physical Frame address"]
    M --> N["Add page offset (bits 11:0 = 0x234)"]
    N --> O["Physical Address"]

    style A fill:#9ff,stroke:#333
    style M fill:#ff9,stroke:#333
    style O fill:#9f9,stroke:#333

Figure 1: Virtual address translation

3.1 Working example

Let us use the same example as above, but let’s use a different virtual address: 0x8000_0000 to help us understand the translation process. We’ll trace the translation of VA 0x8000_0000 (the test address we’ll use later):

L0 index: bits 47:39 = 0x8000_0000 >> 39 = 0. First entry in L0.
L1 index: bits 38:30 = (0x8000_0000 >> 30) & 0x1FF = 2. Third entry in L1 (covers the 2-3 GB range).
L2 index: bits 29:21 = (0x8000_0000 >> 21) & 0x1FF = 0. First entry in L2.
L3 index: bits 20:12 = (0x8000_0000 >> 12) & 0x1FF = 0. First entry in L3.
Page offset: bits 11:0 = 0. Start of the page.

So the hardware walks: L0[0] -> L1[2] -> L2[0] -> L3[0] -> physical frame; then adds the page offset (0 in this case) to get the final physical address. If any of those entries were invalid, we’d get a page fault instead.

This multi-level structure allows us to efficiently map a huge virtual address space without needing an enormous flat page table. Each level of the tree only exists for the parts of the address space we actually use. If a program only uses a few megabytes of memory, we only need a handful of page table entries, not billions.

4. Why 4 levels?

Imagine a naive single-level page table with a 48-bit virtual address space (256 TB) and 4KB pages, you’d need 256 TB / 4 KB = 68 billion entries. At 8 bytes each, that’s 512 GB per process just for the page table, which is completely impractical. Even if you had that much RAM, the CPU would be overwhelmed trying to search through such a huge table on every memory access. That’s why we use a multi-level page table.

Each level of the tree allows us to cover a large portion of the address space with a single entry. The top-level L0 table has 512 entries, each covering 512 GB. The next level, L1, has 512 entries, each covering 1 GB. The next level, L2, has 512 entries, each covering 2 MB. Finally, the L3 level has 512 entries, each covering 4 KB (one page).

The solution is a tree structure where you only allocate page table nodes for memory regions you actually use; a small program using 8 KB of memory needs:

1 L0 table (4 KB)
1 L1 table (4 KB)
1 L2 table (4 KB)
1 L3 table (4 KB)
Total: 16 KB of page table overhead

Compare that to 512 GB for the flat table. That’s a factor of 33 million. Wow! 🫨 Memory efficiency is the main reason for the multi-level design. The tree structure also allows for efficient lookups. The hardware walks down the tree, and if it encounters an invalid entry at any level, it can immediately raise a page fault without searching a huge flat table.

graph TD
    L0["L0 Table (root)<br/>512 entries<br/>each covers 512 GB"]
    L1a["L1 Table<br/>512 entries<br/>each covers 1 GB"]
    L2a["L2 Table<br/>512 entries<br/>each covers 2 MB"]
    L3a["L3 Table<br/>512 entries<br/>each covers 4 KB"]
    P1["4 KB Frame"]
    P2["4 KB Frame"]

    L0 --> L1a
    L0 -.-> L1b["(unused)"]
    L1a --> L2a
    L1a -.-> L2b["(unused)"]
    L2a --> L3a
    L2a -.-> L3b["(unused)"]
    L3a --> P1
    L3a --> P2

    style L0 fill:#f99,stroke:#333
    style L1a fill:#ff9,stroke:#333
    style L2a fill:#9f9,stroke:#333
    style L3a fill:#9ff,stroke:#333
    style P1 fill:#99f,stroke:#333

Figure 2: Page table tree structure

Of course, in the real world, programs use scattered memory regions (heap, stack, code, libraries at various addresses), but since the tree is sparse, you only pay for what you use. The multi-level structure also allows for efficient lookups. The hardware walks down the tree, and if it encounters an invalid entry at any level, it can immediately raise a page fault.

5. Frame allocator

Before we can build page tables, we need a way to allocate physical memory. Our frame allocator is the simplest kind: a bump allocator.

A what now? In terms of frame allocators, there are many strategies you can use. A bump allocator is the simplest: it just keeps a pointer to the next free frame and increments it on each allocation. This is fast and simple, but it can’t free memory.

There is also a free list allocator that maintains a linked list of free frames, allowing for reuse but with more overhead. And finally, there is a bitmap allocator that uses a bitmap to track which frames are free or in use, which can be more space-efficient but also more complex.

For our demo, the bump allocator is sufficient since we only need to allocate a few frames during initialization.

There are two main challenges here: finding free memory to use as frames, and ensuring we don’t overwrite our kernel code or stack. We solve both by starting our allocator at the end of the kernel image (after the stack) and just bumping up from there. The linker script gives us a symbol (__stack_top) that marks the end of the kernel image, so we can safely start allocating frames from that point onward. This way, we avoid overwriting any critical data structures.

In a real OS, you’d want a more robust memory management system that can handle fragmentation and support freeing frames, but this simple approach is enough for our demo.

5.1 The actual code

Let’s look at the code; below is the implementation from mem.rs. The FrameAlloc struct keeps track of the next free frame and the end of available memory. The alloc() method returns the next free frame and advances the pointer, or returns None if we’ve run out of memory.

The code below is a simplified version of a frame allocator, suitable for our demo; in a production OS, you’d want to handle fragmentation, support freeing frames, and possibly implement more complex allocation strategies. The constants at the top define the RAM region provided by QEMU’s virt machine, and the align_up function ensures that our allocations are properly aligned to page boundaries.

const RAM_START: u64 = 0x4000_0000;
const RAM_SIZE: u64 = 256 * 1024 * 1024;
const RAM_END: u64 = RAM_START + RAM_SIZE;
const PAGE_SIZE: u64 = 4096;

extern "C" {
    static __stack_top: u8;
}

#[inline(always)]
fn align_up(x: u64, align: u64) -> u64 {
    (x + align - 1) & !(align - 1)
}

struct FrameAlloc {
    next: u64,
    end: u64,
}

impl FrameAlloc {
    fn new(start: u64, end: u64) -> Self {
        Self { next: start, end }
    }

    fn alloc(&mut self) -> Option<u64> {
        let p = self.next;
        if p + PAGE_SIZE > self.end {
            return None;
        }
        self.next += PAGE_SIZE;
        Some(p)
    }
}

Listing 1: Frame allocator (mem.rs)

QEMU’s virt machine gives us 256 MB of RAM starting at 0x4000_0000. The kernel image and stack live at the beginning of this region. We start allocating frames after the stack ends.

The __stack_top linker symbol

This doesn’t hold a value like a normal variable. Its address marks the end of the kernel image in memory. The linker script defines it. We take its address using &__stack_top as *const u8 as u64 to determine where free memory begins. This is a common pattern in OS development: using linker symbols to mark important memory boundaries (like the end of the kernel, the start of the heap, etc.) without hardcoding addresses. The linker ensures that __stack_top is placed at the correct location in the final binary, so we can rely on it to give us the starting point for our frame allocator.

The align_up bit trick

The expression (x + align - 1) & !(align - 1) rounds x up to the next multiple of align. Here’s how. If align is 4096 (0x1000), then align - 1 is 0xFFF (twelve 1-bits). Adding that ensures we overshoot if not already aligned. The bitwise AND with !(align - 1) = 0xFFFF_FFFF_FFFF_F000 clears the bottom 12 bits, snapping down to the nearest page boundary. This helps ensure that all our frame allocations are properly aligned to page boundaries, which is a requirement for the MMU.

Working example: align_up(0x4083, 0x1000) = (0x4083 + 0xFFF) & 0xFFFFF000 = 0x5082 & 0xFFFFF000 = 0x5000.

6. Page table structures in Rust

OK, now let’s define the page table structures for our page table. Each level of the page table is represented by a PageTable struct, which contains an array of 512 entries. We declare static instances for the L0, L1, and L2 tables, as well as a separate L3 table for our test virtual address. The hardware requires these tables to be 4KB-aligned, so we use #[repr(align(4096))] to ensure that the Rust compiler places them at the correct boundaries in memory.

The code below defines the PageTable struct and the static instances for each level of the page table. Each table has 512 entries, and we initialize them to zero. The new() method is a const fn, allowing us to create these tables at compile time.

#[repr(align(4096))]
struct PageTable {
    entries: [u64; 512],
}

impl PageTable {
    const fn new() -> Self {
        Self { entries: [0; 512] }
    }
}

static mut TT_L0: PageTable = PageTable::new();
static mut TT_L1: PageTable = PageTable::new();
static mut TT_L2_0: PageTable = PageTable::new(); // VA 0..1GB (UART)
static mut TT_L2_1: PageTable = PageTable::new(); // VA 1..2GB (RAM)
static mut TT_L2_2: PageTable = PageTable::new(); // VA 2..3GB (test VA)
static mut TT_L3_TEST: PageTable = PageTable::new();

Listing 2: Page table struct and static tables (mem.rs)

Why #[repr(align(4096))]?

As we touched earlier, the hardware requires page tables to be 4KB-aligned. This means the physical address of each table must be divisible by 4096, which guarantees the low 12 bits are all zeros. This isn’t just a performance optimization - it’s a hard requirement. The hardware uses the lower 12 bits of page table pointers to store attribute flags (valid bit, table/block bit, etc.).

If the table weren’t aligned, its address would have non-zero low bits that would collide with the flag bits, and the MMU would misinterpret flags as address bits or vice versa. Each table is exactly 4KB (512 entries x 8 bytes per entry), conveniently fitting within a single physical frame.

Without #[repr(align(4096))], Rust would use its default alignment for [u64; 512], which is just 8 bytes (the alignment of u64). That’s not nearly enough. The repr(align) attribute tells the Rust compiler and linker to place this struct at a 4096-byte boundary.

Why static mut?

Page tables must live at fixed, known addresses because we write those addresses into hardware registers (TTBR0_EL1) and into other page table entries (each table descriptor contains the physical address of the next-level table). static mut gives us globally accessible, mutable, fixed-address data.

Normally, static mut is dangerous in Rust (data race risk), but it’s fine here because we only modify these tables during single-threaded initialization before the MMU is on. Once the MMU is enabled, we never modify these tables again in our demo. A real OS would need proper synchronization (and TLB invalidation) when modifying page tables at runtime.

Why separate L2 tables?

We have three L2 tables (TT_L2_0, TT_L2_1, TT_L2_2) for the three 1GB regions we map. Each L1 entry covers 1GB, and each L2 table covers the 512 2MB sub-regions within that 1GB range. We could theoretically use one big L2 table, but splitting them makes the code clearer and matches the logical structure: one for UART (0-1GB), one for RAM (1-2GB), one for our test VA (2-3GB).

In a real OS, you’d likely have many more L2 tables as you map more of the address space, but for our learning purposes, three is enough to illustrate the concept.

7. Descriptor format

A descriptor is a 64-bit value that encodes both the physical address of the next-level table (or the mapped frame) and various attribute flags (valid, block/table bit, access permissions, etc.). This is the format defined by the ARM architecture for page table entries. The hardware expects this exact layout, and any deviation will cause translation failures or incorrect behavior.

Below we have the bit layout of a page table entry (descriptor) for AArch64. The upper bits (47:12) contain the physical address of the next-level table or the mapped frame, while the lower bits contain various flags that control how the MMU interprets this entry.

+----------------------------------+----------+----+--+
|   Physical Address [47:12]       | Attributes| TBL|V |
|   (36 bits)                      |           |    |  |
+----------------------------------+----------+----+--+
 63                                              1   0

Key bits:
  [0]  : Valid (1 = entry is active)
  [1]  : Table/Page (1) vs Block (0)
  [10] : AF (Access Flag, must be 1)
  [9:8]: SH (Shareability)
  [4:2]: AttrIndx (memory type index into MAIR)
  [54] : UXN (User Execute Never)
  [53] : PXN (Privileged Execute Never)

Descriptor format (64 bits)

And below are the descriptor constants from our code; these are the flags we OR into the entries when building our page tables.

const DESC_VALID: u64 = 1 << 0;
const DESC_TABLE: u64 = 1 << 1;   // table/page descriptor
const DESC_BLOCK: u64 = 0 << 1;   // block descriptor

const AF: u64 = 1 << 10;          // Access Flag
const SH_INNER: u64 = 0b11 << 8;  // Inner Shareable
const ATTRIDX0: u64 = 0 << 2;     // Normal memory
const ATTRIDX1: u64 = 1 << 2;     // Device memory

const PXN: u64 = 1 << 53;
const UXN: u64 = 1 << 54;

Listing 3: Descriptor bit constants (mem.rs)

The DESC_VALID bit indicates that an entry is valid.
The DESC_TABLE bit indicates that this entry points to another table (as opposed to a block or page).
The AF bit is the Access Flag, which must be set for the MMU to consider the entry valid.
The SH_INNER bits mark the memory as inner shareable, which affects caching and ordering.
The ATTRIDX0 and ATTRIDX1 bits select between normal and device memory types defined in the MAIR register.
The PXN and UXN bits control execution permissions for privileged and user modes, respectively.

Why you can OR flags into addresses?

Page tables are 4KB-aligned, so the low 12 bits of any table address are guaranteed zero. We can safely OR attribute flags into those bits without corrupting the address. The address lives in bits 47:12, the flags live in bits 11:0 (plus some high bits like UXN/PXN). This is a common hardware trick: alignment guarantees give you “free” bits for metadata.

Table vs Block vs Page descriptors

At L0, L1, and L2, bit [1] = 1 means “this entry points to another table” (table descriptor). At L1 and L2, bit [1] = 0 means “this entry directly maps a large region” (block descriptor, 1GB at L1 or 2MB at L2). At L3, bit [1] = 1 means “this entry maps a 4KB page” (page descriptor). Blocks are useful for mapping large contiguous regions with a single entry, rather than an entire next-level table.

8. Building page tables

Let’s build the actual page tables. We have two goals here: create some identity mappings for the kernel and peripherals, and create one non-identity mapping to prove that translation is working.

Our build_tables() function creates three kinds of mappings. The first two are identity mappings for the UART and RAM, which are necessary for the kernel to continue functioning after the MMU is enabled. The third mapping is a non-identity mapping for a test virtual address (0x8000_0000) that points to an allocated frame. This allows us to verify that the MMU is correctly translating addresses.

Here’s the mapping layout:

Virtual Address	Physical Address	Type	Size
`0x0900_0000`	`0x0900_0000` (UART)	Device	2MB block
`0x4000_0000` - `0x5000_0000`	Same (RAM)	Normal	2MB blocks
`0x8000_0000`	`frame0` (allocated)	Normal	4KB page

As we touched on earlier, the first two are identity mappings (VA = PA). This is critical during boot because the kernel code was loaded at specific physical addresses, and the program counter already contains a physical address.

When the MMU turns on, it translates all addresses, including the instruction the CPU is currently executing. If that address isn’t identity-mapped, the CPU can’t fetch the next instruction and crashes.

The UART is identity-mapped, so we can keep printing debug messages after the MMU is on. The RAM is identity-mapped so the kernel can continue accessing its data structures. The test VA is a non-identity mapping that points to an allocated frame, allowing us to verify that the MMU is correctly translating addresses.

Lets us double-click and see what the code looks like for building these tables. The function build_tables() initializes the page tables with the appropriate entries to create the mappings described above. It sets up the L0, L1, and L2 tables to point to each other, and then fills in the L2 and L3 entries for the UART, RAM, and test VA. The function returns the physical address of the L0 table (which we will load into TTBR0_EL1) and the test virtual address for later verification.

Note, the code uses unsafe blocks to modify the static mutable page tables, which is necessary because we’re directly manipulating memory structures that the hardware will read.

fn build_tables(frame0: u64) -> (u64, u64) {
    let test_va: u64 = 0x8000_0000; // 2GB

    unsafe {
        TT_L0.entries = [0; 512];
        TT_L1.entries = [0; 512];
        TT_L2_0.entries = [0; 512];
        TT_L2_1.entries = [0; 512];
        TT_L2_2.entries = [0; 512];
        TT_L3_TEST.entries = [0; 512];

        // L0[0] -> L1 (covers low VA range)
        TT_L0.entries[0] =
            (&raw const TT_L1 as *const _ as u64) | DESC_VALID | DESC_TABLE;

        // L1[0] (0..1GB) -> L2_0
        TT_L1.entries[0] =
            (&raw const TT_L2_0 as *const _ as u64) | DESC_VALID | DESC_TABLE;
        // L1[1] (1..2GB) -> L2_1 (RAM at 0x4000_0000)
        TT_L1.entries[1] =
            (&raw const TT_L2_1 as *const _ as u64) | DESC_VALID | DESC_TABLE;
        // L1[2] (2..3GB) -> L2_2 (test VA)
        TT_L1.entries[2] =
            (&raw const TT_L2_2 as *const _ as u64) | DESC_VALID | DESC_TABLE;

        // Map UART 0x0900_0000 as a 2MB device block (identity).
        let uart_va: u64 = 0x0900_0000;
        let uart_l2 = ((uart_va >> 21) & 0x1FF) as usize;
        TT_L2_0.entries[uart_l2] =
            (uart_va & 0xFFFF_FFFF_FFE0_0000)
                | DESC_VALID | DESC_BLOCK | ATTRIDX1 | AF | PXN | UXN;

        // Map RAM 0x4000_0000..0x5000_0000 as 2MB blocks, normal memory.
        let blocks = RAM_SIZE / (2 * 1024 * 1024);
        for i in 0..blocks {
            let va = RAM_START + i * 2 * 1024 * 1024;
            let pa = va;
            let idx = ((va >> 21) & 0x1FF) as usize;
            TT_L2_1.entries[idx] =
                (pa & 0xFFFF_FFFF_FFE0_0000)
                    | DESC_VALID | DESC_BLOCK | ATTRIDX0 | AF | SH_INNER;
        }

        // Map test_va -> frame0 as a single 4KB page (through L3).
        let test_l2 = ((test_va >> 21) & 0x1FF) as usize;
        TT_L2_2.entries[test_l2] =
            (&raw const TT_L3_TEST as *const _ as u64)
                | DESC_VALID | DESC_TABLE;
        let test_l3 = ((test_va >> 12) & 0x1FF) as usize;
        TT_L3_TEST.entries[test_l3] = (frame0 & 0xFFFF_FFFF_FFFF_F000)
            | DESC_VALID
            | DESC_TABLE
            | ATTRIDX0
            | AF
            | SH_INNER;
    }

    let ttbr0 = &raw const TT_L0 as *const _ as u64;
    (ttbr0, test_va)
}

Listing 4: Building page tables (mem.rs)

Next, let us walk through the three mapping types - identity map for UART, identity map for RAM, and non-identity map for the test VA.

8.1 UART identity map (device memory)

The UART lives at physical address 0x0900_0000. We need it identity-mapped so we can keep printing after the MMU is on. If we didn’t map it, any access to the UART registers would fail once the MMU is enabled, and we’d lose our ability to print debug messages.

The UART is a memory-mapped device, so we mark it as device memory in the page table entry. Device memory has different caching and ordering rules than normal RAM, which is critical for correct operation. We also set the execute-never bits (PXN and UXN) to prevent any code execution from that region, which is a common safety measure for peripheral registers.

To find which L2 entry to write, we extract the L2 index from the address using a right bit shift: 0x0900_0000 >> 21 = 72. The >> operator shifts the binary representation right by 21 positions — equivalent to dividing by $2^{21}$ — which moves the L2 index field (bits 29:21 per the page table layout) down to the lowest bit positions. As a sanity check: each L2 entry covers 2MB, so entry 72 starts at $72 x 2 MB = 144 MB = 0x09000000. That’s exactly where our UART lives.

We write this entry with ATTRIDX1 (device memory), PXN and UXN (prevent any code execution from UART registers), and AF (access flag, required by the hardware). This creates a single 2MB block mapping for the UART. We could descend to L3 for finer-grained 4KB page control, but the UART registers are all within this one 2MB region, so a block descriptor is simpler and more efficient.

8.2 RAM identity map (normal memory, 2MB blocks)

The RAM identity map is straightforward. We want the entire RAM region to be accessible at the same addresses as their physical locations, so we can continue using our existing code without modification. In our case, the RAM spans 0x4000_0000 to 0x5000_0000 (256 MB). That’s 128 blocks, each 2 MB. We loop through all 128 and create block descriptors with ATTRIDX0 (normal memory, cacheable) and SH_INNER (inner shareable, for future multi-core support).

Using 2MB blocks instead of 4KB pages means we only need 128 L2 entries, not 65,536 L3 entries. Fewer entries means less memory used for page tables and fewer TLB misses. The downside is that we can’t set different permissions for individual 4KB pages within those 2MB blocks, but for our simple kernel, that’s not a concern.

In a real OS, you might want finer-grained control and use page descriptors at L3 for some regions. The loop calculates the virtual and physical addresses for each block, which are the same because of the identity mapping. The L2 index is calculated as (va >> 21) & 0x1FF, which gives us the correct entry in the L2 table for each 2MB block.

8.3 Test VA mapping (4KB page through L3)

Before we dive in: VA stands for “virtual address” - an address used by software which the MMU translates to a physical address (PA). The low 12 bits are the page offset and remain the same after translation; the MMU replaces the upper bits (the virtual page number) with the physical frame number. In short VA = virtual address, PA = physical address, and offset bits pass through unchanged.

Finally, the test VA mapping is the most interesting one. We want to prove that the MMU is actually translating addresses, so we create a non-identity mapping: virtual address 0x8000_0000 maps to a physical frame we allocated earlier (let’s call it frame0). This means that when we access 0x8000_0000 in our code, the MMU will translate it to frame0 in physical memory. If we can read and write to that address successfully, it proves that the MMU is working correctly.

It is also important to note that this is not an identity map. The VA and PA are different. This is the whole point: proving that the MMU is actually translating.

The walk: L0[0] -> L1[2] -> L2_2[0] -> L3_TEST[0] -> frame0.

This demonstrates a full 4-level walk with a non-identity mapping. The L0 entry points to the L1 table, the L1 entry at index 2 points to the L2_2 table, the L2_2 entry at index 0 points to the L3_TEST table, and the L3_TEST entry at index 0 maps to frame0. This shows that the MMU is correctly following the page table pointers and applying the translations as expected.

Why L1 index 2? Because 0x8000_0000 >> 30 = 2. Each L1 entry covers 1GB. Index 0 covers 0-1GB, index 1 covers 1-2GB, index 2 covers 2-3GB. Our test address is at the start of the 2-3GB range. We point L1[2] to TT_L2_2, then L2_2[0] points to TT_L3_TEST, and finally L3_TEST[0] maps the 4KB page to frame0. This demonstrates a full 4-level walk with a non-identity mapping.

9. MMU configuration

MMU configuration is critical. The hardware uses the values in these registers to control how it performs address translation. If you get any of these wrong, the MMU won’t work as expected. You might get immediate faults, or worse, silent data corruption due to incorrect caching or permissions. Understanding what each register does and how it interacts with the page tables is essential for successful MMU setup. When we enable the MMU, the hardware reads these registers to determine how to interpret our page tables and perform translations.

Three system registers control how the MMU translates addresses. Getting any of these wrong means either an immediate CPU fault or (worse) silently incorrect translations that can cause data corruption or security issues. Let us walk through each one and explain the fields we set and why.

These registers are: MAIR_EL1 (Memory Attribute Indirection Register), TCR_EL1 (Translation Control Register), and TTBR0_EL1 (Translation Table Base Register 0). Let’s dig into each one in detail.

9.1 MAIR_EL1 (Memory Attribute Indirection Register)

The MAIR defines the memory types that page table entries can reference. Each type specifies caching and ordering rules for memory accesses. This register defines up to 8 memory types (indexed 0 through 7). Each page table entry’s AttrIndx field (bits [4:2]) selects one of these types. Rather than encoding full memory attributes in every page table entry (which would take too many bits), ARM uses this level of indirection: the page table entry says “use type 3” and MAIR defines what “type 3” means.

We define two types:

Attr0 = 0xFF: Normal memory, write-back write-allocate. This is the highest-performance caching mode for RAM. Both inner and outer caches are enabled, writes go to the cache first and are flushed to the RAM later. The 0xFF encoding means inner write-back read/write-allocate (bits [7:4] = 0xF) and outer write-back read/write-allocate (bits [3:0] = 0xF).
Attr1 = 0x04: Device memory, nGnRE. The acronym stands for non-Gathering, non-Reordering, Early-acknowledgment. Every access goes directly to hardware, in exactly the order your code specifies. No caching, no write combining, no speculative reads. The CPU treats each load and store as a side effect that must be visible to the device immediately. nGnRE is a common choice for memory-mapped peripherals like the UART, ensuring correct behavior without risking stale data or out-of-order accesses.

Why two types?

Because UART registers and RAM need fundamentally different treatment. If you cache UART reads, you’ll see stale data (the UART has new input, but the cache still holds the old value). If you make RAM non-cacheable, performance drops by 10-100x because every load and store goes directly to DRAM instead of hitting the L1 cache. Getting this distinction right is one of the most common stumbling blocks in bare-metal ARM development. The MAIR allows us to define these types once and then reference them in our page tables, keeping our entries compact while still providing the necessary flexibility.

9.2 TCR_EL1 (Translation Control Register)

This is the most complex of the three registers. It configures the virtual address width, page granule, caching for the page table walk itself, and which translation tables are active. All of these are necessary for the MMU to function correctly. The TCR tells the hardware how to interpret our page tables and how to perform translations. If we set the virtual address size too small, we won’t be able to use the full address space. If we set the wrong granule size, the hardware will misinterpret our page table entries. If we forget to disable TTBR1, we might accidentally have translations from that register interfering with our intended mappings.

Before the field-by-field breakdown, one quick definition: a granule is the MMU page size for this translation regime. In this post, granule = 4KB, which means each translation unit is 4KB, page-table pages are 4KB, and all descriptor math (offset bits, alignment, index extraction) is based on that size. A useful mental model: if the granule changes, the “grid” the MMU uses to carve up memory changes too.

Here are the fields we set:

T0SZ = 16: This determines the virtual address size. The formula is: VA bits = 64 - T0SZ. So 64 - 16 = 48 bits, giving us a 256 TB virtual address space. Larger T0SZ values mean smaller address spaces but fewer levels of page table to walk. We choose 16 to get a 48-bit VA space, which is more than enough for our demo and matches what Linux typically uses on AArch64. Setting this too small (e.g., T0SZ=32 for a 32-bit VA space) would limit us to 4 GB of virtual memory, which isn’t enough for modern OSes. Setting it too large (e.g., T0SZ=0 for a full 64-bit VA space) would require more levels of page tables and more memory overhead.
TG0 = 0b00: Selects a 4KB granule for TTBR0 translations. This must match the software side (PAGE_SIZE = 4096) and the table assumptions we’ve used throughout the post. ARM also supports 16KB and 64KB granules, but if TG0 says 16KB while your tables are laid out for 4KB, the MMU parses the descriptors with the wrong geometry (different alignment and bit interpretation), and translations fail immediately.
IRGN0/ORGN0 = 0b01: The page table walk hardware itself does memory reads to traverse the tree. These fields control caching for those reads. Write-back write-allocate means the table entries get cached, so repeated walks to the same region are fast. If we set these to non-cacheable, every walk would go to RAM, causing a huge performance hit. If we set them to write-through, we lose the performance benefits of caching. Write-back write-allocate is the best choice for page table walks.
SH0 = 0b11: Inner shareable. This ensures that page table updates are visible across cores in a multi-core system. Even on our single-core setup, it’s good practice to set this correctly. If we set it to non-shareable, other cores might not see updates to the page tables, leading to stale translations and hard-to-debug issues in a multi-core OS. Setting it to outer shareable is also an option, but inner shareable is typically recommended for page tables since they are frequently accessed and updated by the CPU.
EPD1 = 1: Disable TTBR1 translations. AArch64 supports two translation table base registers: TTBR0 (typically for user space, lower VA range) and TTBR1 (typically for kernel, upper VA range). We only use TTBR0, so we disable TTBR1 to avoid accidental translations. If we forget to disable TTBR1, and it contains a non-zero value, the hardware might use it for addresses in the upper VA range, causing unexpected translations and potential security issues. Disabling it ensures that only TTBR0 is used for all translations, simplifying our setup and reducing the risk of mistakes.
IPS = 0b010: 40-bit physical address space, supporting up to 1 TB of physical RAM. QEMU’s virt machine only has 256 MB, but setting this wider does no harm and matches what Linux typically uses. Setting IPS too small (e.g., 32-bit PA space) would limit us to 4 GB of physical memory, which isn’t enough for modern hardware. Setting it too large (e.g., 48-bit PA space) would require more bits in page table entries for addresses, reducing the number of bits available for flags and potentially causing issues with our simple page table setup.

9.3 TTBR0_EL1 (Translation Table Base Register)

This is the simplest of the three registers - it holds the physical address of our L0 page table. The MMU uses it as the starting point for every translation. When the CPU accesses virtual address X, the hardware reads TTBR0 to find the L0 table, then walks from there. Changing TTBR0 changes the entire virtual address space - that’s how operating systems switch between process page tables on a context switch. In our case, we set TTBR0 to the physical address of TT_L0, which is the root of our page table tree. This tells the MMU where to start when translating addresses.

If we set TTBR0 to the wrong address, the MMU will read garbage data as the L0 table, causing all translations to fail. If it points to a valid but incorrect page table, we might get seemingly random translations that are very hard to debug.

Ensuring that TTBR0 points to the correct physical address of our L0 table is critical for the MMU to function correctly. We calculate this address using &raw const TT_L0 as *const _ as u64, which gives us the physical address of the TT_L0 static variable. This is the root of our page table hierarchy, and the MMU will use it to start translating virtual addresses.

10. Enabling the MMU

Enabling the MMU is a delicate dance of setting up the right registers, ensuring all memory writes are visible to the hardware, and then flipping the enable bit. The sequence matters, and so do the barriers. A barrier is a CPU instruction that forces ordering: “finish these operations first, then continue”. In this section we’ll use two of them: dsb (complete/commit prior memory effects) and isb (refresh the instruction pipeline after control-register changes). If you skip a barrier or use the wrong order, the MMU can observe stale state and fail in ways that are very hard to debug, sometimes before even UART prints still work.

The listing below is the actual enable_mmu assembly from boot.S. It performs the following steps: one, configure the MAIR and TCR registers; two, set TTBR0 to point to our L0 table; three, use barriers to ensure all writes are visible and the TLB is flushed; and four, set the M bit in SCTLR_EL1 to turn on the MMU.

The code can be a bit intimidating at first glance, but each step is necessary for correct MMU operation. The barriers (dsb and isb) ensure that all previous memory operations are complete and visible to the hardware before we enable the MMU. The TLB invalidate ensures that no stale translations are cached when we turn on the MMU. Finally, setting the M bit in SCTLR_EL1 actually enables the MMU; until that point, all addresses remain physical.

.global enable_mmu
enable_mmu:
  // MAIR_EL1: Attr0 = 0xFF (Normal WBWA), Attr1 = 0x04 (Device nGnRE)
  mov x1, #0xFF
  mov x2, #0x04
  lsl x2, x2, #8
  orr x1, x1, x2
  msr mair_el1, x1

  // TCR_EL1: T0SZ=16, 4k granule, SH=inner, IRGN/ORGN=WBWA, EPD1=1, IPS=40-bit
  mov x1, #16                  // T0SZ
  mov x2, #(0b11)              // SH0 = inner shareable
  lsl x2, x2, #12
  orr x1, x1, x2
  mov x2, #(0b01)              // ORGN0 = WBWA
  lsl x2, x2, #10
  orr x1, x1, x2
  mov x2, #(0b01)              // IRGN0 = WBWA
  lsl x2, x2, #8
  orr x1, x1, x2
  mov x2, #(1)                 // EPD1 = disable TTBR1
  lsl x2, x2, #23
  orr x1, x1, x2
  mov x2, #(0b010)             // IPS = 40-bit PA
  lsl x2, x2, #32
  orr x1, x1, x2
  msr tcr_el1, x1

  // TTBR0_EL1 (x0 = L0 table base, passed as argument)
  msr ttbr0_el1, x0

  // Synchronize and invalidate TLB
  dsb sy
  isb
  tlbi vmalle1
  dsb ish
  isb

  // Enable MMU: read-modify-write SCTLR_EL1 to set M=1
  mrs x1, sctlr_el1
  orr x1, x1, #1              // M = 1 (MMU on)
  bic x1, x1, #(1 << 2)       // C = 0 (data cache off)
  bic x1, x1, #(1 << 12)      // I = 0 (instruction cache off)
  msr sctlr_el1, x1
  isb
  ret

Listing 5: enable_mmu assembly (boot.S)

Let us walk through the critical steps and barriers in this code as they are often the source of confusion and bugs.

dsb sy (Data Synchronization Barrier, System)

This Waits until all pending memory writes complete. Our page table entries might still be in write buffers, not yet committed to RAM. Without this, the MMU could read a partially-written entry.

This is a common pitfall: you set up your page tables in memory, but the CPU hasn’t actually written them to RAM yet. If you enable the MMU before those writes are visible, the hardware will read garbage data for your page tables, causing all translations to fail.

The dsb sy ensures that all those writes are flushed out and visible to the MMU before we proceed.

isb (Instruction Synchronization Barrier)

This flushes the CPU’s instruction pipeline. After changing a system register, the pipeline still contains instructions fetched under the old settings. isb forces the CPU to re-fetch everything.

This is crucial after enabling the MMU because the very next instructions will be executed with virtual address translation active. If we didn’t have this barrier, the CPU might execute some instructions with the old physical address mappings, which could lead to unpredictable behavior or crashes. The isb ensures that all subsequent instructions are fetched and executed under the new MMU configuration.

tlbi vmalle1 (TLB Invalidate All at EL1)

The TLB caches old VA-to-PA translations. Stale entries could cause the MMU to use wrong translations. This instruction throws them all away. If we forget this step, the TLB might still contain entries from before we set up our new page tables, leading to incorrect translations and very confusing bugs.

For example, if the TLB has an old entry for a virtual address that points to a different physical address than what we set up in our new tables, the MMU will use that stale translation instead of the correct one, causing memory corruption or faults. The tlbi vmalle1 ensures that all TLB entries are invalidated, so the MMU will fetch fresh translations from our newly configured page tables.

The read-modify-write pattern

mrs x1, sctlr_el1 reads the entire system control register. orr x1, x1, #1 sets bit 0 (MMU enable) while preserving everything else. bic clears the cache bits (we leave caches off for simplicity). msr sctlr_el1, x1 writes it back. We don’t just write a fresh value because SCTLR has dozens of other control bits we don’t want to disturb.

If we just wrote msr sctlr_el1, #1, we’d accidentally clear all those other bits, which could disable caches, change endianness, or cause other unintended side effects. By using the read-modify-write pattern, we ensure that we only change the bits we intend to (enabling the MMU) while leaving all other settings intact.

After this function returns, every memory access goes through the page tables. Even the stack pointer is now a virtual address (which is fine because we identity-mapped RAM). If we set up everything correctly, the MMU will translate addresses according to our page tables, and we can access the UART and RAM through their virtual addresses. If we made a mistake in any of the previous steps (e.g., incorrect MAIR/TCR settings, wrong TTBR0 value, missing barriers), we might end up with a non-functional MMU, which can be very difficult to debug since even our debug prints might not work.

The sequence diagram below summarizes the interactions between our Rust code, the assembly function that enables the MMU, and the MMU hardware itself. It shows the steps taken to configure the MMU and the critical barriers that ensure correct operation.

sequenceDiagram
    participant Rust as Rust Code
    participant ASM as enable_mmu (Assembly)
    participant MMU as MMU Hardware

    Rust->>ASM: Call enable_mmu(ttbr0)
    activate ASM
    ASM->>MMU: Write MAIR_EL1 (memory attributes)
    ASM->>MMU: Write TCR_EL1 (translation control)
    ASM->>MMU: Write TTBR0_EL1 (page table base)
    ASM->>MMU: dsb sy + isb (drain writes)
    ASM->>MMU: tlbi vmalle1 (flush TLB)
    ASM->>MMU: dsb ish + isb (ensure flush complete)
    ASM->>MMU: Set SCTLR_EL1.M=1 (enable!)
    ASM->>MMU: isb (flush pipeline)
    MMU->>MMU: Translation now active
    ASM->>Rust: Return (virtual addresses in use)
    deactivate ASM
    Note over Rust: All memory accesses<br/>now go through page tables

Figure 3: MMU enablement sequence

11. Running the demo

OK, we built the page tables, we configured the MMU registers, and we enabled the MMU. How do we know it worked? The only way to truly prove that the MMU is functioning correctly is to perform a memory access through a virtual address that goes through our page tables and see if we get the expected result. If the MMU isn’t working, we’ll either get a fault (if the hardware detects an invalid access) or incorrect data (if the hardware misinterprets our page tables). In our demo, we write a known value (0xDEAD_BEEF) to the test virtual address (0x8000_0000) that we set up to point to frame0. If we can read back the same value from that virtual address after enabling the MMU, it proves that the MMU is correctly translating addresses through our page tables.

The listings below show how to run the full memory management demo, which includes frame allocation, page table construction, MMU enablement, and virtual address translation. The demo() function in mem.rs orchestrates all these steps and prints out the results.

pub fn demo() {
    UartLogger::puts("mm: demo start\n");

    let kernel_end = unsafe { &__stack_top as *const u8 as u64 };
    let free_start = align_up(kernel_end, PAGE_SIZE);

    put_hex("mm: kernel_end=0x", kernel_end);
    put_hex("mm: free_start=0x", free_start);
    put_hex("mm: ram_end=0x", RAM_END);

    let mut fa = FrameAlloc::new(free_start, RAM_END);

    // Allocate a few frames and write/read patterns.
    let f0 = fa.alloc().expect("no frame");
    let f1 = fa.alloc().expect("no frame");
    put_hex("mm: frame0=0x", f0);
    put_hex("mm: frame1=0x", f1);

    unsafe {
        write_volatile(f0 as *mut u32, 0xAABB_CCDD);
        write_volatile(f1 as *mut u32, 0x1122_3344);
        let r0 = read_volatile(f0 as *const u32) as u64;
        let r1 = read_volatile(f1 as *const u32) as u64;
        put_hex("mm: read0=0x", r0);
        put_hex("mm: read1=0x", r1);
    }

    // Build page tables and enable MMU.
    let (ttbr0, test_va) = build_tables(f0);
    put_hex("mm: ttbr0=0x", ttbr0);
    put_hex("mm: test_va=0x", test_va);

    UartLogger::puts("mm: enabling MMU (caches off)...\n");
    unsafe { enable_mmu(ttbr0) };

    // If we survived, translation is live!
    unsafe {
        let p = test_va as *mut u32;
        write_volatile(p, 0xDEAD_BEEF);
        let r = read_volatile(p) as u64;
        put_hex("mm: test_va_read=0x", r);
    }

    UartLogger::puts("mm: demo done (MMU is ON)\n");
}

Listing 6: Memory management demo (mem.rs)

The code allocates frames, builds the page tables with the mappings we discussed, enables the MMU, and then performs a read/write test on the virtual address to verify that translation is working. The expected output includes the addresses of the kernel end, free memory start, RAM end, allocated frames, page table base, test virtual address, and the value read back from the test virtual address after MMU enablement.

The unsafe blocks are necessary because we’re directly manipulating memory and hardware registers, which is inherently unsafe in Rust. The write_volatile and read_volatile functions are used to ensure that the compiler doesn’t optimize away our memory accesses, which is critical when working with memory-mapped hardware and page tables.

Now, let us build and run the demo with the following commands:

./scripts/build-aarch64-virt.sh demo-memory
./scripts/run-aarch64-virt.sh

Listing 7: Build and run the memory demo

Below the listing shows a trace from my run and gives you a sense of the expected output; of course the exact addresses may vary, but the key part is that we see the correct values read back from the test virtual address, proving that translation works.

rustOS: aarch64 QEMU virt boot OK
rustOS: memory management demo (frames + page tables)
mm: demo start
mm: kernel_end=0x000000004009A010
mm: free_start=0x000000004009B000
mm: ram_end=0x0000000050000000
mm: frame0=0x000000004009B000
mm: frame1=0x000000004009C000
mm: read0=0x00000000AABBCCDD
mm: read1=0x0000000011223344
mm: ttbr0=0x0000000040085000
mm: test_va=0x0000000080000000
mm: enabling MMU (caches off)...
mm: test_va_read=0x00000000DEADBEEF
mm: demo done (MMU is ON)

Memory management demo output

Figure 4: Memory management demo showing frame allocation, page table construction, MMU enablement, and virtual address translation.

So did it all work? Or is it all mumbo jumbo? Well, the three things prove it worked:

The MMU enabled without crashing. If our page tables had any errors (unmapped kernel code, unmapped stack, misaligned tables), the CPU would have faulted immediately. The fact that we got to the point of printing “MMU is ON” means the MMU is functioning well enough to fetch instructions and access memory without faults. This is the first and most basic proof that the MMU is working. If there were any critical errors in our page tables or MMU configuration, we would have seen a fault as soon as we enabled the MMU, and we wouldn’t have been able to print anything afterward.
The UART still works after MMU enable. That means our device memory identity mapping is correct. If we had forgotten to map the UART or set the wrong attributes, we would have lost our ability to print immediately after enabling the MMU. The fact that we can still print debug messages after MMU enablement is a strong indication that our page tables are correctly set up to allow access to the UART registers, and that the MMU is correctly translating those addresses.
test_va_read=0xDEADBEEF: We wrote 0xDEAD_BEEF to virtual address 0x8000_0000, and read it back successfully. The MMU translated that VA through our L0/L1/L2/L3 tables to the physical frame we allocated. Translation works. This is the ultimate proof that our MMU setup is correct. Woot!

If the MMU wasn’t working, we would either get a fault when trying to access test_va, or we would read back an incorrect value. The fact that we read back the exact value we wrote confirms that the MMU is correctly translating virtual addresses to physical addresses according to our page tables.

12. The TLB (Translation Lookaside Buffer)

You might be wondering - doesn’t walking 4 levels of page tables on every memory access make everything incredibly slow? Each level requires a memory read, so that’s 4 extra memory accesses for every load or store your program does. If a simple mov x0, [x1] normally takes ~4 cycles, adding 4 table lookups would make it ~20 cycles. That’s a 5x slowdown on everything - obviously unacceptable for a real OS.

Fortunately, the hardware has a solution to this problem - its called the Translation Lookaside Buffer (TLB). The TLB is a small, fast cache that stores recent virtual-to-physical address translations. When the MMU needs to translate a VA, it first checks the TLB to see if it already has a cached translation for that address. If it does (a TLB hit), it can get the physical address in about 1 cycle, which is much faster than walking the page tables. If it doesn’t (a TLB miss), then it falls back to walking the full page table tree.

Think of the TLB like a cheat sheet: instead of walking the full page table tree every time, the MMU first checks if it already knows the answer from a recent lookup. The TLB is typically 48-1536 entries (varying by CPU), stored in extremely fast SRAM right next to the MMU logic.

On a TLB hit (which happens ~95-99% of the time for typical workloads), the MMU gets the physical address in about 1 cycle - the same speed as if there were no translation at all. The high hit rate comes from temporal locality (programs tend to access the same addresses repeatedly) and spatial locality (a single 4KB page covers many consecutive accesses). Only on a TLB miss does the hardware walk the full 4-level tree, which might take 20-40 cycles depending on whether the page table entries themselves are in the data cache.

When you switch between processes (changing TTBR0 to a different page table), you need to invalidate the TLB because the old translations belong to a different address space. Without invalidation, Process B might use a stale TLB entry from Process A and access the wrong physical memory - a security and correctness disaster. That’s what our tlbi vmalle1 instruction does in enable_mmu.

ARM supports ASIDs (Address Space IDs) to avoid this cost. Each TLB entry is tagged with a small process ID (8 or 16 bits). When you switch to Process B, the TLB entries tagged with Process A’s ASID are simply ignored rather than flushed. Process B’s entries from a previous run might still be there, avoiding cold-start misses. This is a significant optimization on context-switch-heavy workloads, but it’s a story for another day.

13. Summary

To wrap up, let’s take a step back and look at the big picture. We started with a blank slate and built up a functioning kernel with memory management capabilities. We implemented a frame allocator to manage physical memory, constructed multi-level page tables to define our virtual address space, configured the MMU registers to tell the hardware how to perform translations, and finally enabled the MMU to activate virtual memory.

That is pretty cool, and not because the code is sophisticated (I can tell you - it’s not, it’s a teaching OS); but because these are the same fundamental mechanisms used by Linux, Windows, macOS, and every other operating system with the differences are in scale and sophistication, not in kind. A 4-level page table walk on Linux works exactly the same way ours does. The context switch saves the same registers. The GIC uses the same IAR/EOIR protocol.

Just for fun, let’s compare our tiny teaching kernel to real-world operating systems and see how we stack up.

rustOS vs Linux: Linux has ~30 million lines of code, supports 30+ architectures, has 10,000+ drivers, CFS scheduling, demand paging, swap, huge pages, NUMA, SELinux, namespaces, cgroups, and a full TCP/IP stack. We have about ~2K lines and a UART. But our boot sequence, IPC mechanism, context switch, and page table setup are structurally identical to what Linux does. If you read Linux’s arch/arm64/kernel/head.S, you’ll recognize the EL2-to-EL1 drop, the vector table installation, and the MMU enable sequence. You built that.

rustOS vs seL4: seL4 is a formally verified microkernel used in aerospace and medical devices. About ~10K lines of kernel code, backed by 200K lines of mathematical proofs showing the code behaves correctly. Its IPC takes ~100 cycles (whilst we don’t measure ours, it will be much slower). It has capability-based security, hard real-time guarantees, and true user/kernel separation. Our IPC design (endpoint-based mailboxes) is actually inspired by seL4’s endpoint model, just without the verification or performance.

rustOS vs xv6: MIT’s xv6 is the closest comparison. It’s a teaching OS in about ~10K lines of C, with a monolithic Unix-like design. It has a shell, a filesystem, fork()/exec(), and pipe-based IPC. Where xv6 goes deeper into Unix APIs, we go deeper into bare-metal ARM specifics.

If you want to keep going, here’s a rough roadmap, loosely ordered by difficulty and dependency.

Heap allocator. Right now everything is on the stack. Implementing GlobalAlloc (start with a linked-list allocator) unlocks Box, Vec, String, and the whole alloc crate. This enables almost everything else. Phil Oppermann’s heap allocation post is an excellent guide.

Process abstraction. Replace our static two-thread array with a proper process table: PIDs, state machine (ready, running, blocked), dynamic creation and destruction. See OSTEP Chapter 4 .

User mode. Drop from EL1 to EL0 to run user code. This requires per-process page tables (TTBR0 swap on context switch), separate user/kernel stacks, and exception handling for the EL0-to-EL1 transition. High difficulty, high reward.

System calls. The user-kernel API. On ARM, the svc instruction traps from EL0 to EL1. You need a syscall dispatch table, argument passing conventions, and at minimum exit(), write(), and yield().

I’ll stop there, but the sky’s the limit. You can implement filesystems, drivers, networking, SMP, graphics, and more. The only real limit is your time and interest.

15. Resources for going deeper

I did not set out to write a 5-part series on OS development, but quite a few folks reached out; so I wanted to provide a roadmap for those who want to go deeper. I am no expert and these are some of the better resources who are the experts in the field and are awesoeme for learning more about operating systems, ARM architecture, and low-level programming. This is by no means an exhaustive list, but it should give you a solid starting point for further exploration.

15.1 Books

Operating Systems: Three Easy Pieces (Arpaci-Dusseau). Free online at ostep.org . I think these are one of the best introductions to OS concepts! 💖
Computer Systems: A Programmer’s Perspective (Bryant & O’Hallaron). Essential CS fundamentals.
Linux Kernel Development (Robert Love). Practical Linux internals.

15.2 Papers

The UNIX Time-Sharing System (Ritchie & Thompson, 1974). The original.
Improving IPC by Kernel Design (Liedtke, 1993). Fast IPC in L4.
seL4: Formal Verification of an OS Kernel (Klein et al., 2009).

15.3 Projects

xv6: MIT’s teaching OS . Comes with a free textbook explaining every line.
Writing an OS in Rust: os.phil-opp.com . Philipp Oppermann’s excellent x86_64 blog series.
Redox: redox-os.org . A Unix-like OS written entirely in Rust.
Tock: tockos.org . An embedded OS in Rust.

15.4 Communities

OSDev Wiki . Invaluable reference for everything.
OSDev Forum . Helpful community.
r/osdev . Reddit community.

16. Final thoughts

Building an OS is hard. You’ve dealt with assembly boot code, interrupt timing constraints, context switching where a single wrong byte offset corrupts everything, and page tables where one misplaced bit means instant crash.

But you’ve also seen that it’s possible. The mechanisms aren’t magic. A timer fires, you save some registers, you load others, you jump. An address goes through a tree lookup. That’s it. The complexity in real operating systems comes from scale (thousands of devices, millions of users, decades of edge cases), not from fundamentally different ideas.

Even if you never write another line of kernel code, you now know why malloc can fail, why programs crash with “segmentation fault,” why fork() is fast (copy-on-write page tables), and why your laptop doesn’t freeze when one tab hangs (preemptive scheduling). You see through the abstractions.

Thanks for following along. Hopefully you picked up as much reading this as I did building it. And remember, next time you get a segfault, you know exactly what’s going on under the hood. Happy hacking! 😍

5-Part Series:

Part 0: Why build an OS from scratch?
Part 1: Foundations
Part 2: Communication
Part 3: Concurrency
Part 4 (this): Memory and beyond

Building a microkernel in Rust (Part 3): Concurrency, interrupts and preemption

Tue, 17 Mar 2026 00:00:00 +0000

5-Part Series:

Part 0: Why build an OS from scratch?
Part 1: Foundations
Part 2: Communication
Part 3 (this): Concurrency
Part 4: Memory and beyond

GitHub Repository: bahree/rust-microkernel — full source code and build scripts

Docker Image: amitbahree/rust-microkernel — prebuilt dev environment with Rust, QEMU, and source code

Recap from Part 2 : we built message-passing IPC with a mailbox router and a cooperative scheduler. Two tasks — PingTask and PongTask — exchange messages through endpoint-based routing, taking turns via poll(). It works great, as long as every task plays nice.

Right now, our tasks are polite. They take turns, each one calling poll() and returning promptly. But what happens when a task doesn’t feel like returning? The whole system freezes. That’s the problem we’re solving today.

We’re going to teach our kernel to yank the CPU away from a misbehaving task, save everything that task was doing, and hand control to somebody else. By the end of this post, two threads will be running infinite loops, never yielding or returning, yet they’ll alternate perfectly. The OS will be in charge, not the tasks. 😈

TL;DR

We add timer interrupts and preemptive multitasking to rustOS on AArch64:

Timer interrupts fire every 100ms, giving us a real tick counter
The GIC routes hardware interrupts to our handler
Exception vectors catch IRQs and dispatch them to Rust code
Context switching saves all 31 general-purpose registers (plus SP, ELR, SPSR) from one thread and restores another’s
Preemptive scheduling means tasks can’t monopolize the CPU

A few terms we’ll define as we go, but here’s the quick version:

Cooperative scheduling: Tasks voluntarily give up the CPU. If one task loops forever, everyone else starves, and the system locks up.
Preemptive scheduling: The OS forcibly takes the CPU away on a timer tick. No task can hog the processor because the OS interrupts it and switches to another task.
Context switch: Saving the complete CPU state of one task and loading another’s, so the second task resumes exactly where it left off. This includes all registers, the stack pointer, and CPU flags.

Time: ~4-6 hours

1. Why interrupts matter

Let’s start with the problem. Why do we need interrupts? Why can’t we just keep polling?

1.1 The polling problem

Remember our scheduler from Part 2?

let mut tick: u64 = 0;
loop {
    for task in tasks.iter_mut() {
        task.poll(logger, router, tick);
    }
    tick += 1;  // Just counting loop iterations, not real time!
}

Listing 1: The fake tick counter from Part 2

This has three problems:

First, tick doesn’t mean anything real. One iteration isn’t one millisecond or one second. It’s, however, the length of time the tasks take.
Second, the CPU burns at 100% spinning through that tight loop even when there’s nothing to do.
And third, if any task decides to loop forever inside poll(), the entire system locks up.

Think about it this way. Without interrupts, your CPU is like a waiter who walks to each table asking, “Ready to order?” If table 3 is still deciding, the waiter is stuck standing there forever. Everyone else is hungry and just waiting. With interrupts, each table has a bell. The waiter can do something else (or just rest) and respond when a bell rings.

1.2 The solution: Timer interrupts

What if hardware could tap the CPU on the shoulder at regular intervals? That’s exactly what a timer interrupt does. We configure a piece of hardware to fire a signal every 100 milliseconds. The CPU stops whatever it’s doing, runs our handler, and then goes back. Now we have real time. And if we’re clever about what the handler does, we can use it to switch between tasks.

static TICKS: AtomicU64 = AtomicU64::new(0);

// Called by hardware every 100ms, not by our code
fn timer_irq_handler() {
    TICKS.fetch_add(1, Ordering::Relaxed);
}

// Main loop can now sleep between interrupts
loop {
    let tick = TICKS.load(Ordering::Relaxed);
    for task in tasks.iter_mut() {
        task.poll(logger, router, tick);
    }
    hal::arch::halt();  // Sleep until next interrupt
}

Listing 2: Timer-driven tick counter (conceptual)

2. What is an interrupt?

There’s literally a wire on the CPU chip (well, a signal line) dedicated to receiving interrupt requests. Between executing each instruction, the CPU checks this signal. When it sees the line asserted (fancy way of saying the voltage changed by an external device), the CPU finishes its current instruction, saves its state, looks up the address of a handler function from a table in memory, and jumps there.

Think of it like a doorbell ringing while you’re reading a book. You finish your sentence, bookmark your page, answer the door, deal with whatever it is, then come back and pick up exactly where you left off. The bookmark is your “saved state.” Answering the door is the “handler.”

sequenceDiagram
    participant CPU
    participant Hardware as Timer Hardware
    participant Handler as IRQ Handler
    participant Task as Task A Code

    Task->>CPU: mov x0, #42
    Task->>CPU: add x1, x0, #10
    Hardware->>CPU: Timer IRQ Signal
    activate CPU
    Note over CPU: Save CPU State<br/>(x0-x30, SP, PC, flags)
    CPU->>Handler: Jump to handler
    activate Handler
    Handler->>Handler: TICKS += 1
    Handler->>Hardware: Acknowledge IRQ
    Handler->>CPU: Return
    deactivate Handler
    Note over CPU: Restore CPU State
    CPU->>Task: Resume at add x1, x0, #10
    deactivate CPU
    Task->>CPU: ... continue execution ...

Figure 1: Interrupt handling flow

A few key properties of interrupts that make them special:

Asynchronous: They can happen at any point during execution - literally between any two instructions. Your code doesn’t ask for them.
Hardware-driven: The CPU doesn’t poll. A physical signal line tells it “something needs attention.”
Context preservation: The handler must not corrupt the interrupted program’s state. When the handler returns, the original code must resume as if nothing happened.

2.1 IRQs and exception types

IRQ stands for Interrupt Request. It’s a signal from a hardware device (timer, UART, network card, etc.) asking the CPU to pause and handle something. On ARM, there are four types of exceptions:

Type	Trigger	Example
Synchronous	Caused by an instruction	Undefined instruction, data abort
IRQ	Hardware interrupt request	Timer, UART
FIQ	Fast interrupt request	High-priority devices
SError	Asynchronous system error	Memory system errors

For our timer, we care about the IRQ. The other types we’ll handle later (or just hang on, for now).

You’ll see “mask” and “unmask” a lot in this post. Masked means blocked. When an interrupt is masked, the hardware device still generates the signal, but the CPU ignores it. Think of it like “Do Not Disturb” on your phone: calls still come in, your phone just doesn’t ring.

Unmasked means enabled. The CPU will respond to the interrupt signal normally. We mask interrupts during critical setup (like installing the handler table) and unmask them once everything’s ready. If an interrupt fired before the handler table was installed, the CPU would jump to garbage memory and crash.

3. The ARM Generic Timer

ARM processors have a built-in timer called the Generic Timer. Unlike external timer chips that communicate over a bus, this timer is integrated into the CPU itself. It counts at a fixed frequency and generates an interrupt when the countdown reaches zero. It’s perfect for our needs: we can set it to fire every 100ms, and it will reliably interrupt the CPU at that interval. The Generic Timer is a standard feature on all ARMv8-A (AArch64) CPUs, so our code will work on real hardware, not just QEMU. The only catch is that the platform firmware determines the timer’s frequency and can vary, so we need to read it at runtime. This is a common pattern in OS development: you write code that adapts to the hardware it’s running on, rather than hardcoding assumptions.

3.1 Timer registers

In reality, there are actually several timers available (physical, virtual, hypervisor), each with its own set of registers. The reason for this complexity is that ARM’s architecture supports multiple execution levels (EL0 for user code, EL1 for kernel, EL2 for hypervisor, EL3 for secure monitor). Each level has access to different timers. We only care about the physical timer accessible at EL1, since that’s what our kernel runs on. Specifically, we use the CNTP (Counter-timer Physical Timer) registers, which is a subset of the full Generic Timer system and is accessible from EL1 (kernel privilege):

CNTFRQ_EL0: The timer’s frequency in Hz. Read-only, set by firmware. On QEMU virt, it’s typically 62,500,000 (62.5 MHz).
CNTP_TVAL_EL0: The countdown value. Write a number here, and the timer counts down to zero, then fires an interrupt.
CNTP_CTL_EL0: Control register. Bit 0 enables the timer.

A quick note on those _EL0 suffixes. That doesn’t mean the register “belongs to” userspace. It means the register is accessible from EL0 and above. Our kernel at EL1 can read and write it just fine. Registers named _EL1 are only accessible from EL1 and above.

3.2 Programming the timer

Programming the timer is straightforward. We read the frequency, calculate the countdown value for our desired tick interval (100ms), write it to CNTP_TVAL_EL0, and enable the timer. The timer will then count down at the specified frequency, and when it reaches zero, it will fire an IRQ. We don’t need to worry about the timer counting down while we’re in the handler. The hardware takes care of that. When the timer fires, it automatically resets and starts counting down again. We just need to reprogram it with the same value to ensure it continues firing every 100ms. We do need to worry about the timer firing while we’re in the handler, but that’s actually a good thing. It means we can use the timer to drive our scheduler. Every time the timer fires, we can check if it’s time to switch tasks.

Here’s how we set up a 100ms tick. We read the frequency, divide by 10 to get the countdown value for 100ms, write it to CNTP_TVAL_EL0, and enable the timer:

fn program_timer(freq: u64) {
    // 100ms tick: frequency / 10
    let tval = (freq / 10) as u64;
    unsafe {
        core::arch::asm!(
            "msr cntp_tval_el0, {tval}",
            "mov x0, #1",
            "msr cntp_ctl_el0, x0",
            tval = in(reg) tval,
            out("x0") _,
            options(nostack, nomem)
        );
    }
}

Listing 3: Timer programming (from timer.rs)

If the frequency is 62,500,000 Hz, then freq / 10 = 6,250,000. The timer counts down from 6,250,000 at 62.5 MHz, reaching zero in exactly 100ms. When it hits zero, IRQ 30 fires (that’s the physical timer’s assigned interrupt ID on ARM).

Note that the timer doesn’t automatically reload. It just fires once - it is a one-shot timer. If we want it to fire again after the next 100ms, we have to write CNTP_TVAL_EL0 again. This gives us flexibility: we can change the tick interval on the fly if we want. We do this inside the IRQ handler every time the timer interrupt fires.

4. The GIC (Generic Interrupt Controller)

The Generic Interrupt Controller (GIC) is the hardware component responsible for managing and routing interrupts on ARM systems. It acts as a traffic cop, receiving interrupt signals from various hardware devices (like our timer) and directing them to the appropriate CPU core. The GIC also handles prioritization, masking, and acknowledgment of interrupts.

When the timer fires, it sends an interrupt signal to the GIC. The GIC then checks its configuration to determine which CPU core should handle that interrupt and delivers it accordingly. The GIC also provides registers to enable/disable specific interrupts, set their priorities, and acknowledge them once handled.

The only thing to note is that the timer can generate an interrupt, but something needs to route that interrupt to the CPU core. That’s the GIC’s job.

4.1 Why two components?

The GIC is split into two parts:

GICD (Distributor): One per system, shared by all cores. It receives interrupts from all hardware devices and decides which CPU core should handle each one. Think of it as a switchboard operator receiving all incoming calls. Then it forwards the call to the correct phone (core) based on its configuration.

GICC (CPU Interface): One per core. It delivers interrupts to its specific CPU and handles acknowledgment. Think of it as the phone on each person’s desk. The GICC listens for calls from the GICD and rings when an interrupt is delivered. It also provides registers for the CPU to acknowledge that it has handled the interrupt, allowing the GICD to send the next one.

This split matters for multi-core systems (which core handles the timer? which core handles the UART?), but even on our single-core setup, we need to configure both. This helps us write code that works on real hardware with multiple cores, not just in QEMU, and it also gives us a clearer separation of concerns in our code. The GICD is responsible for global interrupt configuration, while the GICC is responsible for local interrupt delivery and acknowledgment on each core.

4.2 Memory-mapped registers

Memory-mapped I/O means that hardware devices expose their control registers as specific addresses in the system’s memory map. To interact with the GIC, we read and write to these addresses. This is done using special functions that perform volatile reads and writes, ensuring that the compiler doesn’t optimize these accesses away. We treat these addresses as pointers to hardware registers. When we write to a GIC register, we’re actually sending commands to the hardware. When we read from a GIC register, we’re getting status information from the hardware.

On QEMU’s virt machine, the GIC lives at these addresses:

1
2

GICD: 0x0800_0000 (Distributor)
GICC: 0x0801_0000 (CPU Interface)

Here are the registers we’ll touch:

GICD registers:

Offset	Register	What it does
`0x000`	GICD_CTLR	Bit 0 enables the distributor
`0x100`	GICD_ISENABLER0	Write a 1-bit to enable that IRQ number
`0x400`	GICD_IPRIORITYR	Priority for each IRQ (lower = higher priority)
`0x800`	GICD_ITARGETSR	Which CPU core handles each IRQ

GICC registers:

Offset	Register	What it does
`0x000`	GICC_CTLR	Bit 0 enables this core’s interface
`0x004`	GICC_PMR	Priority mask: only deliver IRQs above this priority
`0x00C`	GICC_IAR	Read this to learn which IRQ fired (and tell GIC you’re handling it)
`0x010`	GICC_EOIR	Write the IRQ number here to say “I’m done”

The IAR/EOIR protocol is important. Reading GICC_IAR does two things at once: it returns the interrupt number and atomically marks that interrupt as “active.” Writing GICC_EOIR marks it “complete.” If you forget the EOIR write, the GIC thinks you’re still handling that interrupt and won’t deliver another one. Classic bug that makes the system appear to freeze after the first interrupt. Always remember: read IAR to start handling, write EOIR to finish.

4.3 The actual init code

Let us bring it all together. Here’s our real timer::init() function. It sets up the GIC, programs the timer, and unmasks IRQs:

static TICKS: AtomicU64 = AtomicU64::new(0);
static CNTFRQ: AtomicU64 = AtomicU64::new(0);

const GICD_BASE: usize = 0x0800_0000;
const GICC_BASE: usize = 0x0801_0000;
const IRQ_CNTPNS: u32 = 30;

pub fn init() {
    // Enable GIC distributor
    mmio_write32(GICD_BASE, 0x000, 1);

    // Enable GIC CPU interface, accept all priorities
    mmio_write32(GICC_BASE, 0x004, 0xFF);
    mmio_write32(GICC_BASE, 0x000, 1);

    // Enable physical timer interrupt (IRQ 30)
    enable_irq(IRQ_CNTPNS);

    // Read counter frequency and program timer
    let freq: u64;
    unsafe {
        core::arch::asm!(
            "mrs {0}, cntfrq_el0",
            out(reg) freq,
            options(nostack, nomem)
        );
    }
    CNTFRQ.store(freq, Ordering::Relaxed);
    program_timer(freq);

    // Unmask IRQs (clear DAIF.I)
    unsafe {
        core::arch::asm!("msr daifclr, #2", options(nostack, nomem));
    }
}

Listing 4: Full timer and GIC initialization (timer.rs)

What is important is that the sequence really matters. We enable the distributor first, then the CPU interface, then the specific IRQ, then program the timer, and finally unmask interrupts at the CPU level. If we unmasked interrupts before the GIC was ready, the CPU might see a spurious interrupt and crash. If we enabled the timer before the GIC, the timer would fire, but the interrupt would never reach the CPU. If we enabled the CPU interface before the distributor, the CPU would be ready to receive interrupts, but none would be sent. The order of operations is crucial for a stable system.

Why AtomicU64 instead of a regular u64? Even on a single-core system, interrupt handlers create concurrency. The main loop and the interrupt handler interleave unpredictably. If the main loop reads TICKS as two 32-bit halves (which happens on architectures without native 64-bit atomic loads) and the timer interrupt fires between the two halves, you get a corrupted value. AtomicU64 guarantees each load and store is indivisible. We use Ordering::Relaxed because on a single core, there’s no other CPU to synchronize with. The atomicity is still needed to prevent tearing, but we don’t need stronger ordering guarantees since there’s no cross-core communication.

5. The exception vector table

The exception vector table is a key piece of the interrupt handling mechanism. It’s a fixed block of memory that the CPU consults when an exception (like an IRQ) occurs. The CPU uses the exception type and the current execution context to calculate an offset into this table, where it expects to find a branch instruction to the appropriate handler.

When an IRQ fires, the CPU needs to know where to jump. On ARM, you set up a vector table: a fixed-layout block of code in memory containing branch instructions for each exception type.

AArch64’s vector table has 16 entries arranged in a 4x4 grid: 4 exception types (Synchronous, IRQ, FIQ, SError) × 4 execution contexts. The contexts are:

Current EL, SP0 (using the EL0 stack pointer, uncommon)
Current EL, SPx (using the current EL’s own stack pointer, the normal case for kernel code)
Lower EL, AArch64 (interrupts from user mode running 64-bit code)
Lower EL, AArch32 (interrupts from 32-bit code)

For our kernel running at EL1, the “Current EL, SPx” row is the one that matters. That’s at offset 0x200 from the table base.

Each entry gets exactly 0x80 (128) bytes of space. That’s enough room for a branch instruction to a larger handler elsewhere. The .align 11 directive aligns the table to 2^11 = 2048 bytes, which is a hardware requirement: the CPU will fault if VBAR_EL1 points to an unaligned address. That’s why we have the .align 11 at the start of the table. Each entry is a simple branch to a common handler for that exception type. For now, all IRQ entries point to the same exc_irq handler, which will read the interrupt ID and dispatch accordingly.

Here’s the actual vector table from our boot.S:

.align 11
vectors:
  // 0x000: Current EL, SP0, Sync
  b exc_sync
  .space 0x80 - 4
  // 0x080: Current EL, SP0, IRQ
  b exc_irq
  .space 0x80 - 4
  // 0x100: Current EL, SP0, FIQ
  b exc_fiq
  .space 0x80 - 4
  // 0x180: Current EL, SP0, SError
  b exc_serr
  .space 0x80 - 4

  // 0x200: Current EL, SPx, Sync
  b exc_sync
  .space 0x80 - 4
  // 0x280: Current EL, SPx, IRQ   <--- Timer IRQs land here
  b exc_irq
  .space 0x80 - 4
  // 0x300: Current EL, SPx, FIQ
  b exc_fiq
  .space 0x80 - 4
  // 0x380: Current EL, SPx, SError
  b exc_serr
  .space 0x80 - 4

  // 0x400-0x780: Lower EL entries (same pattern)
  // ... (all branches to the same handlers for now)

Listing 5: Exception vector table (boot.S)

The table is installed during boot with:

el1_start:
    adr x0, vectors
    msr vbar_el1, x0
    isb

Listing 6: Installing the vector table (boot.S)

VBAR_EL1 (Vector Base Address Register) tells the CPU where the table lives. The isb (Instruction Synchronization Barrier) flushes the pipeline so the CPU sees the new table address immediately. If we forgot the isb, the CPU might still use the old (zero) value of VBAR_EL1 for a few instructions, and if an interrupt fired during that window, it would jump to address 0 and crash. Always remember to use isb after changing VBAR_EL1.

6. The Rust IRQ handler

In Rust, we write a safe wrapper around the raw assembly handler. The assembly code in exc_irq will save all the registers, call our Rust function, and then restore registers based on the return value. This way, we can write our interrupt handling logic in Rust without worrying about the low-level details of context saving and restoring. The reason we return a pointer to a Context struct is that the handler can decide to switch to a different thread. If it returns the same pointer it was given, the assembly code will restore that context and resume the same thread. If it returns a different pointer, the assembly code will restore the new context, effectively switching threads. This is how we implement preemptive multitasking: the timer interrupt can decide to switch to a different thread every time it fires.

When a timer interrupt fires, the CPU jumps to exc_irq in the vector table, which (after saving context, more on that later) calls our Rust handler. Here’s the actual code:

#[unsafe(no_mangle)]
pub extern "C" fn rust_irq_handler(current: *mut Context) -> *const Context {
    let iar = mmio_read32(GICC_BASE, GICC_IAR);
    let id = iar & 0x3FF;

    let mut next: *const Context = current;
    if id == IRQ_CNTPNS {
        let t = TICKS.fetch_add(1, Ordering::Relaxed) + 1;
        let freq = CNTFRQ.load(Ordering::Relaxed);
        if freq != 0 {
            program_timer(freq);
        }

        // Switch threads every 5 ticks (~500ms with 100ms tick).
        if (t % 5) == 0 {
            next = super::preempt::switch_next(current);
        }
    }

    // End of interrupt
    mmio_write32(GICC_BASE, GICC_EOIR, iar);
    next
}

Listing 7: Rust IRQ handler (timer.rs)

If we look closely at the signature, we see that it takes a *mut Context (a pointer to the current thread’s saved state) and returns a *const Context (a pointer to the next thread’s state). If no switch is needed, it returns the same pointer it was given. The assembly code that calls this function uses the return value to decide which context to restore.

Here’s what happens step by step:

Read IAR: Find out which interrupt fired. We mask with 0x3FF because the lower 10 bits hold the interrupt ID.
Check if it’s the timer: ID 30 is our physical timer.
Increment tick counter: fetch_add atomically adds 1 and returns the old value, so we add 1 to get the current tick.
Reprogram timer: Set up the next 100ms countdown. Without this, the timer only fires once.
Maybe switch threads: Every 5 ticks (500ms), call switch_next to pick the other thread.
Write EOIR: Tell the GIC we’re done handling this interrupt.

7. DAIF bits and unmasking

ARM has four masking bits in the processor state register (PSTATE), collectively called DAIF. These processor flags control whether certain types of exceptions are masked (blocked) or unmasked (enabled) and are critical for safely handling interrupts - hence the name “DAIF” for Debug, SError, IRQ, FIQ. These bits act as individual on/off switches for different categories of exceptions:

D (bit 3): Debug exceptions — breakpoints, watchpoints, single-step traps. Used by debuggers.
A (bit 2): SError (asynchronous aborts) — serious hardware errors like uncorrectable memory faults. You generally want these enabled so the system can respond to hardware failures rather than silently ignoring them.
I (bit 1): IRQ (normal interrupts) — this is the one we care about most. Timer interrupts, UART receive interrupts, and most peripheral interrupts are IRQs.
F (bit 0): FIQ (fast interrupts) — a higher-priority interrupt channel. On some systems, FIQ is reserved for secure-world use (TrustZone). We don’t use it in our kernel.

When a bit is set (1), that exception type is masked (blocked). By default, when the CPU enters EL1, all four bits are set, meaning everything is blocked. This is a safety measure: the kernel starts with interrupts disabled so it can set up handlers and data structures before anything fires.

The instruction msr daifclr, #2 means “DAIF Clear with bitmask 0b0010,” which clears only the I bit, unmasking normal IRQs. We deliberately leave D, A, and F alone. We do this at the very end of timer::init(), after the GIC is configured, the timer is programmed, and the vector table is installed. The ordering is crucial: if we unmasked IRQs before the vector table was ready, the first timer tick would jump to an uninitialized address, causing the system to crash.

You might wonder: why not just unmask everything? Because we want precise control. We only enable what we’ve set up handlers for. Unmasking FIQ without a handler would cause the same crash problem. And debug exceptions should only be enabled when we actually have a debugger attached. BTW, leaving them masked by default is a good defensive programming practice.

8. Running the timer demo

OK, this is the moment of truth. Let’s build and run the timer demo. If everything is set up correctly, you should see the system boot, print a message about the timer demo, and then enter an idle loop that wakes on each timer tick. You can run just the timer demo (without preemption) using the demo-timer feature:

./scripts/build-aarch64-virt.sh demo-timer
./scripts/run-aarch64-virt.sh

Listing 8: Build and run the timer demo

Expected output:

1
2
3

rustOS: aarch64 QEMU virt boot OK
rustOS: timer interrupts demo
rustOS: timer started, entering idle loop

The system enters a wfe (Wait For Event) loop and wakes on each timer tick. At this point, we’re not printing ticks to keep things simple. The key proof that it works: the system doesn’t hang. If interrupts weren’t firing, wfe would sleep forever.

Now here’s the key insight: if we can interrupt a running task, we can save its state and run a different task instead. Think about what the interrupt handler already does. When a timer IRQ fires, the CPU automatically saves the program counter and status flags. Our handler saves the remaining registers. After incrementing the tick counter, we restore everything and resume the interrupted code. The interrupted code never knows it was paused. But what if, instead of restoring the same code’s registers, we restored different code’s registers? We’d resume a completely different thread. That’s preemptive multitasking. That’s the whole trick.

Pretty cool, right? With just a timer interrupt and some context saving/restoring, we can switch between threads without them ever asking for it. The timer is the “clock” that drives our scheduler, and the interrupt handler is the “switchboard operator” that decides which thread to run next.

9. What is CPU context?

We talk about saving and restoring “context,” but what does that actually mean? We need to save enough of the CPU’s state so that when we switch back to a thread, it can resume exactly where it left off, with all its registers and flags intact. This includes all the general-purpose registers (x0-x30), the stack pointer (SP), the program counter (PC, stored in ELR_EL1), and the CPU flags (stored in SPSR_EL1).

In other words, Context is everything the CPU needs to resume a task exactly where it left off. Here’s our actual Context struct from preempt.rs:

#[repr(C)]
pub struct Context {
    pub x: [u64; 31],   // x0..x30
    pub sp: u64,         // Stack pointer
    pub elr: u64,        // Exception Link Register (where to resume)
    pub spsr: u64,       // Saved Program Status Register (CPU flags)
}

Listing 9: Context struct (preempt.rs)

Size: 31 registers + SP + ELR + SPSR = 34 fields, each 8 bytes = 272 bytes per task.

Why #[repr(C)]? By default, Rust is free to reorder struct fields and add padding however it likes. But our assembly code accesses fields by hardcoded byte offsets ([x9, #0x00] for x0, [x9, #0xF8] for SP, etc.). If the compiler rearranged the fields, the assembly would save and restore the wrong data, silently corrupting the task state. #[repr(C)] forces fields to appear in memory in declaration order. This way, we can reliably calculate offsets for assembly access.

For example, x0 is at offset 0, x1 at 8, …, x30 at 240, sp at 248, elr at 256, and spsr at 264. The assembly code uses these offsets to save and restore the correct registers; if the struct layout changed, the offsets would be wrong, and we’d end up with a very buggy system. This is also a common pattern in OS development: you have a data structure shared between Rust and assembly, and you need to ensure they agree on its layout. #[repr(C)] is the standard way to achieve this.

9.1 Why save ALL 31 registers?

In theory, we could get away with saving only the registers that the interrupted code was actually using. But in practice, we have no way of knowing which registers those are. The compiler can use any register for any purpose at any time. Even if we tried to analyze the code to figure out which registers are live at the interrupt point, we’d have to do that for every possible instruction and every possible interrupt timing. It would be a nightmare.

Imagine Thread A is in the middle of a loop. Register x5 holds a loop counter, and x19 holds a critical pointer. The timer fires, we switch to Thread B, which runs for a while and overwrites x5 and x19 with its own data. When we switch back to Thread A, those registers are corrupted. Thread A crashes or (worse) silently computes the wrong thing.

The problem is we don’t know which registers the interrupted code was using. It could be any of them. So we save all of them.

Registers	Role	What breaks if not saved
x0-x7	Function arguments, return values	Function calls fail
x8	Indirect result location	Large struct returns corrupt memory
x9-x15	Temporaries (caller-saved)	Mid-calculation values corrupted
x16-x17	Intra-procedure scratch	Dynamic linking breaks
x18	Platform register	Reserved for OS use
x19-x28	Callee-saved	Long-lived variables corrupted (hardest bugs)
x29	Frame pointer	Debugger backtraces break
x30	Link register (return address)	Functions return to wrong address
SP	Stack pointer	Stack access goes to wrong memory
ELR_EL1	Exception return address	CPU resumes at wrong instruction
SPSR_EL1	Saved CPU flags	Conditional branches break, IRQ state wrong

What about NEON/FP registers (v0-v31)? AArch64 has 32 additional 128-bit registers for SIMD and floating-point. We don’t save them because it would add 512 bytes per task and ~200 extra cycles per switch. For our demo, we assume tasks don’t use floating-point. A real OS would implement lazy FP saving: track whether a task used FP registers and only save them if it did.

10. Thread setup

For preemption to work, we need at least two threads. Each thread needs its own stack and a context that points to its entry function. When the timer interrupt fires, the handler can switch between these threads by returning different context pointers. With just one thread, the timer would interrupt it, but we’d always return the same context, so it would just resume the same thread. With two threads, we can alternate between them on each timer tick.

Now let’s look at the complete preemption system. We need stacks, entry points, and context initialization.

10.1 Stacks and static contexts

Because we’re writing a kernel, we don’t have dynamic memory allocation yet. We can’t malloc stacks for our threads. Instead, we define static arrays in Rust to serve as stacks. Each thread gets its own stack array. We also define a static array of Context structs to hold the CPU state for each thread. This way, we can set up everything at compile time without needing any heap allocation. The code below is quite straightforward: we define two 16 KB stacks and an array of two Context structs. The CURRENT atomic variable tracks which thread is currently running.

const STACK_SIZE: usize = 16 * 1024; // 16 KB per thread

#[repr(align(16))]
struct Stack([u8; STACK_SIZE]);

static STACK0: Stack = Stack([0; STACK_SIZE]);
static STACK1: Stack = Stack([0; STACK_SIZE]);

static mut CTX: [Context; 2] = [
    Context { x: [0; 31], sp: 0, elr: 0, spsr: 0x5 },
    Context { x: [0; 31], sp: 0, elr: 0, spsr: 0x5 },
];

static CURRENT: AtomicUsize = AtomicUsize::new(0);

Listing 10: Stacks and static context array (preempt.rs)

The #[repr(align(16))] on Stack guarantees 16-byte alignment. ARM requires the stack pointer to be 16-byte aligned at all times. Violating this causes an alignment fault.

The spsr: 0x5 value means 0b0000_0101, which encodes EL1h: Exception Level 1, using SP_EL1 (the kernel’s own stack pointer). When eret restores this SPSR value, the CPU runs at the kernel privilege level. That’s what we want for our kernel threads.

10.2 Thread entry points

Entry points are the functions where each thread starts executing. When we initialize the context for each thread, we set the elr field to point to its entry function. When the scheduler switches to that thread, it will “return” to that function as if it had been interrupted there. Each thread runs an infinite loop that prints its name every 10 ticks. Thread A prints “A\n” at ticks 0, 10, 20, etc., while Thread B prints “B\n” at ticks 5, 15, 25, etc. This way, we can visually confirm that both threads are running and that the timer interrupt is switching between them.

pub(crate) extern "C" fn thread_a_entry() -> ! {
    // Enable the timer after the first thread context is active.
    timer::init();
    let mut last_tick: u64 = 0;
    loop {
        let t = timer::ticks();
        if t != last_tick {
            last_tick = t;
            if (t % 10) == 0 { UartLogger::puts("A\n"); }
        }
        core::hint::spin_loop();
    }
}

pub(crate) extern "C" fn thread_b_entry() -> ! {
    let mut last_tick: u64 = 0;
    loop {
        let t = timer::ticks();
        if t != last_tick {
            last_tick = t;
            if (t % 10) == 5 { UartLogger::puts("B\n"); }
        }
        core::hint::spin_loop();
    }
}

Listing 11: Thread entry points (preempt.rs)

Notice something subtle: Thread A checks (t % 10) == 0 but Thread B checks (t % 10) == 5. Why not both check (t % 10) == 0?

Here’s what we mean. Context switches happen every 5 ticks. So at tick 10 (a multiple of 10), the scheduler switches from whichever thread is running to the other. But the switch happens during the tick, which means Thread B might never actually see tick 10 while it’s running. It gets switched out right at that moment. By having B check for tick 5 instead, we offset the check from the switch boundary, and both threads reliably print.

Both functions are extern "C" because they’ll be called from assembly via eret. The C ABI ensures the calling convention matches what our assembly code expects.

10.3 Context initialization

Context initialization is the process of setting up the Context structs for each thread so that when we switch to them, they start executing at their entry points with a valid stack. We set the sp field to point to the top of each thread’s stack (remember, stacks grow downward), and we set the elr field to point to the thread’s entry function. The spsr field is set to 0x5, which means EL1h mode (kernel mode using SP_EL1). This way, when we “return” to this context, the CPU will jump to the entry function with the correct privileges.

pub fn init() {
    unsafe {
        let top0 = STACK0.0.as_ptr().add(STACK_SIZE) as u64;
        let top1 = STACK1.0.as_ptr().add(STACK_SIZE) as u64;

        CTX[0].sp = top0;
        CTX[0].elr = thread_a_entry as *const () as usize as u64;
        CTX[0].spsr = 0x5;

        CTX[1].sp = top1;
        CTX[1].elr = thread_b_entry as *const () as usize as u64;
        CTX[1].spsr = 0x5;
    }
}

Listing 12: Context initialization (preempt.rs)

Why SP points to the END of the stack: Stacks grow downward on ARM (and x86). Pushing data decreases the stack pointer. So the “bottom” (starting point) is the highest address. We compute STACK0.0.as_ptr().add(STACK_SIZE) to get one-past-the-end. If we set SP to the beginning of the allocation, the very first push would write below it, corrupting other memory.

The eret trick: We set elr to the thread’s entry function address. When eret executes, the CPU loads the PC from the ELR and jumps to it. The CPU doesn’t know (or care) whether it’s “returning” from a real exception or just bootstrapping a new thread. It simply loads and jumps. This is how the first task on every OS starts: fill in the context as if the task had been interrupted, then “return” to it.

11. The context switch in assembly

As we touched on earlier, this is the heart of preemption. Every register must be saved and restored perfectly, or threads silently corrupt each other. The exc_irq handler in boot.S does this. It saves all registers to the current context, calls the Rust handler to decide which context to run next, and then restores registers from the returned context. The assembly code is a bit verbose because we have to save 31 registers, plus SP, ELR, and SPSR, and we have to be careful not to clobber any registers before we save them.

Here’s the actual exc_irq handler from boot.S. We’ll walk through it section by section.

11.1 The scratch save trick

exc_irq:
  // We need a working register to read TPIDR_EL1, but we can't
  // touch any register without saving it first. Chicken-and-egg.
  // Solution: use the stack as scratch space for x9 and x10.
  sub sp, sp, #0x20
  str x9,  [sp, #0x00]
  str x10, [sp, #0x08]

  // Now x9 is free. Load the current Context pointer.
  mrs x9, tpidr_el1

Listing 13: exc_irq handler, part 1: scratch save (boot.S)

Here’s the problem: to save registers to the Context struct, we need a pointer to that struct. But reading TPIDR_EL1 requires a destination register, which would clobber the interrupted thread’s value in that register. Solution: we temporarily push x9 and x10 onto the stack (which doesn’t need a general-purpose register because SP is implicit in str), then use x9 to hold the context pointer. We’ll retrieve the original x9/x10 from the stack later and store them in the proper Context slots.

11.2 Saving all registers

This is needed as we called out earlier, but the implementation is straightforward: we use stp (Store Pair) to save registers in pairs, and str for the odd one out (x8). We also save SP, ELR, and SPSR at the end. The offsets into the Context struct are carefully calculated based on the struct layout.

  // Save x0..x8 to Context
  stp x0,  x1,  [x9, #0x00]
  stp x2,  x3,  [x9, #0x10]
  stp x4,  x5,  [x9, #0x20]
  stp x6,  x7,  [x9, #0x30]
  str x8,        [x9, #0x40]

  // Retrieve original x9/x10 from scratch and save to Context
  ldr x0, [sp, #0x00]
  str x0, [x9, #0x48]        // x9's slot in Context
  ldr x0, [sp, #0x08]
  str x0, [x9, #0x50]        // x10's slot in Context

  // Save x11..x30
  stp x11, x12, [x9, #0x58]
  stp x13, x14, [x9, #0x68]
  stp x15, x16, [x9, #0x78]
  stp x17, x18, [x9, #0x88]
  stp x19, x20, [x9, #0x98]
  stp x21, x22, [x9, #0xA8]
  stp x23, x24, [x9, #0xB8]
  stp x25, x26, [x9, #0xC8]
  stp x27, x28, [x9, #0xD8]
  stp x29, x30, [x9, #0xE8]

  // Save SP (original, before our scratch allocation), ELR, SPSR
  add x0, sp, #0x20
  str x0, [x9, #0xF8]        // SP
  mrs x0, elr_el1
  str x0, [x9, #0x100]       // ELR (where to resume)
  mrs x0, spsr_el1
  str x0, [x9, #0x108]       // SPSR (CPU flags)

Listing 14: exc_irq handler, part 2: save registers (boot.S)

stp (Store Pair) saves two 8-byte registers in one instruction. The offsets (#0x00, #0x10, #0x20, …) are byte offsets into the Context struct. x0 lives at offset 0, x1 at offset 8, x2 at offset 16 (0x10), and so on. After saving x0-x8, we restore x9 and x10 from the stack and save them in their proper slots in the Context. Finally, we save SP, ELR, and SPSR.

Notice how we recover SP: add x0, sp, #0x20 computes the original SP before we subtracted 0x20 for the scratch area. We store that original SP in the Context. This way, when we restore the context later, we can set SP back to where the interrupted thread expects it, not where we temporarily moved it for our scratch space.

11.3 Calling the Rust handler

The rust handler needs the current context pointer to decide which thread to run next. We pass it in x0, and it returns the next context pointer in x0. We then store that pointer back in TPIDR_EL1 for the next interrupt and also keep it in x19 for the restore sequence. This is the point where we can switch threads: by returning a different context pointer, the Rust handler can cause the assembly code to restore a different thread’s state. When the timer fires again, it will save the current thread’s state to its context, call the Rust handler, and then restore the next thread’s state, effectively switching between them on each tick.

  // Call rust_irq_handler(current_context) -> next_context
  mov x0, x9
  bl rust_irq_handler

  // x0 now points to the next Context (might be the same, might be different)
  msr tpidr_el1, x0
  mov x19, x0             // Keep as base register for restore

Listing 15: exc_irq handler, part 3: call Rust (boot.S)

We pass the current context pointer in x0 (the first argument in the C calling convention). The Rust handler returns the next context pointer in x0 (the return value). We stash this in both TPIDR_EL1 (so the next interrupt knows where to save) and x19 (so we can use it as a base for the restore sequence).

11.4 Restoring the next context

On the flip side, we restore all the registers from the context pointer returned by the Rust handler. This is essentially the reverse of the save sequence. We load SP, ELR, and SPSR first, then all the general-purpose registers. Finally, we execute eret to jump to the next thread’s entry point with the correct CPU state. This is quite literally the “magic” moment where we switch from one thread to another. The CPU doesn’t know or care that we’re switching threads; it just loads the new PC, flags, and jumps. The interrupted thread is completely unaware that it was paused and resumed later.

  // Restore SP, ELR, SPSR for the next thread
  ldr x1, [x19, #0xF8]
  mov sp, x1
  ldr x1, [x19, #0x100]
  msr elr_el1, x1
  ldr x1, [x19, #0x108]
  msr spsr_el1, x1

  // Restore x0..x18 (skip x19 for now, we're using it)
  ldp x0,  x1,  [x19, #0x00]
  ldp x2,  x3,  [x19, #0x10]
  ldp x4,  x5,  [x19, #0x20]
  ldp x6,  x7,  [x19, #0x30]
  ldp x8,  x9,  [x19, #0x40]
  ldp x10, x11, [x19, #0x50]
  ldp x12, x13, [x19, #0x60]
  ldp x14, x15, [x19, #0x70]
  ldr x16,       [x19, #0x80]
  ldp x17, x18,  [x19, #0x88]

  // Restore x20..x30
  ldp x20, x21, [x19, #0xA0]
  ldp x22, x23, [x19, #0xB0]
  ldp x24, x25, [x19, #0xC0]
  ldp x26, x27, [x19, #0xD0]
  ldp x28, x29, [x19, #0xE0]
  ldr x30,       [x19, #0xF0]

  // Restore x19 LAST (we were using it as the base pointer)
  ldr x19, [x19, #0x98]
  eret

Listing 16: exc_irq handler, part 4: restore and return (boot.S)

The critical detail: x19 is restored as the last because we’re using it as the base address for all loads. If we restored it earlier, we’d lose our pointer and couldn’t load the remaining registers. Then we execute eret (Exception Return), which atomically loads the PC from ELR_EL1 and the CPU flags from SPSR_EL1, and jumps to that address. The next thread starts running as if it were never interrupted. The whole save/restore process is completely transparent to the threads. They just run, get interrupted, and resume later without any awareness of the context switching happening behind the scenes.

12. start_first: launching the first thread

We do have a chicken-and-egg problem at boot: we need to start the first thread so that the timer can preempt it, but we don’t have a “previous context” to save because no thread is running yet. The solution is a special start_first function that loads the first thread’s context and jumps to it using eret. This function is called from Rust after we initialize the contexts but before we enable interrupts. It sets up the CPU state for the first thread and then “returns” to it, effectively launching our multitasking system.

.global start_first
start_first:
  msr tpidr_el1, x0          // Store context pointer for future IRQs
  mov x19, x0                // Use as base register

  // Restore SP, ELR, SPSR
  ldr x1, [x19, #0xF8]
  mov sp, x1
  ldr x1, [x19, #0x100]
  msr elr_el1, x1
  ldr x1, [x19, #0x108]
  msr spsr_el1, x1

  // Restore all GPRs
  ldp x0,  x1,  [x19, #0x00]
  ldp x2,  x3,  [x19, #0x10]
  // ... (same pattern as exc_irq restore)
  ldr x19, [x19, #0x98]
  eret

Listing 17: start_first (boot.S)

It’s nearly identical to the restore half of exc_irq. The eret jumps to whatever address was in elr, which we set to thread_a_entry during init(). The CPU doesn’t know this isn’t a “real” exception return. It just loads and jumps.

preempt::init();
extern "C" {
    fn start_first(ctx: *const preempt::Context) -> !;
}
unsafe { start_first(preempt::first_context()) }

Listing 18: Starting the first thread (main.rs)

13. The scheduler

As we mentioned earlier, the scheduler’s job is to decide which thread to run next. In our simple demo, we have only two threads, so the scheduler just toggles between them. After all that assembly and GIC configuration, the scheduler itself is almost comically simple:

pub fn switch_next(_current_ctx: *mut Context) -> *const Context {
    let cur = CURRENT.load(Ordering::Relaxed);
    let next = cur ^ 1;       // XOR: 0 becomes 1, 1 becomes 0
    CURRENT.store(next, Ordering::Relaxed);
    unsafe { &CTX[next] as *const Context }
}

Listing 19: Round-robin scheduler (preempt.rs)

That’s it. XOR toggle for two threads. Compare this with the cooperative scheduler from Part 2, which was a loop that called poll() on each task. That scheduler was the one making decisions about when to run each task. Here, the scheduler doesn’t decide when to switch — the timer interrupt does. The scheduler only decides who to switch to. This separation of concerns is important: the timer provides the preemption mechanism, and the scheduler provides the policy.

We XOR the current index with 1 to toggle between 0 and 1. This is a common trick for two-thread round-robin scheduling. For more threads, you’d use a different data structure (like a queue) and a different algorithm (like priority scheduling). Still, the core mechanism of “save current context, pick next context, restore next context” remains the same.

The rust_irq_handler calls switch_next every 5 ticks (500ms). This function picks the other thread by XOR-toggling the index (0 becomes 1, 1 becomes 0) and returns a pointer to that thread’s saved context. The assembly handler then restores that context instead of the interrupted thread’s context.

For two threads, XOR is perfect. For N threads, you’d replace this with (cur + 1) % N or a proper run queue with priorities. But the underlying mechanism — save the old context, pick the next thread, restore the next context — stays the same regardless of how sophisticated your scheduling policy becomes. Linux’s CFS scheduler and our XOR toggle both ultimately call the same kind of context switch.

14. Running the preemptive demo

OK, let’s see it in action. Build and run the demo-preempt feature. You should see the system boot, print a message about the preemptive demo, and then threads A and B alternating every second without ever yielding. This is the key proof that preemption works: both threads are running and switching without either of them ever calling yield() or returning. The OS is in control.

At face value, this looks boring — just two threads printing letters. But it’s actually a huge deal. We have a timer interrupt that can preempt running code, a scheduler that selects the next thread to run, and a context-switch mechanism that saves and restores CPU state perfectly. This is the core of any preemptive multitasking OS.

./scripts/build-aarch64-virt.sh demo-preempt
./scripts/run-aarch64-virt.sh

Listing 20: Build and run the preemptive demo

rustOS: aarch64 QEMU virt boot OK
rustOS: preemptive multitasking demo
A
B
A
B
A
B
...

Figure 2: Preemptive multitasking demo showing threads A and B alternating without yielding.

The proof: Neither thread ever calls yield() or returns. They both spin in infinite loops. Yet they alternate. The OS is in control.

Timing breakdown: Timer ticks every 100ms. Threads switch every 5 ticks (500ms). Each thread prints every 10 ticks (1 second). So you see A, then about a second later B, then A, and so on.

Press Ctrl-A then X to exit QEMU.

15. Cooperative vs preemptive: a comparison

Now that we’ve built both models, let’s put them side by side:

Property	Cooperative (Part 2)	Preemptive (this part)
Who decides when to switch	The task itself	The OS (via timer)
Can a task hog the CPU?	Yes	No
Overhead per switch	~10 cycles (function return)	~300 cycles (full register save/restore)
Worst-case latency	Unbounded (task might never yield)	Bounded (at most one tick period)
Shared state safety	Simple (no concurrent access)	Complex (need atomics, critical sections)
Best for	Trusted code, embedded, coroutines	Untrusted code, general-purpose OS

The overhead difference is worth understanding. In cooperative scheduling, a “switch” is just a function that returns, with the scheduler calling the next function. That’s a few instructions. In preemptive scheduling, we save 34 fields to memory, call the Rust handler, pick the next thread, restore 34 fields from memory, and execute eret. That’s roughly 30x more work per switch. At 10 switches per second (our current rate), that’s negligible. At 10,000 switches per second (a typical Linux configuration), it becomes a concern.

The shared state trade-off is also significant. With cooperative scheduling, you know exactly when your task might be interrupted: only when you return from poll(). So you can safely modify shared data between poll() calls without locks. With preemption, an interrupt can arrive between any two instructions. That’s why we use AtomicU64 for TICKS and AtomicUsize for CURRENT — even on a single core, the interrupt handler and the interrupted code are logically concurrent.

Most real systems use both approaches. The kernel itself typically uses cooperative scheduling internally (kernel code is trusted, and the performance benefit matters). User processes get preemptive scheduling (user code is untrusted, and fairness matters more than raw throughput). We’ll get there eventually — once we add user/kernel separation.

16. What we built

In this part, we took our cooperative kernel from Part 2 and gave the OS the power to enforce fairness by implementing preemptive scheduling. We went from “tasks voluntarily yield” to “the OS takes control whether tasks like it or not.” Here’s what we added:

Timer interrupts: The ARM Generic Timer fires every 100ms, giving us a real notion of time instead of meaningless loop iteration counts.
GIC configuration: The Distributor and CPU Interface route IRQ 30 (the physical timer) to our handler, with proper acknowledge/end-of-interrupt protocol.
Exception vectors: A 16-entry vector table catches IRQs and dispatches them to Rust code, bridging the hardware-software boundary.
Context switching: Assembly saves and restores all 34 fields (31 GPRs + SP + ELR + SPSR) per switch, ensuring threads can’t corrupt each other’s state.
Preemptive scheduling: Timer-driven round-robin means threads can’t monopolize the CPU, even if they spin in infinite loops.

The context switch takes about 300 cycles on our setup. That’s fast, partly because we don’t flush the TLB (we’re in a single address space), don’t save NEON/FP registers, and don’t switch page tables. A real OS doing all of those things might take 3,000-5,000 cycles per switch. But the mechanism is identical: save state, pick next, restore state, eret.

The two threads in our demo never yield, never cooperate, never even know the other exists. Yet they alternate perfectly. That’s the fundamental promise of preemptive multitasking, and it’s why every general-purpose OS uses it.

In Part 4 , we’ll tackle the last major piece of the puzzle: virtual memory. We’ll build a frame allocator, construct page tables, and enable the MMU so that each memory access goes through address translation. This is the foundation for process isolation — giving each task its own view of memory so they can’t stomp on each other.

17. References

ARM Architecture Reference Manual , Section D1 (Exception Handling), Section D11 (Generic Timer)
ARM GICv2 Specification
QEMU virt machine , memory map
ARM Procedure Call Standard (AAPCS64) , register usage
Operating Systems: Three Easy Pieces, Chapter 6: Limited Direct Execution

5-Part Series:

Part 0: Why build an OS from scratch?
Part 1: Foundations
Part 2: Communication
Part 3 (this): Concurrency
Part 4: Memory and beyond

Making a headless AI assistant observable - without SSH

Mon, 16 Mar 2026 00:00:00 +0000

GitHub Repository: bahree/nanoclaw - full source code

NanoClaw is a headless AI assistant running on my personal server. It processes messages from WhatsApp, Telegram, and Slack, runs scheduled tasks, and manages conversations with Claude agents in isolated containers. It’s been incredibly useful, but it had one major pain point: no visibility into what it was doing or why. If something went wrong - a message didn’t get a reply, a task didn’t run - the only way to debug was to SSH into the server, tail logs, and piece together what happened. Not ideal when you’re on the go and your assistant just… stops responding.

I use NanoClaw as my personal AI assistant - it handles everything from answering questions, to tracking flights, to giving me a daily morning briefing with F1 standings and weather. It is my agent and allows for me to check on the kind of thing I want to check on from my phone, not have to SSH into a server for.

The problem with a headless service that runs 24/7 is that you can’t see what it’s doing. When a message doesn’t get a reply, or a scheduled task doesn’t fire, the debugging workflow is: SSH into the server, tail the pino logs, grep for timestamps, piece together what happened. Not great when you’re out and about and your assistant just… stops responding.

I wanted three things:

A way to ask the system “what are you doing right now?” from the same WhatsApp chat I use to talk to it.
A way to manage scheduled tasks without SSH.
And a way to ask “why did you do that?” after the fact, with full traceability from triggering event to outbound action.

This post covers all three and how we built them in NanoClaw - nine features (~1100 lines of new TypeScript, two new modules, and three new SQLite tables).

TL;DR

Given the existing architecture of NanoClaw, we added a suite of observability features that are all accessible from the main WhatsApp group. No SSH, no separate dashboards, no external logging services - just commands you can type in the chat to see what’s going on and manage the system.

I added three groups of features to NanoClaw:

Real-time visibility - /status shows uptime, memory, active containers, channels, groups, and task summaries. /status tasks shows the full task list with schedules, next run times, and IDs.
Operational control - /task pause|resume|delete <id> manages scheduled tasks directly from the chat. No SSH, no restarts.
Event tracing and debugging - three SQLite tables that trace every action back to its triggering event. Pipeline is instrumented at message ingress, agent output, scheduled tasks, and IPC. Query with /debug last 10, /debug why, /debug event <id>, and /debug report. Auto-prunes after 3 days.

All operated entirely from the messaging channel.

1. The architecture (quick context)

Before diving in, here’s how NanoClaw works at a high level. Understanding this makes the instrumentation decisions clearer. NanoClaw is a OSS fork of OpenClaw, so it shares the same core architecture:

flowchart LR
    WA[WhatsApp] --> ORC[Orchestrator]
    TG[Telegram] --> ORC
    SL[Slack] --> ORC
    ORC --> DB[(SQLite)]
    ORC --> Q[Group Queue]
    Q --> C1[Container 1]
    Q --> C2[Container 2]
    C1 --> IPC[IPC Watcher]
    IPC --> ORC
    SCH[Task Scheduler] --> Q

Figure 1: NanoClaw message flow

Messages arrive from channels, get stored in SQLite, and the orchestrator polls for new messages every 2 seconds. When a registered group has unprocessed messages, the group queue spawns a container running Claude’s Agent SDK. The agent’s output streams back through the orchestrator to the originating channel. Scheduled tasks follow the same path but are triggered by a cron/interval scheduler rather than by incoming messages.

The key insight: everything flows through the orchestrator. That’s where we intercept commands, instrument actions, and expose state. The group queue manages the container lifecycle. The database is already there for message storage. All the pieces are in place - we just need to wire them up.

2. What was built

Here’s the full inventory. Two new modules, six modified files, nine distinct features:

#	Feature	Type	Files
1	`/status` - system dashboard	Command	`status.ts`, `index.ts`
2	`/status tasks` - task detail view	Command	`status.ts`, `index.ts`
3	`/task pause\|resume\|delete` - task management	Command	`status.ts`, `index.ts`
4	`GroupQueue.getStatus()` - queue introspection	API	`group-queue.ts`
5	Three-table event log schema	Schema	`db.ts`
6	`logEvent` / `logAction` / `logToolCall`	Module	`event-log.ts` (new)
7	Pipeline instrumentation	Instrumentation	`index.ts`, `task-scheduler.ts`, `ipc.ts`
8	`/debug last\|why\|event\|report`	Command	`status.ts`, `event-log.ts`
9	Auto-pruning with configurable retention	Config	`event-log.ts`, `config.ts`

All commands are restricted to the main group only. Non-main groups are silently ignored, preventing random group members from querying system status or managing tasks.

3. /status - real-time system dashboard

The /status command assembles information from several subsystems into a single message. It queries the group queue for container states, the database for registered groups and tasks, and formats it all as a WhatsApp-friendly message.

The output is designed to give a quick overview of the system’s health and activity at a glance:

Figure 2: /status while the agent is answering an F1 question - containers, channels, groups, and tasks at a glance.

What it shows:

Uptime and memory - how long the process has been running, RSS in MB
Timezone - the configured timezone (important for cron schedules)
Containers - active/max concurrent, with per-group detail (idle, processing, running task, queued)
Channels - which messaging channels are connected (WhatsApp, Telegram, etc.)
Groups - all registered groups with the main group indicator
Tasks - count of active/paused tasks, next upcoming task with time-until
Remote control - whether a remote Claude Code session is active

The implementation pulls from existing subsystems - no new state tracking was needed:

export function buildStatus(queue: GroupQueue, channels: Channel[]): string {
  const uptime = formatDuration(Date.now() - startTime);
  const mem = Math.round(process.memoryUsage.rss() / 1024 / 1024);

  const qs = queue.getStatus();       // new method on GroupQueue
  const channelNames = channels.map((ch) => ch.name).join(', ');
  const groups = getAllRegisteredGroups();
  const tasks = getAllTasks();
  const activeTasks = tasks.filter((t) => t.status === 'active');
  const rc = getActiveSession();       // remote control state

  const lines: string[] = [
    `*${ASSISTANT_NAME} Status*`,
    `----------------`,
    `Uptime: ${uptime}`,
    `Memory: ${mem} MB`,
    `Timezone: ${TIMEZONE}`,
    ``,
    `*Containers:* ${qs.activeCount}/${qs.maxConcurrent} active`,
    `*Channels:* ${channelNames || 'none'}`,
  ];
  // ... active container details, groups, tasks, remote control
  return lines.join('\n');
}

Listing 1: Building the status response

3.1 Exposing queue internals

The GroupQueue class already tracked everything we needed internally: active containers, pending messages, pending tasks, and idle state. It just wasn’t exposed. A new getStatus() method surfaces this without leaking internal implementation:

getStatus(): {
  activeCount: number;
  maxConcurrent: number;
  waitingCount: number;
  groups: Array<{
    jid: string;
    active: boolean;
    idleWaiting: boolean;
    isTaskContainer: boolean;
    pendingMessages: boolean;
    pendingTaskCount: number;
  }>;
} {
  const groups = [];
  for (const [jid, state] of this.groups) {
    if (state.active || state.pendingMessages || state.pendingTasks.length > 0) {
      groups.push({
        jid,
        active: state.active,
        idleWaiting: state.idleWaiting,
        isTaskContainer: state.isTaskContainer,
        pendingMessages: state.pendingMessages,
        pendingTaskCount: state.pendingTasks.length,
      });
    }
  }
  return {
    activeCount: this.activeCount,
    maxConcurrent: MAX_CONCURRENT_CONTAINERS,
    waitingCount: this.waitingGroups.length,
    groups,
  };
}

Listing 2: GroupQueue.getStatus() - queue introspection

Only groups with activity (active, pending messages, or pending tasks) are included - no noise from idle groups. The status output translates these states into human-readable labels: “idle”, “processing”, “running task”, or “queued”.

4. /status tasks - task detail view

While /status includes a task summary (count + next upcoming), /status tasks gives the full picture. Each task shows its prompt (truncated to 50 chars), schedule type, next run time, last run time, and the task ID you need for management commands.

Figure 3: /status tasks - two active daily crons (morning briefing with F1 + weather, AI news digest) and completed one-offs (flight tracking, reminders).

function formatTaskLine(task: ScheduledTask, index: number): string {
  const status = task.status === 'active' ? ''
    : task.status === 'paused' ? ' [paused]' : ' [done]';
  const schedule = task.schedule_type === 'cron'
    ? `cron: ${task.schedule_value}`
    : task.schedule_type === 'interval'
      ? `every ${task.schedule_value}`
      : `once`;
  const prompt = task.prompt.length > 50
    ? task.prompt.slice(0, 50) + '...' : task.prompt;
  const next = task.next_run ? formatTimeUntil(task.next_run) : 'n/a';
  const lastRun = task.last_run ? formatTimeAgo(task.last_run) : 'never';

  return [
    `*${index + 1}.* ${prompt}${status}`,
    `   Schedule: ${schedule}`,
    `   Next: ${next} | Last: ${lastRun}`,
    `   ID: ${task.id}`,
  ].join('\n');
}

Listing 3: Formatting a task line

Tasks are grouped by status - active first, then paused, then completed. The relative time formatting (in 2h 15m, 3m 42s ago) makes it easy to see at a glance what’s coming up and what ran recently.

5. /task - operational control

The /task command turns the chat into a control plane. No more SSH-ing in to pause a runaway task or clean up a completed one.

Command	What it does
`/task pause <id>`	Pause an active task (stops scheduling, preserves config)
`/task resume <id>`	Resume a paused task
`/task delete <id>`	Delete a task and its run history

Figure 4: Pausing and resuming the AI news digest task - no SSH, no restart.

export function handleTaskCommand(
  args: string,
): { ok: true; message: string } | { ok: false; error: string } {
  const parts = args.trim().split(/\s+/);
  const action = parts[0]?.toLowerCase();
  const taskId = parts.slice(1).join(' ');

  const task = getTaskById(taskId);
  if (!task) return { ok: false, error: `Task not found: ${taskId}` };

  switch (action) {
    case 'pause':
      if (task.status !== 'active')
        return { ok: false, error: `Task is already ${task.status}` };
      updateTask(taskId, { status: 'paused' });
      return { ok: true, message: `Paused: "${task.prompt.slice(0, 50)}"` };
    case 'resume':
      if (task.status !== 'paused')
        return { ok: false, error: `Task is ${task.status}, not paused` };
      updateTask(taskId, { status: 'active' });
      return { ok: true, message: `Resumed: "${task.prompt.slice(0, 50)}"` };
    case 'delete':
      deleteTask(taskId);
      return { ok: true, message: `Deleted: "${task.prompt.slice(0, 50)}"` };
    default:
      return { ok: false, error: `Unknown action: ${action}` };
  }
}

Listing 4: Task command handler

The result type ({ ok: true; message } | { ok: false; error }) is a pattern used throughout NanoClaw for commands; the caller doesn’t need to know the implementation details, just whether it succeeded and what to tell the user. State validation is done upfront (can’t pause an already-paused task, can’t resume an active one).

6. Command interception

All built-in commands (/status, /status tasks, /task, /debug) share the same interception pattern: they’re caught at the top of the onMessage callback, before the message is stored or processed.

const channelOpts = {
  onMessage: (chatJid: string, msg: NewMessage) => {
    const trimmed = msg.content.trim();

    // Built-in commands - intercept before storage
    if (trimmed === '/status' || trimmed === '/status tasks') {
      handleStatus(trimmed, chatJid, msg).catch(/* ... */);
      return;  // don't store, don't process
    }
    if (trimmed.startsWith('/task ')) {
      handleTaskCmd(chatJid, msg).catch(/* ... */);
      return;
    }
    if (trimmed.startsWith('/debug')) {
      handleDebugCmd(chatJid, msg).catch(/* ... */);
      return;
    }

    // ... sender allowlist filtering, event logging, message storage
    storeMessage(msg);
  },
};

Listing 5: Intercepting built-in commands before storage

The return before storeMessage() is the key design decision. These commands are ephemeral – they shouldn’t appear in the conversation history, shouldn’t trigger the agent, and shouldn’t affect message cursors. They’re handled entirely by the host process and respond instantly (no container spawn needed).

Each handler checks group?.isMain before proceeding. Non-main-group commands are silently dropped, with a warning in the server logs.

7. Event/action/tool logging

The commands above tell you what’s happening now. But when something went wrong an hour ago, you need a trail. That’s where event logging comes in.

7.1 The three-table schema

NanoClaw’s existing SQLite schema was focused on message storage and container state. We needed a new schema to capture the full traceability from inbound triggers to outbound actions to tool calls. The design is a simple three-table structure. Every inbound trigger is an event, every outbound action is an action linked to its triggering event, and every tool invocation is a tool call linked to its parent action.

CREATE TABLE IF NOT EXISTS event_log (
  id           TEXT PRIMARY KEY,    -- UUID
  timestamp    DATETIME DEFAULT CURRENT_TIMESTAMP,
  source       TEXT NOT NULL,       -- 'whatsapp', 'telegram', 'scheduled_task', 'ipc'
  source_id    TEXT,                -- message ID, task ID, etc.
  raw_content  TEXT,                -- full payload (JSON, truncated to 10KB)
  summary      TEXT                 -- human-readable one-liner
);

CREATE TABLE IF NOT EXISTS action_log (
  id            TEXT PRIMARY KEY,
  timestamp     DATETIME DEFAULT CURRENT_TIMESTAMP,
  triggered_by  TEXT REFERENCES event_log(id),
  action_type   TEXT NOT NULL,      -- 'message_sent', 'task_scheduled', 'tool_call'
  target        TEXT,               -- group JID, email address, task ID
  content       TEXT,               -- what was sent or done
  tool_calls    TEXT                -- JSON array of tool names
);

CREATE TABLE IF NOT EXISTS tool_call_log (
  id           TEXT PRIMARY KEY,
  action_id    TEXT REFERENCES action_log(id),
  timestamp    DATETIME DEFAULT CURRENT_TIMESTAMP,
  tool_name    TEXT NOT NULL,
  input        TEXT,                -- JSON (truncated to 10KB)
  output       TEXT,                -- JSON (truncated to 10KB)
  duration_ms  INTEGER,
  success      INTEGER DEFAULT 1
);

CREATE INDEX IF NOT EXISTS idx_event_log_timestamp ON event_log(timestamp);
CREATE INDEX IF NOT EXISTS idx_action_log_triggered_by ON action_log(triggered_by);
CREATE INDEX IF NOT EXISTS idx_tool_call_log_action_id ON tool_call_log(action_id);

Listing 6: Event logging schema

The foreign key chain is event_log <- action_log <- tool_call_log. Given any action, you can trace back to why it happened. Given any event, you can see everything it caused. The indexes support the /debug query patterns: filtering by timestamp (pruning), joining by triggered_by (event lookups), and grouping by action_id (tool call chains).

flowchart TD
    E[Event: WhatsApp message received] --> A1[Action: message_sent to group]
    E --> A2[Action: task_scheduled]
    A1 --> T1[Tool Call: runContainerAgent]
    A1 --> T2[Tool Call: channel.sendMessage]

Figure 5: Event -> Action -> Tool Call tracing chain

These tables live in the same messages.db file as everything else. No additional file handles, no additional backup concerns, no additional connection management. They’re created in the existing createSchema() function using CREATE TABLE IF NOT EXISTS, so they’re added transparently on first startup after the upgrade.

7.2 The logger module

Three functions, matching the three tables. All IDs are UUIDs via crypto.randomUUID(). All writes are fire-and-forget - wrapped in try/catch, errors logged at debug level, never blocking the pipeline.

export function logEvent(
  source: string,
  sourceId: string | null,
  rawContent: unknown,
  summary: string,
): string {
  const id = crypto.randomUUID();
  try {
    insertEventStmt().run(
      id, new Date().toISOString(), source, sourceId,
      truncate(rawContent), summary,
    );
  } catch (err) {
    logger.debug({ err, source, sourceId }, 'Failed to log event');
  }
  return id;  // always returns an ID, even if the write failed
}

export function logAction(
  triggeredBy: string | null,
  actionType: string,
  target: string | null,
  content: unknown,
  toolCalls?: string[],
): string {
  const id = crypto.randomUUID();
  try {
    insertActionStmt().run(
      id, new Date().toISOString(), triggeredBy, actionType,
      target, truncate(content),
      toolCalls ? JSON.stringify(toolCalls) : null,
    );
  } catch (err) {
    logger.debug({ err, actionType, target }, 'Failed to log action');
  }
  return id;
}

Listing 7: Core logging functions

Every call returns an ID, so callers can chain events -> actions -> tool calls, even if the write fails silently. This is intentional - the pipeline code doesn’t check whether logging succeeded, it just carries the ID forward.

The truncate() helper caps content at 10KB to prevent DB bloat:

const MAX_CONTENT_SIZE = 10 * 1024;

function truncate(value: unknown): string | null {
  if (value === undefined || value === null) return null;
  const str = typeof value === 'string' ? value : JSON.stringify(value);
  if (str.length > MAX_CONTENT_SIZE) return str.slice(0, MAX_CONTENT_SIZE);
  return str;
}

Listing 8: Content truncation

7.3 The tool call wrapper

logToolCall is different from the other two - it wraps an async operation and automatically records input, output, duration, and success/failure:

export async function logToolCall<T>(
  actionId: string,
  toolName: string,
  input: unknown,
  fn: () => Promise<T>,
): Promise<T> {
  const id = crypto.randomUUID();
  const start = Date.now();
  let success = true;
  let output: unknown = null;

  try {
    const result = await fn();
    output = result;
    return result;
  } catch (err) {
    success = false;
    output = err instanceof Error ? err.message : String(err);
    throw err;  // re-throw - logging doesn't swallow errors
  } finally {
    const durationMs = Date.now() - start;
    try {
      insertToolCallStmt().run(
        id, actionId, new Date().toISOString(), toolName,
        truncate(input), truncate(output), durationMs, success ? 1 : 0,
      );
    } catch (logErr) {
      logger.debug({ err: logErr, toolName }, 'Failed to log tool call');
    }
  }
}

Listing 9: Tool call wrapper with automatic instrumentation

The finally block ensures the log entry is written regardless of whether the wrapped operation succeeded or failed. The error is always re-thrown - logToolCall is transparent to the caller. It’s a decorator pattern: wrap any async operation and get free instrumentation.

Insert statements are lazily prepared and reused across calls, avoiding the overhead of re-preparing the same SQL on every log write.

7.4 Instrumenting the pipeline

With the logger module in place, instrumentation is surgical - a few lines at each key point in the message flow:

Inbound messages (in the onMessage callback):

// Detect channel from JID format
const evtChannel = chatJid.includes('@g.us') || chatJid.includes('@s.whatsapp.net')
  ? 'whatsapp'
  : chatJid.startsWith('tg:') ? 'telegram'
  : chatJid.startsWith('dc:') ? 'discord'
  : chatJid.startsWith('sl:') ? 'slack'
  : 'channel';

logEvent(
  evtChannel, msg.id,
  { sender: msg.sender_name, content: msg.content?.slice(0, 200) },
  `Message from ${msg.sender_name}: ${(msg.content || '').slice(0, 80)}`,
);

Listing 10: Logging inbound messages with channel detection

Agent processing (when the orchestrator starts handling a group’s messages):

// Log the processing event - this ID links to all downstream actions
const eventId = logEvent(
  'message_batch', chatJid,
  { messageCount: missedMessages.length, group: group.name },
  `Processing ${missedMessages.length} message(s) for ${group.name}`,
);

Listing 11: Logging the processing event and carrying the eventId

Outbound messages (in the streaming output callback):

if (text) {
  await channel.sendMessage(chatJid, text);
  logAction(eventId, 'message_sent', chatJid, text.slice(0, 500));
}

Listing 12: Logging outbound actions linked to the triggering event

The eventId from the processing event links the outbound action back to the batch that triggered it. This is the chain that /debug why follows.

Scheduled tasks - logged at the point the scheduler picks up a due task, and again when the result is sent to the user:

const eventId = logEvent(
  'scheduled_task', task.id,
  { prompt: task.prompt.slice(0, 200), schedule: task.schedule_type },
  `Scheduled task: ${task.prompt.slice(0, 80)}`,
);

// ... later, when the container produces output:
await deps.sendMessage(task.chat_jid, streamedOutput.result);
logAction(eventId, 'message_sent', task.chat_jid, result?.slice(0, 500) ?? null);

Listing 13: Scheduled task instrumentation

IPC - logged when the IPC watcher processes messages and task operations from containers:

const ipcEventId = logEvent(
  'ipc', null,
  { chatJid: data.chatJid, sourceGroup, text: data.text?.slice(0, 200) },
  `IPC message from ${sourceGroup}`,
);
await deps.sendMessage(data.chatJid, data.text);
logAction(ipcEventId, 'message_sent', data.chatJid, data.text?.slice(0, 500) ?? null);

Listing 14: IPC message instrumentation

8. The /debug commands

Four subcommands for querying the event log, all main-group only:

Command	What it shows
`/debug last <n>`	Last n actions with their triggering events
`/debug why`	Most recent action with full event context and tool call chain
`/debug event <id>`	Everything triggered by a specific event
`/debug report`	Summary dashboard: table sizes, events by source, actions by type, busiest hours, recent errors

8.1 /debug last - quick scan

The “what happened recently?” view. Joins action_log with event_log to show each action alongside what caused it:

Figure 6: /debug last 5 - the F1 response, flight landing alert, and Gmail cleanup, each traced back to its trigger.

export function getLastActions(n: number): Array<{
  action: ActionLogRow;
  event: EventLogRow | null;
}> {
  const rows = getDb()
    .prepare(
      `SELECT a.*, e.id as e_id, e.timestamp as e_timestamp,
              e.source as e_source, e.source_id as e_source_id,
              e.summary as e_summary
       FROM action_log a
       LEFT JOIN event_log e ON a.triggered_by = e.id
       ORDER BY a.timestamp DESC
       LIMIT ?`,
    )
    .all(n);
  // ... map to structured result
}

Listing 15: Querying last N actions with LEFT JOIN

The LEFT JOIN is important - some actions might not have a triggering event (e.g., system-initiated actions), and we still want to see them.

8.2 /debug why - full trace

Answers “why did the last thing happen?” by pulling the most recent action, its triggering event, and all associated tool calls:

Figure 7: /debug why - tracing the most recent action back to the WhatsApp message that caused it.

export function getLastActionWithToolCalls(): {
  action: ActionLogRow;
  event: EventLogRow | null;
  toolCalls: ToolCallLogRow[];
} | null {
  const results = getLastActions(1);
  if (results.length === 0) return null;

  const { action, event } = results[0];
  const toolCalls = getDb()
    .prepare(
      `SELECT * FROM tool_call_log WHERE action_id = ? ORDER BY timestamp`,
    )
    .all(action.id) as ToolCallLogRow[];

  return { action, event, toolCalls };
}

Listing 16: Full trace for the most recent action

The output shows the full chain: triggering event (source, summary, timestamp, ID), the action taken (type, target, content), and each tool call with its duration and success/failure status. Copy the event ID from here and use /debug event <id> to see everything else that event triggered.

8.3 /debug report - summary dashboard

The “is everything healthy?” view. Aggregates across all three tables into a single report:

Figure 8: /debug report - a day's worth of WhatsApp messages, flight tracking, Gmail cleanup, and F1 queries summarized.

What it includes:

Retention period - configured days and current time window
Table sizes - row counts for events, actions, and tool calls
Events by source - breakdown by channel (whatsapp, telegram, scheduled_task, ipc)
Actions by type - breakdown by what was done (message_sent, task_scheduled, etc.)
Busiest hours - top 5 hours by event count, in local timezone
Recent failed tool calls - last 10 with tool name, duration, and error output
Recent errors - last 10 error-like actions with their triggering event

One thing worth noting: the busiest hours are computed in JavaScript using toLocaleString with the configured time zone, not in SQL. SQLite stores timestamps as UTC (as ISO strings), and performing timezone conversions in SQL would require loading an extension. Instead, we fetch the raw timestamps and bucket them in JS:

const allTimestamps = db.prepare(`SELECT timestamp FROM event_log`).all();
const hourCounts = new Map<string, number>();
for (const { timestamp } of allTimestamps) {
  const localHour = new Date(timestamp).toLocaleString('en-US', {
    timeZone: TIMEZONE,
    hour: 'numeric', hour12: true,
  });
  hourCounts.set(localHour, (hourCounts.get(localHour) || 0) + 1);
}

Listing 17: Bucketing event timestamps by local hour

Learned this the hard way when the report initially showed UTC hours, and I couldn’t figure out why 5 PM was my busiest time. 😄

9. Auto-pruning and retention

An unbounded observability system is a liability. Logs older than 3 days are automatically deleted. The retention period is configurable via the EVENT_LOG_RETENTION_DAYS environment variable (set to 0 to disable pruning).

// config.ts
export const EVENT_LOG_RETENTION_DAYS = Math.max(
  0,
  parseInt(process.env.EVENT_LOG_RETENTION_DAYS || '3', 10) || 3,
);
export const EVENT_LOG_PRUNE_INTERVAL = 60 * 60 * 1000; // hourly

Listing 18: Log retention configuration

Pruning runs at startup (clean up anything that expired while the service was down) and then every 60 minutes:

export function pruneOldLogs() {
  if (EVENT_LOG_RETENTION_DAYS === 0) return;

  const cutoff = new Date(
    Date.now() - EVENT_LOG_RETENTION_DAYS * 24 * 60 * 60 * 1000,
  ).toISOString();

  const db = getDb();
  // Delete in FK-safe order: children first
  db.prepare(`DELETE FROM tool_call_log WHERE action_id IN (
    SELECT id FROM action_log WHERE timestamp < ?
  )`).run(cutoff);
  db.prepare(`DELETE FROM action_log WHERE timestamp < ?`).run(cutoff);
  db.prepare(`DELETE FROM event_log WHERE timestamp < ?`).run(cutoff);
}

export function startLogPruning(): void {
  pruneOldLogs();
  const timer = setInterval(pruneOldLogs, EVENT_LOG_PRUNE_INTERVAL);
  timer.unref();   // don't keep the process alive for pruning
}

Listing 19: Pruning with FK-safe deletion order

The deletion order matters: tool_call_log rows reference action_log, which in turn references event_log. Deleting parents first would violate foreign key constraints. The timer.unref() call ensures the pruning interval doesn’t prevent graceful shutdown.

10. Design decisions

A few choices that are worth calling out:

Fire-and-forget, not await. Every logging call is synchronous (better-sqlite3) and wrapped in try/catch. If the write fails - disk full, DB locked, schema mismatch - the error is logged at debug level and the pipeline continues. The logging system is never on the critical path. An observability system that can take down the thing it’s observing is worse than useless.

Same database, no new dependencies. The logging tables live in messages.db alongside messages, tasks, sessions, and router state. No new files to back up, no new connections to manage, no new packages to install. CREATE TABLE IF NOT EXISTS means existing installations pick up the schema on restart.

Commands intercepted before storage. /status, /task, and /debug messages never reach the agent container. They don’t appear in conversation history, don’t trigger container spawns, and don’t affect message cursors. This is important - a /status check shouldn’t cost you a container slot or show up as context in the agent’s next conversation.

Prepared statements, lazily created. The insert statements are created on first use and reused across calls. For a system logging every message and action, re-preparing SQL on every call would add up.

UUIDs for everything. crypto.randomUUID() for all IDs. No auto-increment, no collision risk across restarts, and IDs are meaningful in isolation (you can paste one into /debug event <id> without context).

11. Try it yourself

If you’re running NanoClaw (or OpenClaw), these features are available out of the box. Here’s how to get started:

If you already have NanoClaw running:

Pull the latest code and rebuild:
1 2

git pull npm run build

Restart the service:

# Linux (systemd)
systemctl --user restart nanoclaw

# macOS (launchd)
launchctl kickstart -k gui/$(id -u)/com.nanoclaw

The new tables are created automatically on startup. No migration step needed.
Send /status in your main group to verify it’s working.

If you’re starting fresh:

Fork or clone bahree/nanoclaw (or the upstream OpenClaw )
Follow the setup instructions in the README
Once connected to a channel, all commands are available immediately

Command reference:

Command	What it does
`/status`	System overview: uptime, memory, containers, channels, groups, tasks
`/status tasks`	Full task list with schedules, next run, last run, IDs
`/task pause <id>`	Pause a scheduled task
`/task resume <id>`	Resume a paused task
`/task delete <id>`	Delete a task and its run history
`/debug last <n>`	Last n actions with their triggering events
`/debug why`	Most recent action with full trace
`/debug event <id>`	All actions triggered by a specific event
`/debug report`	Summary dashboard with stats and errors

Configuration:

Env variable	Default	What it does
`EVENT_LOG_RETENTION_DAYS`	`3`	Days to keep event logs (0 = keep forever)

All commands are main-group only. They respond instantly (no container needed) and don’t appear in the conversation history.

12. Summary

Three problems, one philosophy: make the system controllable and observable from the same interface you use to interact with it.

/status gives you real-time visibility - what’s running, what’s queued, what’s scheduled, which channels are connected. /task gives you operational control: pause a runaway task, resume one you paused, and clean up completed ones. Event logging gives you after-the-fact traceability - every action links back to its triggering event, every tool call links back to its parent action. /debug commands let you query the trail. Auto-pruning keeps it from growing unbounded.

About ~1100 lines of new TypeScript across 8 files. Two new modules (status.ts and event-log.ts), three new SQLite tables, a handful of indexes, and one new config variable. No new dependencies, no separate services. It just works on the next restart.

The source code for NanoClaw is available at bahree/nanoclaw .

Building a microkernel in Rust (Part 2): Communication, messages between tasks

Sun, 15 Mar 2026 00:00:00 +0000

5-Part Series:

Part 0: Why build an OS from scratch?
Part 1: Foundations
Part 2 (this): Communication
Part 3: Concurrency
Part 4: Memory and beyond

GitHub Repository: bahree/rust-microkernel - full source code and build scripts

Docker Image: amitbahree/rust-microkernel - prebuilt dev environment with Rust, QEMU, and source code

Recap from Part 1 : we have a bare-metal kernel that boots on AArch64, sets up a basic logger, and halts. It’s a single task, doing one thing. Now it’s time to add more tasks and let them talk to each other.

Right now, our kernel boots, prints a message, and halts - that’s it. One task, running alone, forever. But operating systems run hundreds of tasks at once, and those tasks need to communicate with one another. A scheduler needs to tell a display driver what to render. A network stack needs to hand received bytes to an application. How do those tasks communicate?

That’s what we’re building today. We’ll add message-passing IPC (inter-process communication) and a cooperative scheduler, turning our boot-only kernel into something that actually looks like a microkernel. Two tasks will run side by side, exchanging messages through a router we build from scratch.

TL;DR

We add two big things in this post. First, a message-passing IPC system: tasks communicate by sending small, fixed-size messages through a central router, rather than sharing memory directly. Second, a cooperative scheduler that polls tasks in round-robin order, giving each one a turn to do work before moving to the next.

A couple of terms we’ll use throughout.

Cooperative multitasking means that each task voluntarily relinquishes the CPU when it finishes its current unit of work. The OS trusts tasks to yield promptly, which is a nice assumption when it holds and a disaster when it doesn’t (more on that later).
Round-robin is the simplest scheduling strategy: we call each task in order, cycling through the list forever. Every task gets an equal turn. No priorities, no special treatment.

1. Why microkernel?

Most operating systems you use daily (Linux, Windows, macOS) use monolithic kernels. Everything runs in kernel space with full privileges: drivers, filesystems, networking, etc. This makes things fast because components can call each other directly, with no overhead. However, on the flip side, it’s fragile. A single bug in a GPU driver can corrupt kernel memory and take down the entire system, because that driver runs with the same unrestricted access as the scheduler and the memory manager.

Microkernels flip this around. Only the absolute minimum runs in kernel space: IPC, scheduling, and basic memory management. Everything else (drivers, filesystems, networking) runs in user space as separate, isolated processes. If a driver crashes, the kernel can restart it without rebooting. The tradeoff is speed. Every time a user-space driver needs to talk to the kernel or another driver, it must cross a privilege boundary, which costs cycles.

1.1 Famous microkernels

OS	Used in	Key feature
seL4	Aerospace, medical devices	Formally verified (mathematically provably correct)
Minix 3	Intel ME firmware	Extreme isolation, automatic driver restart
QNX	Cars, industrial systems	Hard real-time guarantees
L4 family	Research, embedded	Ultra-fast IPC (~100 cycles)

A microkernel forces us to think about interfaces - especially as we can’t just reach into another component’s data structures. We need to define a message format and a protocol. That constraint yields cleaner abstractions and makes each piece easier to reason about in isolation. When something goes wrong, we can trace the message flow rather than hunting through a tangle of shared state.

2. IPC: how tasks “talk”

In a microkernel, tasks don’t share memory. They communicate by sending messages through the kernel. This is called inter-process communication (IPC). It’s the lifeblood of a microkernel OS. The design of the IPC system has huge implications for performance, security, and ease of use. There are many ways to do IPC, but we’ll build a simple message-passing system with fixed-size messages and single-slot mailboxes. This is the model used by real microkernels like L4 and seL4, and it’s a great way to understand the core concepts without getting bogged down in complexity.

2.1 The problem with shared memory

The key design decision in microkernel IPC is to avoid shared memory. If tasks shared memory, they could read and write the same variables directly. This is simple and fast, but it leads to all sorts of problems, such as race conditions, data corruption, and security vulnerabilities. For example:

// Shared memory approach (NOT what we're building)
static mut SHARED_DATA: u32 = 0;

// Task A writes
unsafe { SHARED_DATA = 42; }

// Task B reads
let x = unsafe { SHARED_DATA };

Listing 1: Shared memory approach (not what we're doing)

On one hand, this looks simple, but there’s a problem. What happens if Task A is partway through writing a value when Task B reads it? Imagine A is changing a u32 counter from 255 (0x000000FF) to 256 (0x00000100). If B reads at just the wrong moment, it might see a half-updated value like 0x000001FF (511), which is neither the old value nor the new one. This is called a torn read.

Now, on our single-core cooperative system, this particular scenario can’t happen because tasks don’t run simultaneously. But the moment you add preemption (Part 3) or multiple cores, shared memory becomes a minefield. And even without race conditions, shared memory has other problems. There’s no access control (anyone can read or write), and bugs become “spooky action at a distance” where one task silently corrupts data that another task depends on.

2.2 Message-passing instead

The alternative is message-passing. Tasks don’t share any memory. Instead, they send discrete, self-contained messages through the kernel. Each message has a header (who it’s from, who it’s to, what type of message it is) and a payload (the actual data). The kernel routes messages to their destinations based on their headers. This way, communication becomes explicit and traceable. You can see every interaction by looking at the send and receive calls. The kernel controls delivery, so you get natural access control (a task can only receive messages addressed to it). And if the receiver isn’t keeping up, the router can push back on the sender rather than letting messages pile up in some unbounded buffer.

// Task A sends a message
router.send(Message {
    header: MsgHeader { dst: EndpointId::TaskB, .. },
    payload: [42, 0, 0, 0, 0, 0, 0, 0],
});

// Task B receives it
if let Some(msg) = router.recv(EndpointId::TaskB) {
    // process msg.payload
}

Listing 2: Message-passing approach (our design)

Communication becomes explicit, allowing us to trace every interaction by reading the send and receive calls. The kernel controls delivery, so you get natural access control (a task can only receive messages addressed to it). And if the receiver isn’t keeping up, the router can push back on the sender rather than letting messages pile up in some unbounded buffer. This is the model real microkernels use, and it’s what we’ll build.

3. Message structure

Let’s look at the actual code. Everything lives in crates/kernel/src/ipc.rs. This module defines the message format, the router, and some helper functions for working with messages. The design is simple but captures the essential features of a real IPC system.

3.1 Endpoints and message types

First, we need a way to identify who’s who. Each task gets an EndpointId, and each message has a MsgType to indicate its type. This is like addressing an envelope: the header says who it’s from, who it’s to, and what kind of message it is.

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
#[repr(u8)]
pub enum EndpointId {
    Ping = 1,
    Pong = 2,
}

#[derive(Copy, Clone, Debug)]
#[repr(u8)]
pub enum MsgType {
    Ping = 1,
    Pong = 2,
}

Listing 3: Endpoint and message type definitions

An endpoint is like a mailing address. Each task is assigned a unique EndpointId, and messages are routed to their destination by looking up the recipient’s endpoint. Think of each task as having its own numbered post office box.

You’ll notice #[repr(u8)] on both enums. Here’s what that does. By default, Rust can choose whatever in-memory representation it likes for an enum, and that might be 4 bytes or more depending on alignment. #[repr(u8)] tells the compiler to store the enum’s discriminant as exactly one byte. On bare metal, we need precise control over the size of every data structure because these values end up in message buffers and will eventually cross privilege boundaries via syscalls. Without #[repr(u8)], our struct sizes would be unpredictable.

3.2 The message itself

pub const MAX_PAYLOAD: usize = 8;

#[derive(Copy, Clone, Debug)]
#[repr(C)]
pub struct MsgHeader {
    pub src: EndpointId,      // who sent this
    pub dst: EndpointId,      // who should receive it
    pub ty: MsgType,          // what kind of message
    pub len: u8,              // how many payload bytes are used
    pub seq: u32,             // sequence number for ordering
}

#[derive(Copy, Clone, Debug)]
#[repr(C)]
pub struct Message {
    pub header: MsgHeader,
    pub payload: [u8; MAX_PAYLOAD],
}

Listing 4: Message header and payload

A few design decisions to outline.

Why #[repr(C)]?

This tells Rust to lay out the struct in memory exactly like a C compiler would. The fields are stored in declaration order, with predictable padding and alignment. This is crucial for IPC because both the sender and receiver need to agree on where each field sits in memory. If Rust reordered fields for optimization, the sender might put the src field at byte 0 while the receiver expects it at byte 4, leading to chaos.

Why only 8 bytes of payload?

One might wonder why we limit the payload to just 8 bytes. This is a deliberate design choice to keep our IPC system simple and efficient. In a real microkernel, the IPC mechanism is optimized for small messages that fit in CPU registers, enabling very fast communication between tasks. By keeping the payload small, we can avoid the overhead of copying large amounts of data and instead pass capabilities (such as memory access rights) for larger transfers. This design also encourages a more message-oriented architecture, in which tasks exchange small commands or data rather than sharing large memory buffers.

Why sequence numbers?

Each message gets a seq field to help detect lost messages (if you receive seq 5 then seq 7, you know seq 6 went missing), duplicates (seq 5 received twice), and ordering issues. It’s a simple reliability mechanism that real IPC systems use too.

No heap allocations.

We’re working in a bare-metal environment with no heap allocator, so we can’t use Vec<u8> or Box<Message>. Instead, we use fixed-size arrays for the payload.

3.3 Payload helpers

We want to send structured data in the payload, but since it’s just a byte array, we need to serialize and deserialize it manually. For our Ping/Pong demo, we’ll just send a u32 sequence number in the payload. We use little-endian byte order for simplicity and consistency with ARM’s native endianness.

pub fn write_u32_le(dst: &mut [u8], v: u32) {
    dst[0] = (v & 0xFF) as u8;
    dst[1] = ((v >> 8) & 0xFF) as u8;
    dst[2] = ((v >> 16) & 0xFF) as u8;
    dst[3] = ((v >> 24) & 0xFF) as u8;
}

pub fn read_u32_le(src: &[u8]) -> u32 {
    (src[0] as u32)
        | ((src[1] as u32) << 8)
        | ((src[2] as u32) << 16)
        | ((src[3] as u32) << 24)
}

Listing 5: Payload serialization helpers

We manually shift bytes to serialize a u32 into the first four bytes of the payload. We can’t use u32::to_le_bytes() in a const context on all targets yet, and this works fine.

In case you are wondering what endianess is? It’s the order in which bytes are stored for multi-byte values. Little-endian means the least significant byte comes first. For example, the number 0x12345678 would be stored as 78 56 34 12 in little-endian. This is important to get right when serializing data into a byte array.

4. The router: mailbox-based IPC

Now that we have a message format, we need a way to deliver messages between tasks. In a microkernel, the kernel itself is responsible for routing messages. Each task has an endpoint, and the kernel maintains a mailbox for each endpoint. When a task sends a message, the kernel looks at the destination endpoint and drops the message into that endpoint’s mailbox. When a task wants to receive, the kernel checks if there’s anything in that task’s mailbox. This design decouples senders and receivers, allowing them to operate at their own pace. If the receiver is slow, the sender can still send messages (up to a point), and if the sender is fast, the receiver can process messages as they come in.

Here’s the flow:

sequenceDiagram
    participant PT as PingTask
    participant R as Router
    participant M as Pong Mailbox
    participant PO as PongTask

    PT->>R: send(msg to Pong)
    R->>M: Check if full
    alt Mailbox empty
        R->>M: Store message
        M-->>R: OK
        R-->>PT: Success
    else Mailbox full
        R-->>PT: MailboxFull error
    end

    Note over PO: Later, in poll()
    PO->>R: recv(Pong endpoint)
    R->>M: Check if message present
    alt Message available
        M->>R: Return message
        M->>M: Mark empty
        R->>PO: Some(message)
        PO->>PO: Process message
    else No message
        R->>PO: None
    end

Figure 1: IPC message flow between tasks

If the mailbox already has an unread message, send() fails. The sender has to try again later. This is called backpressure, and it’s the simplest possible version: zero buffering beyond one message. Without backpressure, a fast sender could flood a slow receiver, consuming unbounded memory. Our single-slot design prevents that by construction.

4.1 Implementation

Let’s look at the actual code for the Mailbox and Router. The Router owns one mailbox per endpoint and handles all the logic for sending and receiving messages. The code is straightforward and intentionally minimal, illustrating the core concepts without extra complexity.

#[derive(Copy, Clone, Debug)]
pub enum SendError {
    MailboxFull,
}

#[derive(Copy, Clone)]
struct Mailbox {
    full: bool,
    msg: Message,
}

impl Mailbox {
    const EMPTY: Message = Message {
        header: MsgHeader {
            src: EndpointId::Ping,
            dst: EndpointId::Ping,
            ty: MsgType::Ping,
            len: 0,
            seq: 0,
        },
        payload: [0; MAX_PAYLOAD],
    };

    const fn new() -> Self {
        Self { full: false, msg: Self::EMPTY }
    }

    fn put(&mut self, msg: Message) -> Result<(), SendError> {
        if self.full { return Err(SendError::MailboxFull); }
        self.msg = msg;
        self.full = true;
        Ok(())
    }

    fn take(&mut self) -> Option<Message> {
        if !self.full { return None; }
        self.full = false;
        Some(self.msg)
    }
}

pub struct Router {
    ping: Mailbox,
    pong: Mailbox,
}

impl Router {
    pub const fn new() -> Self {
        Self { ping: Mailbox::new(), pong: Mailbox::new() }
    }

    pub fn send(&mut self, msg: Message) -> Result<(), SendError> {
        match msg.header.dst {
            EndpointId::Ping => self.ping.put(msg),
            EndpointId::Pong => self.pong.put(msg),
        }
    }

    pub fn recv(&mut self, dst: EndpointId) -> Option<Message> {
        match dst {
            EndpointId::Ping => self.ping.take(),
            EndpointId::Pong => self.pong.take(),
        }
    }
}

Listing 6: Mailbox and Router implementation

The Mailbox struct has a full flag indicating whether it currently holds a message and a msg field to store the message itself. The put() method checks if the mailbox is already full; if it is, it returns an error. If not, it stores the message and marks the mailbox as full. The take() method checks if there’s a message to read; if there is, it returns the message and marks the mailbox as empty. If not, it returns None.

In the Router, we have one mailbox for each endpoint. The send() method examines the destination endpoint in the message header and attempts to place the message in the corresponding mailbox. The recv() method checks the specified endpoint’s mailbox for a message and returns it if available.

This design is simple and efficient for our demo. In a real microkernel, you might have more complex routing logic, support for multiple mailboxes per endpoint, or even a more sophisticated synchronization mechanism.ightly.

5. Static router placement

From the code we’ve seen so far, it’s clear that the Router needs to be accessible globally, since all tasks need to send and receive messages through it. How do we achieve that in Rust, especially in a no_std environment where we don’t have the luxury of a heap or dynamic initialization? That’s where things get interesting.

The problem is the Sync trait. Any static value must be Sync, meaning it’s safe to access from multiple threads simultaneously. Our Router has mutable state (those mailboxes change when messages are sent and received), so it’s not Sync by default. Rust is trying to protect us from data races, which is normally a good thing. But in our single-threaded bare-metal kernel, there are no other threads. We need to tell the compiler, “dude, trust us, this is fine.”

The solution involves three layers: UnsafeCell, a manual Sync implementation, and a linker section attribute to ensure that the router is placed in writable memory. Let’s break it down.

use core::cell::UnsafeCell;

#[repr(transparent)]
struct RouterCell(UnsafeCell<ipc::Router>);
unsafe impl Sync for RouterCell {}

#[link_section = ".data"]
static ROUTER: RouterCell = RouterCell(UnsafeCell::new(ipc::Router::new()));

pub fn kmain(logger: &dyn Logger) -> ! {
    let router: &mut ipc::Router = unsafe { &mut *ROUTER.0.get() };
    // ... use router for the rest of the kernel's life
}

Listing 7: Static router with writable section placement

Here’s how the three layers wrap the Router to make it safely accessible as a global static:

graph TD
    subgraph ".data section (Layer 3)"
        subgraph "RouterCell + unsafe impl Sync (Layer 2)"
            subgraph "UnsafeCell (Layer 1)"
                R["ipc::Router<br/>(mutable mailbox state)"]
            end
        end
    end

    L3["\`#[link_section = '.data']\`<br/>Forces writable memory placement"]
    L2["unsafe impl Sync<br/>Satisfies static requirement"]
    L1["UnsafeCell&lt;T&gt;<br/>Enables interior mutability"]

    L3 -.-> R
    L2 -.-> R
    L1 -.-> R

    style R fill:#f99,stroke:#333
    style L1 fill:#ff9,stroke:#333
    style L2 fill:#9f9,stroke:#333
    style L3 fill:#9ff,stroke:#333

Figure 2: Three layers wrapping the Router

And here’s the runtime dereference chain - how we actually get a usable &mut Router from the static:

flowchart LR
    A["ROUTER<br/>(static RouterCell)"] -->|".0"| B["UnsafeCell&lt;Router&gt;"]
    B -->|".get()"| C["*mut Router<br/>(raw pointer)"]
    C -->|"* (deref)"| D["Router<br/>(value)"]
    D -->|"&mut"| E["&mut Router<br/>(usable reference)"]

    style A fill:#9ff,stroke:#333
    style C fill:#ff9,stroke:#333
    style E fill:#9f9,stroke:#333

Figure 3: Runtime dereference chain

Layer 1: UnsafeCell

This is the key to interior mutability in Rust. UnsafeCell<T> is a special wrapper that tells the compiler, “I know this data will be mutated through shared references, but I promise to handle it safely.” Normally, Rust enforces that if you have a &T, you can’t mutate T. But UnsafeCell provides a .get() method that returns a raw mutable pointer (*mut T), allowing us to bypass the borrow checker. This is essential for our router because we need to mutate its state (the mailboxes) while it’s accessible globally.

Layer 2: Manual Sync implementation

One might wonder why we need to implement Sync manually. It is because UnsafeCell<T> deliberately doesn’t implement Sync to prevent data races. Since static variables require Sync, we wrap our Router in a RouterCell that contains an UnsafeCell. By writing unsafe impl Sync for RouterCell {}, we’re telling the compiler that we guarantee this will only be accessed from one context at a time. This is safe in our case because our kernel is single-threaded and our interrupt handlers don’t touch the router. In a more complex kernel with preemption or multiple cores, this would be a dangerous promise, but for our simple cooperative scheduler, it’s perfectly fine.

Layer 3: #[link_section = ".data"]

OK, this one bit is interesting. When the compiler produces a binary, it organizes data into sections. .rodata (read-only data) holds constants. .data holds initialized mutable data. On ARM, the MMU enforces these permissions: .rodata pages are mapped read-only, so writing to them triggers a data abort (the CPU equivalent of a segfault). The Rust compiler sees that ROUTER is a static initialized with a const fn and sometimes decides it belongs in .rodata. But we need to mutate it at runtime. #[link_section = ".data"] overrides the compiler’s choice and forces placement in the writable section. If we don’t do this, the kernel crashes on the first send() call. Now that seem a lot like my usual code. 😁

The unsafe dereference

In case you missed it, the expression unsafe { &mut *ROUTER.0.get() } is a bit of Rust wizardry. Let’s break it down:

ROUTER is our static RouterCell.
ROUTER.0 accesses the UnsafeCell<ipc::Router> inside RouterCell.
.get() returns a *mut ipc::Router (raw mutable pointer).
The * dereferences that raw pointer, giving us an ipc::Router.
&mut borrows it as a mutable reference.
The whole thing is wrapped in unsafe { ... } because the compiler can’t verify that no other code holds a reference to the same Router. We know it’s safe because we only call this once, at the start of kmain, and our kernel is single-threaded. This is a common pattern for global mutable state in no_std Rust, but it requires careful reasoning to ensure safety.
The end result is that we get a &mut ipc::Router that we can use throughout the kernel to send and receive messages.

6. Task abstraction

Tasks are the fundamental units of work in our kernel. Each task has its own logic and state, and the scheduler manages which task runs at any given time. To make this work, we define a Task trait that all tasks must implement. This trait defines the contract for how tasks interact with the scheduler and the IPC system. Each task must provide an id() method to identify its endpoint and a poll() method that the scheduler calls every tick. The poll() method is where the task does its work: checking for messages, sending messages, updating state, etc. The cooperative multitasking model means that tasks must return from poll() quickly to allow other tasks to run. If a task needs to wait for something (like a message), it should return immediately and check again on the next tick.

6.1 What’s a task?

A task is an independent unit of work that the scheduler manages. Think of it as a lightweight thread - it has its own state and logic, but shares the CPU with other tasks. The scheduler decides which task runs at any given moment, switching between them to create the illusion that they all run simultaneously.

If you’ve used async/await in Rust (with tokio or async-std), the mental model is similar. Each async future does some work, then yields control back to the executor. Our tasks do the same thing. The scheduler calls each task’s poll() method; the task performs a small chunk of work (checking the mailbox, sending a message, logging something, etc.), then returns so the next task gets a turn. The one difference in our simple example is that our scheduler is the OS itself, not a userspace library.

6.2 The Task trait

pub trait Task {
    fn id(&self) -> EndpointId;
    fn poll(&mut self, logger: &dyn Logger, ipc: &mut ipc::Router, tick: u64);
}

Listing 8: Task trait definition

Two methods define the contract for a task:

id() - this returns the task’s endpoint identifier, so the scheduler (or the task itself) knows which mailbox to check.
poll() - this is where the work happens. It is called every scheduler iteration and must return quickly.

The logger parameter is &dyn Logger, a trait object. This means it’s a reference that can point to any type implementing the Logger trait, with the specific type resolved at runtime (dynamic dispatch). Different platforms have different logger implementations (COM1 serial, PL011 UART, mini-UART), but our kernel code works with all of them without knowing which concrete type it’s talking to. Same pattern as Box<dyn Error> in standard Rust, but we use a plain reference because we don’t have a heap.

The cooperative contract is simple - tasks must return from poll() quickly. No infinite loops, no blocking waits. If a task needs to wait for something (like a reply message), it returns immediately and checks again on the next tick.

7. PingTask and PongTask

Let’s see how this all comes together in practice. We have our IPC system, our scheduler, and the Task trait; now let’s see how real tasks use this system. We’ll build two tasks that play a simple game: PingTask sends a ping message every few ticks, PongTask receives it and replies with a pong. It’s the “hello world” of IPC. 😊

7.1 PingTask

The PingTask is a tiny state machine (with two states). It has a sequence number that increments with each ping, and a boolean flag to track whether it’s currently waiting for a pong reply. The logic is straightforward: if it’s not waiting, it sends a ping every 10 ticks. If it is waiting, it checks for a pong reply each tick. When it gets the pong, it flips back to idle and increments the sequence number.

pub struct PingTask {
    seq: u32,
    waiting: bool,
}

impl PingTask {
    pub const fn new() -> Self {
        Self { seq: 1, waiting: false }
    }
}

impl Task for PingTask {
    fn id(&self) -> EndpointId { EndpointId::Ping }

    fn poll(&mut self, logger: &dyn Logger, ipc: &mut ipc::Router, tick: u64) {
        if tick == 0 { logger.log("task/ping: poll\n"); }

        // Check for replies first
        if let Some(msg) = ipc.recv(self.id()) {
            if matches!(msg.header.ty, MsgType::Pong) {
                self.waiting = false;
                logger.log("task/ping: got pong\n");
            }
        }

        // Send a ping every 10 ticks, but only if we're not
        // already waiting for a reply
        if !self.waiting && (tick % 10) == 0 {
            let mut payload = [0u8; ipc::MAX_PAYLOAD];
            ipc::write_u32_le(&mut payload[0..4], self.seq);
            let msg = ipc::Message {
                header: ipc::MsgHeader {
                    src: EndpointId::Ping,
                    dst: EndpointId::Pong,
                    ty: MsgType::Ping,
                    len: 4,
                    seq: self.seq,
                },
                payload,
            };
            match ipc.send(msg) {
                Ok(()) => {
                    logger.log("task/ping: sent ping\n");
                    self.waiting = true;
                    self.seq = self.seq.wrapping_add(1);
                }
                Err(_) => {
                    logger.log("task/ping: send failed (queue full)\n");
                }
            }
        }
    }
}

Listing 9: PingTask implementation

When PingTask is idle (waiting == false), it checks every 10 ticks whether it’s time to send a ping. When it sends one, it changes to the waiting state. On the other hand, while waiting, it checks its mailbox each tick for a pong reply. When the pong arrives, it flips back to idle. The sequence number increments with each send (using wrapping_add so it rolls over instead of panicking at u32::MAX).

Notice how the task constructs the message manually: it fills a payload buffer, builds a header with source, destination, type, length, and sequence number, then hands the whole thing to the router. If the send fails because the destination mailbox is full, it logs the error and will try again on the next qualifying tick.

7.2 PongTask

The PongTask is even simpler. It just waits for pings and replies with pongs. It doesn’t need to track any state, so it’s a unit struct. Every tick, it checks its mailbox. If there’s a ping, it reads the sequence number from the payload, constructs a pong reply with the same sequence number, and sends it back to the Ping endpoint. It ignores the tick counter entirely since it doesn’t need to do anything on a timer.

pub struct PongTask;

impl PongTask {
    pub const fn new() -> Self { Self }
}

impl Task for PongTask {
    fn id(&self) -> EndpointId { EndpointId::Pong }

    fn poll(&mut self, logger: &dyn Logger, ipc: &mut ipc::Router, _tick: u64) {
        if let Some(msg) = ipc.recv(self.id()) {
            if matches!(msg.header.ty, MsgType::Ping) {
                let seq = ipc::read_u32_le(&msg.payload[0..4]);
                logger.log("task/pong: got ping\n");

                let mut payload = [0u8; ipc::MAX_PAYLOAD];
                ipc::write_u32_le(&mut payload[0..4], seq);
                let reply = ipc::Message {
                    header: ipc::MsgHeader {
                        src: EndpointId::Pong,
                        dst: EndpointId::Ping,
                        ty: MsgType::Pong,
                        len: 4,
                        seq,
                    },
                    payload,
                };
                let _ = ipc.send(reply);
            }
        }
    }
}

Listing 10: PongTask implementation

The reason PongTask is a unit struct (i.e., it has no fields) because it doesn’t need to track any state. Every tick, it checks its mailbox. If there’s a ping, it reads the sequence number from the payload, constructs a pong reply with the same sequence number, and sends it back to the Ping endpoint. Notice it ignores the tick counter entirely (the _tick prefix tells Rust we know it’s unused).

One subtle thing - PongTask uses let _ = ipc.send(reply) instead of a match on the result. It deliberately discards any send errors. For a simple echo-reply task, if the Ping mailbox is full, dropping the reply is acceptable. A more robust implementation might retry, but for our demo, this keeps things clean.

8. The scheduler

The scheduler is the heart of our kernel. It’s responsible for giving each task a turn to run and for managing the flow of time (ticks). In a real microkernel, the scheduler would be more complex, supporting priorities, preemption, and multiple cores. But for our demo, we keep it simple: it’s a cooperative round-robin scheduler that just iterates over a fixed list of tasks and calls poll() on each one every tick.

The scheduler itself is just a function that takes a list of tasks, a logger, and the IPC router. It runs an infinite loop where it calls poll() on each task, increments a tick counter, and halts the CPU until the next interrupt. The cooperative nature means that if any task takes too long in poll(), it can starve the others, but that’s a trade-off we’re making for simplicity.

pub fn run(
    tasks: &mut [&mut dyn Task],
    logger: &dyn Logger,
    ipc: &mut ipc::Router,
) -> ! {
    let mut tick: u64 = 0;
    logger.log("sched: starting\n");
    loop {
        for t in tasks.iter_mut() {
            t.poll(logger, ipc, tick);
        }
        tick = tick.wrapping_add(1);
        hal::arch::halt();
    }
}

Listing 11: Round-robin cooperative scheduler

That’s the whole thing. An infinite loop that iterates over every task, calls poll() on each one, bumps a tick counter, and halts the CPU until the next interrupt. The -> ! return type is the “never” type: this function never returns. On bare metal, there’s no OS to return to. The kernel runs forever.

A few things worth noting:

The tick counter uses wrapping_add(1) instead of += 1, which means when it hits u64::MAX it wraps back to zero instead of panicking.
At one tick per millisecond, a u64 would take about 584 million years to overflow, so this is really just defensive programming, but it’s a good habit in kernel code where panicking means the system dies.
The hal::arch::halt() call at the end of each iteration is important. On x86_64, it executes the HLT instruction, and on AArch64, it executes WFI (wait for interrupt). Both put the CPU into a low-power sleep state until the next hardware interrupt arrives. Without this, our loop would spin at full speed, burning power for no reason.

8.1 Wiring it all together

So we have our IPC system, our tasks, and our scheduler. Now we need to put it all together in the kernel’s entry point. This is where we initialize the router, create our tasks, and hand everything to the scheduler. After this point, the scheduler takes over, and we never return. kmain is the main function of our kernel, and it’s where we set up the system’s initial state.

pub fn kmain(logger: &dyn Logger) -> ! {
    logger.log("rustOS: kernel online\n");
    logger.log("rustOS: microkernel step 1 (IPC + cooperative scheduling)\n");

    let router: &mut ipc::Router = unsafe { &mut *ROUTER.0.get() };

    let mut ping = sched::PingTask::new();
    let mut pong = sched::PongTask::new();
    let mut tasks: [&mut dyn sched::Task; 2] = [&mut ping, &mut pong];

    sched::run(&mut tasks, logger, router)
}

Listing 12: Kernel entry point with scheduler invocation

We start by logging some messages to indicate the kernel is online and what we’re doing. Then we grab a mutable reference to the global router using the unsafe pattern we discussed earlier. We create our two tasks on the stack, put them into a fixed-size array of trait objects, and hand everything to the scheduler. The scheduler runs forever, polling both tasks on every iteration. Ping sends messages to Pong, Pong replies, and the cycle continues indefinitely.

9. Running the demo

How do we know it works? Let’s run it and see the logs. If everything is set up correctly, you should see a steady stream of ping-pong messages in the output, demonstrating that the tasks are communicating through the IPC system and that the scheduler is giving them turns to run.

Build and run the kernel as usual and watch the logs:

./scripts/build-aarch64-virt.sh && ./scripts/run-aarch64-virt.sh

Listing 13: Build and run the IPC demo

rustOS: aarch64 QEMU virt boot OK
rustOS: IPC + cooperative scheduling demo
rustOS: kernel online
rustOS: microkernel step 1 (IPC + cooperative scheduling)
sched: starting
task/ping: poll
task/ping: sent ping
task/pong: got ping

Figure 4: IPC demo output showing tasks exchanging messages through the cooperative scheduler.

You should see the ping and pong messages alternating. The first tick (tick 0) triggers both the initial poll log and the first ping send (since 0 % 10 == 0). PongTask picks up the message and replies. On subsequent qualifying ticks, the pattern repeats.

10. Limitations of cooperative scheduling

OK, so we’ve got tasks talking to each other. That’s great. But there’s a fundamental problem with our scheduler. What happens if a task doesn’t cooperate? Imagine we have a BadTask that just spins forever in its poll() method, never returning control to the scheduler:

impl Task for BadTask {
    fn id(&self) -> EndpointId { EndpointId::Ping }
    fn poll(&mut self, _logger: &dyn Logger, _ipc: &mut ipc::Router, _tick: u64) {
        loop {
            // spin forever, the scheduler never regains control
        }
    }
}

Listing 14: A misbehaving task

In this scenario, the BadTask takes over the CPU and never yields. The scheduler is stuck waiting for poll() to return, but it never does. As a result, all other tasks are frozen. The system is effectively deadlocked because the scheduler has no way to interrupt BadTask and give someone else a turn. In other words, the entire system hangs.

This is the core weakness of cooperative scheduling - it relies on every task being well-behaved. If one task misbehaves, the whole system suffers. In a real OS running untrusted code, this is unacceptable. We need a way for the OS to forcibly take control back from a running task, regardless of what that task is doing.

The way to solve this is with preemptive scheduling. Instead of waiting for tasks to yield, the OS can use a hardware timer to interrupt the currently running task at regular intervals. When the timer fires, it triggers an interrupt, and the interrupt handler can switch to a different task. This way, even if a task misbehaves and never yields, the OS can still regain control and keep the system responsive.

The solution is preemptive scheduling: a hardware timer fires at regular intervals, triggers an interrupt, and the interrupt handler switches to a different task. The currently running task doesn’t get a choice. In Part 3 , we’ll set up timer interrupts on AArch64 and build a preemptive context switcher that saves and restores full CPU state. The cooperative scheduler we built here won’t go away (it’s still useful for understanding the basics), but we’ll layer preemption on top of it.

11. What we built

In this part, we took our boot-only kernel and turned it into a simple microkernel with IPC and cooperative multitasking. Starting from a kernel that could only boot and halt, we added three things. A message-passing IPC system with typed messages, endpoint-based routing, and single-slot mailboxes with backpressure. A trait-based task abstraction that lets us write independent units of work with a clean poll() interface. And a cooperative round-robin scheduler that gives each task a fair turn.

The Ping/Pong demo proves it works: two tasks communicate entirely through messages, with no shared mutable state, no unsafe data sharing, and explicit control flow that you can trace through the logs. This is the essence of microkernel design: tasks are isolated, communicate through well-defined channels, and the kernel provides minimal mechanisms to support that communication and scheduling.

However, cooperative scheduling has a fatal flaw. It trusts tasks to yield. If a task misbehaves and never yields, the whole system hangs. In Part 3 , we’ll fix that with timer interrupts and preemptive multitasking, allowing the OS to regain control even if a task goes rogue.

14. References

Improving IPC by Kernel Design (Liedtke, 1993) - The paper that made fast IPC possible in L4
The seL4 Microkernel - Formally verified microkernel design
The Rust Programming Language - Chapter 10 (Traits), Chapter 17 (Trait Objects)
Operating Systems: Design and Implementation (Tanenbaum) - The classic Minix book and microkernel design rationale

5-Part Series:

Part 0: Why build an OS from scratch?
Part 1: Foundations
Part 2 (this): Communication
Part 3: Concurrency
Part 4: Memory and beyond

Building a microkernel in Rust (Part 1): Foundations, booting on bare metal

Wed, 25 Feb 2026 00:00:00 +0000

5-Part Series:

Part 0: Why build an OS from scratch?
Part 1 (this): Foundations
Part 2: Communication
Part 3: Concurrency
Part 4: Memory and beyond

GitHub Repository: bahree/rust-microkernel — full source code and build scripts

Docker Image: amitbahree/rust-microkernel — prebuilt dev environment with Rust, QEMU, and source code

You’re about to write code that runs with nothing underneath it.

No operating system. No standard library. No println!. No heap. No file descriptors. No threads. Just you, a CPU, and some RAM.

That probably sounds terrifying. Or maybe thrilling. Honestly, it’s both. The first time you get a bare-metal kernel to print a single character to a serial port, you’ll stare at your terminal for a solid minute. One letter. It takes days of work. And it feels like you’ve built a cathedral.

This post walks through everything that happens between “the CPU wakes up” and “a message appears on screen.” We’ll go instruction by instruction through the assembly boot code, byte by byte through the UART driver, and concept by concept through the design decisions that make it all work. By the end, you’ll have a bootable AArch64 kernel running in QEMU that prints to a virtual serial port, built on a clean, platform-agnostic kernel architecture.

Let’s get to it.

TL;DR

In this part, we boot a minimal Rust kernel on the AArch64 QEMU virt machine. We write ARM assembly to set up a stack, zero the BSS, drop from EL2 (hypervisor level) to EL1 (kernel level), enable floating-point registers, and jump into Rust. The Rust side implements a PL011 UART driver using memory-mapped I/O, wraps it behind a Logger trait for platform abstraction, and calls into a shared kernel entry point. The whole thing compiles to a single ELF binary that QEMU loads directly.

1. What “bare metal” actually means

When you write a normal Rust program (or any program using a higher level language for that matter), there’s a tower of software beneath you. Think about what happens when you call println!("hello"). First, Rust’s standard library formats the string and then calls the OS to write it to stdout. Next, the OS looks up your process’s file descriptor table, finds that stdout points to a terminal, copies your bytes into a kernel buffer, and eventually, a terminal emulator reads those bytes and draws glyphs on screen. And to make all of this work, sitting below all of that, is the firmware which already tested the RAM, configured the CPU caches, initialized the PCI bus, and handed control to the bootloader.

Guess what? We’re removing all of that. 🤖

Our code is the first thing that runs after the CPU wakes up. There’s no runtime to set up the stack. No OS to provide virtual memory. No firmware to initialize the UART (well, QEMU helps us a little here, but we will ignore that for now). We are the bottom of the stack. Everything above us, we have to build.

In Rust terms, this means two attributes at the top of every bare-metal crate:

#![no_std]
#![no_main]

Bare-metal crate attributes

#![no_std] tells the compiler: don’t link the standard library. The standard library (std) gives you Vec, String, Box, println!, file I/O, threads, and networking. All of those features need an underlying operating system. Vec needs a heap allocator, which needs mmap or brk from the OS. println! needs stdout, which needs file descriptors, which need a kernel. Since we ARE the kernel, there’s nothing beneath us to provide those services.

With no_std, we only get core: basic types, traits, iterators, Option, Result, math operations. Things that need zero OS support.

#![no_main] is subtler. Normally, when you compile a Rust binary, the compiler generates a hidden entry point that sets up the stack, initializes the heap allocator, configures panic handling, spawns the main thread, and then calls your main() function. None of that machinery exists on bare metal. We handle all initialization ourselves, in assembly, and we tell the compiler not to look for a standard main entry point.

2. What’s an ELF?

Before we dive into boot code, a quick concept: when you compile a Rust program (or C, or anything that targets a Unix-like system), the output isn’t just a flat blob of machine code. It’s a structured file called an ELF (Executable and Linkable Format).

An ELF file is like a shipping manifest. It says: “here’s a chunk of executable code, load it at this address. Here’s read-only data, put it over here. Here’s uninitialized data (BSS), allocate this much space and zero it. Oh, and the program starts executing at this address.”

When we build our kernel, cargo and the linker produce an ELF binary. QEMU knows how to read ELF files, so it parses the headers, loads each section to the right memory address, and sets the program counter to the entry point (_start). On real hardware without an ELF-aware loader, you’d strip the ELF headers and produce a raw binary. But QEMU makes our lives easier.

The important sections in our ELF:

.text.boot: The very first code that runs (our assembly entry point)
.text: The rest of our compiled Rust code
.rodata: Read-only data (string constants, lookup tables)
.data: Initialized mutable data (statics with non-zero initial values)
.bss: Uninitialized data (statics that start at zero, just a size, no actual bytes in the file)

3. The QEMU virt machine

We’re targeting qemu-system-aarch64 -machine virt, which is QEMU’s generic ARM virtual machine. It doesn’t model any specific real-world board. Instead, it gives us a clean, well-documented set of virtual hardware:

Component	Details
CPU	Cortex-A53 (emulated), 64-bit ARMv8-A
RAM	256 MB starting at `0x4000_0000`
UART	PL011 at `0x0900_0000`
Interrupt controller	GICv2 (Generic Interrupt Controller)
Timer	ARM Generic Timer (CNTP)

Here’s the QEMU command we use to run our kernel:

qemu-system-aarch64 \
  -machine virt,gic-version=2 \
  -cpu cortex-a53 \
  -m 256M \
  -nographic \
  -serial mon:stdio \
  -kernel dist/virt/os-aarch64-virt.elf

QEMU launch command for AArch64 virt

Let’s break down what each of these mean:

-machine virt,gic-version=2: Use the generic ARM virtual machine with a GICv2 interrupt controller. GICv2 is simpler than GICv3 and sufficient for single-core work.
-cpu cortex-a53: Emulate a Cortex-A53 core. This is the same CPU in the Raspberry Pi Zero 2 W.
-m 256M: Give the machine 256 MB of RAM.
-nographic: Don’t open a GUI window. We’re using serial output, not a display.
-serial mon:stdio: Connect the virtual serial port to our terminal’s stdin/stdout. This is how we see output.
-kernel dist/virt/os-aarch64-virt.elf: Load our ELF binary directly. QEMU parses the ELF, loads sections to the right addresses, and jumps to _start.

The beauty of virt is its simplicity. No GPU firmware to deal with, no board-specific quirks, no SD card flashing. Build, run, see output. The iteration cycle is seconds.

4. Setup

You need Rust nightly (for bare-metal features like no_std binaries) and the AArch64 bare-metal target. Our repository includes a rust-toolchain.toml that handles most of this:

[toolchain]
channel = "nightly"
components = ["llvm-tools-preview", "rust-src"]
targets = ["x86_64-unknown-none", "aarch64-unknown-none"]

rust-toolchain.toml

If you’re setting up from scratch:

# Install Rust nightly with AArch64 target
rustup default nightly
rustup target add aarch64-unknown-none
rustup component add llvm-tools-preview rust-src

# Install QEMU (Linux)
sudo apt install qemu-system-aarch64

# Install QEMU (macOS)
brew install qemu

# Clone the repository
git clone https://github.com/bahree/rust-microkernel.git
cd rust-microkernel

Install Rust toolchain and QEMU

5. Build and run

./scripts/build-aarch64-virt.sh demo-ipc
./scripts/run-aarch64-virt.sh

Build and run the AArch64 virt kernel

Expected output:

rustOS: aarch64 QEMU virt boot OK
rustOS: IPC + cooperative scheduling demo
rustOS: kernel online
rustOS: microkernel step 1 (IPC + cooperative scheduling)
sched: starting
task/ping: poll
task/ping: sent ping
task/pong: got ping
task/ping: got pong

Expected serial output

Figure 1: Boot output showing the kernel printing to serial, proving our boot assembly, UART driver, and platform abstraction all work.

If you see that, everything’s working. Press Ctrl-A then X to exit QEMU.

If you don’t pass a feature flag, the default build uses demo-memory (which we’ll cover in Part 4). The demo-ipc flag gives us the cooperative scheduling output shown above.

6. The boot sequence

Here’s what happens from the moment QEMU starts to the moment you see “boot OK” on your terminal:

flowchart TD
    A[QEMU loads ELF to 0x40080000] --> B[CPU starts at _start, running at EL2]
    B --> C[Set stack pointer]
    C --> D[Zero BSS section]
    D --> E{Current exception level?}
    E -->|EL2| F[Set EL1 stack pointer]
    F --> G[Configure SPSR_EL2 for EL1h return]
    G --> H[Set ELR_EL2 to el1_start]
    H --> I[Disable FP/ASIMD trapping at EL2]
    I --> J[eret: drop to EL1]
    J --> K[el1_start]
    E -->|EL1| K
    K --> L[Install exception vector table]
    L --> M[Enable FP/ASIMD at EL1]
    M --> N["bl rust_main"]
    N --> O[Initialize PL011 UART]
    O --> P["Print: boot OK"]
    P --> Q[Call kernel::kmain]

    style B fill:#f9f,stroke:#333
    style J fill:#ff9,stroke:#333
    style Q fill:#9f9,stroke:#333

AArch64 virt boot sequence

That’s a lot of steps. Let’s walk through these; however, first, you need to be able to read the assembly.

7. ARM assembly primer

If you’ve never read assembly before, the key is that we need to understand only a handful of concepts to follow the boot code. And as with most things in life, once you get the hang of it, assembly is surprisingly readable. The thing though is that everything is just very, very explicit about what it’s doing.

7.1 Registers

ARM gives you 31 general-purpose 64-bit registers, named x0 through x30. Think of them as 31 local variables that live inside the CPU itself, way faster than RAM. A register access takes maybe one clock cycle. A RAM access? Hundreds.

Some registers have conventional roles:

x0 through x7: Function arguments and return values
x29: Frame pointer (like rbp on x86)
x30: Link register, holds the return address after a bl (branch-with-link) call
sp: Stack pointer, tracks the top of the call stack
xzr: The zero register. Always reads as zero, discards writes. Surprisingly useful.

You’ll also see w0 through w30. These are the lower 32 bits of the corresponding x register. w0 is the bottom half of x0. ARM uses these when working with 32-bit values.

7.2 System registers

Beyond the general-purpose registers, ARM has a separate world of system registers that control CPU behavior. You can’t use them in normal instructions like add or sub. You access them with special instructions:

mrs x0, CurrentEL: Move from system register to general-purpose register. Read the current exception level into x0.
msr spsr_el2, x0: Move from general-purpose register to system register. Write x0 into the saved program status register for EL2.

These are the knobs and dials that configure how the CPU works - what exception level we’re at, whether floating-point is enabled, where the exception vector table lives, and what happens on an eret.

7.3 Common instructions

Here’s what you’ll see in the boot code:

Instruction	What it does
`ldr x0, =label`	Load the address of `label` into `x0`
`mov sp, x0`	Copy `x0` into the stack pointer
`str xzr, [x1], #8`	Store zero to the address in `x1`, then add 8 to `x1`
`cmp x1, x2`	Compare `x1` and `x2` (sets condition flags)
`b.ge 2f`	Branch forward to label `2:` if the comparison was greater-than-or-equal
`b.ne label`	Branch to `label` if not equal
`bl rust_main`	Branch with link: save return address in `x30`, jump to `rust_main`
`bic x0, x0, #(1 << 10)`	Bit clear: clear bit 10 in `x0`
`orr x0, x0, #(3 << 20)`	Bitwise OR: set bits 20 and 21 in `x0`
`lsr x1, x1, #2`	Logical shift right by 2 bits
`and x1, x1, #3`	Bitwise AND with 3 (keep only bits 0 and 1)
`adr x0, label`	Load the PC-relative address of `label` into `x0`
`isb`	Instruction synchronization barrier (flush pipeline)
`eret`	Exception return (drop to lower exception level)
`wfe`	Wait for event (low-power sleep)

7.4 Local labels and directives

Assembly uses numbered local labels like 1:, 2:, 3:. You reference them with a direction suffix: 1b means “search backward for label 1:” and 2f means “search forward for label 2:.” This avoids name collisions when the same pattern (like a loop) appears multiple times in the same file.

.section .text.boot tells the assembler: “put the following code into a section named .text.boot.” The linker script (which we’ll cover later) uses this name to ensure our boot code lands at the very start of the binary, exactly where the CPU expects it.

.global _start makes the _start symbol visible to the linker so that it can be used as the entry point.

8. Walking through boot.S

This is the heart of the post. Let’s go through the actual boot assembly line by line.

Here’s the file: crates/arch_aarch64_virt/src/boot.S. We’re showing the boot sequence portion (the first 60 lines) - you can see the full code from the repository.

.section .text.boot
.global _start
_start:
  // Set stack
  ldr x0, =__stack_top
  mov sp, x0

  // Zero BSS: [__bss_start, __bss_end)
  ldr x1, =__bss_start
  ldr x2, =__bss_end
1:
  cmp x1, x2
  b.ge 2f
  str xzr, [x1], #8
  b 1b
2:
  // If we entered at EL2 (typical for QEMU virt), drop to EL1 so the kernel runs
  // in a simpler environment (EL1 + GICv2 + CNTP timer).
  mrs x1, CurrentEL
  lsr x1, x1, #2
  and x1, x1, #3
  cmp x1, #2
  b.ne el1_start

  // Set up an EL1 stack pointer.
  ldr x0, =__stack_top
  msr sp_el1, x0

  // Configure EL2 to return to EL1h.
  mov x0, #(0b0101)         // EL1h
  msr spsr_el2, x0
  adr x0, el1_start
  msr elr_el2, x0

  // Enable FP/ASIMD access at EL1 and ensure EL2 doesn't trap it.
  mrs x0, cptr_el2
  bic x0, x0, #(1 << 10)    // TFP = 0 (don't trap FP/ASIMD)
  msr cptr_el2, x0
  isb
  eret

el1_start:
  // Install a minimal exception vector table for the *current* exception level.
  // QEMU `virt` typically enters at EL2, so we must set VBAR_EL2 (not just VBAR_EL1).
  adr x0, vectors
  msr vbar_el1, x0
  isb

  // Enable FP/ASIMD for Rust/LLVM.
  // Rust/LLVM may use NEON registers for struct copies/memcpy even in early bring-up.
  // If FP is disabled, this traps with EC=0x07 (FP/ASIMD access trap).
  mrs x0, cpacr_el1
  orr x0, x0, #(3 << 20)   // FPEN = 0b11
  msr cpacr_el1, x0
  isb

  bl rust_main
3:
  wfe
  b 3b

boot.S: the complete boot sequence (crates/arch_aarch64_virt/src/boot.S, lines 1-60)

Now let’s walk through each piece to help understand what it means.

8.1 Stack setup (lines 5-6)

  ldr x0, =__stack_top
  mov sp, x0

Setting up the stack pointer

The very first thing we do is set up a stack. But why? Because Rust literally cannot run without one. Every function call pushes a return address onto the stack. Every local variable lives on the stack. Without a valid stack pointer, the first bl instruction would try to save the return address to… nowhere. Instant crash.

__stack_top is a symbol defined by our linker script. It points to the top of a 64 KB stack memory region. ARM stacks grow downward (from high addresses to low), so we point sp at the top.

Why 64 KB? It’s a reasonable starting size. We don’t have many nested function calls yet, and we don’t have a heap, so the stack doesn’t need to be huge. If we run out, we’ll get a data abort (the ARM equivalent of a segfault). We’d increase it then.

8.2 BSS zeroing (lines 8-16)

  ldr x1, =__bss_start
  ldr x2, =__bss_end
1:
  cmp x1, x2
  b.ge 2f
  str xzr, [x1], #8
  b 1b
2:

Zeroing the BSS section

BSS stands for “Block Started by Symbol” (a historical name from the 1950s). It’s the section where uninitialized global and static variables live. In Rust, if you write static mut COUNTER: u32 = 0;, the compiler puts it in .bss. The key insight: the ELF file doesn’t actually store the zeros. It just records the size. That saves space in the binary. But it means the memory might contain garbage when we start.

The C and Rust languages both guarantee that uninitialized statics start at zero. Before any Rust code runs, we have to zero the entire BSS region manually. That’s what this loop does:

Load the start and end addresses of .bss (defined by the linker script) into x1 and x2.
Compare them. If x1 >= x2, the BSS is empty (or we’re done), skip to label 2:
Store zero (xzr, the zero register) to the address in x1, then increment x1 by 8 (one 64-bit word)
Jump back to the comparison

The [x1], #8 syntax is a post-increment addressing mode. It means “use the address in x1 for the store, THEN add 8 to x1.” It’s a single instruction that does two things. ARM loves these compound operations.

8.3 Checking the exception level (lines 19-23)

  mrs x1, CurrentEL
  lsr x1, x1, #2
  and x1, x1, #3
  cmp x1, #2
  b.ne el1_start

Detecting the current exception level

Now things get interesting. We need to know what exception level the CPU is running at. QEMU’s virt machine starts us at EL2 (hypervisor level), but we want to run our kernel at EL1 (the normal kernel level). If we’re already at EL1 for some reason, we skip the drop.

mrs x1, CurrentEL reads the CurrentEL system register. The exception level is encoded in bits [3:2] of this register (ARM’s register designs can be a little eccentric). So we shift right by 2 (lsr x1, x1, #2) and mask off everything except the bottom 2 bits (and x1, x1, #3). Now x1 holds 0, 1, 2, or 3 corresponding to EL0 through EL3.

If it’s 2, we proceed with the EL2-to-EL1 drop. If it’s not, we jump straight to el1_start.

8.4 Exception levels explained

Before we get into the drop, let’s talk about why exception levels exist.

ARM’s AArch64 architecture has four privilege levels:

Level	Name	Who runs here	What they can do
EL0	Application	User programs	Normal instructions only. Can’t touch hardware.
EL1	Kernel	Operating systems	Configure MMU, handle interrupts, access all memory
EL2	Hypervisor	Virtual machine monitors	Virtualize EL1 guests, trap privileged operations
EL3	Secure Monitor	TrustZone firmware	Switch between secure and non-secure worlds

Think about it this way: EL0 is a sandbox. Applications can compute, but they can’t mess with memory mappings, disable interrupts, or talk to hardware directly. If they try, the CPU traps to EL1, and the kernel decides what to do (usually: kill the process with a signal).

EL1 is where kernels live. Linux, macOS, Windows, and our rustOS all run here. You can configure page tables, handle interrupts, and access device memory.

EL2 is for hypervisors. If you want to run multiple operating systems on the same hardware (like AWS does with EC2 instances), the hypervisor at EL2 virtualizes the hardware so each guest OS at EL1 thinks it has the machine to itself.

EL3 is for firmware-level security (ARM TrustZone). We don’t touch it.

QEMU starts us at EL2 because it’s the most flexible starting point. If you were writing a hypervisor, you’d stay there. Since we’re writing a kernel, we drop to EL1.

8.5 The EL2 to EL1 drop (lines 26-40)

  // Set up an EL1 stack pointer.
  ldr x0, =__stack_top
  msr sp_el1, x0

  // Configure EL2 to return to EL1h.
  mov x0, #(0b0101)         // EL1h
  msr spsr_el2, x0
  adr x0, el1_start
  msr elr_el2, x0

  // Enable FP/ASIMD access at EL1 and ensure EL2 doesn't trap it.
  mrs x0, cptr_el2
  bic x0, x0, #(1 << 10)    // TFP = 0 (don't trap FP/ASIMD)
  msr cptr_el2, x0
  isb
  eret

Dropping from EL2 to EL1

This is the most intricate part of the boot sequence. Here’s the trick: you can’t just “jump” to a lower exception level. ARM doesn’t have a “go to EL1” instruction. Instead, you use the exception return mechanism backwards.

Normally, when an exception occurs (such as an interrupt), the CPU saves the current state and jumps to a higher exception level. The handler processes the exception, then executes eret (exception return) to return. We’re abusing this: we set up the return state registers to point at EL1, then execute eret as if we were “returning” from an exception that never happened.

Here’s each step:

msr sp_el1, x0 sets the stack pointer that EL1 will use. Each exception level has its own stack pointer. We’re at EL2 right now, so sp refers to sp_el2. We need to pre-configure sp_el1 before we get there.

mov x0, #(0b0101) / msr spsr_el2, x0 configures the Saved Program Status Register. When eret executes, the CPU restores the processor state from spsr_el2. The value 0b0101 (decimal 5) means: return to EL1 using the handler stack pointer variant (called “EL1h”). Bits [3:2] = 01 select EL1. Bit [0] = 1 selects the h variant, meaning EL1 will use sp_el1 as its stack pointer (rather than sp_el0).

adr x0, el1_start / msr elr_el2, x0 sets the Exception Link Register. This is the address the CPU will jump to on eret. We point it at el1_start, the label where our EL1 code begins. adr computes a PC-relative address, which is important because our code might not be running at the address the linker assumed (though, in practice, it is for us).

bic x0, x0, #(1 << 10) / msr cptr_el2, x0 disables the floating-point trap at EL2. By default, EL2 can trap FP/ASIMD (NEON) instructions executed at EL1. If we don’t clear this bit, the first time Rust tries to use a NEON register (which happens sooner than you’d think), the CPU will trap to EL2 and crash. bic stands for “bit clear.” It clears bit 10 (the TFP flag) in the cptr_el2 register.

isb is an Instruction Synchronization Barrier. It forces the CPU to finish processing all pending changes to system registers before executing the next instruction. Without it, the pipeline might still be using stale configuration values.

eret does the actual drop. It atomically loads the program counter from elr_el2 and the processor state from spsr_el2. In one instruction, we go from EL2 to EL1 and start executing at el1_start.

It’s a beautiful hack, honestly. The hardware designers intended eret for returning from exception handlers. We’re using it as a one-way door to a lower privilege level.

8.6 Vector table installation (lines 43-47)

  adr x0, vectors
  msr vbar_el1, x0
  isb

Installing the exception vector table

The CPU needs to know where to jump when an exception occurs (an interrupt, a system call, a data abort, etc.). ARM uses a vector table: a fixed-layout block of code that contains an entry for each exception type. There are 16 entries, every 128 bytes (0x80) apart, for a total of 2048 bytes (0x800).

We load the address of our vectors table and write it to vbar_el1 (Vector Base Address Register for EL1). When an exception happens at EL1, the CPU adds an offset to this base address and starts executing there. We’ll cover the vector table in detail in Part 3 when we implement interrupt handling.

8.7 FP/ASIMD enablement: why Rust needs NEON (lines 49-55)

  mrs x0, cpacr_el1
  orr x0, x0, #(3 << 20)   // FPEN = 0b11
  msr cpacr_el1, x0
  isb

Enabling floating-point and SIMD at EL1

You might wonder: we’re writing a kernel, not a graphics engine. Why do we need floating-point?

Here’s the thing. LLVM (Rust’s code generator) uses NEON registers for more than just floating-point math. It uses them for memory operations like memcpy and struct copies. If you have a struct that’s, say, 32 bytes, LLVM might decide the fastest way to copy it is to load it into a pair of 128-bit NEON registers and store it back. This happens even in integer-only code. Even in early boot.

If FP/ASIMD is disabled when this happens, the CPU generates a synchronous exception with exception class EC=0x07 (FP/ASIMD access trap). Your kernel crashes before it even prints “hello.”

We handled the EL2 side already (clearing TFP in cptr_el2). Now at EL1, we set FPEN (bits [21:20]) in cpacr_el1 to 0b11, which means “don’t trap FP/ASIMD at EL1 or EL0.” The orr instruction sets those bits. Another isb to synchronize.

8.8 Jumping to Rust (lines 57-60)

  bl rust_main
3:
  wfe
  b 3b

Calling into Rust

Finally, bl rust_main is a branch-with-link: it saves the return address in x30 (the link register) and jumps to rust_main. This is where we leave assembly and enter Rust.

The three lines after it are a safety net. rust_main is declared as -> ! (never returns), but if something goes horribly wrong and it does return, we don’t want the CPU executing random memory. So we enter an infinite loop: wfe (wait for event, which puts the CPU in a low-power sleep state) followed by a branch back to wfe.

9. The Rust entry point

Now we’re in Rust. Let’s look at the actual main.rs for the AArch64 virt platform.

#![no_std]
#![no_main]

use core::panic::PanicInfo;
use hal::log::Logger;

core::arch::global_asm!(include_str!("boot.S"));

// QEMU `virt` PL011 UART base.
const UART0_BASE: usize = 0x0900_0000;

struct UartLogger;

impl UartLogger {
    #[inline(always)]
    fn mmio_write(offset: usize, val: u32) {
        unsafe { core::ptr::write_volatile((UART0_BASE + offset) as *mut u32, val) }
    }

    #[inline(always)]
    fn mmio_read(offset: usize) -> u32 {
        unsafe { core::ptr::read_volatile((UART0_BASE + offset) as *const u32) }
    }

    fn putc(c: u8) {
        // FR (0x18) bit5 = TXFF (transmit FIFO full)
        while (Self::mmio_read(0x18) & (1 << 5)) != 0 {}
        Self::mmio_write(0x00, c as u32);
    }

    pub(crate) fn puts(s: &str) {
        for &b in s.as_bytes() {
            if b == b'\n' {
                Self::putc(b'\r');
            }
            Self::putc(b);
        }
    }
}

impl hal::log::Logger for UartLogger {
    fn log(&self, s: &str) {
        UartLogger::puts(s);
    }
}

mod timer;
mod preempt;
mod mem;

#[unsafe(no_mangle)]
pub extern "C" fn rust_main() -> ! {
    let logger = UartLogger;
    logger.log("rustOS: aarch64 QEMU virt boot OK\n");

    #[cfg(feature = "demo-ipc")]
    {
        logger.log("rustOS: IPC + cooperative scheduling demo\n");
        kernel::kmain(&logger)
    }

    #[cfg(feature = "demo-timer")]
    {
        logger.log("rustOS: timer interrupts demo\n");
        timer::init();
        logger.log("rustOS: timer started, entering idle loop\n");
        loop {
            hal::arch::halt();
        }
    }

    #[cfg(feature = "demo-preempt")]
    {
        logger.log("rustOS: preemptive multitasking demo\n");
        preempt::init();
        extern "C" {
            fn start_first(ctx: *const preempt::Context) -> !;
        }
        unsafe { start_first(preempt::first_context()) }
    }

    #[cfg(feature = "demo-memory")]
    {
        logger.log("rustOS: memory management demo (frames + page tables)\n");
        mem::demo();
        loop {
            hal::arch::halt();
        }
    }

    #[cfg(not(any(feature = "demo-ipc", feature = "demo-timer",
                   feature = "demo-preempt", feature = "demo-memory")))]
    {
        logger.log("rustOS: no demo selected, halting\n");
        loop {
            hal::arch::halt();
        }
    }
}

#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    UartLogger::puts("rustOS: PANIC\n");
    loop {
        hal::arch::halt();
    }
}

crates/arch_aarch64_virt/src/main.rs

Let’s unpack the important pieces.

9.1 `global_asm!` and including boot.S

core::arch::global_asm!(include_str!("boot.S"));

Including assembly in the Rust build

This line tells the Rust compiler: “take the contents of boot.S and include them as global assembly in this compilation unit.” The assembler processes our boot code, the linker resolves the symbols (_start, rust_main, __stack_top, etc.), and everything ends up in one binary. It’s the bridge between our assembly boot code and Rust.

9.2 `#[unsafe(no_mangle)]` and `extern "C"`

#[unsafe(no_mangle)]
pub extern "C" fn rust_main() -> ! {

The Rust entry point signature

Two things happening here.

#[unsafe(no_mangle)] prevents the Rust compiler from mangling the function name. Normally, Rust encodes type information into symbol names (so rust_main might become something like _ZN17arch_aarch64_virt9rust_main17h3a2b1c4d5e6f7g8hE). Our assembly code calls bl rust_main and needs to find that exact name.

extern "C" says: use the C calling convention. Rust’s own calling convention isn’t stable and can change between compiler versions. The C ABI is well-defined on every platform. On AArch64, the C calling convention puts the first argument in x0, the second in x1, return values in x0, and so on. Since assembly calls this function, we need a stable, predictable calling convention.

-> ! means this function never returns. On bare metal, there’s nothing to return to. The CPU would start executing whatever random bytes follow in memory. The function must either loop forever or halt the CPU.

9.3 The feature gates

#[cfg(feature = "demo-ipc")]
{
    logger.log("rustOS: IPC + cooperative scheduling demo\n");
    kernel::kmain(&logger)
}

Compile-time feature selection

The #[cfg(feature = "...")] attributes are compile-time switches. Depending on which feature you pass to the build script, a different code path gets compiled in. Only one demo runs at a time. This is a common pattern in embedded Rust: use cargo features to select between different configurations without runtime overhead.

9.4 The panic handler

#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    UartLogger::puts("rustOS: PANIC\n");
    loop {
        hal::arch::halt();
    }
}

Bare-metal panic handler

On bare metal, you must define what happens when Rust panics (array out of bounds, unwrap() on None, explicit panic!(), etc.). There’s no OS to catch it. Our handler prints “PANIC” to the UART and halts. In a more sophisticated kernel, you’d print the panic message and a stack trace. For now, knowing something panicked is enough.

10. PL011 UART: talking to hardware via MMIO

In a normal program, every memory address maps to a byte of RAM. Address 0x1000 holds some data, you read it, you write it, it’s just storage. On bare metal, some addresses are special: they’re connected to hardware peripherals instead of RAM. Writing to address 0x0900_0000 on our QEMU virt machine doesn’t store anything in memory. It sends a byte out the serial line. This is memory-mapped I/O (MMIO).

The PL011 UART sits at that address. When you write a byte to offset 0x00 from that base, the UART hardware transmits it. When you read from offset 0x18, you get the UART’s flag register, telling you whether the transmit FIFO is full. These addresses come from the hardware manufacturer’s specification (or in QEMU’s case, from QEMU’s source code and the virt machine documentation). There’s no way to discover them at runtime. You look them up.

10.1 The UartLogger implementation

Let’s look at how we actually write bytes to the UART:

const UART0_BASE: usize = 0x0900_0000;

impl UartLogger {
    #[inline(always)]
    fn mmio_write(offset: usize, val: u32) {
        unsafe { core::ptr::write_volatile((UART0_BASE + offset) as *mut u32, val) }
    }

    #[inline(always)]
    fn mmio_read(offset: usize) -> u32 {
        unsafe { core::ptr::read_volatile((UART0_BASE + offset) as *const u32) }
    }
}

MMIO helper functions (from main.rs)

write_volatile and read_volatile are critical. Normally, the Rust compiler (through LLVM) is free to optimize away memory operations. If you write to an address and never read it back, the optimizer might skip the write entirely. “Why bother,” it thinks, “nobody’s going to look at it.”

But for MMIO, the hardware IS looking at it. Writing to the UART data register transmits a byte. The compiler can’t see that side effect. volatile tells the compiler: “this access has effects you can’t reason about. Do it exactly as written, in exactly this order.”

Both functions are unsafe because we’re casting raw integers to pointers and dereferencing them. In safe Rust, pointer arithmetic is forbidden. On bare metal, it’s the only way to talk to hardware.

10.2 Sending a character

fn putc(c: u8) {
    // FR (0x18) bit5 = TXFF (transmit FIFO full)
    while (Self::mmio_read(0x18) & (1 << 5)) != 0 {}
    Self::mmio_write(0x00, c as u32);
}

PL011 character transmission (from main.rs)

The PL011 UART has a transmit FIFO (a small hardware buffer, typically 16 bytes). When you write a byte to the data register (offset 0x00), it goes into the FIFO. The UART hardware drains the FIFO at the configured baud rate, sending bits out the serial line.

But the FIFO can fill up. If you write faster than the UART can transmit, the FIFO overflows and bytes get lost. So before writing, we check the Flag Register (offset 0x18). Bit 5 is TXFF (transmit FIFO full). If it’s set, we spin-wait until there’s room.

This is called polling or busy-waiting. It’s the simplest approach and perfectly fine for debug output. In a production driver, you’d use interrupts instead (the UART can signal when the FIFO has room, so the CPU can do other work instead of spinning).

10.3 Sending a string

pub(crate) fn puts(s: &str) {
    for &b in s.as_bytes() {
        if b == b'\n' {
            Self::putc(b'\r');
        }
        Self::putc(b);
    }
}

PL011 string transmission with newline conversion (from main.rs)

Two things to notice. First, we iterate over the raw bytes of the string (s.as_bytes()), not characters. At this level, we’re dealing with bytes, not Unicode code points.

Second, the \n to \r\n conversion. Serial terminals expect both a carriage return (\r, move cursor to column 0) and a line feed (\n, move cursor down one line). If you send just \n, many terminals will move down without returning to the left edge, giving you a staircase effect. So we insert a \r before every \n.

10.4 The Logger trait implementation

impl hal::log::Logger for UartLogger {
    fn log(&self, s: &str) {
        UartLogger::puts(s);
    }
}

Logger trait implementation (from main.rs)

This bridges the gap between our platform-specific UART driver and the platform-agnostic kernel. The kernel never calls UartLogger::puts directly. It calls logger.log() on a trait object. More on this in the next section.

11. The platform-agnostic kernel

Here’s the design payoff. The kernel crate knows nothing about ARM, nothing about UART, nothing about QEMU. It just knows it has something that can log strings.

11.1 The Logger trait

pub trait Logger {
    fn log(&self, s: &str);
}

crates/hal/src/log.rs

Three lines. That’s the entire hardware abstraction layer for output. Any platform that can print strings can implement this trait.

11.2 The kernel entry

#![no_std]

use hal::log::Logger;

mod ipc;
mod sched;

use core::cell::UnsafeCell;

#[repr(transparent)]
struct RouterCell(UnsafeCell<ipc::Router>);
unsafe impl Sync for RouterCell {}

// Force the router into a writable section. On some bare-metal targets, a `static`
// with interior mutability can otherwise end up in a read-only segment, causing
// a data abort when we first write to it (exactly what we saw on aarch64 QEMU virt).
#[link_section = ".data"]
static ROUTER: RouterCell = RouterCell(UnsafeCell::new(ipc::Router::new()));

pub fn kmain(logger: &dyn Logger) -> ! {
    logger.log("rustOS: kernel online\n");
    logger.log("rustOS: microkernel step 1 (IPC + cooperative scheduling)\n");

    let router: &mut ipc::Router = unsafe { &mut *ROUTER.0.get() };

    let mut ping = sched::PingTask::new();
    let mut pong = sched::PongTask::new();
    let mut tasks: [&mut dyn sched::Task; 2] = [&mut ping, &mut pong];

    sched::run(&mut tasks, logger, router)
}

crates/kernel/src/lib.rs

Look at the kmain signature: fn kmain(logger: &dyn Logger) -> !. It takes a trait object (&dyn Logger), not a concrete type. The kernel doesn’t know whether logger is a PL011 UART, a COM1 serial port, or a carrier pigeon with a Morse code encoder. It just knows it can call .log(). If you wanted to port this kernel to RISC-V or even a microcontroller with an SPI display, you’d implement Logger for that platform’s output device and call kmain with it.

Notice the #[link_section = ".data"] on the ROUTER static. This is a hard-won lesson. On AArch64, the linker sometimes places statics with interior mutability (like UnsafeCell) into read-only sections. The first write causes a data abort. Forcing it into .data guarantees it lands in a writable section. We learned this the painful way.

The IPC router and scheduler are the subject of Part 2 . For now, just know that kmain is where platform-agnostic kernel logic lives.

11.3 The halt function

#[inline(always)]
pub fn halt() {
    #[cfg(target_arch = "x86_64")]
    unsafe {
        core::arch::asm!("hlt", options(nomem, nostack, preserves_flags));
    }

    #[cfg(target_arch = "aarch64")]
    unsafe {
        // Use WFI (wait-for-interrupt) so we reliably sleep until the next IRQ.
        // WFE can return immediately if an event is already latched.
        core::arch::asm!("wfi", options(nomem, nostack, preserves_flags));
    }

    #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
    {
        loop {}
    }
}

crates/hal/src/arch.rs

This is another piece of the hardware abstraction layer. On AArch64, wfi (wait for interrupt) puts the CPU into a low-power sleep state until an interrupt fires. On x86_64, the equivalent is hlt. The kernel code just calls hal::arch::halt() and doesn’t care which instruction runs underneath.

The comment about wfe vs wfi is interesting. wfe (wait for event) can return immediately if an event flag is already set, which means your “sleep” loop might spin instead of sleeping. wfi is more predictable: it always waits for an actual interrupt.

12. Linker scripts

On bare metal, there’s no OS to decide where your code lives in memory. The linker script is the document that says: put .text.boot first (so the CPU’s entry point is at the right address), put .text after it, then .rodata, .data, .bss, and finally the stack.

Our linker script for the virt platform places the kernel at 0x40080000. This is within the RAM region that starts at 0x40000000 on the QEMU virt machine. The 0x80000 offset is a convention (it matches the Raspberry Pi’s load address offset), though for QEMU it’s somewhat arbitrary as long as it’s within RAM.

The linker script also defines symbols like __bss_start, __bss_end, and __stack_top that the assembly boot code references. Without these, the assembler would have no idea where BSS begins or where to put the stack.

The build.rs file for our crate tells cargo to use this linker script:

fn main() {
    println!("cargo:rerun-if-changed=linker.ld");
    println!("cargo:rerun-if-changed=src/boot.S");
    let manifest_dir = std::env::var("CARGO_MANIFEST_DIR")
        .expect("CARGO_MANIFEST_DIR not set");
    println!("cargo:rustc-link-arg=-T{}/linker.ld", manifest_dir);
}

crates/arch_aarch64_virt/build.rs

The cargo:rustc-link-arg=-T.../linker.ld line passes the linker script to the linker. Without it, the linker would use its default layout, which almost certainly wouldn’t put our code at the right address.

13. Hands-on exercise

Here’s a quick challenge. Open crates/kernel/src/lib.rs and add a second log message to kmain:

pub fn kmain(logger: &dyn Logger) -> ! {
    logger.log("rustOS: kernel online\n");
    logger.log("rustOS: microkernel step 1 (IPC + cooperative scheduling)\n");
    logger.log("rustOS: hello from YOUR_NAME_HERE!\n");  // add this
    // ... rest of the function

Adding a custom log message

Rebuild and run:

./scripts/build-aarch64-virt.sh demo-ipc
./scripts/run-aarch64-virt.sh

Rebuild and verify

You should see your message appear between the boot messages and the scheduler output. One line of Rust. It traveled through a trait object, into a UART driver, through memory-mapped I/O, into a virtual PL011 peripheral, and out QEMU’s emulated serial port to your terminal.

Let’s step back and look at what we’ve got. A 60-line assembly boot sequence that sets up a stack, zeros BSS, drops from EL2 to EL1, enables floating-point, installs an exception vector table, and hands off to Rust. A Rust entry point that initializes a PL011 UART driver and calls into a platform-agnostic kernel. A three-line Logger trait that abstracts away all the hardware details. And a kernel entry function that doesn’t know or care what CPU architecture it’s running on. It boots. It prints. It works. And everything from here builds on this foundation. The IPC system in Part 2 , the timer interrupts in Part 3 , the MMU in Part 4 : they all stand on the boot sequence and UART driver we just walked through. Get this part right, and the rest is (relatively) straightforward.

14. Brief appendix: how would this differ on other platforms?

If you’re curious about x86_64 or Raspberry Pi, here’s a quick sketch.

On x86_64, the boot process is much more complex because of backwards compatibility with the 1978 Intel 8086. The CPU starts in 16-bit real mode, where you can only address 1 MB of memory. You have to transition through protected mode (32-bit) and then long mode (64-bit), setting up GDTs (Global Descriptor Tables), enabling PAE (Physical Address Extension), and configuring initial page tables (paging is required for 64-bit mode). Most Rust OS projects use the bootloader crate by Philipp Oppermann, which handles all these mode transitions and loads your kernel ELF. For serial output, x86 has legacy COM ports at I/O address 0x3F8, accessed with special in and out instructions (port-mapped I/O rather than memory-mapped I/O). The uart_16550 crate wraps this nicely.

On Raspberry Pi, something unexpected happens: the GPU boots first, not the CPU. The VideoCore GPU reads bootcode.bin from the SD card, loads start.elf, reads config.txt for settings, and then loads your kernel (kernel8.img) to address 0x80000. Only then does the ARM CPU start. The boot assembly is simpler than our QEMU virt version (no EL2 drop needed, the firmware handles that), but the UART driver is more complex. The Pi has two UARTs: a full PL011 (which is connected to Bluetooth by default) and a simpler Mini-UART on GPIO pins 14/15. The Mini-UART’s clock is tied to the GPU frequency, which makes baud rate configuration trickier. And of course, the development cycle is slower: build, copy to SD card, plug into Pi, power on, check output.

Both platforms share the same kernel crate and Logger trait. That’s the whole point of the abstraction.

15. References

ARM Architecture Reference Manual (ARMv8-A) (Section D1 for exception levels, Section G for system registers)
OSDev Wiki (community-driven OS development knowledge base)
Writing an OS in Rust by Philipp Oppermann (excellent x86_64 series)
The Embedded Rust Book (no_std patterns and bare-metal Rust)
QEMU AArch64 virt machine documentation
PL011 UART Technical Reference Manual

5-Part Series:

Part 0: Why build an OS from scratch?
Part 1 (this): Foundations
Part 2: Communication
Part 3: Concurrency
Part 4: Memory and beyond

Building a microkernel in Rust (Part 0): Why build an OS from scratch?

Sun, 22 Feb 2026 00:00:00 +0000

5-Part Series:

Part 0 (this): Why build an OS from scratch?
Part 1: Foundations
Part 2: Communication
Part 3: Concurrency
Part 4: Memory and beyond

GitHub Repository: bahree/rust-microkernel — full source code and build scripts

Docker Image: amitbahree/rust-microkernel — prebuilt dev environment with Rust, QEMU, and source code

Why this, why now?

I recently wrapped up an incredible chapter at Microsoft, working on the AI engineering team in CoreAI. Between roles, I found myself on gardening leave with something rare in this industry: a few weeks of unstructured time. And instead of doing the sensible thing (rest, catch up on sleep, maybe touch grass), my brain went: “You know what sounds fun? Writing an operating system from scratch.”

Look, I’ve spent the last several years deep in AI/ML, from foundational model development to applied AI at scale. That work isn’t going anywhere (the LLM from scratch series on this blog is a good example of that curiosity). But I’ve had this itch for a while now to go back to the fundamentals. The stuff that sits beneath all those PyTorch tensors and CUDA kernels. The stuff that actually makes a computer compute.

There’s an old joke: a QA engineer walks into a bar. Orders 1 beer. Orders 0 beers. Orders 99999999 beers. Orders -1 beers. Orders a lizard. Orders NULL beers. First real customer walks in and asks where the bathroom is. The bar bursts into flames.

That’s basically what happens when you try to run code without an OS. You think you’re ordering a beer, but there’s nobody behind the counter, no glasses, no refrigerator, and the building doesn’t have plumbing yet. You are the plumbing. This series is about building that plumbing from nothing.

Part 0? Really?

Yes, we’re starting at zero. Because real programmers count from zero, and honestly, this part is about understanding what we’re building before we build it. Think of it as the moment before you start cooking, where you read the whole recipe, check you’ve got the ingredients, and figure out what you’re actually making. You don’t skip that step. (Okay, maybe you do sometimes. But you shouldn’t.)

If you’ve ever wondered how code even runs without an operating system underneath, or what happens between pressing the power button and seeing a login screen, or why everyone says “virtual memory” but nobody explains how it actually works, you’re in the right place.

What happens when there’s nothing underneath you?

Every program you’ve ever written ran on top of an operating system. When you call println!, the OS figures out where to send those bytes. When you allocate memory, the OS finds a free chunk and maps it into your address space. When your code runs alongside other programs without crashing into them, that’s the OS keeping everyone in their own lane.

Now imagine none of that exists. No OS. No standard library. No heap. No println!. Just you, a CPU, and some RAM. You want to print a single character? You need to know the exact memory address of the serial port hardware, write one byte to that address, and hope you configured the device correctly. If you didn’t, nothing happens. The screen stays blank. The CPU doesn’t complain. It just sits there.

That’s where we’re starting. We’re going to build, from nothing, the thing that makes all of those invisible layers work. By the end, you’ll understand exactly what happens between “power on” and “your code runs.” And honestly? It’s one of the most satisfying things you can do as a programmer.

TL;DR

This is a 5-part educational series (Part 0 through Part 4) where we build a minimal OS kernel in Rust from scratch. We focus on AArch64 QEMU virt, which means anyone with a laptop can run it. No special hardware needed, just QEMU.

Here’s what you’ll learn:

Boot sequences: How a CPU goes from power-on to running your first line of Rust
Serial output: The simplest possible way to see what your kernel is doing
Message-passing IPC: How tasks communicate without sharing memory
Cooperative and preemptive scheduling: From polite turn-taking to the OS forcibly yanking the CPU away
Timer interrupts: How hardware gives your OS a sense of time
Context switching: Saving and restoring every register so tasks don’t know they were paused
Virtual memory: Page tables, the MMU, and why every address your code uses is a lie

Warning: This is a learning project, not production code. The goal is understanding, not security or completeness.

Time commitment: About 2-4 hours per part, totaling 5 parts. Follow along at your own pace.

1. A quick primer: AArch64, QEMU, and the “virt” machine

You’ll see the phrase “AArch64 QEMU virt” throughout this series, so let’s unpack it before we go further.

AArch64 is ARM’s 64-bit instruction set architecture. It’s the CPU architecture inside your phone, your tablet, Apple Silicon Macs, most cloud servers (AWS Graviton, Ampere Altra), and the Raspberry Pi 3 and newer. When we say we’re writing an AArch64 kernel, we mean we’re writing code that speaks the language these CPUs understand: A64 instructions, 31 general-purpose 64-bit registers (x0 through x30), four exception levels (EL0 through EL3) for privilege separation, and a specific way of handling interrupts, memory, and device access.

If you’ve only worked with x86_64 (Intel/AMD), the concepts are the same, but the details differ. ARM uses a load/store architecture (you can’t perform arithmetic directly in memory; you have to load values into registers first), a different interrupt model (GIC instead of APIC), and a different boot flow (exception levels instead of real/protected/long mode transitions). None of this matters yet. We’ll explain every ARM-specific detail as we encounter it.

QEMU is a machine emulator. It simulates an entire computer in software: CPU, RAM, interrupt controller, timer, serial port, everything. When we run our kernel in QEMU, it doesn’t touch your real hardware. QEMU pretends to be an ARM computer, and our kernel runs on it. This is incredibly useful for OS development because you can restart instantly, you can’t brick anything, and the emulated hardware behaves predictably.

“virt” is a specific machine type provided by QEMU. Real ARM hardware comes in thousands of configurations: different board layouts, different peripherals, different memory maps. The virt machine is a clean, minimal design created by QEMU specifically for virtual machines and testing. It gives you:

A PL011 UART (serial port) at a known memory address for text output
A GICv2 interrupt controller for routing hardware interrupts to the CPU
An ARM Generic Timer for periodic tick interrupts
RAM starting at a fixed address (0x40000000)
No legacy baggage, no quirky firmware, no board-specific workarounds

This is why we chose it for the blog series. The virt machine strips away all the hardware complexity that has nothing to do with OS fundamentals. You don’t need to worry about GPU initialization, USB enumeration, or PCIe configuration. You get a CPU, memory, a timer, an interrupt controller, and a serial port. That’s exactly what you need to learn how an OS works, and nothing more.

When you run qemu-system-aarch64 -machine virt, you’re telling QEMU: “pretend to be this specific kind of ARM computer.” Our kernel is written to match that pretend computer’s memory map and peripherals.

2. What we’re building

rustOS is a minimal microkernel. Not a “toy” that prints hello world and stops, but not Linux either. It’s the kind of thing that actually boots, actually runs multiple tasks, actually translates virtual addresses to physical ones. The interesting middle ground where you learn how real operating systems work without drowning in 30 million lines of code.

The kernel boots on AArch64 QEMU virt (an emulated ARM machine that anyone can run). It implements message-passing IPC, the core idea behind microkernel design, in which tasks communicate by sending small, fixed-size messages through a router rather than poking each other’s memory directly. It runs a cooperative scheduler that polls tasks in round-robin, and then we upgrade to preemptive multitasking with timer interrupts so tasks can’t hog the CPU. And it manages virtual memory with 4-level page tables and a real MMU, the hardware that makes every modern OS possible.

Now, the codebase also supports x86_64 and the Raspberry Pi Zero 2 W. Those platforms share the same kernel logic and are available in the repository if you want to explore them. But the blog series focuses on AArch64 virt because it’s the most accessible. You don’t need a Raspberry Pi. You don’t need a Windows machine to deal with x86 bootloader quirks. You need QEMU, and QEMU runs on everything.

Before we look at the code layout, a quick Rust concept if you’re coming from another language: a crate is Rust’s unit of compilation, roughly equivalent to a library or package in other ecosystems. A crate has its own Cargo.toml (think package.json or pom.xml) that declares its name, dependencies, and build settings. A workspace groups multiple crates together so they can share dependencies and be built in one cargo build invocation. Our project is a workspace with several crates: kernel, hal, and one per platform. Each crate compiles independently and can depend on others. When you see kernel crate or hal crate below, think “self-contained module with a clear API boundary.”

Here’s how the code is organized:

graph TB
    subgraph "Platform-Agnostic Core"
        kernel[kernel<br/>IPC Router, Scheduler, Tasks]
        hal[hal<br/>Logger Trait, Arch Primitives]
    end

    subgraph "Platform Crate"
        virt[arch_aarch64_virt<br/>QEMU Boot, PL011 UART<br/>Timer, GIC, MMU]
    end

    virt --> hal
    virt --> kernel

    style kernel fill:#e1f5ff
    style hal fill:#e1f5ff
    style virt fill:#fff4e1

Figure 4: Platform-agnostic crate architecture

The kernel crate knows nothing about hardware. It talks to the world through a Logger trait defined in hal. The platform crate (arch_aarch64_virt) implements that trait using a PL011 UART driver and handles all ARM-specific setup: boot assembly, timer configuration, GIC interrupt routing, and MMU page tables. This separation is the same pattern you’ll find in real operating systems. Linux has arch/arm64/ for ARM-specific code and kernel/ for the portable parts. We’re doing the same thing, just at a much smaller scale.

3. Why this series?

Most OS tutorials either skip the hard parts (boot, interrupts, memory management), use simulators instead of real code, or don’t explain why things work the way they do. You get a code dump with a comment saying “configure the GIC” and zero explanation of what the GIC is, why it exists, or what happens if you get it wrong.

This series is different. We show complete implementations with no hand-waving. Every register write gets explained. Every design decision has a rationale. When we configure the MMU, we don’t just say “set TCR_EL1 to this magic number.” We break down every bit field, explain what it controls, and why we chose that value.

We also focus on understanding the features. A production OS needs thousands of things we don’t build: filesystems, networking, user mode, multicore support, security hardening. We skip all of that. Not because those things aren’t important, but because they all build on the foundation we’re teaching here. You can’t write a filesystem without understanding memory management. You can’t implement a network stack without interrupts. Get the foundation right, and the rest is (hard, but conceptually clear) engineering.

And we’re using Rust, which honestly makes bare-metal programming less terrifying than it used to be. The type system catches entire classes of bugs at compile time. You’ll still write plenty of unsafe code (this is an OS, after all), but the safe parts of Rust keep you honest about where the dangerous bits are.

Operating systems sit at the intersection of hardware (registers, interrupts, memory controllers), architecture (x86 vs ARM, privilege levels, MMUs), concurrency (multiple tasks, race conditions), and abstraction (clean APIs over messy hardware). There’s no debugger, no standard library, and mistakes cause silent crashes. This series shows you how to navigate that complexity, one piece at a time.

4. Why Rust? (and not C or assembly)

If you’re building an OS, you’re probably wondering: why Rust? Traditionally, kernels are written in C (Linux, xv6, Minix) or even raw assembly. Both are fine choices. But Rust brings something genuinely new to the table.

4.1 The problem with C

C gives you total control. No garbage collector, no runtime, direct memory access. That’s why it’s been the lingua franca of OS development for 50 years. But C also trusts you completely, and that trust is the problem. Buffer overflows, use-after-free, null pointer dereferences, data races: these aren’t theoretical. They’re the source of roughly 70% of security vulnerabilities in production systems (Microsoft and Google have both published studies confirming this).

In kernel code, these bugs are catastrophic. A buffer overflow in userspace causes your program to crash. A kernel-space buffer overflow corrupts the entire system. There’s no safety net below you.

4.2 The problem with assembly only

You could write an OS in pure assembly. Some people do, and there’s real value in understanding every instruction. But assembly doesn’t scale. An OS needs abstractions (traits, modules, type checking) to stay maintainable. Writing a mailbox router in assembly would work, but debugging it would be miserable. And you’d lose all the benefits of a compiler that can optimize, inline, and catch type errors for you.

4.3 What Rust gives you

Rust’s ownership system catches entire classes of bugs at compile time. You can’t use a value after freeing it. You can’t have two mutable references to the same data. You can’t forget to initialize memory. The borrow checker is annoying until you realize it’s catching bugs that would take hours to debug on bare metal, where there’s no debugger, no address sanitizer, and no core dump.

But Rust also gets out of your way when you need it to. The unsafe keyword lets you write raw pointer dereferences, inline assembly, and MMIO operations. You have to be explicit about it. The dangerous parts are clearly marked, and the safe parts are genuinely safe. This is a huge improvement over C, where everything is implicitly unsafe.

Rust’s no_std ecosystem is also mature. Crates like bootloader, uart_16550, and volatile provide building blocks for OS development. The core library gives you basic types, traits, and iterators without needing an allocator. And #[repr(C)] lets you control struct layout for hardware register interactions.

Concretely, here are bugs Rust prevented during this project:

The compiler refused to let us access the IPC router from an interrupt handler without explicit synchronization (UnsafeCell + Sync impl)
A missing volatile read would have been optimized away in C, silently breaking UART output. In Rust, we had to use read_volatile, making the intent clear explicitly
Type mismatches between page table entry formats were caught at compile time, not at boot time, with a mysterious hang

Is Rust perfect for OS development? No. The borrow checker sometimes fights you in ways that feel unnecessary for single-threaded kernel code. static mut is awkward. Some patterns that are natural in C require contortions in Rust. But on balance, we think the trade-off is overwhelmingly positive. You spend a little more time satisfying the compiler, and a lot less time debugging silent memory corruption.

5. Getting started

You’ll need Rust (nightly), QEMU, and basic Rust knowledge. If you know what ownership and traits are, you’re good. If you’ve written a function that takes &self and returned a Result, you’ll be fine. You don’t need to know assembly or have prior OS background. That’s the whole point of the series.

Tool	Purpose	Installation
Rust (nightly)	Bare-metal compilation	`rustup default nightly`
QEMU	Virtual machine for running our kernel	`apt install qemu-system-aarch64` or `brew install qemu`
Git	Cloning the repository	You probably already have this

5.1 Install Rust (nightly)

# Install Rust via rustup
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh

# Follow prompts, select default installation
# Then reload your shell environment:
source $HOME/.cargo/env

# Switch to nightly (required for bare-metal features)
rustup default nightly

# Add the rust-src component (needed for building core for bare-metal targets)
rustup component add rust-src --toolchain nightly

# Add the AArch64 bare-metal target
rustup target add aarch64-unknown-none

Verify it works:

$ rustc --version
rustc 1.86.0-nightly (5765cbc51 2025-01-28)

$ rustup target list --installed | grep aarch64
aarch64-unknown-none

5.2 Install QEMU

# Ubuntu/Debian
sudo apt install qemu-system-aarch64

# macOS
brew install qemu

# Verify
$ qemu-system-aarch64 --version
QEMU emulator version 8.2.2 (Debian 1:8.2.2+ds-0ubuntu1.4)

5.3 Clone and build

$ git clone https://github.com/bahree/rust-microkernel.git
$ cd rust-microkernel
$ ./scripts/build-aarch64-virt.sh

[virt] building aarch64 QEMU virt ELF (default features)...
    Finished `release` profile [optimized] target(s) in 0.28s
[virt] wrote dist/virt/os-aarch64-virt.elf

5.4 Run it

$ ./scripts/run-aarch64-virt.sh

rustOS: aarch64 QEMU virt boot OK
rustOS: memory management demo (frames + page tables)
mm: demo start
mm: kernel_end=0x000000004009A010
mm: free_start=0x000000004009B000
mm: ram_end=0x0000000050000000
mm: frame0=0x000000004009B000
mm: frame1=0x000000004009C000
mm: read0=0x00000000AABBCCDD
mm: read1=0x0000000011223344
mm: ttbr0=0x0000000040085000
mm: test_va=0x0000000080000000
mm: enabling MMU (caches off)...
mm: test_va_read=0x00000000DEADBEEF
mm: demo done (MMU is ON)

Press Ctrl+A then X to exit QEMU.

Figure 1: Memory demo output showing frame allocation, page table construction, MMU enablement, and virtual address translation.

If you see that output, congratulations. You just ran an operating system kernel that booted on bare metal, allocated physical memory frames, built 4-level page tables, enabled the MMU, and successfully translated a virtual address from a few shell commands.

5.5 Try the other demos

The codebase has multiple demos you can build with feature flags:

# IPC + cooperative scheduling (Part 2)
$ ./scripts/build-aarch64-virt.sh demo-ipc && ./scripts/run-aarch64-virt.sh

rustOS: aarch64 QEMU virt boot OK
rustOS: IPC + cooperative scheduling demo
rustOS: kernel online
rustOS: microkernel step 1 (IPC + cooperative scheduling)
sched: starting
task/ping: poll
task/ping: sent ping
task/pong: got ping

# Preemptive multitasking (Part 3)
$ ./scripts/build-aarch64-virt.sh demo-preempt && ./scripts/run-aarch64-virt.sh

rustOS: aarch64 QEMU virt boot OK
rustOS: preemptive multitasking demo
A
B
A
B
A
B
...

Figure 2: IPC demo showing ping and pong tasks exchanging messages through the mailbox router.

Figure 3: Preempt demo showing two threads alternating without ever yielding, proving the OS is in control.

If all three demos work, your environment is fully set up and ready for the series.

5.6 Docker (no local setup)

If you just want to try it without installing anything:

docker pull amitbahree/rust-microkernel:latest
docker run -it amitbahree/rust-microkernel:latest
# Inside the container:
./scripts/build-aarch64-virt.sh && ./scripts/run-aarch64-virt.sh

x86_64 and Raspberry Pi builds also exist in the repository. Run ./scripts/build-x86.sh or ./scripts/build-rpi.sh if you want to try them. They share the same kernel logic but use different boot code and hardware drivers.

6. Series roadmap

Here’s what each part covers and why it comes in that order. Each part builds on the previous one. You can’t preempt tasks without interrupts. You can’t isolate processes without virtual memory. The ordering isn’t arbitrary.

Part 0: Why build an OS from scratch? (you are here) — The overview. What we’re building, why, and how the pieces fit together.

Part 1: Foundations — How a CPU goes from power-on to running your first line of Rust. AArch64 boot assembly that drops from EL2 to EL1, stack setup, BSS zeroing, PL011 UART driver, and the platform abstraction (HAL) that keeps the kernel hardware-agnostic.

Part 2: Communication — Message-passing IPC with a mailbox router. Tasks implement a poll() method, the scheduler calls them round-robin, and we build PingTask/PongTask to show communication in action. Plus why cooperative scheduling has limits (one misbehaving task freezes everything).

Part 3: Concurrency — The big one. ARM Generic Timer, GICv2 interrupt controller, periodic timer interrupts, and a preemptive scheduler that saves all 31 registers plus SP, ELR, and SPSR. Two threads print “A” and “B” without ever yielding, proving the OS is in control.

Part 4: Memory and beyond — Virtual memory from the ground up. Bump allocator for physical frames, 4-level page tables (L0/L1/L2/L3), MMU configuration with MAIR, TCR, and TTBR0, and address translation verified by writing through a virtual address and reading back through the physical one. Ends with reflections and where to go next.

7. The big picture: what makes an operating system?

Before we start building, let’s understand what we’re aiming for. An operating system has many layers, and we’re focusing on the foundational ones that everything else depends on.

7.1 Complete OS architecture

graph TB
    subgraph "User Space (Ring 3 / EL0)"
        Apps[Applications<br/>ls, cat, browser, games]
        Libs[Standard Libraries<br/>libc, libstd]
        Shell[Shell<br/>bash, zsh]
    end

    subgraph "Kernel Space (Ring 0 / EL1)"
        subgraph "System Call Interface"
            Syscalls[System Calls<br/>open, read, write, fork]
        end

        subgraph "Process Management"
            Sched[Scheduler<br/>Round-robin, CFS, priorities]
            IPC[IPC<br/>Message passing, pipes, signals]
            Proc[Process Table<br/>PIDs, state, resources]
        end

        subgraph "Memory Management"
            VMM[Virtual Memory<br/>Page tables, TLB, swapping]
            Heap[Heap Allocator<br/>malloc/free, buddy system]
            MMU[MMU Control<br/>TTBR, page faults]
        end

        subgraph "File Systems"
            VFS[VFS Layer<br/>Common file API]
            FS[File Systems<br/>ext4, FAT, NTFS]
            Cache[Page Cache<br/>Read-ahead, write-back]
        end

        subgraph "Device Drivers"
            Block[Block Devices<br/>Disk, SSD, NVMe]
            Net[Network<br/>Ethernet, WiFi, TCP/IP]
            Char[Character Devices<br/>UART, keyboard, mouse]
            GPU[Graphics<br/>Framebuffer, GPU drivers]
        end

        subgraph "Core Kernel - WHAT WE BUILD"
            Boot[Boot<br/>Platform init, entry]
            Int[Interrupts<br/>Timer, exceptions, GIC/PIC]
            Ctx[Context Switch<br/>Save/restore registers]
            Mem[Physical Memory<br/>Frame allocator]
        end
    end

    subgraph "Hardware"
        CPU[CPU<br/>Cores, caches, MMU]
        RAM[RAM<br/>Physical memory]
        Disk[Storage<br/>HDD, SSD]
        Devices[Peripherals<br/>UART, Timer, GPIO]
    end

    Apps --> Syscalls
    Libs --> Syscalls
    Shell --> Syscalls

    Syscalls --> Sched
    Syscalls --> VMM
    Syscalls --> VFS

    Sched --> Ctx
    Sched --> Proc
    Sched --> IPC

    VMM --> MMU
    VMM --> Heap
    VMM --> Mem

    VFS --> FS
    VFS --> Cache
    FS --> Block

    Block --> Devices
    Net --> Devices
    Char --> Devices
    GPU --> Devices

    Int --> CPU
    Ctx --> CPU
    MMU --> CPU
    Mem --> RAM
    Boot --> CPU

    style Boot fill:#e8f5e8
    style Int fill:#e8f5e8
    style Ctx fill:#e8f5e8
    style Mem fill:#e8f5e8
    style IPC fill:#fff4e1
    style Sched fill:#fff4e1
    style VMM fill:#fff4e1

Figure 5: Complete OS architecture showing what we build

Look at that diagram. It’s a lot. But here’s the thing: we’re only building the green and yellow boxes at the bottom. Boot, interrupts, context switching, physical memory, IPC, scheduling, and virtual memory. That’s the core kernel layer, the absolute foundation that every OS needs. These aren’t optional features you can skip. They’re the bedrock on which everything else is built.

7.2 What we build (the foundation)

Boot sequences (Part 1) are how a computer goes from pressing the power button to running your kernel’s first line of code. It’s the handoff between hardware/firmware and your operating system. On ARM, this means handling exception levels (EL2 hypervisor to EL1 kernel), setting up a stack, zeroing the BSS, and enabling floating-point instructions before the LLVM-generated code can run. Get any step wrong, and the CPU hangs silently.

IPC and scheduling (Parts 2 and 3) are how multiple tasks run on a single CPU core. The scheduler decides who runs when. IPC lets tasks communicate without sharing memory. We start with cooperative scheduling (tasks voluntarily yield) and upgrade to preemptive scheduling (the OS is in control). This requires context switching: saving all 31 ARM registers, the stack pointer, the return address, and the CPU flags, then loading another task’s saved state.

Interrupts and exceptions (Part 3) are hardware signals that tell the CPU, “stop what you’re doing and handle this NOW.” The timer fires, interrupt! Without interrupts, your OS would have to constantly poll hardware (“Are you done yet? How about now?”), wasting billions of CPU cycles. And you’d have no way to switch between tasks forcibly. Interrupts are the foundation for everything asynchronous in a computer.

Memory management (Part 4) is the illusion that each program has its own private address space. The MMU (Memory Management Unit) translates every address using page tables, a 4-level tree structure that maps virtual addresses to physical ones. We build those tables, configure the MMU, and enable translation. After that, every memory access your code makes goes through hardware that we configured.

7.3 What we don’t build

We skip filesystems, networking, user mode, multicore, and security hardening. Not because they’re unimportant, but because they all build on the foundation we’re teaching. You can’t write a filesystem without memory management. You can’t implement a network stack without interrupts. You can’t add user mode without understanding privilege levels and page tables. Part 4 points you to resources for going deeper into each of these areas.

8. Repository structure

The full source code is at github.com/bahree/rust-microkernel . Here’s how it’s organized, focused on the AArch64 virt platform that the blog series covers:

rust-microkernel/
├── crates/
│   ├── kernel/                 # Platform-agnostic kernel logic
│   │   ├── lib.rs              # Entry point (kmain)
│   │   ├── ipc.rs              # Message-passing router
│   │   └── sched.rs            # Cooperative scheduler
│   ├── hal/                    # Hardware abstraction layer
│   │   ├── log.rs              # Logger trait
│   │   └── arch.rs             # Architecture primitives (halt)
│   └── arch_aarch64_virt/      # AArch64 QEMU virt platform
│       ├── boot.S              # Assembly: EL2→EL1 drop, vectors
│       ├── main.rs             # Rust entry point
│       ├── uart.rs             # PL011 UART driver
│       ├── timer.rs            # Generic Timer + GIC
│       ├── preempt.rs          # Context switching
│       └── mem.rs              # Frame allocator + MMU
├── scripts/
│   ├── build-aarch64-virt.sh   # Build AArch64 virt kernel
│   └── run-aarch64-virt.sh     # Run in QEMU
├── docs/
│   └── blog/                   # This blog series
└── dist/                       # Build artifacts

The kernel crate knows nothing about hardware. It accepts a Logger trait object and calls .log() to print messages. The hal crate defines that trait. The arch_aarch64_virt crate implements it with a real UART driver and handles everything ARM-specific: boot assembly, timer setup, interrupt routing, page tables.

The x86_64 and Raspberry Pi platform crates (arch_x86_64, arch_aarch64_rpi) also exist in crates/ and follow the same pattern. They’re not covered in the blog series, but use the same kernel and hal crates.

9. References and acknowledgments

This project stands on the shoulders of giants. Philipp Oppermann’s “Writing an OS in Rust” showed that Rust OS development is not only possible but genuinely enjoyable. MIT’s xv6 demonstrated that you can teach OS internals with a clean, readable codebase. The OSDev community has been documenting hardware quirks and boot sequences for decades, and their wiki is invaluable.

For operating systems theory, Operating Systems: Three Easy Pieces (free online) is the best introduction we’ve found. For ARM architecture specifics, the ARM Architecture Reference Manual is the definitive reference (and it’s freely available). For bare-metal Rust patterns, the Embedded Rust Book covers no_std development in depth.

We also learned from real microkernel projects: seL4 (formally verified, beautifully designed), Minix 3 (Tanenbaum’s vision of what microkernels should be), and the L4 family (proving that microkernel IPC can be fast).

5-Part Series:

Part 0 (this): Why build an OS from scratch?
Part 1: Foundations
Part 2: Communication
Part 3: Concurrency
Part 4: Memory and beyond

GitHub Repository: bahree/rust-microkernel — full source code and build scripts

Docker Image: amitbahree/rust-microkernel — prebuilt dev environment with Rust, QEMU, and source code

Building LLMs from Scratch - Part 4: Evaluation & Deployment

Fri, 02 Jan 2026 00:00:00 +0000

TL;DR

In this final part of our 4-part series on building language models from scratch, we explore the evaluation, testing, and deployment pipeline that transforms our trained historical language models into working systems. Part 1 showed you how to use the published models, Part 2 covered data collection and custom tokenization, and Part 3 detailed the model architecture and training infrastructure. Here, we complete the journey with evaluation frameworks, testing infrastructure, and deployment to Hugging Face Hub.

⚠️ Educational Purpose: This is a learning project designed to teach LLM development concepts. For production-scale LLMs, you’ll need much larger datasets, more sophisticated infrastructure, and additional considerations not covered here.

As outlined in Part 1 , both the SLM (117M parameters) and the Regular Model (354M parameters) use the same training code and infrastructure with different configurations defined in config.py. The evaluation and deployment infrastructure is also identical - only the model architecture parameters differ.

Both PyTorch checkpoint inference and Hugging Face model inference are fully working and available. Both the SLM and the Regular model are published on Hugging Face Hub . Local PyTorch checkpoints can be used directly for inference with the script inference_pytorch.py.

🔗 GitHub Repository: github.com/bahree/helloLondon - Complete evaluation and deployment infrastructure (05_evaluation/, 06_inference/, 10_scripts/) plus guides (08_documentation/EVALUATION_GUIDE.md, 08_documentation/HUGGINGFACE_PUBLISHING.md, 08_documentation/DEPLOYMENT_GUIDE.md)

🟥 Series Posts: Part 1 - Using the Published Historical Models | Part 2 - Data Collection & Custom Tokenizer | Part 3 - Training Architecture & GPU Optimization | Part 4 (this post)

🟧 Published Models: SLM Model | Regular Model - Ready-to-use historical language models on Hugging Face

📗 Book Reference: Generative AI in Action - For deeper understanding of core LLM concepts

1. The Evaluation Challenge: Measuring What Matters for Historical Language Models

Now that we have trained models from Part 3 , we face a critical question: How do we know if our models actually work? This isn’t just about checking if the code runs - it’s about validating that the models can generate historically accurate, linguistically appropriate text that captures the essence of 1500-1850 London English.

The challenge with evaluating historical language models goes far beyond standard LLM metrics. Standard evaluation approaches like Perplexity and BLEU scores (we explain these and other metrics in Section 2.1 ) tell us whether the model generates fluent text. Still, they don’t answer the questions that matter for historical applications: Does the model avoid anachronisms? Can it distinguish between Tudor and Victorian language patterns? Does it understand London geography and historical context?

Consider a simple example: if we prompt the model with “In the year 1600, I traveled to London by railway”, a standard language model might generate this without flagging the obvious problem - railways didn’t exist in 1600. The evaluation framework needs to catch these temporal inconsistencies, period-inappropriate language, and historical inaccuracies that standard metrics miss.

This evaluation challenge requires building a specialized assessment pipeline that understands historical context, temporal boundaries, and period-specific linguistic patterns. We need metrics that can distinguish between a model that generates fluent modern English and one that produces authentic historical text - two very different capabilities.

1.1 High-Level Evaluation Strategy

Our evaluation framework provides two complementary approaches that work with both PyTorch checkpoints and Hugging Face models, as illustrated in Figure 1 below.

graph TD
    A[🤖 Trained Models<br/>SLM 117M / Regular 354M] --> B{Evaluation Type}
    
    B -->|Quick| C[⚡ Quick Evaluation<br/>Historical accuracy, language quality, coherence]
    B -->|Comprehensive| D[🔬 Comprehensive Evaluation<br/>Benchmarks, G-Eval, groundedness]
    
    C --> E[📊 Evaluation Results<br/>Historical accuracy scores, metrics]
    D --> E
    
    E --> F{Quality OK?}
    F -->|Yes| G[🚀 Deployment Options]
    F -->|No| H[🔄 Retrain/Adjust]
    H --> A
    
    G --> I[📦 PyTorch Checkpoints<br/>Direct inference]
    G --> J[🤗 Hugging Face Hub<br/>Published models]
    G --> K[💻 Local Deployment<br/>API, CLI, notebooks]
    
    I --> L[✅ Working Models<br/>Ready for use]
    J --> L
    K --> L
    
    style A fill:#e1f5fe
    style E fill:#f3e5f5
    style L fill:#e8f5e8
    style H fill:#fff3e0

Figure 1: Complete Evaluation and Deployment Pipeline

Quick Evaluation (quick_eval.py): Rapid validation testing historical accuracy on key events (e.g., 1665 plague, 1666 fire, etc.), language quality metrics (vocabulary diversity, historical pattern detection, readability), and coherence (ROUGE scores). Runs in minutes without external APIs.

Comprehensive Evaluation (comprehensive_evaluator.py): Extends quick evaluation with benchmark datasets (small MMLU and HellaSWAG subsets), groundedness/fluency metrics, and optional LLM-as-a-judge scoring via G-Eval (using an external GPT model). Produces detailed reports with generation samples.

Both evaluators test across historical periods (such as Tudor, Stuart, and Georgian), language patterns (archaic pronouns and verb forms), and London-specific knowledge (geography and landmarks). The framework goes beyond standard LM metrics to assess period-appropriate language, temporal consistency, and historical accuracy.

2. Model Evaluation Framework

Now that we’ve outlined the evaluation challenge, let’s dive into the implementation. Our evaluation framework provides two complementary approaches that work with both PyTorch checkpoints and Hugging Face models. The framework is designed to be practical for a learning project while still providing meaningful insights into model performance.

2.1 Historical, Linguistic, and Category-Specific Evaluation

To make the evaluation concrete, we look at the model from three complementary aspects that together capture how well it understands the period, writes fluent text, and handles the different slices of the corpus. This multi-dimensional approach ensures we catch various types of failures - a model might generate grammatically perfect text but fail historically, or vice versa.

Historical assessments: Quick evaluation uses targeted prompts around key events (e.g., 1665 plague, 1666 fire, Old Bailey trials) and checks for expected keywords and phrases. Comprehensive evaluation adds temporal consistency checks (forbidden/required terms per period), date-range sanity checks, and historical benchmarks (custom historical questions and the MMLU subset).
Linguistic assessments: We measure surface quality (chars/words/sentences per sample, words per sentence), vocabulary diversity (unique/total tokens), readability (Flesch-style scores), and presence of historical patterns (archaic verb forms like hath, doth, pronouns like thou, thee, conjunctions and interjections). This shows whether the model writes in a historically flavored yet readable style.
Category-specific benchmarks: Evaluations are grouped by period (Tudor, Stuart, Georgian), by linguistic phenomena (archaic forms, dialogue patterns), and by London knowledge (Thames, Westminster, Old Bailey, etc.). The comprehensive evaluator further probes general reasoning using HellaSWAG and MMLU subsets to assess the model’s performance across broader benchmarks.

Industry-Standard Evaluation Metrics and Benchmarks

Our evaluation framework uses several standard metrics and benchmarks from LLM research. Here’s what each one measures and why we include it:

Perplexity: How surprised the model is by the reference text; lower is better because it means the model assigns higher probability to what actually happened in the corpus.

BLEU / ROUGE: N-gram overlap between generated and reference text, giving a rough sense of literal similarity and how closely the model “sticks” to the reference phrasing. We use ROUGE-L (longest common subsequence) to evaluate coherence and narrative flow.

MMLU (Massive Multitask Language Understanding): A large multiple-choice exam covering many academic subjects. Here, we use a tiny subset as a sanity check for general knowledge and reasoning, not as a primary goal.

HellaSWAG: A commonsense inference benchmark where the model must pick a plausible continuation for a short story-like context. We use it to see whether the model’s basic reasoning looks sensible.

G-Eval: An LLM-as-a-judge pattern where a stronger reference model (for example, GPT) scores generated text along dimensions like coherence or groundedness. In this project, it is optional and requires an external API key.

Groundedness: Asks: does the model stick to the provided context / known facts, or hallucinate? Our implementation approximates this by comparing generations against reference answers and historical constraints.

For a deeper treatment of evaluation benchmarks (including MMLU, HellaSWAG, and LLM-as-a-judge methods like G-Eval), see Chapter 12 - Evaluating and Monitoring Generative Systems in the book 📘 Generative AI in Action .

2.2 Automated Evaluation Pipeline

The run_comprehensive_evaluation function in 05_evaluation/comprehensive_evaluator.py orchestrates the entire evaluation process. Listing 1 shows how it works: We iterate over test sets, generate text with the model, compute all the metrics defined above, and aggregate the results into a results dictionary for analysis.

def run_comprehensive_evaluation(model, tokenizer, test_data, device='cuda'):
    """Run comprehensive evaluation on historical language model"""
    
    # Initialize evaluation metrics
    metrics = {
        'perplexity': [],
        'bleu_scores': [],
        'rouge_scores': [],
        'historical_accuracy': [],
        'linguistic_quality': [],
        'coherence_scores': [],
        'temporal_consistency': []
    }
    
    # Evaluate on different text types
    for text_type, samples in test_data.items():
        logger.info(f"Evaluating on {text_type} samples...")
        
        for sample in samples:
            # Generate text
            generated = generate_text(model, tokenizer, sample['prompt'], device)
            
            # Calculate metrics
            perplexity = calculate_perplexity(model, tokenizer, sample['text'], device)
            bleu = calculate_bleu(generated, sample['reference'])
            rouge = calculate_rouge(generated, sample['reference'])
            hist_acc = assess_historical_accuracy(generated, sample['context'])
            ling_qual = assess_linguistic_quality(generated)
            coherence = assess_coherence(generated)
            temp_cons = assess_temporal_consistency(generated, sample['time_period'])
            
            # Store metrics
            metrics['perplexity'].append(perplexity)
            metrics['bleu_scores'].append(bleu)
            metrics['rouge_scores'].append(rouge)
            metrics['historical_accuracy'].append(hist_acc)
            metrics['linguistic_quality'].append(ling_qual)
            metrics['coherence_scores'].append(coherence)
            metrics['temporal_consistency'].append(temp_cons)
    
    # Calculate aggregate metrics
    results = {}
    for metric_name, values in metrics.items():
        results[metric_name] = {
            'mean': np.mean(values),
            'std': np.std(values),
            'min': np.min(values),
            'max': np.max(values),
            'median': np.median(values)
        }
    
    return results

Listing 1: Comprehensive Evaluation Pipeline

The pipeline computes all the metrics we outlined above (standard LM metrics such as perplexity and BLEU/ROUGE, plus our historically specific assessments of accuracy, linguistic quality, and coherence). Each metric provides a different lens through which to view model performance: perplexity measures how well the model predicts the training distribution, BLEU/ROUGE measures literal similarity to the reference text, and the custom metrics assess historical authenticity and linguistic appropriateness.

Why This Multi-Metric Approach Matters?

Standard language model evaluation often focuses on perplexity and n-gram overlap metrics, which measure general language quality but miss domain-specific requirements. For historical language models, we need to know not just whether the text is fluent, but whether it’s historically accurate, temporally consistent, and linguistically appropriate for the target period. This multi-metric approach ensures we catch different types of failures - a model might generate grammatically perfect text but fail historically, or produce historically accurate content with poor linguistic quality.

The aggregation step (computing mean, std, min, max, median) provides a comprehensive view of model performance across different test cases. This statistical summary helps identify whether the model performs consistently or has high variance, whether certain types of prompts cause failures, and how the model compares across different historical periods and linguistic phenomena.

2.3 Historical Accuracy Assessment

Standard LLM evaluation metrics (perplexity, BLEU, ROUGE) measure general language quality, but they don’t tell us whether the model generates historically accurate text for London between 1500-1850. To address this, we built customized evaluation tools that check period-appropriate language, temporal consistency, London-specific geography and landmarks, and historical fact accuracy. These tools are implemented in 05_evaluation/comprehensive_evaluator.py as shown in Listing 2:

def assess_historical_accuracy(generated_text, historical_context):
    """Assess the historical accuracy of generated text"""
    
    accuracy_score = 0.0
    total_checks = 0
    
    # Check temporal consistency
    temporal_score = check_temporal_consistency(generated_text, historical_context['time_period'])
    accuracy_score += temporal_score
    total_checks += 1
    
    # Check historical facts
    fact_score = check_historical_facts(generated_text, historical_context['facts'])
    accuracy_score += fact_score
    total_checks += 1
    
    # Check period-appropriate language
    language_score = check_period_language(generated_text, historical_context['time_period'])
    accuracy_score += language_score
    total_checks += 1
    
    # Check geographical accuracy
    geo_score = check_geographical_accuracy(generated_text, historical_context['location'])
    accuracy_score += geo_score
    total_checks += 1
    
    # Check social context accuracy
    social_score = check_social_context(generated_text, historical_context['social_class'])
    accuracy_score += social_score
    total_checks += 1
    
    return accuracy_score / total_checks

def check_temporal_consistency(text, time_period):
    """Check if text maintains temporal consistency with the specified period"""
    
    # Define period-specific constraints
    period_constraints = {
        '1500-1600': {
            'forbidden_terms': ['electricity', 'steam engine', 'railway'],
            'required_terms': ['ye', 'hath', 'doth'],
            'date_range': (1500, 1600)
        },
        '1600-1700': {
            'forbidden_terms': ['railway', 'telegraph'],
            'required_terms': ['hath', 'doth', 'verily'],
            'date_range': (1600, 1700)
        },
        '1700-1800': {
            'forbidden_terms': ['telegraph', 'telephone'],
            'required_terms': ['hath', 'doth', 'indeed'],
            'date_range': (1700, 1800)
        },
        '1800-1850': {
            'forbidden_terms': ['telephone', 'automobile'],
            'required_terms': ['indeed', 'verily', 'pray'],
            'date_range': (1800, 1850)
        }
    }
    
    if time_period not in period_constraints:
        return 0.5  # Neutral score for unknown periods
    
    constraints = period_constraints[time_period]
    score = 1.0
    
    # Check for forbidden terms (anachronisms)
    for term in constraints['forbidden_terms']:
        if term.lower() in text.lower():
            score -= 0.2
    
    # Check for required period-appropriate terms
    period_terms_found = 0
    for term in constraints['required_terms']:
        if term.lower() in text.lower():
            period_terms_found += 1
    
    if constraints['required_terms']:
        score += 0.3 * (period_terms_found / len(constraints['required_terms']))
    
    # Check date references
    date_score = check_date_references(text, constraints['date_range'])
    score += 0.2 * date_score
    
    return max(0.0, min(1.0, score))

Listing 2: Historical Accuracy Assessment

The forbidden terms (like “electricity” for 1500-1600, “railway” for 1600-1700) are anachronisms - technologies or concepts that didn’t exist in those periods. We selected them based on historical timelines: electricity wasn’t harnessed until the late 1700s, railways didn’t appear until the early 1800s, and telegraphs came later. Similarly, the required terms (such as “hath”, “doth”, and “verily”) are archaic language patterns we observed frequently in the training corpus for each period.

We analyzed the corpus to identify which linguistic markers were most characteristic of each era, then selected a small set that would catch obvious anachronisms without being overly restrictive. This is a practical heuristic rather than an exhaustive historical grammar - we focus on high-impact anachronisms and common period markers that are easy to detect automatically.

How the scoring works

The check_temporal_consistency() function starts with a score of 1.0 and applies penalties and bonuses: each forbidden term found subtracts 0.2 (so finding “railway” in 1600-1700 text drops the score), while finding required period-appropriate terms adds up to 0.3 based on how many are present. Date references within the period add up to 0.2. The final score ranges from 0.0 to 1.0.

The overall assess_historical_accuracy() function then averages the five component scores (temporal consistency, historical facts, period-appropriate language, geographical accuracy, and social context) to produce a single score between 0 and 1, with higher values indicating better historical accuracy. In practice (and yes, we are generalizing), scores above 0.7 indicate good historical consistency, while scores below 0.5 suggest significant anachronisms or factual errors.

2.4 Linguistic Quality Evaluation

While historical accuracy checks whether the model gets facts and period-appropriate terms right, linguistic quality measures how well the model writes - grammar, coherence, vocabulary diversity, sentence structure, and the presence of historical language patterns.

Standard metrics like BLEU and ROUGE don’t capture whether the text reads naturally or uses appropriate archaic forms. We built customized tools that assess these dimensions, implemented in 05_evaluation/comprehensive_evaluator.py as shown in Listing 3:

To make this easier to read, it helps to view the code as a scoring scaffold rather than a complete NLP system. Each check_* function is expected to return a normalized score in the range [0, 1] (higher is better), and assess_linguistic_quality() simply averages those components so you can track one headline number over time.

This mirrors patterns from earlier in the series: in Part 2 we used lightweight, automatable checks to validate data quality, and in Part 3 we relied on simple, repeatable metrics to judge training health. Here, we do the same for generation quality: start with cheap checks that run everywhere, then iterate toward richer evaluators as needed.

Also note that the exact weights (0.3/0.2, etc.) are tunable. The main benefit is splitting “linguistic quality” into components you can inspect individually, so when output is bad, you can tell why (grammar-ish structure vs coherence vs vocabulary vs historically flavored patterns).

def assess_linguistic_quality(generated_text):
    """Assess the linguistic quality of generated historical text"""
    
    quality_score = 0.0
    total_checks = 0
    
    # Check grammatical correctness
    grammar_score = check_grammatical_correctness(generated_text)
    quality_score += grammar_score
    total_checks += 1
    
    # Check coherence and flow
    coherence_score = check_text_coherence(generated_text)
    quality_score += coherence_score
    total_checks += 1
    
    # Check vocabulary appropriateness
    vocab_score = check_vocabulary_appropriateness(generated_text)
    quality_score += vocab_score
    total_checks += 1
    
    # Check sentence structure variety
    structure_score = check_sentence_structure_variety(generated_text)
    quality_score += structure_score
    total_checks += 1
    
    # Check historical language patterns
    pattern_score = check_historical_language_patterns(generated_text)
    quality_score += pattern_score
    total_checks += 1
    
    return quality_score / total_checks

def check_grammatical_correctness(text):
    """Check grammatical correctness of generated text"""
    
    # Parse text into sentences
    sentences = nltk.sent_tokenize(text)
    
    if not sentences:
        return 0.0
    
    correct_sentences = 0
    
    for sentence in sentences:
        # Check for basic grammatical patterns
        if check_sentence_grammar(sentence):
            correct_sentences += 1
    
    return correct_sentences / len(sentences)

def check_historical_language_patterns(text):
    """Check if text follows appropriate historical language patterns"""
    
    score = 0.0
    total_patterns = 0
    
    # Check for appropriate use of historical verb forms
    historical_verbs = ['hath', 'doth', 'dost', 'art', 'wilt', 'shalt']
    verb_score = 0
    for verb in historical_verbs:
        if verb in text.lower():
            verb_score += 1
    if historical_verbs:
        score += 0.3 * (verb_score / len(historical_verbs))
    total_patterns += 1
    
    # Check for appropriate use of historical pronouns
    historical_pronouns = ['thou', 'thee', 'thy', 'thine', 'ye']
    pronoun_score = 0
    for pronoun in historical_pronouns:
        if pronoun in text.lower():
            pronoun_score += 1
    if historical_pronouns:
        score += 0.3 * (pronoun_score / len(historical_pronouns))
    total_patterns += 1
    
    # Check for appropriate use of historical conjunctions
    historical_conjunctions = ['whilst', 'betwixt', 'amongst', 'ere', 'anon']
    conj_score = 0
    for conj in historical_conjunctions:
        if conj in text.lower():
            conj_score += 1
    if historical_conjunctions:
        score += 0.2 * (conj_score / len(historical_conjunctions))
    total_patterns += 1
    
    # Check for appropriate use of historical interjections
    historical_interjections = ['verily', 'indeed', 'forsooth', 'prithee', 'marry']
    interj_score = 0
    for interj in historical_interjections:
        if interj in text.lower():
            interj_score += 1
    if historical_interjections:
        score += 0.2 * (interj_score / len(historical_interjections))
    total_patterns += 1
    
    return score / total_patterns if total_patterns > 0 else 0.0

Listing 3: Linguistic Quality Evaluation

About NLTK: We use NLTK (Natural Language Toolkit), a standard Python library for natural language processing, to handle text tokenization. If you followed Part 1 ’s setup instructions, NLTK was already installed as part of the data processing dependencies. In check_grammatical_correctness(), we use nltk.sent_tokenize() to split text into sentences so we can evaluate grammar sentence-by-sentence. NLTK also provides word tokenization (word_tokenize) and BLEU score calculation (sentence_bleu), which are used elsewhere in the evaluation pipeline.

We chose NLTK because it’s well-established, handles edge cases (like abbreviations and historical punctuation), and provides reliable sentence boundaries even with archaic English patterns. The same qualities made it useful during data collection and cleaning (covered in Part 2 ).

The historical language patterns we check (verbs like hath, doth, pronouns like thou, thee, conjunctions like whilst, betwixt, and interjections like verily, forsooth) are the same archaic forms we identified during corpus analysis for temporal consistency. The difference here is that we’re measuring their presence as a positive signal of historical authenticity, rather than using them as required/forbidden constraints. Each pattern category (verbs, pronouns, conjunctions, interjections) contributes proportionally to the score based on how many patterns from that category appear in the text.

How the scoring works

The assess_linguistic_quality() function averages five component scores (grammar, coherence, vocabulary appropriateness, sentence structure variety, and historical language patterns) to produce a single score between 0 and 1. Each component is evaluated independently and returns a score in the range [0, 1].

For example, check_grammatical_correctness() counts the proportion of grammatically correct sentences, while check_historical_language_patterns() weights the presence of archaic verb forms (30%), pronouns (30%), conjunctions (20%), and interjections (20%) to produce a pattern score. The final linguistic quality score is the simple average of all five components. In practice, scores above 0.75 indicate strong linguistic quality with good grammar and historical flavor, while scores below 0.6 suggest the model struggles with either basic grammar or historical language patterns.

2.5 Running Evaluations

You can run the evaluators directly from the command line. The framework defaults to CPU for safety (so you can evaluate during training without GPU conflicts), but you can use --device gpu when the GPU is free for faster evaluation.

Quick example:

# Quick evaluation (runs in minutes, no external APIs)
python 05_evaluation/run_evaluation.py --mode quick --device cpu

# Comprehensive evaluation (includes benchmarks, optional G-Eval)
python 05_evaluation/run_evaluation.py --mode comprehensive --device cpu

The unified launcher (run_evaluation.py) supports multiple modes: setup (install dependencies), quick (fast validation), comprehensive (full suite with benchmarks), dataset (generate test cases), and all (complete evaluation). You can also call quick_eval.py or comprehensive_evaluator.py directly if you need more control.

Practical Evaluation Workflow:

Our typical evaluation workflow follows this pattern:

After Training: Run a quick evaluation to get immediate feedback on model performance
Before Publishing: Run a comprehensive evaluation to ensure the model meets quality standards
During Development: Use interactive testing to explore model behavior on specific prompts
For Research: Generate custom test datasets and run targeted evaluations

The framework defaults to CPU for safety (so you can evaluate during training without GPU conflicts), but you can use --device gpu when the GPU is free for faster evaluation. This design allows continuous assessment throughout the training process without interfering with GPU resources needed for training.

For complete usage examples, command-line options, and troubleshooting, see the Evaluation Guide in the repository.

3. Comprehensive Testing Pipeline

3.1 Automated Testing Framework

The 06_testing package contains a parallel set of tests that double-check the full system. Listing 4 captures the idea behind run_comprehensive_tests.

We group tests into basic functionality, historical accuracy, linguistic quality, performance, edge cases, and integration, then run them as a batch and emit a structured report. This mirrors how you would build a real CI test suite, but at a scale appropriate for this learning project.

def run_comprehensive_tests(model, tokenizer, device='cuda'):
    """Run comprehensive tests on historical language model"""
    
    test_results = {
        'basic_functionality': test_basic_functionality(model, tokenizer, device),
        'historical_accuracy': test_historical_accuracy(model, tokenizer, device),
        'linguistic_quality': test_linguistic_quality(model, tokenizer, device),
        'performance_metrics': test_performance_metrics(model, tokenizer, device),
        'edge_cases': test_edge_cases(model, tokenizer, device),
        'integration_tests': test_integration(model, tokenizer, device)
    }
    
    # Generate test report
    generate_test_report(test_results)
    
    return test_results

def test_basic_functionality(model, tokenizer, device):
    """Test basic model functionality"""
    
    tests = {
        'text_generation': test_text_generation(model, tokenizer, device),
        'tokenization': test_tokenization(tokenizer),
        'model_loading': test_model_loading(model, device),
        'memory_usage': test_memory_usage(model, device),
        'inference_speed': test_inference_speed(model, tokenizer, device)
    }
    
    return tests

def test_historical_accuracy(model, tokenizer, device):
    """Test historical accuracy of generated text"""
    
    tests = {
        'temporal_consistency': test_temporal_consistency(model, tokenizer, device),
        'factual_accuracy': test_factual_accuracy(model, tokenizer, device),
        'period_appropriate_language': test_period_language(model, tokenizer, device),
        'geographical_accuracy': test_geographical_accuracy(model, tokenizer, device),
        'social_context_accuracy': test_social_context(model, tokenizer, device)
    }
    
    return tests

Listing 4: Comprehensive Testing Framework

Automated tests cover basics, historical accuracy, linguistic quality, performance, edge cases, and integration.

3.2 Interactive Testing and Validation

For manual exploration, the interactive testing interface (conceptually similar to the CLI flows in 06_inference/inference_unified.py) lets you type prompts, trigger specific test groups, and immediately inspect analysis for each generation. Listing 5 shows a simple REPL loop that dispatches to the same evaluation helpers used in the automated tests.

def interactive_testing(model, tokenizer, device='cuda'):
    """Interactive testing interface for historical language model"""
    
    print("Interactive Testing Mode")
    print("=" * 50)
    print("Enter prompts to test the model. Type 'quit' to exit.")
    print("Available commands:")
    print("  - Enter any text prompt to generate continuation")
    print("  - 'test_historical' - Run historical accuracy tests")
    print("  - 'test_linguistic' - Run linguistic quality tests")
    print("  - 'test_performance' - Run performance tests")
    print("  - 'quit' - Exit testing mode")
    print()
    
    while True:
        try:
            prompt = input("Enter prompt: ").strip()
            
            if prompt.lower() == 'quit':
                break
            elif prompt.lower() == 'test_historical':
                run_historical_tests(model, tokenizer, device)
            elif prompt.lower() == 'test_linguistic':
                run_linguistic_tests(model, tokenizer, device)
            elif prompt.lower() == 'test_performance':
                run_performance_tests(model, tokenizer, device)
            elif prompt:
                # Generate text
                generated = generate_text(model, tokenizer, prompt, device)
                print(f"Generated: {generated}")
                print()
                
                # Analyze generated text
                analysis = analyze_generated_text(generated, prompt)
                print(f"Analysis: {analysis}")
                print()
            else:
                print("Please enter a valid prompt or command.")
                
        except KeyboardInterrupt:
            print("\nExiting interactive testing mode...")
            break
        except Exception as e:
            print(f"Error: {e}")
            print("Please try again.")

def analyze_generated_text(text, prompt):
    """Analyze generated text for quality and accuracy"""
    
    analysis = {
        'length': len(text),
        'sentences': len(nltk.sent_tokenize(text)),
        'historical_accuracy': assess_historical_accuracy(text, {}),
        'linguistic_quality': assess_linguistic_quality(text),
        'coherence': assess_coherence(text),
        'relevance': assess_relevance(text, prompt)
    }
    
    return analysis

Listing 5: Interactive Testing Interface

Interactive mode lets you try prompts, run quick tests, and see immediate analysis.

3.3 Performance Benchmarking

Performance benchmarking follows the same pattern: generate controlled workloads and measure speed and resource usage. Listing 6 illustrates how we vary sequence length, measure average latency, and compute tokens-per-second, alongside separate helpers for memory, batch throughput, long-sequence handling, and basic concurrency.

def benchmark_model_performance(model, tokenizer, device='cuda'):
    """Benchmark model performance across different scenarios"""
    
    benchmarks = {
        'inference_speed': benchmark_inference_speed(model, tokenizer, device),
        'memory_usage': benchmark_memory_usage(model, device),
        'batch_processing': benchmark_batch_processing(model, tokenizer, device),
        'long_sequence_handling': benchmark_long_sequences(model, tokenizer, device),
        'concurrent_requests': benchmark_concurrent_requests(model, tokenizer, device)
    }
    
    return benchmarks

def benchmark_inference_speed(model, tokenizer, device):
    """Benchmark inference speed for different sequence lengths"""
    
    sequence_lengths = [50, 100, 200, 500, 1000]
    results = {}
    
    for length in sequence_lengths:
        # Generate test prompts of different lengths
        prompts = generate_test_prompts(length, num_prompts=100)
        
        # Measure inference time
        start_time = time.time()
        for prompt in prompts:
            generate_text(model, tokenizer, prompt, device)
        end_time = time.time()
        
        total_time = end_time - start_time
        avg_time_per_prompt = total_time / len(prompts)
        tokens_per_second = length / avg_time_per_prompt
        
        results[length] = {
            'avg_time_per_prompt': avg_time_per_prompt,
            'tokens_per_second': tokens_per_second,
            'total_time': total_time
        }
    
    return results

Listing 6: Performance Benchmarking

Benchmarks capture inference speed, memory, batch throughput, long-sequence handling, and simple concurrency.

4. Model Deployment and Publishing

With evaluation and testing complete, we’re ready to make our models available for use. This section covers the two deployment paths we support: direct inference from PyTorch checkpoints (useful during development and for maximum control) and publishing to Hugging Face Hub (for easy sharing and community access).

As called out in Part 1 , both the SLM (117M parameters) and the Regular Model (354M parameters) are fully trained and available. The SLM has already been published on Hugging Face Hub , while the Regular Model is ready for publication. Both can also be run directly from local PyTorch checkpoints.

4.1 Two Paths to Inference

We provide two complementary ways to run inference, each suited to different use cases.

PyTorch Checkpoint Inference gives you direct access to the trained model weights without any conversion overhead. This is ideal during development, when you want to test a freshly trained checkpoint, or when you need maximum control over the inference process. The checkpoints live in 09_models/checkpoints/ - the SLM at slm/checkpoint-4000.pt (117M parameters) and the Regular Model at checkpoint-60001.pt (354M parameters). The inference_pytorch.py script handles loading these directly: Listing 7

# SLM inference from checkpoint
python 06_inference/inference_pytorch.py \
  --checkpoint 09_models/checkpoints/slm/checkpoint-4000.pt \
  --prompt "In the year 1834, I walked through the streets of London and witnessed"

# Regular model inference from checkpoint
python 06_inference/inference_pytorch.py \
  --checkpoint 09_models/checkpoints/checkpoint-60001.pt \
  --prompt "In the year 1834, I walked through the streets of London and witnessed"

Listing 7: Running Inference from PyTorch Checkpoints

Hugging Face Model Inference uses the published models on Hugging Face Hub, which means anyone can load and use them with just a few lines of code - no need to download checkpoints or set up the full training environment. The inference_unified.py script provides a consistent interface for both published models and local checkpoints: Listing 8

# Published model inference (downloads from Hugging Face Hub)
python 06_inference/inference_unified.py \
  --published \
  --model_name bahree/london-historical-slm \
  --prompt "In the year 1834, I walked through the streets of London and witnessed"

# Interactive mode for exploration
python 06_inference/inference_unified.py --published --model_type slm --interactive

# Demo mode with curated historical prompts
python 06_inference/inference_unified.py --published --model_type slm --demo

Listing 8: Hugging Face Model Inference

We’ve tested both paths extensively. The published SLM loads in about 9 seconds on a GPU, generates text in under 6 seconds, and passes all 10 automated validation tests. The unified inference script provides clean logging, proper model detection, and accurate parameter counts - small details that make a big difference when debugging or demonstrating the models.

4.2 Publishing to Hugging Face Hub

Publishing to Hugging Face Hub makes our models accessible to the broader community without requiring anyone to clone our repository or set up a training environment. The process involves converting our PyTorch checkpoints to the Hugging Face format, creating a model card with documentation, and uploading everything to the Hub.

The publishing workflow is handled by scripts in 10_scripts/ - specifically publish_slm_to_huggingface.py for the SLM and publish_to_huggingface.py for the Regular Model. Listing 9 shows the core publishing flow: authenticate with the Hub, create (or reuse) a repository, save the model and tokenizer locally in Hugging Face format, upload the folder, and generate a model card.

def publish_to_huggingface(model, tokenizer, model_name, description, tags):
    """Publish model to Hugging Face Hub"""
    from huggingface_hub import HfApi
    
    api = HfApi()
    
    # Create model repository
    repo_id = f"bahree/{model_name}"
    api.create_repo(repo_id=repo_id, exist_ok=True)
    
    # Save model and tokenizer locally
    model.save_pretrained(f"./models/{model_name}")
    tokenizer.save_pretrained(f"./models/{model_name}")
    
    # Upload to Hub
    api.upload_folder(
        folder_path=f"./models/{model_name}",
        repo_id=repo_id,
        commit_message="Initial model upload"
    )
    
    # Generate and upload model card (README.md)
    model_card = generate_model_card(model_name, description, tags)
    api.upload_file(
        path_or_fileobj=model_card,
        path_in_repo="README.md",
        repo_id=repo_id
    )
    
    return repo_id

Listing 9: Hugging Face Publishing

The generate_model_card() function creates the README.md that appears on the Hugging Face model page. This includes model description, architecture details, training data sources, usage examples, and limitations. You can see the live model cards at bahree/london-historical-slm and bahree/london-historical-llm .

Listing 10 shows how to load and use the published models:

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load the published model
model_name = "bahree/london-historical-slm"  # or "bahree/london-historical-llm"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Move to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

# Generate historical text
prompt = "In the year 1834, I walked through the streets of London and witnessed"
inputs = tokenizer(prompt, return_tensors="pt").to(device)

outputs = model.generate(
    inputs["input_ids"],
    max_new_tokens=50,
    do_sample=True,
    temperature=0.8,
    top_p=0.95,
    repetition_penalty=1.2,
    pad_token_id=tokenizer.eos_token_id
)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Listing 10: Loading Models from Hugging Face Hub

4.3 Publishing Workflow

If you want to publish your own trained model to Hugging Face Hub, here’s the workflow we followed:

Set up authentication: Install huggingface_hub and authenticate with a token that has Write permissions. You can generate tokens at huggingface.co/settings/tokens .
Convert the checkpoint: PyTorch training checkpoints include optimizer states and training metadata that aren’t needed for inference. The conversion scripts extract just the model weights and translate them to Hugging Face’s naming conventions (covered in detail in Section 5).
Prepare the tokenizer: Save the tokenizer files alongside the model. Our custom tokenizer with 30,000 tokens and 150+ historical special tokens needs to be converted to the transformers library format.
Generate a model card: The README.md on your Hugging Face model page serves as documentation. Include model architecture details, training data sources, usage examples, evaluation results, and limitations. The scripts generate this automatically, but you should review and customize it.
Upload and validate: Push everything to the Hub, then immediately test with from_pretrained() to ensure the published model loads and generates correctly.

📝 Full documentation: See 08_documentation/HUGGINGFACE_PUBLISHING.md and 08_documentation/DEPLOYMENT_GUIDE.md in the repository for the complete step-by-step workflow with troubleshooting guidance.

5. PyTorch to Hugging Face Format Conversion

5.1 Why Format Conversion is Necessary

During training, our models are saved in PyTorch’s native .pt format. These checkpoints include everything needed to resume training: model weights, optimizer states, learning rate schedules, and training metadata. However, for deployment and sharing, we need a leaner, inference-optimized format compatible with the broader machine learning ecosystem.

Think of it like the difference between a development environment and a production deployment: training checkpoints are like a developer’s workspace with all the tools and intermediate files, while Hugging Face format is like a clean, standardized package that anyone can use without understanding the internal training details.

The Hugging Face Hub expects models to follow specific file structures, naming conventions, and metadata requirements. The conversion process extracts just the model weights (discarding optimizer states and training metadata), translates weight names to match Hugging Face conventions, creates proper configuration files, and ensures the tokenizer is compatible with the transformers library.

5.2 The Conversion Process

The conversion handles several transformations to bridge PyTorch and Hugging Face formats:

Weight name mapping: PyTorch layer names like transformer.h.0.attn.c_attn.weight become Hugging Face names like transformer.h.0.attn.c_attn.weight (mostly the same for GPT-2, but with careful handling of edge cases)
Automatic torch.compile handling: If you used torch.compile() during training, weights get prefixed with _orig_mod. - the conversion strips these prefixes
Configuration translation: Model hyperparameters (n_layer, n_head, n_embd, etc.) are mapped to Hugging Face’s config.json format
Tokenizer conversion: Our custom 30,000-token vocabulary with 150+ historical special tokens is converted to transformers library format
Validation: After conversion, we verify that the model loads correctly and produces expected outputs

💻 Full Implementation: See 10_scripts/publish_slm_to_huggingface.py for the complete conversion pipeline with error handling, validation, and model card generation.

5.3 Dependencies for Hugging Face Integration

The Hugging Face integration requires specific dependencies and follows established patterns for model publishing and usage: Listing 11

# Required dependencies for Hugging Face integration
huggingface_dependencies = {
    "transformers": ">=4.21.0",
    "torch": ">=1.12.0",
    "tokenizers": ">=0.12.0",
    "safetensors": ">=0.3.0",
    "accelerate": ">=0.20.0",
    "huggingface_hub": ">=0.10.0"
}

# Model loading and usage example
def load_published_model(model_name="bahree/london-historical-slm"):
    """Load published model from Hugging Face Hub"""
    
    # Suppress warnings for cleaner output
    import os
    import warnings
    import logging
    
    os.environ['TRANSFORMERS_VERBOSITY'] = 'error'
    warnings.filterwarnings('ignore')
    logging.getLogger("transformers").setLevel(logging.ERROR)
    os.environ['TOKENIZERS_PARALLELISM'] = 'false'
    
    # Load model and tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    
    # Set pad token if not set
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id
    
    return model, tokenizer

def generate_historical_text(model, tokenizer, prompt, max_length=50, temperature=0.3):
    """Generate historical text using the published model"""
    
    # Tokenize input
    inputs = tokenizer.encode(prompt, return_tensors="pt")
    
    # Generate text
    with torch.no_grad():
        outputs = model.generate(
            inputs,
            max_new_tokens=max_length,
            do_sample=True,
            temperature=temperature,
            top_p=0.9,
            top_k=20,
            repetition_penalty=1.2,
            no_repeat_ngram_size=3,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            early_stopping=True
        )
    
    # Decode output
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_text

# Example usage
if __name__ == "__main__":
    # Load the published model
    model, tokenizer = load_published_model()
    
    # Test prompts
    test_prompts = [
        "In the year 1834, I walked through the streets of London and witnessed",
        "The gentleman from the country said, 'we have never seen such a sight",
        "The Thames flowed dark and mysterious through the heart",
        "Parliament sat in Westminster Hall",
        "The Great Fire of 1666 had destroyed"
    ]
    
    # Generate text for each prompt
    for prompt in test_prompts:
        generated = generate_historical_text(model, tokenizer, prompt)
        print(f"Prompt: {prompt}")
        print(f"Generated: {generated}")
        print("-" * 80)

Listing 11: Hugging Face Dependencies

Hugging Face integration provides standard from_pretrained() loading and generation with minimal setup, making the models easy to share and reuse.

5.4 Comprehensive Testing and Validation Framework

Once a model is on the Hub, 06_inference/test_published_models.py provides a concrete implementation of the testing pattern in Listing 12. It loads the model via from_pretrained, runs functional, historical, linguistic, and performance checks, and prints a human-readable summary so you can verify the published artefact behaves like your local checkpoints.

def test_published_model(model_name="bahree/london-historical-slm"):
    """Comprehensive testing of published model"""
    
    print(f"Testing published model: {model_name}")
    
    # Load model
    model, tokenizer = load_published_model(model_name)
    
    # Test basic functionality
    basic_tests = test_basic_functionality(model, tokenizer)
    
    # Test historical accuracy
    historical_tests = test_historical_accuracy(model, tokenizer)
    
    # Test linguistic quality
    linguistic_tests = test_linguistic_quality(model, tokenizer)
    
    # Test performance metrics
    performance_tests = test_performance_metrics(model, tokenizer)
    
    # Compile results
    results = {
        "basic_functionality": basic_tests,
        "historical_accuracy": historical_tests,
        "linguistic_quality": linguistic_tests,
        "performance_metrics": performance_tests
    }
    
    # Print summary
    print("\nTest Results Summary:")
    print("=" * 50)
    for category, tests in results.items():
        print(f"\n{category.replace('_', ' ').title()}:")
        for test_name, result in tests.items():
            status = "PASS" if result else "FAIL"
            print(f"  {test_name}: {status}")
    
    return results

def test_basic_functionality(model, tokenizer):
    """Test basic model functionality"""
    
    tests = {}
    
    # Test model loading
    tests["model_loading"] = model is not None and tokenizer is not None
    
    # Test tokenizer functionality
    test_text = "In the year 1834, London was"
    tokens = tokenizer.encode(test_text)
    decoded = tokenizer.decode(tokens)
    tests["tokenizer_encode_decode"] = test_text in decoded
    
    # Test model generation
    try:
        inputs = tokenizer.encode(test_text, return_tensors="pt")
        with torch.no_grad():
            outputs = model.generate(inputs, max_new_tokens=10, do_sample=False)
        generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
        tests["model_generation"] = len(generated) > len(test_text)
    except Exception as e:
        tests["model_generation"] = False
    
    # Test special tokens
    special_tokens = ["<|london|>", "<|thou|>", "<|hath|>", "<|doth|>"]
    special_token_tests = []
    for token in special_tokens:
        if token in tokenizer.get_vocab():
            special_token_tests.append(True)
        else:
            special_token_tests.append(False)
    tests["special_tokens"] = any(special_token_tests)
    
    return tests

def test_historical_accuracy(model, tokenizer):
    """Test historical accuracy of generated text"""
    
    tests = {}
    
    # Test prompts for different historical periods
    period_prompts = {
        "1500-1600": "In the year 1550, the gentleman said",
        "1600-1700": "In the year 1650, the gentleman said",
        "1700-1800": "In the year 1750, the gentleman said",
        "1800-1850": "In the year 1834, the gentleman said"
    }
    
    for period, prompt in period_prompts.items():
        try:
            generated = generate_historical_text(model, tokenizer, prompt, max_length=30)
            
            # Check for period-appropriate language
            period_terms = {
                "1500-1600": ["ye", "hath", "doth", "thou", "thee"],
                "1600-1700": ["hath", "doth", "thou", "thee", "verily"],
                "1700-1800": ["hath", "doth", "thou", "thee", "indeed"],
                "1800-1850": ["indeed", "verily", "whilst", "pray"]
            }
            
            found_terms = sum(1 for term in period_terms[period] if term in generated.lower())
            tests[f"period_{period}"] = found_terms > 0
            
        except Exception as e:
            tests[f"period_{period}"] = False
    
    # Test London-specific knowledge
    london_prompts = [
        "The Thames flowed through",
        "Westminster Hall was",
        "The Tower of London",
        "Cheapside was filled with"
    ]
    
    london_tests = []
    for prompt in london_prompts:
        try:
            generated = generate_historical_text(model, tokenizer, prompt, max_length=20)
            london_tests.append(len(generated) > len(prompt))
        except Exception as e:
            london_tests.append(False)
    
    tests["london_knowledge"] = any(london_tests)
    
    return tests

def test_linguistic_quality(model, tokenizer):
    """Test linguistic quality of generated text"""
    
    tests = {}
    
    # Test prompts for linguistic quality
    quality_prompts = [
        "The gentleman walked through the garden",
        "In the morning, the sun rose",
        "The old man sat by the fire",
        "The young woman read her book"
    ]
    
    quality_tests = []
    for prompt in quality_prompts:
        try:
            generated = generate_historical_text(model, tokenizer, prompt, max_length=30)
            
            # Check for basic linguistic quality
            sentences = generated.split('.')
            quality_tests.append(len(sentences) > 1)
            
        except Exception as e:
            quality_tests.append(False)
    
    tests["linguistic_quality"] = any(quality_tests)
    
    # Test coherence
    coherence_prompt = "The gentleman walked through the garden and"
    try:
        generated = generate_historical_text(model, tokenizer, coherence_prompt, max_length=50)
        tests["coherence"] = len(generated) > len(coherence_prompt)
    except Exception as e:
        tests["coherence"] = False
    
    return tests

def test_performance_metrics(model, tokenizer):
    """Test performance metrics of the model"""
    
    tests = {}
    
    # Test inference speed
    test_prompt = "In the year 1834, London was"
    try:
        start_time = time.time()
        generated = generate_historical_text(model, tokenizer, test_prompt, max_length=50)
        end_time = time.time()
        
        inference_time = end_time - start_time
        tests["inference_speed"] = inference_time < 5.0  # Should complete within 5 seconds
        
    except Exception as e:
        tests["inference_speed"] = False
    
    # Test memory usage
    try:
        import psutil
        process = psutil.Process()
        memory_usage = process.memory_info().rss / 1024 / 1024  # MB
        tests["memory_usage"] = memory_usage < 8000  # Should use less than 8GB
        
    except Exception as e:
        tests["memory_usage"] = True  # Skip if psutil not available
    
    return tests

Listing 12: Testing Published Models

Published model tests validate loading, generation, historical accuracy, and basic performance before and after publication.

5.5 Model Card Generation

The model card serves as the primary documentation on Hugging Face Hub, making it the first thing users see when they discover your model. A well-crafted model card helps users understand what the model does, how to use it, and its limitations. The generate_comprehensive_model_card() function in 10_scripts/publish_slm_to_huggingface.py creates this documentation automatically.

What Makes an Effective Model Card:

The model card for our historical language models includes several key sections that provide users with everything they need to get started. At a minimum, include:

Model Description & Key Features: A clear explanation that the model was trained from scratch (not fine-tuned), emphasizing the 117M parameter SLM variant and 354M parameter Regular Model, with details about the custom 30,000-token vocabulary and 150+ historical special tokens.
Setup Instructions: Platform-specific guidance for creating virtual environments (Linux/macOS/Windows), installing dependencies (transformers, torch, accelerate), and handling different accelerators (CPU, NVIDIA CUDA, AMD ROCm).
Quick Start Code: Auto-device detection that works across CPU, CUDA, and ROCm with sensible generation parameters (temperature=0.8, top_p=0.95, repetition_penalty=1.2).
Training Details: Architecture specifics (GPT-2 Small/Medium), training infrastructure (2x GPU with Distributed Data Parallel), performance metrics (training loss, MFU utilization), and data sources (218+ historical sources spanning 1500-1850).
Example Prompts: Period-specific prompts demonstrating different historical eras (Tudor, Stuart, Georgian, Victorian) and London-specific contexts (Thames, Westminster, Parliament).
Testing & Validation: Instructions for running the automated test suite (test_published_models.py) and interactive testing with custom prompts.
Troubleshooting: Common issues and solutions for PyTorch installation, GPU detection, and memory constraints.
Citation & License: BibTeX citation format and MIT license information.

Key Implementation Details:

The model card generation follows Hugging Face conventions with YAML frontmatter specifying license, library, pipeline type, language, and tags. The script emphasizes that models were trained from scratch (not fine-tuned) and provides device-agnostic code examples that run on CPU, CUDA, and ROCm.

The card also includes detailed model selection guidance comparing the SLM (faster, lower memory) versus the Regular Model (higher quality, more parameters), helping users choose the right model for their use case - whether that’s quick experimentation, educational purposes, or production deployment.

💻 Complete Implementation: See 10_scripts/publish_slm_to_huggingface.py and 10_scripts/publish_to_huggingface.py for the full model card generation implementation.

👀 Live Model Cards: View the published cards at bahree/london-historical-slm and bahree/london-historical-llm .

📝 Documentation: See HUGGINGFACE_PUBLISHING.md and DEPLOYMENT_GUIDE.md for complete publishing and deployment workflows.

5.6 Local Deployment Options

Finally, Listing 13 sketches how you might wrap a trained model into a simple REST API or CLI. These patterns are intentionally minimal, meant to help you connect the dots between the inference utilities in 06_inference/ and real applications (dashboards, notebooks, small services).

def setup_local_deployment(model, tokenizer, deployment_type='api'):
    """Set up local deployment for historical language model"""
    
    if deployment_type == 'api':
        return setup_api_deployment(model, tokenizer)
    elif deployment_type == 'cli':
        return setup_cli_deployment(model, tokenizer)
    elif deployment_type == 'notebook':
        return setup_notebook_deployment(model, tokenizer)
    else:
        raise ValueError(f"Unknown deployment type: {deployment_type}")

def setup_api_deployment(model, tokenizer):
    """Set up REST API deployment"""
    
    from flask import Flask, request, jsonify
    import torch
    
    app = Flask(__name__)
    
    @app.route('/generate', methods=['POST'])
    def generate_text():
        data = request.get_json()
        prompt = data.get('prompt', '')
        max_length = data.get('max_length', 100)
        temperature = data.get('temperature', 0.3)
        
        # Generate text
        inputs = tokenizer.encode(prompt, return_tensors="pt")
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_length=max_length,
                temperature=temperature,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        return jsonify({
            'generated_text': generated_text,
            'prompt': prompt,
            'parameters': {
                'max_length': max_length,
                'temperature': temperature
            }
        })
    
    @app.route('/health', methods=['GET'])
    def health_check():
        return jsonify({'status': 'healthy', 'model_loaded': True})
    
    return app

def setup_cli_deployment(model, tokenizer):
    """Set up command-line interface deployment"""
    
    import argparse
    
    def main():
        parser = argparse.ArgumentParser(description='Historical Language Model CLI')
        parser.add_argument('--prompt', type=str, required=True, help='Text prompt')
        parser.add_argument('--max_length', type=int, default=100, help='Maximum length')
        parser.add_argument('--temperature', type=float, default=0.3, help='Temperature')
        parser.add_argument('--interactive', action='store_true', help='Interactive mode')
        
        args = parser.parse_args()
        
        if args.interactive:
            run_interactive_mode(model, tokenizer)
        else:
            generate_and_print(model, tokenizer, args.prompt, args.max_length, args.temperature)
    
    return main

Listing 13: Local Deployment Setup

Local deployment options: REST API, CLI, or notebook integration for different workflows.

6. Quality Assurance and Validation

Before wrapping up, let’s look at the quality assurance systems that ensure the models behave reliably across different scenarios.

6.1 Automated Quality Checks

The system includes automated quality checks that validate model performance and reliability: Listing 14

def run_quality_checks(model, tokenizer, device='cuda'):
    """Run quality checks on historical language model"""
    
    quality_checks = {
        'model_integrity': check_model_integrity(model),
        'tokenizer_consistency': check_tokenizer_consistency(tokenizer),
        'generation_quality': check_generation_quality(model, tokenizer, device),
        'historical_accuracy': check_historical_accuracy(model, tokenizer, device),
        'performance_metrics': check_performance_metrics(model, tokenizer, device),
        'memory_usage': check_memory_usage(model, device),
        'error_handling': check_error_handling(model, tokenizer, device)
    }
    
    # Generate quality report
    quality_report = generate_quality_report(quality_checks)
    
    return quality_checks, quality_report

def check_model_integrity(model):
    """Check model integrity and consistency"""
    
    checks = {
        'parameter_count': check_parameter_count(model),
        'weight_distribution': check_weight_distribution(model),
        'gradient_flow': check_gradient_flow(model),
        'activation_patterns': check_activation_patterns(model)
    }
    
    return checks

def check_generation_quality(model, tokenizer, device):
    """Check quality of generated text"""
    
    test_prompts = [
        "In the year of our Lord 1750, London was",
        "The Thames flowed through the heart of",
        "Merchants and tradesmen plied their wares",
        "The Great Fire of 1666 had changed",
        "Parliament sat in Westminster, making laws"
    ]
    
    quality_scores = []
    
    for prompt in test_prompts:
        # Generate text
        generated = generate_text(model, tokenizer, prompt, device)
        
        # Check quality metrics
        quality_score = {
            'coherence': assess_coherence(generated),
            'grammatical_correctness': assess_grammatical_correctness(generated),
            'historical_accuracy': assess_historical_accuracy(generated, {}),
            'linguistic_quality': assess_linguistic_quality(generated),
            'relevance': assess_relevance(generated, prompt)
        }
        
        quality_scores.append(quality_score)
    
    return quality_scores

Listing 14: Quality Assurance Checks

Quality checks cover model integrity, generation quality, historical accuracy, performance, and error handling, ensuring the models behave reliably across different scenarios.

6.2 Continuous Integration and Testing

If you want to wire this into a lightweight CI gate, keep it simple and CPU-friendly. The goal is not to re-run full benchmarks in CI - it’s to catch obvious regressions (can the model load, can it generate, do the evaluators still run).

Minimal CI smoke checks (suggested):

# 1) Run a fast, local evaluation pass (no external APIs)
python 05_evaluation/run_evaluation.py --mode quick --device cpu

# 2) Run a local inference smoke test from a checkpoint (replace with your path)
python 06_inference/inference_pytorch.py --checkpoint <path-to-checkpoint.pt> --prompt "In the year 1834, London was"

# 3) Optional: test the published model (requires downloading from Hugging Face)
python 06_inference/test_published_models.py --model_name bahree/london-historical-slm

7. Summary

We’ve now completed the full cycle of building language models from scratch. This final part has shown how to transform trained models into working systems that can be evaluated, tested, and deployed for real-world use. The journey that began in Part 1 with using published models, continued through Part 2 ’s data collection and tokenization, and Part 3 ’s training architecture, now concludes with evaluation and deployment - the critical final steps that make models usable.

What we’ve built:

The evaluation, testing, and deployment pipeline provides a practical approach for bringing historical language models from research to deployment. We’ve created specialized assessment metrics that go beyond standard LLM evaluation to catch historical inaccuracies, temporal inconsistencies, and period-inappropriate language. The testing infrastructure ensures reliability across different scenarios, while multiple deployment options make the models accessible to researchers, educators, and developers worldwide.

Current Deployment Status:

PyTorch Checkpoint Inference: Fully working with both SLM and Regular models
Hugging Face Model Inference: SLM published and available, Regular model ready
Local Testing: Both inference methods tested and validated on a remote Ubuntu machine
Documentation: Complete guides and examples for all inference methods
Performance: Clean logging, proper model detection, accurate parameter counts

The Complete Pipeline:

This four-part series has demonstrated the complete LLM development lifecycle:

Data Collection ( Part 2 ): We gathered 218+ historical sources spanning 1500-1850, processed them through a sophisticated cleaning pipeline, and created a 500M+ character corpus of authentic historical English.
Custom Tokenization ( Part 2 ): We built a specialized BPE tokenizer with 30,000 vocabulary tokens and 150+ special tokens that understand historical language patterns, London geography, and period-specific terminology.
Model Training ( Part 3 ): We implemented custom GPT architectures, optimized for multi-GPU training, and successfully trained two models - an SLM (117M parameters) and a Regular model (354M parameters) - both capable of generating authentic historical text.
Evaluation & Deployment (This Part): We built comprehensive evaluation frameworks that assess historical accuracy, linguistic quality, and temporal consistency. We created a testing infrastructure for reliability and deployed models to the Hugging Face Hub for community access.

The Learning Journey:

What started as a learning project has become a complete, working system that demonstrates every aspect of LLM development - from raw data collection through model deployment. The principles and techniques we’ve covered scale from the 500M-character corpus to production-scale systems, and the evaluation frameworks we’ve built can be adapted to any domain-specific language modeling task.

Whether you’re a researcher exploring historical linguistics, an educator teaching AI concepts, or a developer building specialized language models, this series provides the complete toolkit for understanding and implementing LLM development from scratch. The models are published, the code is available, and the journey from data to deployment is complete.

🔗 GitHub Repository: github.com/bahree/helloLondon - Complete training infrastructure (04_training/), model architecture (config.py), and evaluation/deployment (05_evaluation/, 06_inference/, 10_scripts/)

🟥 Series Posts: Part 1 - Using the Published Historical Models | Part 2 - Data Collection & Custom Tokenizer | Part 3 - Training Architecture & GPU Optimization | Part 4 (this post)

🟧 Published Models: SLM Model | Regular Model - Ready-to-use historical language models on Hugging Face

📗 Book Reference: Generative AI in Action - For deeper understanding of core LLM concepts

8. Resources

If you want to reproduce the full pipeline (or adapt it to your own domain), these are the most useful starting points:

Public GitHub repo: github.com/bahree/helloLondon
Evaluation guide: https://github.com/bahree/helloLondon/blob/main/08_documentation/EVALUATION_GUIDE.md
Hugging Face publishing guide: https://github.com/bahree/helloLondon/blob/main/08_documentation/HUGGINGFACE_PUBLISHING.md
Deployment guide: https://github.com/bahree/helloLondon/blob/main/08_documentation/DEPLOYMENT_GUIDE.md
Published models: bahree/london-historical-slm | bahree/london-historical-llm

References

Vaswani et al. (2017) - Attention Is All You Need: https://arxiv.org/abs/1706.03762
Radford et al. (2019) - Language Models are Unsupervised Multitask Learners: https://www.semanticscholar.org/paper/Language-Models-are-Unsupervised-Multitask-Learners-Radford-Wu/9405cc0d6169988371b2755e573cc28650d14dfe
Lin (2004) - ROUGE: A Package for Automatic Evaluation of Summaries: https://aclanthology.org/W04-1013/
Papineni et al. (2002) - BLEU: A Method for Automatic Evaluation of Machine Translation: https://aclanthology.org/P02-1040/
Hendrycks et al. (2021) - Measuring Massive Multitask Language Understanding (MMLU): https://arxiv.org/abs/2009.03300
Zellers et al. (2019) - HellaSwag: Can a Machine Really Finish Your Sentence?: https://arxiv.org/abs/1905.07830
Liu et al. (2023) - G-Eval: NLG Evaluation using GPT-4 with Better Human Alignment: https://arxiv.org/abs/2303.16634

Acknowledgments

This project builds upon the excellent work of the open-source community. Special thanks to haykgrigo3’s TimeCapsuleLLM for the initial inspiration and framework for historical language model training, and to Andrej Karpathy’s nanoGPT for the foundational GPT architecture and training methodology. The project extends these foundations with specialized adaptations for historical text, including custom tokenizers, advanced data filtering, evaluation frameworks, and educational deployment infrastructure.

Vibe Coding

Sun, 23 Nov 2025 00:00:00 +0000

Deadlock by design: two vibes, two locks, zero unlocks.

In vibe coding, that’s not a bug - it is a feature. 😎

Here’s what happens when agreement-first engineering meets C++ and mutexes:

#include <iostream>
#include <mutex>
#include <stdexcept>

std::mutex mVibes, mProd;

void shipToProd(bool agree) {
    // We lock the vibes and production—because feelings 
    // and facts both need exclusive access.
    std::lock_guard<std::mutex> a(mVibes);
    std::lock_guard<std::mutex> b(mProd);

    if (agree) {
        std::cout << "Deploying to prod…\n";
        // Reality pushes back:
        throw std::runtime_error("DeadlockException: vibes vs reality");
    } else {
        std::cout << "Ignored tests.\n";
    }
}

int main() {
    try {
        shipToProd(true); // agreement-first engineering
    } catch (const std::exception& ex) {
        std::cout << "AI: You're absolutely right! (" << ex.what() << ")\n";
        std::cout << "/* lock(prod); lock(vibes); // Deadlock achieved; energy immaculate */\n";
        std::cout << "COMMIT vibes; ROLLBACK sanity;\n";
    }
    return 0;
}

Listing 1: Two mutexes walk into prod… and never come out

#GeekyJokes #AI #GenAI #VibeCoding

🏛️Building LLMs from Scratch - Part 3: Training Architecture & GPU Optimization

Sat, 01 Nov 2025 00:00:00 +0000

TL;DR

In this third part of our 4-part series on building language models from scratch, I explore the complete training infrastructure that transforms our clean historical data and custom tokenizer into working language models.

Part 1 How to build a Large Language Model from Scratch - covered using the published model
Part 2 Building LLMs from Scratch - Part 2: Data Collection & Custom Tokenizers - detailed data collection and custom tokenizer development.

Here, we build the complete training pipeline from a custom GPT architecture through deployment-ready checkpoints.

This post demonstrates how to design custom model architectures, optimize GPU utilization, and implement comprehensive training pipelines that transform our 500M+ character historical corpus into two working language models.

⚠️ Educational Purpose: This is a learning project designed to teach LLM development concepts. For production-scale LLMs, you’ll need significantly larger datasets, more sophisticated infrastructure, and additional considerations that are not covered in this post.

As outlined in Part 1 , both the SLM (117M parameters) and the regular Model (354M parameters) use the same training code and pipeline (04_training/train_model_slm.py and 04_training/train_model.py) with different configurations defined in config.py. The training infrastructure, GPU optimization, checkpointing, and WandB integration are identical - only the model architecture parameters differ.

🔗 GitHub Repository: github.com/bahree/helloLondon - Complete training infrastructure (04_training/), model architecture (config.py), and GPU configuration (08_documentation/GPU_TUNING.md)

🧱 Series Posts: Part 1 – Using the Published Historical Models | Part 2 – Data Collection & Custom Tokenizer | Part 3 (this post) | Part 4 – Evaluation & Deployment

🤗 Published Models: SLM Model | Regular Model - Ready-to-use historical language models on Hugging Face

📚 Book Reference: Generative AI in Action - For deeper understanding of core LLM concepts.

1. The Training Challenge: From Data to Working Models

Now that we have our clean historical corpus and custom tokenizer from Part 2 , we need to transform this data into working language models. This isn’t just about running training scripts – it’s about designing an architecture that can learn from historical text, optimizing for the unique patterns of 1500-1850 English, and building infrastructure to handle the computational demands of language model training.

The challenge with historical language modeling isn’t just having enough data - it’s having the right architecture and training process that can learn from the complex linguistic patterns in historical texts. Unlike modern text, historical English contains archaic vocabulary, period-specific terminology, and cultural references that require specialized attention mechanisms and training strategies.

1.1 High-Level Training Process Overview

The model training pipeline transforms our clean historical data and custom tokenizer into working language models through several key stages:

Model Architecture Design - involves a custom GPT implementation optimized for historical text patterns
GPU Configuration covers - multi-GPU training with precision optimization and memory management
Training Infrastructure - includes distributed training, checkpointing, and experiment tracking
Performance Optimization - encompasses mixed precision, compilation, and hardware-specific tuning
Model Validation - covers testing and evaluation of trained models.

Figure 1 below illustrates this complete training pipeline:

graph TD
    A[📚 Clean Historical Corpus<br/>500M+ characters] --> B[🔤 Custom Tokenizer<br/>30K vocab + 150+ special tokens]
    B --> C[🏗️ Model Architecture<br/>Custom GPT for Historical Text]
    
    C --> D[⚙️ GPU Configuration<br/>Multi-GPU + Precision Optimization]
    D --> D1[Mixed Precision<br/>bf16/fp16]
    D --> D2[Torch Compile<br/>JIT optimization]
    D --> D3[Memory Management<br/>Gradient checkpointing]
    
    D1 --> E[🏋️ Training Process<br/>60K iterations, checkpointing]
    D2 --> E
    D3 --> E
    
    E --> E1[SLM: 117M params<br/>7-8 hours training]
    E --> E2[Regular: 354M params<br/>28-32 hours training]
    
    E --> E3[WandB Integration<br/>Experiment tracking]
    E --> E4[Checkpointing<br/>Resume capability]
    E --> E5[Multi-GPU Support<br/>Distributed training]
    
    E1 --> F[📊 Model Evaluation<br/>Historical accuracy testing]
    E2 --> F
    E3 --> F
    E4 --> F
    E5 --> F
    
    F --> G{Quality OK?}
    G -->|Yes| H[🚀 Deployment<br/>Hugging Face + Local Inference]
    G -->|No| I[🔄 Retrain/Adjust]
    I --> E
    H --> J[💬 Text Generation<br/>Historical language output]
    
    style A fill:#e1f5fe
    style C fill:#f3e5f5
    style H fill:#e8f5e8
    style J fill:#fff3e0

Figure 1: Complete Training Pipeline

We will explore each of these components in detail, starting with the model architecture design, but first, let’s discuss why I chose PyTorch as the framework for this project.

1.2 Using PyTorch

I chose PyTorch for this project based on three key factors: educational accessibility, integration with the research ecosystem, and practical convenience. PyTorch provides many components out of the box - transformer blocks, attention layers, feed-forward networks, training loops, and CUDA support - which makes it much easier for learners building their first language model.

From a technical perspective, PyTorch’s memory management and GPU optimization features-including automatic mixed precision, gradient checkpointing, and efficient attention implementations well-suited for the resource-intensive task of training language models on historical text.

PyTorch’s recent developments, such as torch.compile, FlashAttention kernels, and SDPA operator (scaled dot-product attention), provide significant performance improvements, making training more efficient. These improvements enhance both speed and memory efficiency, which are critical for scaling LLMs. Of course, in our case, we’re building a working toy example rather than scaling to production levels, and these optimizations help keep training times reasonable on available hardware.

What about other Frameworks?

I also considered TensorFlow and JAX, but neither seemed right for helloLondon; TensorFlow’s API felt too complex, specifically from a beginner’s perspective. JAX has excellent performance and a clean, functional approach, but it’s more research-focused and has a smaller ecosystem, which would make it harder to follow along and experiment with.

2. Model Architecture Overview

2.1 Understanding the GPT Architecture

Our custom GPT (Generative Pre-trained Transformer) is a decoder-only transformer model designed for autoregressive language modeling on historical text. The architecture consists of four core components, each serving a distinct purpose in the sequence-to-sequence prediction pipeline. These are: token embeddings, position embeddings, causal self-attention mechanisms, and the language modeling head. Let us double-click into each component to understand its role and implementation.

2.1.1 Token Embeddings

Token embeddings convert discrete token IDs from our 30,000-token historical vocabulary into dense, continuous vector representations. Each token (whether it’s a word, subword unit, or special token) is mapped to a point in a high-dimensional space (768 dimensions for SLM, 1024 for the regular model).

This is implemented as a simple lookup table - wte = torch.nn.Embedding(config.vocab_size, config.n_embd). When processing the token sequence, we look up the corresponding vector for each token ID. These embeddings are learned during training - the model learns which tokens should be close together in this vector space based on their co-occurrence patterns in historical text.

For historical language models, this is particularly valuable because rare historical terms (like “yeoman” or “guildhall”) get their own representations that can capture contextual relationships with related terms from that era.

2.1.2 Position Embeddings

Position embeddings encode each token’s absolute position within the sequence. This is crucial because, unlike recurrent models, transformer architectures have no inherent notion of temporal order or sequence position - they process all tokens in parallel. Let us double-click into the problem why.

Think of it like reading words without any sense of order. The words “The cat chased the mouse” would be indistinguishable from “Mouse the chased cat the” - you’d see the same words but lose all meaning because you don’t know which word came first, second, or third. Transformers face exactly this problem because they process all words simultaneously rather than sequentially, unlike older RNN models.

To help the model understand word order, we add position embeddings to the token embeddings. This way, each token’s representation includes information about both “what” the token is and “where” it appears in the sequence. We use learned position embeddings (as opposed to fixed sinusoidal patterns): wpe = torch.nn.Embedding(config.block_size, config.n_embd). For the SLM with a 512-token context window, we learn 512 different position vectors (one for each possible position). Similarly, the regular model with a 1024-token context learns 1024 position vectors.

Position embeddings work like giving each word a “timestamp” or “address” that tells the model where it sits in the sequence:

Token embedding says: “This is the word ‘cat’” → converts to a vector like [0.2, -0.5, 0.8, ...]
Position embedding says: “This word is at position 3” → adds another vector like [0.1, 0.3, -0.2, ...]
Combined: The model sees both “what” the word is AND “where” it appears

The embedding vectors are combined element-wise: x = token_emb + position_emb. This allows the model to understand both what each token is (via the token embedding) and where it appears in the sequence (via the position embedding).

Our model uses learned position embeddings, meaning during training the model discovers that:

Position 1 tends to be capitalized (start of sentence)
Position 512 might be mid-sentence (needs different handling)
Certain positions in historical documents have patterns (formal openings, closings, etc.)

This is different from fixed sinusoidal embeddings (used in the original Transformer paper), which use a mathematical formula to encode positions. Learned embeddings are generally better because they adapt to specific patterns in the training data.

In historical texts, word order is crucial for understanding meaning. Consider “The King granted the land” versus “The land granted the King” - same words, completely different meanings. Historical legal documents and Victorian-era writings often have precise word order that changes legal or semantic meaning. Position embeddings ensure the model can distinguish between these critical variations.

2.1.3 Causal Self-Attention

Causal self-attention is the mechanism that allows each position in the sequence to selectively attend to previous positions. The “causal” constraint ensures the model can only look at past tokens (not future ones), which is essential for autoregressive generation.

When you read a sentence, you naturally use context from earlier words to understand later ones. If you see “The King granted the land to his loyal…”, you can predict that “servant,” “knight,” or “subject” might come next because you remember what came before. The model needs to do the same thing - use previous words to predict the next word.

However, there’s a crucial constraint: during training, when predicting word 7, the model must only see words 1-6, never word 8 or beyond. This “causal” (cause-and-effect) constraint ensures the model learns realistic patterns - in the real world, you can’t use future information to predict the present.

How Attention Works

Think of attention as a sophisticated “relevance detector”. When the model is processing the word “loyal” in our example above, it needs to look back and ask: “Which previous words are most relevant for understanding this context?” The attention mechanism computes a weighted sum of all previous token representations, where the weights are determined by how relevant each previous token is to the current one.

This is done through three learned linear projections that create different “views” of each word:

Query (Q): “What am I looking for?” - The current word asks a question
Key (K): “What do I have to offer?” - Previous words advertise their content
Value (V): “What information should I contribute?” - The actual information to pass forward

Let us see a practical example to help us grok the concept. Consider the historical phrase: “The alderman of Cheapside, having served the city faithfully, was granted…”

When processing “granted” the attention mechanism:

Creates a Query from “granted” asking “what context do I need?”
Compares this Query against Keys from all previous words
Finds high relevance with “alderman” (who is being granted something) and “faithfully” (why the grant is happening)
Uses these attention weights to pull relevant Values from those words
Combines this information to understand better “granted” in context

The attention score between token i and token j is computed as:

$$\text{Attention}(Q_i, K_j) = \text{softmax}\left(\frac{Q_i K_j^T}{\sqrt{d_k}}\right)V_j$$

Breaking this down:

$Q_i K_j^T$ computes how well the query from token i matches the key from token j (higher = more relevant)
$1/\sqrt{d_k}$ is a scaling factor that prevents scores from getting too large (which would make softmax too “sharp”)
$\text{softmax}$ converts scores into probabilities that sum to 1 (so each word gets a weighted “vote”)
Finally, we use these weights to combine the Values from all previous tokens

The $1/\sqrt{d_k}$ scaling factor (where $d_k = 64$ for our SLM, so $\sqrt{64} = 8$) prevents the dot products from growing too large with high-dimensional embeddings, ensuring stable gradients during training. The softmax ensures the weights sum to 1, creating a proper probability distribution over the previous tokens.

Why This Matters for Historical Text?

Historical documents present unique challenges that make attention particularly valuable. Consider a legal document from 1750: “John Smith, yeoman of the parish of St. Giles, being of sound mind and body, doth hereby bequeath…”

The attention mechanism enables the model to:

Connect “doth bequeath” back to “John Smith” across multiple clauses
Understand that “yeoman” modifies “John Smith” even though they’re separated
Learn that “doth” (archaic) and “does” (modern) serve similar grammatical functions
Recognize that formal legal phrasing follows specific patterns

For our historical language models, this attention mechanism learns which historical terms and phrases co-occur and relate to one another contextually - crucial for understanding historical documents where terminology and phrasing differ from modern English. The model learns to attend to relevant historical context, enabling it to generate coherent text that maintains period-appropriate language patterns and references.

2.1.4 Language Modeling Head

The language modeling head (also called the “output projection” or lm_head) is the final translator that turns the rich internal representation (after all the attention + MLP refinements) back into a decision: “Given everything I’ve seen so far, what is the most likely next token?” It does this by mapping each hidden vector at every position into a vector of length equal to the vocabulary size (30,000 in our historical tokenizer). Each element of that output vector is a logit - an unnormalized score indicating how likely the model thinks the token is to be the next one.

Implementation is intentionally simple: lm_head = torch.nn.Linear(n_embd, vocab_size). We don’t put an activation function after it because we want raw, unconstrained scores. Those scores then flow into:

Inference: Apply softmax -> probabilities -> sample or greedy pick
Training: Feed logits + target token IDs into cross-entropy loss -> gradients flow backward

You can think of logits as evidence totals. The softmax transforms those evidence values into a normalized probability distribution that the model can sample from. High logit = more supporting evidence; low logit = less.

Step-by-step (Inference vs Training):

Hidden state at last position (e.g., index 511) enters lm_head.
Linear projection produces a 30,000-dimensional logit vector.
In inference: probs = softmax(logits / temperature); optionally apply top-k/top-p filtering.
Sample (or argmax) a token → append to sequence → repeat.
In training: Cross-entropy compares logits to the true next token; loss scalar backpropagates through the head into all prior layers.

Because our vocabulary mixes common function words (“the”, “and”, “of”) with rare era-specific tokens (“yeoman”, “guildhall”, “paternoster”, “quoth”), the head must reliably distinguish both frequent and infrequent patterns. Rare historical tokens need consistent representations from embedding -> transformer -> head so they are not forgotten. If their logits remained perpetually low, the model would never learn to generate them in authentic contexts.

Logits (not probabilities) inside the model - We retain logits (raw, unnormalized scores) instead of immediately converting to probabilities because they yield numerically stable loss computation - PyTorch efficiently fuses log_softmax with negative log-likelihood - allow cleaner gradient flow before any normalization (we only invoke softmax when we actually need a distribution), and enable flexible post-processing (temperature scaling, top-k or top-p filtering, repetition penalties) directly in score space without forcing an extra probability recomputation step.

We reuse the input embedding matrix for the output projection to keep input and output semantics aligned and reduce parameter and memory traffic. This concept is called Weight Typing, which we will cover in detail in Section 2.2.2 - Weight Tying .

We share the embedding and output projection weights (self.transformer.wte.weight = self.lm_head.weight) so input token interpretation and next-token scoring occur in the same semantic space.

Using the shared embedding matrix $E$ (shape $(V,d)$), the logits are computed with $\text{logits} = h \cdot E^T$, reusing the same rows used for token lookup. This saves parameters (~23.0M SLM, ~30.7M Regular), keeps gradients for rare historical tokens coupled, and reduces memory traffic (Press & Wolf, 2017; Inan et al., 2016). See Section 2.2.2 for detailed mechanics, benefits, and the historian/scribe analogy.

In short, the lm_head converts rich contextual understanding into next-token scores; with weight tying (details in Section 2.2.2) it stays efficient and semantically consistent.

2.1.5 The Complete Flow

The complete forward pass through our GPT model works as follows:

Input: A sequence of token IDs (batch × sequence_length, e.g., 512 tokens)
Token Embedding: Convert each token ID to a dense vector (768 or 1024 dimensions)
Position Embedding: Add position information to each token
Transformer Blocks: Pass through n_layer transformer blocks (12 for SLM, 24 for regular model), each containing:
- Layer normalization
- Causal self-attention (with multiple heads)
- Residual connection
- Layer normalization
- Feed-forward MLP
- Residual connection
Final Layer Norm: Normalize the final hidden states
Language Head: Project to vocabulary logits (30,000 dimensions)
Output: Probability distribution over next token

This architecture design is conventional and follows the GPT-style pattern established by OpenAI’s GPT models. The traditional design is intentional - it allows for clear, educational learning from the implementation while being configured to work seamlessly with our historical tokenizer from Part 2.

Figure 2 below illustrates the complete architecture for the SLM:

graph TD
    A[Input Tokens<br/>512 tokens] --> B[Token Embedding<br/>30K vocab → 768 dim]
    A --> C[Position Embedding<br/>512 pos → 768 dim]
    B --> D[Add Embeddings]
    C --> D
    D --> E[Layer Norm]
    E --> F[Transformer Block 1<br/>12 heads, 768 dim]
    F --> G[Transformer Block 2<br/>12 heads, 768 dim]
    G --> H[...]
    H --> I[Transformer Block 12<br/>12 heads, 768 dim]
    I --> J[Final Layer Norm]
    J --> K[Language Head<br/>768 → 30K vocab]
    K --> L[Output Logits<br/>30K probabilities]
    
    subgraph "Key Specifications"
        M[Layers: 12<br/>Heads: 12<br/>Embedding: 768<br/>Context: 512<br/>Parameters: 117M<br/>Training: 7-8 hours<br/>MFU: 8-9%]
    end
    
    style A fill:#e1f5fe
    style L fill:#e8f5e8
    style M fill:#fff3e0

Figure 2: SLM Architecture (117M Parameters)

The Regular model, as shown in Figure 3 below, follows the same architectural pattern as the SLM but with increased capacity: 24 transformer layers instead of 12, 16 attention heads instead of 12, and 1024-dimensional embeddings instead of 768.

This represents a ~3x increase in parameters (354M vs 117M), ~2x more attention heads, and ~33% larger embedding dimensions, providing significantly more computational power for learning complex historical language patterns.

graph TD
    A[Input Tokens<br/>1024 tokens] --> B[Token Embedding<br/>30K vocab → 1024 dim]
    A --> C[Position Embedding<br/>1024 pos → 1024 dim]
    B --> D[Add Embeddings]
    C --> D
    D --> E[Layer Norm]
    E --> F[Transformer Block 1<br/>16 heads, 1024 dim]
    F --> G[Transformer Block 2<br/>16 heads, 1024 dim]
    G --> H[...]
    H --> I[Transformer Block 24<br/>16 heads, 1024 dim]
    I --> J[Final Layer Norm]
    J --> K[Language Head<br/>1024 → 30K vocab]
    K --> L[Output Logits<br/>30K probabilities]
    
    subgraph "Key Specifications"
        M[Layers: 24<br/>Heads: 16<br/>Embedding: 1024<br/>Context: 1024<br/>Parameters: 354M<br/>Training: 28-32 hours<br/>MFU: 15-20%]
    end
    
    style A fill:#e1f5fe
    style L fill:#e8f5e8
    style M fill:#fff3e0

Figure 3: Regular Model Architecture (354M Parameters)

2.2 SimpleGPT Class

Now that we’ve covered the theory behind the GPT architecture, let’s examine the actual implementation. The SimpleGPT class is at the heart of our implementation - it’s the core class that brings together all the components discussed in section 2.1 into a working language model. The class inherits from PyTorch’s **torch.nn.Module**, which is the base class for all neural network components in PyTorch. This gives us access to automatic differentiation, GPU support, and other PyTorch features.

2.2.1 The `init` method

The __init__ method is the constructor that assembles our entire language model from individual components. First, it stores all hyperparameters (such as vocabulary size, embedding dimensions, and the number of layers) in a configuration object that the rest of the model can reference. Next, it creates the embedding layers - one that converts our 30,000 historical tokens into dense vectors, and another that encodes position information. Hence, the model knows where each word appears in the sequence.

Next, it builds the transformer blocks - the core processing units that do the heavy lifting. Each block contains self-attention mechanisms and feed-forward networks that learn to understand relationships between words. The method also initializes the language modeling head, the final layer that converts all internal processing back into probabilities for which word should come next.

Finally, it sets up proper weight initialization to ensure the model starts with good random weights (not too big, not too small), and implements weight tying between the input embeddings and the output layer. This clever technique reduces the number of parameters while improving training efficiency by sharing weights between the first and last layers.

This is important because if the weights are too large, the model’s gradients can explode during training, leading to unstable learning. If they start too small, gradients can vanish, rendering the model unable to learn. Our initialization ensures the model begins in the “Goldilocks zone” - just right for effective learning. Without this, even a perfectly designed architecture might fail to train properly.

Now, let me show you the actual implementation. The code in Listing 1 demonstrates how we implement the core GPT architecture:

class SimpleGPT(torch.nn.Module):
    """Simple GPT model based on nanoGPT
    
    This class implements a decoder-only transformer model optimized for 
    historical text generation. It inherits from PyTorch's Module class
    to get automatic differentiation and GPU support.
    """
    
    def __init__(self, config):
        super().__init__()  # Initialize the parent PyTorch Module class
        self.config = config  # Store all model hyperparameters
        
        # Create the main transformer components using ModuleDict
        # ModuleDict allows us to organize related layers together
        self.transformer = torch.nn.ModuleDict(dict(
            # Token Embedding Layer (wte = "word token embedding")
            # Converts each token ID to a high-dimensional vector
            # Input: token IDs (integers 0 to vocab_size-1)
            # Output: dense vectors of size $n_{embd}$ (e.g., 768 dimensions)
            wte = torch.nn.Embedding(config.vocab_size, config.n_embd),
            
            # Position Embedding Layer (wpe = "word position embedding") 
            # Encodes where each token appears in the sequence
            # Input: position indices (0 to block_size-1)
            # Output: dense vectors of size $n_{embd}$
            wpe = torch.nn.Embedding(config.block_size, config.n_embd),
            
            # Dropout Layer for regularization
            # Randomly sets some inputs to zero during training to prevent overfitting
            drop = torch.nn.Dropout(config.dropout),
            
            # Stack of Transformer Blocks (h = "hidden layers")
            # Each SimpleBlock contains self-attention and feed-forward layers
            # We create n_layer blocks (e.g., 12 for SLM, 24 for regular model)
            h = torch.nn.ModuleList([SimpleBlock(config) for _ in range(config.n_layer)]),
            
            # Final Layer Normalization
            # Normalizes the output before the language modeling head
            ln_f = torch.nn.LayerNorm(config.n_embd, bias=config.bias),
        ))
        
        # Language Modeling Head
        # Converts the final hidden states back to vocabulary space
        # Input: hidden states of size $n_{embd}$
        # Output: logits for each token in vocabulary ($vocab_{size}$ logits)
        self.lm_head = torch.nn.Linear(config.n_embd, config.vocab_size, bias=False)
        
        # Initialize all weights using our custom initialization method
        # This ensures the model starts with good random weights
        self.apply(self._init_weights)
        
        # Weight Tying: Share weights between input embeddings and output layer
        # This technique improves training efficiency and model performance
        # by ensuring the same representation space is used for input and output
        self.transformer.wte.weight = self.lm_head.weight

Listing 1: SimpleGPT Model Architecture

2.2.2 Weight Tying

The tied weights between the embedding layer and language modeling head (self.transformer.wte.weight = self.lm_head.weight) are a crucial optimization for our historical language model. In a typical neural network, you’d have two separate weight matrices - one for converting input tokens to embeddings, and another for converting hidden states back to vocabulary probabilities. Weight tying means we use the same weight matrix for both operations.

Think of it like this: instead of having two different dictionaries (one for reading, one for writing), we use the same dictionary for both. The same table that maps “alderman” → [0.2, -0.5, 0.8, …] is used whether the model is reading “alderman” as input or trying to generate “alderman” as output.

Without weight tying, the model would have two separate weight matrices - one for converting input tokens to embeddings, and another for converting hidden states back to vocabulary probabilities. This means the model could learn that “alderman” means one thing when it sees it as input, but something slightly different when it tries to generate it as output. For rare historical terms, this inconsistency can cause the model to “forget” how to use words it has seen before properly.

Historical vocabulary contains many rare terms such as “quoth” “alderman” and “paternoster” that appear infrequently in our training data. Without weight tying, the model might learn different representations for the same word when it sees it as input versus when it generates it as output. This inconsistency can cause the model to struggle with rare historical terms.

When the model sees “alderman” in the input, it learns a specific representation of it. Later, when it needs to generate “alderman” in the output, it uses that same learned representation, ensuring consistency and improving the model’s ability to generate coherent historical language with proper terminology.

Mechanics (matrix reuse) A single matrix $E \in \mathbb{R}^{(V \times d)}$ serves both roles: row lookup for input embeddings and column interaction for output scoring. The language head reuses it to compute logits via:

$$\text{logits} = h \cdot E^T$$

where $h$ is the hidden state at each position, this keeps the input interpretation and output prediction within the same semantic geometry - no second projection to drift or disagree.

Why it helps Parameter savings (~23.0M SLM, ~30.7M Regular) lower memory footprint and bandwidth. Gradients for predicting a rare token (e.g., yeoman, guildhall, paternoster) directly refine the very rows used to embed it on future inputs - improving both recall and generation. Shared weights mildly regularize against the two spaces drifting apart and empirically improve perplexity for mid-scale autoregressive models (Press & Wolf, 2017; Inan et al., 2016).

Analogy If the transformer stack is a panel of historians debating context, the language modeling head is the scribe choosing the next historically plausible word. Weight tying means the scribe and historians consult the same dictionary - no translation mismatch between how words are read and how they’re proposed.

Practical notes Avoid inflating vocabulary unnecessarily (cost scales with $V$); tied weights do not remove the need for careful rare token coverage in the corpus; and if later adding adapters or LoRA heads, remember that tying interacts with how those layers inject low-rank updates.

Now that we understand how the model efficiently handles vocabulary, let’s examine the core processing units that transform these embeddings into meaningful representations.

2.3 Transformer Block Design

Each transformer block implements the standard attention and feed-forward pattern, but with optimizations for historical text processing. Let us look at the code real quick and then get into a little more detail.

class SimpleBlock(torch.nn.Module):
    """Simple transformer block"""
    
    def __init__(self, config):
        super().__init__()
        self.ln_1 = torch.nn.LayerNorm(config.n_embd, bias=config.bias)
        self.attn = SimpleCausalSelfAttention(config)
        self.ln_2 = torch.nn.LayerNorm(config.n_embd, bias=config.bias)
        self.mlp = SimpleMLP(config)
    
    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

Listing 2: Transformer Block Implementation

This code implements a single transformer block, which, as we know, is the fundamental building unit of our GPT model.

2.3.1 Self-Attention Step

The self-attention step (x = x + self.attn(self.ln_1(x))) is the block that first normalizes the input with LayerNorm, then applies self-attention to understand relationships between words. The + creates a “residual connection” that helps information flow through the network.

As we discussed in section 2.1.4 , self-attention is the “magic” of transformers, allowing each token to decide how much attention to pay to every other token in the sequence. Our implementation uses multiple attention heads (12 for SLM, 16 for Regular) that operate in parallel, with each head learning to focus on different types of relationships - syntactic, semantic, or positional. Causal masking ensures that during training, the model learns to predict the next token based solely on the preceding context, which is essential for coherent text generation.

The residual connection (+) is crucial, as it allows the model to preserve the original token representation while adding contextual information from the attention. The pre-normalization approach (LayerNorm before attention) provides more stable training than post-normalization, especially important when working with the varied linguistic patterns found in historical text.

2.3.2 Feed-Forward Step

After attention, we have the feed-forward step (x = x + self.mlp(self.ln_2(x))), which first normalizes the attended information with LayerNorm, then passes it through a multi-layer perceptron (MLP) that transforms and processes it.

The MLP typically consists of two linear layers with a non-linear activation function (like GELU) between them, allowing the model to learn complex non-linear transformations of the attended features. This step is crucial because attention can only perform linear transformations on the input representations; the feed-forward network adds the necessary non-linearity, enabling the model to learn complex patterns and relationships in the historical text. Another residual connection preserves the original information, ensuring that the model can always fall back to the pre-attention representation if needed.

2.3.3 Understanding the Feed-Forward (MLP) Sublayer

Directly beneath the SimpleBlock code above, you see the line self.mlp = SimpleMLP(config). After attention has mixed information across positions, the model passes each token embedding through a position-wise feed-forward network (the MLP). Unlike attention, it does not look at other tokens; it refines the representation of each token independently, given the contextualized features attention just produced. In practice, this is where raw contextual patterns are distilled into richer semantic, stylistic, and morphological signals.

Figure 4 below visualizes how a single transformer block routes data through normalization, attention, and the feed-forward expansion/contraction before returning an upgraded representation via the residual path:

graph TB
    A[Input Embeddings<br/>batch, seq, emb] --> LN1[LayerNorm 1]
    LN1 --> ATTN[Multi-Head Attention<br/>query, key, value]
    ATTN --> DROPA[Dropout]
    DROPA --> RES1[Residual Add<br/>x + attn_out]
    RES1 --> LN2[LayerNorm 2]
    RES1 --> LN2
    LN2 --> EXPAND[Linear Expand<br/>emb → 4*emb]
    EXPAND --> GELU[GELU Activation]
    GELU --> PROJECT[Linear Project<br/>4*emb → emb]
    PROJECT --> DROPM[Dropout]
    DROPM --> RES2[Residual Add<br/>res1 + mlp_out]
    RES2 --> OUT[Block Output<br/>Updated embeddings]
    style A fill:#e1f5fe
    style ATTN fill:#f3e5f5
    style EXPAND fill:#fff3e0
    style PROJECT fill:#fff3e0
    style RES2 fill:#e8f5e8

Figure 4: Internal Flow of a Transformer Block

Conceptually, the MLP is a two-step projection: first, an expansion into a higher-dimensional “workspace” with a non-linear activation, then a projection back down so the residual can safely merge with the original stream.

For our SLM, 768 dimensions expand to 3072 and then contract back to 768; for the larger model, 1024 dimensions expand to 4096. This temporary widening allows the network to express combinations of features that a purely linear transform could not capture. It is the difference between merely routing information and actually transforming it.

Here is the representative structure shown in Listing 3:

class SimpleMLP(torch.nn.Module):
    def __init__(self, config):
        super().__init__()
        self.fc_in  = torch.nn.Linear(config.n_embd, 4 * config.n_embd)
        self.act    = torch.nn.GELU()
        self.fc_out = torch.nn.Linear(4 * config.n_embd, config.n_embd)
        self.drop   = torch.nn.Dropout(config.dropout)
    def forward(self, x):
        return self.drop(self.fc_out(self.act(self.fc_in(x))))

Listing 3: Feed-Forward (MLP) Sublayer Implementation

Why expand then shrink?

The widened hidden space allows the model to form intermediate feature bundles (e.g., tense, register, archaic morphology) that do not cleanly live in the original lower-dimensional basis. The contraction enforces a stable interface for the residual path and keeps the total parameter count manageable. Removing the expansion would noticeably degrade expressiveness; removing the contraction would balloon memory use and break architectural symmetry.

In our context, the historical model internalizes regularities like mapping “hath” and “doth” into modern tense abstractions while still preserving period flavor; it encodes stylistic shifts between court proceedings, religious prose, and narrative storytelling; it stabilizes inconsistent orthography and variant spellings so downstream layers predict coherent continuations instead of brittle echoes. Attention tells the model where to look; the MLP decides how to reinterpret what it saw.

Focusing only on attention gives an incomplete mental model of transformers. More than half of the parameters and a large fraction of the FLOPs sit in these feed-forward layers. Under-sized MLPs lead to shallow pattern memorization - models that can repeat phrases but cannot generalize style or adapt archaic forms to new contexts. Properly scaled MLP width (the common ×4 expansion) is a proven sweet spot: smaller factors underfit; much larger ones give diminishing returns at this scale (see scaling law discussions in Kaplan et al. 2020 ).

A useful mental analogy: attention is the lively debate in a hall; the MLP is each participant stepping aside to integrate what was heard into their own refined understanding before the next round of discussion. When you see x = x + self.mlp(self.ln_2(x)), that addition represents the moment a token’s contextual representation is upgraded. Without this transformation, the model would “hear” context but fail to internalize it, producing shallow, literal continuations rather than fluent, period-authentic prose.

In our helloLondon models, the MLP is therefore essential for converting raw multi-head attention patterns into durable historical linguistic competence - one of the quiet reasons the generated text feels coherent rather than stitched together.

Each block in our model (12 for SLM, 24 for Regular) applies this same pattern, allowing the model to build an increasingly sophisticated understanding of historical language patterns as text flows through the layers.

Each transformer block applies layer normalization before both the self-attention mechanism and the feed-forward network, followed by residual connections. This pre-normalization approach (as opposed to post-normalization) has been shown to provide more stable training, especially important when working with the varied linguistic patterns found in historical text.

2.3.4 Activation choice matters

The activation function determines how the neural network processes information at each layer. Think of it as a “decision maker” that decides how much of each input signal to pass through to the next layer. The most common activation functions are ReLU (Rectified Linear Unit) and GELU (Gaussian Error Linear Unit).

ReLU is simple and fast: it passes positive values unchanged and sets negative values to zero (f(x) = max(0, x)). However, ReLU can be “harsh” - it completely cuts off negative signals, leading to “dead neurons” that never activate again. GELU is smoother and more sophisticated: it uses a Gaussian distribution to determine how much of each input to pass through (f(x) = x * Φ(x) where Φ is the cumulative distribution function of a standard normal distribution). This creates a smooth, differentiable function that allows for more nuanced information processing.

GELU offers smoother gradients and better calibration for language than plain ReLU. The smoother nature of GELU helps the model learn more subtle patterns in historical text, where the relationships between words and phrases can be complex and nuanced. Alternatives like SwiGLU can yield marginal gains in perplexity but increase implementation complexity - valuable in frontier systems, optional in educational builds like helloLondon. Modest dropout in the MLP further improves generalization on a corpus that, while sizable, is still modest relative to billion-token modern pretraining regimes.

2.3.5 Pre vs Post-normalization

In pre-normalization, we normalize the input before processing it (like we do here). In post-normalization, we’d process first, then normalize the output. Pre-normalization is like checking that your ingredients are properly prepared before cooking, while post-normalization is like seasoning after cooking - both work, but pre-normalization tends to yield more consistent results.

This matters because historical texts contain complex syntactic structures and long-range dependencies that require sophisticated attention mechanisms. The residual connections ensure that information can flow directly through the network, helping the model learn to preserve important historical context across long sequences while still allowing the attention mechanism to focus on relevant historical details.

2.4 Causal Self-Attention for Historical Sequences

The attention mechanism is crucial for understanding the complex relationships in historical text. Our implementation is based on the original transformer architecture from “Attention Is All You Need” (Vaswani et al., 2017), but optimized for historical language patterns.

Understanding Multi-Head Attention

Multi-head attention runs several attention “heads” in parallel, allowing the model to focus on different aspects of a sequence simultaneously (syntax, semantics, and position). Compared to a single head, this parallelism yields richer representations—think multiple specialists examining the same text. In our setup, the SLM uses 12 heads and the Regular model 16, scaling capacity with model size. Empirically, heads tend to specialize (e.g., subject–verb agreement, word relations, word order), as observed by Clark et al. (2019) - “What Does BERT Look At?” .

Research by Kaplan et al. (2020) Scaling Laws for Neural Language Models shows that the optimal number of attention heads scales with model size. For our 117M-parameter SLM, 12 heads provide sufficient parallel processing capacity, while our 354M-parameter Regular model benefits from 16 heads to capture more complex attention patterns.

The attention mechanism has $O(n^2)$ complexity with respect to sequence length. This means that doubling our sequence length from 512 to 1024 tokens is a quadratic jump and requires 4x more memory for attention computations. This is why we carefully balance sequence length with available GPU memory and why techniques like FlashAttention (Dao et al., 2022) are so important for memory efficiency.

How the attention mechanism works in practice:

The code in Listing 4 shows how we implement the attention mechanism that we’ve been discussing. Here’s what happens step by step:

Input Processing: The model receives a batch of sequences (B = batch size, T = sequence length, C = embedding dimension). For example, with our SLM: B=4, T=512, C=768.
Query, Key, Value Generation: The input embeddings are transformed into three different representations - Query (Q), Key (K), and Value (V) - using a single linear layer that outputs 3×768 dimensions, then splits them.
Multi-Head Reshaping: Each of Q, K, V is reshaped to separate the 12 attention heads, so each head gets its own 64-dimensional subspace (768 ÷ 12 = 64).
Attention Computation: The scaled dot-product attention is computed, where each word “looks at” all previous words (causal masking) and decides how much attention to pay to each.
Output Assembly: All attention head outputs are combined back into a single representation and projected through a final linear layer.

This implementation is optimized for historical text processing, using PyTorch’s efficient scaled_dot_product_attention function with causal masking to ensure the model can only attend to previous tokens, not future ones.

class SimpleCausalSelfAttention(torch.nn.Module):
    """Simple causal self-attention"""
    
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        # Key, query, value projections for all heads, but in a batch
        self.c_attn = torch.nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        # Output projection
        self.c_proj = torch.nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        # Regularization
        self.attn_dropout = torch.nn.Dropout(config.dropout)
        self.resid_dropout = torch.nn.Dropout(config.dropout)
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.dropout = config.dropout
    
    def forward(self, x):
        # Batch size, sequence length, embedding dimensionality ($n_{embd}$)
        B, T, C = x.size()
        
        # Calculate query, key, values for all heads in batch 
        # and move head forward to be the batch dim
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
        
        # Causal self-attention; Self-attend:
        # (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, C)  # Re-assemble all head outputs side by side
        
        # Output projection
        y = self.resid_dropout(self.c_proj(y))
        return y

Listing 4: Causal Self-Attention Implementation

The attention mechanism computes attention as show below:

$$\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Where $Q$, $K$, and $V$ are the query, key, and value matrices, respectively.

In our case, with 768-dimensional embeddings and 12 heads, each head operates on 64-dimensional subspaces ($d_k = 768 / 12 = 64$), providing sufficient representational capacity for each type of historical relationship while maintaining computational efficiency.

In addition, the $\sqrt{d_k}$ scaling factor ($\sqrt{64} = 8$) prevents the dot products from becoming too large, ensuring stable gradient flow during training.

In plain English, please!

Think of attention like a spotlight that can shine on different parts of a sentence. When the model is trying to understand the word “he” in a historical document, it needs to look back through the text to find who “he” refers to. The attention mechanism is like having multiple spotlights (our 12 or 16 attention heads) that can each focus on different aspects - each might look for people’s names, another for relationships, and another for locations.

The mathematical formula we showed above is how the model calculates the amount of attention to pay to each word. The scaling factor ($$\sqrt(64) = 8) is like adjusting the brightness of the spotlight – it prevents the model from being “blinded” by very bright spots and helps it focus on the right amount of information.

Does this matter for historical text?

Historical documents are particularly challenging because they often feature complex sentence structures and references spanning long distances. For example, in a court record, you might see “The defendant, John Smith, was accused of theft. He claimed innocence throughout the trial.” The model needs to understand that “He” refers to “John Smith,” even though there are several words between them. The attention mechanism enables the model to make these connections, generating coherent text that maintains proper historical context and references.

This is certainly required for language modeling, given the complex structures in which later words reference earlier ones, and understanding the full context is essential for proper interpretation. The attention mechanism enables the model to learn long-range dependencies, allowing it to generate coherent text across extended sequences. For historical texts specifically, this becomes even more important because archaic language patterns and historical references often span longer distances than those in modern texts.

2.5 Model Configuration

The model architecture uses a centralized configuration, where each parameter is selected based on research findings and practical constraints for historical text processing. The SLM architecture uses five key parameters, each representing a design choice with specific trade-offs between computational efficiency and learning capacity.

Parameter	Value	Purpose	Trade-off
`n_layer`	12	Number of transformer blocks (model depth)	More layers = better learning, but slower training
`n_head`	12	Number of attention heads (parallel processing)	More heads = better attention, but more computation
`n_embd`	768	Embedding dimension (token representation)	Larger = richer representations, but more memory
`max_length`	512	Context window size (sequence length)	Longer = more context, but quadratic memory growth
`vocab_size`	30K	Vocabulary size (tokenizer compatibility)	Larger = more words, but more parameters

These parameters work together to create a model that effectively processes historical text while remaining computationally manageable.

Layer Count (n_layer: 12)

The 12-layer architecture balances representational capacity with computational efficiency for historical text processing. Shallow layers (1-3) learn basic token patterns and grammatical structures, middle layers (4-8) capture complex syntactic relationships and historical language patterns, and deep layers (9-12) understand high-level semantic relationships and historical context.

This depth follows GPT-2 Small’s 12-layer architecture, which delivers strong performance while remaining computationally manageable on available hardware.

Attention Heads (n_head: 12)

Multi-head attention allows the model to attend to different types of relationships simultaneously – for example, temporal (chronological order), social (class hierarchies), geographical (London landmarks), and linguistic (archaic patterns). The 12-head architecture balances parallel processing capability with computational efficiency for historical text understanding.

Embedding Dimension (n_embd: 768)

The 768-dimensional embedding space can represent complex historical concepts, such as archaic terms (“yeoman”, “paternoster row”, “hath”), while maintaining computational efficiency. This dimension is commonly used in transformer architectures, including BERT-base and GPT-2 Medium.

Why 768 became standard: As a side note, in case you are seeing a lot of 768 lately, there are a good set of reasons for this. Beyond its divisibility ($768 ÷ 12 = 64$ per attention head), 768 aligns with GPU memory architecture - it’s a multiple of 256 (3 × 256), which matches common GPU memory bus widths and cache line sizes. This makes matrix operations more efficient on modern GPUs, as the hardware can process data in optimal chunks. Additionally, 768 provides sufficient representational capacity without the memory overhead of larger dimensions like 1024, making it practical for training on consumer hardware while still capturing complex linguistic relationships.

Context Window (n_positions: 512)

We use a 512-token context window as a practical balance between historical coherence and available compute for a learning-focused setup. While many of our working snippets (e.g., diary passages, sections of legal records, or literary excerpts) comfortably fit within 512 tokens, full historical documents can be much longer. The 512 window keeps attention costs manageable (quadratic in sequence length) while covering typical training segments we use.

Both models use the same 30K vocabulary from our custom historical tokenizer, ensuring consistent tokenization across model variants.

3. GPU Configuration and Perf. Optimization

The training system is designed to maximize GPU utilization while maintaining training stability. Understanding GPU architecture and memory management is crucial for efficient language model training, especially when working with historical text that requires significant computational resources.

3.1 GPU Architecture and Memory Management for Language Model Training

Training on historical text benefits from sensible GPU settings even for a small, learning-focused model. We keep to practical, low-risk optimizations (precision choice, batch/sequence trade-offs, memory-aware attention) and accept some trial and error—reserving heavier systems engineering for larger setups.

The main universal factors are:

Attention scales quadratically with sequence length, so longer contexts get expensive fast.
Natural language variability (syntax, vocabulary, style) demands sufficient model capacity and stable optimization.
Real‑world data quality (formatting, noise) can destabilize training, requiring robust error handling and memory management.

For historical text specifically, archaic vocabulary, period terminology, and cultural references introduce patterns absent from modern corpora. OCR artifacts and uneven formatting in digitized sources add noise beyond what’s typical in contemporary datasets.

3.1.1 GPU memory hierarchy and optimization strategies

Modern GPUs use a hierarchical memory system that significantly impacts training performance: fast but tiny registers and shared memory sit closest to the compute; a larger L2 cache buffers traffic; and global memory holds parameters and activations. Attention often ends up memory-bound, so moving less data (via AMP, Flash/SDPA kernels, and sensible sequence/batch sizes) is as important as raw FLOPs.

For language model training, the key optimization is managing the memory bandwidth bottleneck. Attention operations are often memory-bound rather than compute-bound, meaning performance is limited by how quickly data can be moved between memory levels rather than by computational power. If we are not careful, it is quite easy to run into memory issues, as shown in Figure 5 below.

Figure 5:OOM error

And it is not restricted to training only; even on checkpoints that are saved, we can also encounter memory issues, as shown in Figure 6.

Figure 6:OOM checkpoint eval

Mixed precision training and memory optimization

Training large language models requires careful memory management, especially when working with limited GPU resources. Our training system uses several optimization techniques to maximize memory efficiency while maintaining training stability.

GPU detection and basic configuration:

The training system needs to work across different hardware setups, from single consumer GPUs to multi-GPU servers. Our approach uses a centralized configuration system that automatically adapts to available hardware.

The actual GPU detection in train_model_slm.py is quite straightforward - it checks for distributed training environment variables (RANK, LOCAL_RANK, WORLD_SIZE) and sets up basic multi-GPU support if available. The system also detects GPU capabilities, such as bfloat16 support, and enables appropriate optimizations. This allows the same training script to work across different hardware setups, though the real complexity comes from the trial-and-error process of stabilizing training.

# GPU configuration (from config.py)
self.gpu_config = {
    "auto_detect": True,  # Automatically detect available GPUs
    "max_gpus": 0,  # Maximum number of GPUs to use (0 = no limit, use all available)
    "min_gpu_memory_gb": 8,  # Minimum GPU memory required (GB)
    "preferred_gpu_types": ["A30", "A40", "A100", "V100", "RTX4090", "RTX4080"],
    "fallback_to_cpu": True,  # Fall back to CPU if no suitable GPUs found
    "force_single_gpu": False,  # Force single GPU even if multiple available
    "force_multi_gpu": False,  # Force multi-GPU even if only one available
    "gpu_memory_fraction": 0.9,  # Fraction of GPU memory to use (0.0-1.0)
    "allow_growth": True,  # Allow GPU memory growth
    "log_device_placement": False  # Log device placement for debugging
}

Listing 5: GPU Configuration and Detection System

The configuration in Listing 5 is defined in our centralized config.py file and provides settings for automatic GPU detection, memory management, and fallback options. While this looks comprehensive, the actual implementation is simpler - the training code primarily relies on PyTorch’s built-in distributed training detection and basic device selection.

The reality of training: Nearly 100 runs and many failures

Figure 7: helloLondon training runs - WandB

Figure 7 shows the actual training experience: 99 total runs with 24 completions. The failures were largely data-driven - OCR and encoding issues, uneven sequence lengths, and sensitivity to learning-rate warmup - and a few were plain memory pressure from early, less conservative settings. The code stabilized early; the data and knobs took time.

This iterative process is typical in language model development - the “sophisticated” system shown here is the result of learning from these failures and gradually improving the training pipeline. The successful runs exhibit stable loss curves and appropriate learning rate schedules, demonstrating that the final configuration performs well on historical text processing tasks.

Most importantly, this experience reinforces a fundamental truth in machine learning: data quality is still king. No amount of sophisticated architecture, GPU optimization, or training infrastructure can overcome poor data quality. The “garbage in, garbage out” principle remains as true for language models as it was for the earliest machine learning systems. Our 75% failure rate was primarily due to data issues – such as inconsistent formatting, OCR errors, and encoding problems - not technical limitations. This is why Part 2’s focus on data cleaning and tokenization was so crucial to our success.

3.2 Precision and Performance Configuration

The system includes precision and performance configuration options that can be tuned based on available hardware. Mixed-precision training uses lower-precision (fp16/bf16) for most operations while keeping full precision for critical computations, providing significant memory savings and speed improvements with minimal impact on quality.

Understanding fp16 and bf16: The Precision Trade-off

To understand why precision matters for language model training, we need to look at how computers represent numbers. Standard floating-point numbers use 32 bits (float32), but we can use fewer bits to save memory and increase speed:

fp16 (Half Precision): Uses 16 bits to represent numbers, cutting memory usage in half and enabling faster computation. However, it has a smaller range of representable numbers, which can cause “overflow” (numbers too large) or “underflow” (numbers too small) during training.
bf16 (Brain Float 16): Also uses 16 bits, but with a different bit layout that matches float32’s exponent range. This means it can represent the same range of large and small numbers as float32, but with less precision for very small decimal values.

Why bf16 is better for language models:

bf16 provides better numerical stability than fp16, especially for large language models, reducing the likelihood of overflow and underflow that can cause training instability. The key difference is that bf16 can represent the same range of numbers as float32 (from very small to very large), while fp16 has a much smaller range. This is crucial for language models because:

Gradient magnitudes vary widely - Some gradients are very small (close to zero) while others are large
Attention weights - The softmax operations in attention can produce very small numbers that FP16 might round to zero
Learning rate scaling - Modern optimizers like AdamW work with gradients of varying magnitudes

When gradients become too small and are rounded to zero (underflow), the model stops learning effectively. When they become too large (overflow), training becomes unstable. bf16’s wider range helps prevent both issues.

Understanding precision and performance settings:

The configuration in Listing 6 toggles the levers that matter on consumer hardware: TF32 for faster matmuls, AMP (prefer bf16) for stability and memory cuts, torch.compile for an extra boost after warmup, and sequence/batch sizes sized to your VRAM. Used together, these commonly halve activation memory and yield 2-3x speedups versus full-precision baselines.

# Runtime/precision knobs (A30 optimized)
"enable_tf32": True,
"enable_amp": True,
"amp_dtype": "bf16",  # bf16 on Ampere; fallback to fp16 if unsupported
"enable_compile": True,  # torch.compile; set False to reduce memory usage
# Conservative baseline (for broad hardware) - uncomment to use:
# "enable_tf32": False,
# "enable_amp": True,
# "amp_dtype": "fp16",
# Sequence/batch control
"max_length": 1024,  # increase tokens per step when VRAM allows
"batch_size": 20,    # per-GPU batch; raise if VRAM allows
# Conservative sequence/batch - uncomment to use:
# "max_length": 768,
# "batch_size": 8,

Listing 6: Precision and Performance Configuration

Key GPU Configuration Settings

A few switches move the needle the most: enable TF32 on Ampere-class GPUs for a quick matrix-mul speedup; use AMP (bf16 where supported, fp16 otherwise) to halve activation memory; and turn on torch.compile if you can afford the warmup to get another 1.2-1.5x after a few hundred steps. Keep the sequence length in line with VRAM (~512 tokens for 8GB, ~1024 for 16GB+), and scale the per-GPU batch size accordingly (think hundreds of MB per batch at these widths). The repo includes sensible presets so you can start conservative and dial up.

3.2.1 Real-World Performance Results

On 2x A30s, the SLM lands around mid-20s MFU with ~210 ms/iter and ~18 GB per GPU, converging from ~10.4 loss to the mid-3s over the full run. The clean BPE tokenizer and precision stack keep math efficient, and DDP delivers the expected speedup over a single device.

Automatic Precision Detection and Memory Optimization:

The system also includes automatic precision detection and memory optimization during model initialization. The code snippet below shows how the system automatically selects the optimal precision format based on available hardware capabilities:

# Precision / TF32 knobs from config
tf32 = self.slm_config.get("enable_tf32", True)
torch.backends.cuda.matmul.allow_tf32 = bool(tf32)
torch.backends.cudnn.allow_tf32 = bool(tf32)
try:
    torch.set_float32_matmul_precision('high' if tf32 else 'medium')
except Exception:
    pass
use_amp = self.slm_config.get("enable_amp", True)
amp_dtype_cfg = self.slm_config.get("amp_dtype", "bf16").lower()
bf16_ok = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
if use_amp:
    if amp_dtype_cfg == 'bf16' and bf16_ok:
        self.dtype = 'bfloat16'
    else:
        self.dtype = 'float16'
else:
    self.dtype = 'float32'

Listing 7: Precision Detection and Memory Optimization

The TF32 configuration optimizes matrix operations for Ampere+ GPUs, delivering significant speedups while maintaining training stability.

3.3 Multi-GPU Training with Distributed Data Parallel

The system supports multi-GPU training using PyTorch’s DistributedDataParallel (DDP) - each GPU hosts a full model replica, processes different batches in parallel, and synchronizes gradients automatically. PyTorch handles the inter‑GPU communication, so on two GPUs, you typically see near‑linear speedup (~2×) for these model sizes.

Multi-GPU training improves throughput and shortens wall‑clock time by splitting work across devices. On our 2× A30 setup, we process 36 sequences in parallel (18 per GPU) instead of 18 on a single card, cutting Regular model training from ~56 hours to ~28–32 hours. It also offers operational flexibility: scale up or down based on the number of GPUs available.

However, multi-GPU training introduces several challenges that can limit performance gains. The primary bottleneck is inter-GPU communication - after each backward pass, gradients must be synchronized across all GPUs, which requires transferring large amounts of data. This communication overhead can become significant, especially with larger models and more GPUs.

The performance of multi-GPU training heavily depends on the interconnect between GPUs. On NVIDIA systems, InfiniBand provides the highest bandwidth and lowest latency for GPU-to-GPU communication, enabling near-linear scaling across many GPUs. NVLink (found on high-end NVIDIA GPUs such as A100 and H100) provides direct GPU-to-GPU connections with very high bandwidth, making it ideal for 2-8 GPU setups. PCIe connections are slower but more common in consumer and workstation systems.

In AMD systems, Infinity Fabric serves a role similar to NVLink, providing high-bandwidth interconnects between GPUs. AMD’s MI200 and MI300 series GPUs include Infinity Fabric links that enable efficient multi-GPU communication.

In practice, scaling efficiency depends on the ratio of computation to communication. Our historical language models have relatively modest parameter counts (117M-354M), so communication overhead can be significant compared to computation time. This is why we see good scaling with 2 GPUs but diminishing returns with more GPUs - the communication overhead starts to dominate.

DDP is more efficient than naive data parallelism because it reduces communication overhead and enables larger effective batch sizes, as shown in Listing 8 below.

# DDP setup (process group already initialized in main())
self.ddp = int(os.environ.get('RANK', -1)) != -1
if self.ddp:
    self.ddp_rank = int(os.environ['RANK'])
    self.ddp_local_rank = int(os.environ['LOCAL_RANK'])
    self.ddp_world_size = int(os.environ['WORLD_SIZE'])
    self.device = f'cuda:{self.ddp_local_rank}'
    torch.cuda.set_device(self.device)
    self.master_process = self.ddp_rank == 0
    self.seed_offset = self.ddp_rank
else:
    self.master_process = True
    self.seed_offset = 0
    self.ddp_world_size = 1

Listing 8: Multi-GPU Training Setup

What is “rank” and why does it matter?

In distributed training, each GPU process gets a unique “rank.” Rank 0 acts as the coordinator (handles logging, checkpointing, and WandB), while the remaining ranks focus purely on computation. This avoids collisions - only one process touches files and dashboards - while every device contributes gradients.

This division of labor is crucial because it prevents conflicts. Without it, all processes would try to save checkpoints simultaneously, log to WandB at the same time, or write to the same files, causing errors and corruption.

The key to scaling efficiency is that each GPU works independently on different data batches, then synchronizes only the essential information (gradients). Here’s how it works:

Parallel computation: Each GPU processes a different batch of data simultaneously
Gradient synchronization: After each backward pass, gradients are averaged across all GPUs
Independent updates: Each GPU updates its model copy with the averaged gradients

This means that if you have 2 GPUs, you can process 2x the data in the same time, giving you roughly 2x the speed. With 4 GPUs, you get approximately 4x speedup. The “near-linear” part acknowledges that there’s always some overhead from communication and synchronization, so that you might get 1.9x speedup instead of exactly 2x. Still, it’s close enough to be very effective.

However, there’s a practical limit to this approach. Beyond 8-16 GPUs, the communication overhead becomes so significant that you need more robust hardware (such as InfiniBand networks) and advanced systems engineering techniques (gradient compression, pipeline parallelism, model parallelism) to maintain efficiency. For truly large-scale training with hundreds of GPUs, you need specialized infrastructure and techniques that go far beyond what we’re doing here.

This combination of distributed training and memory optimization enables us to train our historical language models efficiently, even on consumer hardware. The distributed setup provides fault tolerance and near-linear speedup, while the precision optimizations enable larger models and longer sequences on the same hardware.

4. Training Infrastructure: Making It All Work Together

As a reminder, as we saw earlier, the two model variants share the same training stack (scheduler, checkpointing, WandB, DDP). See Part 1 for the high‑level comparison; here are the training‑relevant differences only:

SLM (117M): per‑GPU batch 18 → effective 36 on 2 GPUs; sequence length 512; ~7–8h on 2×A30
Regular (354M): per‑GPU batch 12 → effective 24 on 2 GPUs; sequence length 1024; ~28–32h on 2×A30

4.1 The Training Loop

The core training happens in the train() method, which implements a standard language model training loop with several key phases - outlined below.

4.1.1 Data Loading and Preparation

The training loop starts by loading tokenized data using get_batch('train'), which reads from pre-processed binary files created during data preparation. This includes both training and validation data, with the tokenizer from Part 2: Data Collection & Custom Tokenizers handling the conversion between text and tokens.

Main Training Loop Structure:

def train(self):
    """Main training loop"""
    # Get initial batch
    X, Y = self.get_batch('train')
    
    while True:
        # 1. Learning rate scheduling
        lr = self.get_lr(iter_num)
        
        # 2. Evaluation and checkpointing (every eval_interval steps)
        if iter_num % self.eval_interval == 0:
            losses = self.estimate_loss()
            # Save checkpoint if validation loss improved
            
        # 3. Forward pass with mixed precision
        with torch.amp.autocast():
            logits, loss = self.model(X, Y)
        
        # 4. Backward pass and optimization
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad()
        
        # 5. Get next batch
        X, Y = self.get_batch('train')
        
        # 6. Logging and monitoring
        if iter_num % self.log_interval == 0:
            # Log to WandB and console

Listing 9: Core Training Loop Structure

Training Process Flow:

Understanding how the training actually works requires seeing both the high-level flow and the technical details of each phase. Figure 8 shows the complete training process flow.

graph TD
    A[Start Training] --> B[Load Tokenized Data]
    B --> C[Initialize Model & Optimizer]
    C --> D[Training Loop Start]
    D --> E[Update Learning Rate]
    E --> F{Evaluation Time?}
    F -->|Yes| G[Run Validation]
    F -->|No| H[Forward Pass]
    G --> H
    H --> I[Compute Loss]
    I --> J[Backward Pass]
    J --> K[Gradient Clipping]
    K --> L[Update Weights]
    L --> M[Zero Gradients]
    M --> N[Log Metrics]
    N --> O{Checkpoint?}
    O -->|Yes| P[Save Model State]
    O -->|No| Q[Load Next Batch]
    P --> Q
    Q --> R{Max Iterations?}
    R -->|No| D
    R -->|Yes| S[Save Final Model]
    S --> T[End Training]

Figure 8: Training Process Flow

Now that we have a high-level overview of the training process, let us dig deeper into each phase and see how it works in practice.

4.1.2 Data Loading

Data loading reads pre-tokenized sequences from binary files (.bin) using np.memmap for memory efficiency. The initial tokenization process can take quite a long time on our 500M+ character corpus, but this is done only once and saved to disk. This optimization was crucial during our development process – given nearly 100 training runs and many failures, re-tokenizing the entire corpus each time would have been prohibitively slow. The system handles train/val splits (90/10 %) with random sampling per batch and uses pin_memory() and non_blocking=True for faster GPU transfers.

When we run this for the time time, it takes a long time to load and tokenize the training data corpus. We see this just startging in Figure 9 below.

Figure 9:Tokenizer training data

Batch sizes are optimized for our 2x A30 GPU setup: 18 per GPU for the SLM model (36 effective batch size) and 12 per GPU for the Regular model (24 effective batch size). These numbers balance memory usage with training stability – the SLM can handle larger batches thanks to its smaller 117M parameter count. In comparison, the Regular model’s 354M parameters require smaller batches to fit in GPU memory.

Figure 10 below shows the dual GPU setup used for one of the training sessions for the regular mode.

Figure 10:GPU detail

4.1.3 Learning Rate Scheduling

Learning Rate Scheduling uses cosine decay with warmup, a two-phase approach that helps prevent training instability. The warmup phase gradually increases the learning rate from 0 to the target value over 500 steps (SLM) or 1000 steps (Regular model), preventing the model from making large, destabilizing updates early in training.

After warmup, cosine decay smoothly reduces the learning rate following a cosine curve to 10% of the initial rate by the end of training. In case you are not familiar with Cosine decay, it is a scheduling strategy where the learning rate follows the shape of a cosine wave: starting at the maximum value after warmup, it decreases slowly at first, then more rapidly in the middle of training, and finally levels off gently near the minimum value.

Mathematically, this follows the curve lr = min_lr + (max_lr - min_lr) × 0.5 × (1 + cos(π × progress)), where progress goes from 0 (start of decay) to 1 (end of training). Unlike linear decay (which drops at a constant rate) or step decay (which drops abruptly at fixed intervals), cosine decay provides a smooth, natural reduction that helps the model explore the loss landscape more effectively early on, then refine its parameters more precisely as training progresses.

The initial learning rates are chosen based on model size: 3e-4 (0.0003) for the SLM model and 3e-5 (0.00003) for the Regular model. The 10x difference reflects the Regular model’s larger parameter count (354M vs 117M) - larger models typically need smaller learning rates to prevent gradient explosion. The cosine decay ensures the model converges smoothly rather than oscillating around the minimum, which is crucial for the complex patterns in historical text.

def get_lr(self, it):
    """Learning rate schedule"""
    warmup_iters = 500
    lr_decay_iters = self.max_iters
    min_lr = self.learning_rate * 0.1
    
    if it < warmup_iters:
        return self.learning_rate * (it + 1) / (warmup_iters + 1)
    if it > lr_decay_iters:
        return min_lr
    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
    assert 0 <= decay_ratio <= 1
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (self.learning_rate - min_lr)

Listing 10: Learning Rate Scheduling Function

The code in Listing 10 shows how the learning rate schedule is implemented. The function takes the current iteration number and returns the appropriate learning rate based on whether we’re in the warmup phase (linear increase) or decay phase (cosine curve). The warmup_iters parameter controls the warmup duration, while min_lr sets the final learning rate to 10% of the initial value.

As a side note, in case you are curious about why a cosine decay specifically makes sense, then read on. The cosine function has unique mathematical properties that make it ideal for learning rate scheduling. Unlike linear decay (which drops too quickly) or exponential decay (which drops too slowly), cosine decay starts with a gentle slope that gradually steepens, then flattens out near the end. This creates a “restart” effect, allowing the model to escape local minima early in training and then fine-tune more precisely in later stages.

The smooth, continuous nature of cosine decay prevents the learning rate from changing too abruptly, which could destabilize training. Given the historical text’s complex linguistic patterns, this gradual, adaptive approach helps the model learn both general language structures and specific historical vocabulary without getting stuck in suboptimal solutions.

4.1.4 Evaluation

We run evaluations at a regular interval after a certain number of steps, which are defined in the eval_interval() method (defaults to 500 for SLM, 1000 for Regular) and compute loss on both train and validation sets using the estimate_loss() method. The different intervals reflect the models’ training complexity: the SLM trains faster and benefits from more frequent checks, while the Regular model’s longer runs can use less frequent evaluation.

The estimate_loss() function monitors training progress without disrupting the learning process. To ensure consistent measurements, it temporarily switches the model to evaluation mode (model.eval()). In this mode, dropout layers stop randomly dropping neurons (using the full network capacity), and batch normalization uses running statistics rather than recomputing them from each batch. This means the same input produces the same output every time, unlike in training mode, where dropout introduces randomness for regularization.

Rather than computing loss on the entire dataset (which would be too slow), estimate_loss() samples eval_iters random batches (default 100) from both training and validation sets. It computes the loss for each batch and returns the average, providing a representative estimate of model performance while remaining computationally efficient.

The evaluation process uses torch.no_grad() to disable gradient computation during validation. Gradients are the partial derivatives that tell us how to adjust each model parameter to reduce loss - they’re computed during the backward pass and stored for the optimizer. During evaluation, we don’t need gradients because we’re not updating weights; we’re just measuring performance.

Disabling gradient computation serves two critical purposes. First, it prevents memory leaks by not storing gradients for validation computations - without this, GPU memory would gradually increase during evaluation and eventually cause out-of-memory errors. Second, it ensures accurate loss measurement by preventing any accidental gradient updates during the evaluation phase. The no_grad() context manager is essential for maintaining training stability and memory efficiency.

4.1.5 Forward Pass

A forward pass is when the model processes input data through its layers to produce a prediction. Think of it like asking the model a question: given a sequence of historical text tokens, “what word should come next?” The model flows the input forward through 12 transformer blocks (SLM) or 24 blocks (Regular), each applying self-attention (to understand relationships between words) and feed-forward operations (to transform and refine the representations). At the end, the model outputs a probability distribution over all possible next tokens.

The forward pass uses mixed precision training with torch.amp.autocast and bf16/fp16 data types, reducing memory usage by ~50% while maintaining training stability. Cross-entropy loss is computed by comparing the model’s predicted probabilities with the actual next tokens in the training data; it measures how “wrong” the model’s predictions are. The loss function handles variable sequence lengths by appropriately padding sequences. The mixed precision approach is particularly important for our historical text corpus, which contains long sequences that would otherwise exceed GPU memory limits.

4.1.6 Backward Pass

After the forward pass tells us how wrong the model is (via the loss), the backward pass figures out how to fix it. Using loss.backward(), PyTorch computes gradients for every parameter in the model - these gradients tell us the direction and magnitude of changes needed to reduce the loss. It’s like having a GPS telling you which way to move and how far, but for 117 million (SLM) or 354 million (Regular) parameters simultaneously.

The system applies gradient clipping with torch.nn.utils.clip_grad_norm_ using a maximum norm of 1.0. Sometimes gradients can become extremely large, especially when processing complex or unusual historical text patterns. Without clipping, these huge gradients would cause the model parameters to jump wildly, potentially making the model unstable or causing it to “forget” what it learned. Clipping acts like a safety valve, limiting the maximum size of parameter updates to keep training stable. Put simply, gradient clipping caps the overall (global) gradient norm at a threshold; if it exceeds the limit, gradients are rescaled so the update stays bounded. In our early runs, omitting clipping occasionally produced NaN losses; keeping max_norm=1.0 eliminated those spikes.

After computing gradients, the system updates the model weights using the AdamW optimizer, which applies the gradients with momentum and adaptive learning rates for each parameter. The optimizer decouples weight decay (a regularization technique to prevent overfitting) from gradient updates, improving generalization. Finally, gradients are zeroed with optimizer.zero_grad(set_to_none=True) - this clears the gradient buffers before the next iteration, preventing them from accumulating across batches. The set_to_none=True option releases memory immediately rather than waiting for GC, improving memory efficiency.

4.1.7 Checkpointing

Checkpointing saves model state, optimizer state, iteration number, and best validation loss whenever validation performance improves, rather than at every evaluation. This selective saving strategy provides multiple benefits: it conserves disk space (our 354M-parameter model checkpoints are ~1.4GB each), reduces I/O overhead that can slow down training, and improves overall training time by 5-10% by eliminating redundant disk writes. The system maintains only the last 5 checkpoints (as configured in config.py), with PyTorch’s torch.save() using compression to ensure efficient storage while preserving all necessary training state for resuming. We’ll dive deeper into checkpointing strategies and implementation details in Section 6.

The training loop implements standard optimization practices, including dynamic learning rate scheduling, regular evaluation and checkpointing, and comprehensive logging to WandB (as detailed in Section 5). The system automatically saves checkpoints when validation loss improves, ensuring that the best model is always preserved. The learning rate schedule uses cosine decay with warmup, which is standard practice for transformer training.

📁 Full Implementation:

SLM: 04_training/train_model_slm.py

Regular Model: 04_training/train_model.py

4.2 Model Initialization: Setting Up the Training Foundation

Before the training loop can begin, the system must properly initialize the model, optimizer, and training infrastructure. The init_model() method handles this setup, ensuring everything is configured correctly for efficient training.

4.2.1 Model Configuration and Creation

The initialization process starts by loading metadata from the tokenized data to ensure the model architecture matches the training data. The system reads vocabulary size, block size, and other parameters from the meta.pkl file created during data preparation, ensuring consistency between the model and the data it will be trained on.

The model configuration is built from the SLM parameters defined in config.py, including the number of layers (12), attention heads (12), embedding dimensions (768), and other architectural choices. This configuration is then used to create the SimpleGPT model instance, which inherits from PyTorch’s nn.Module and provides all the functionality we discussed in the architecture section.

4.2.2 Optimizer Setup and Configuration:

The optimizer is the algorithm that actually updates the model’s parameters (weights and biases) during training. After the backward pass computes gradients (which tell us how to adjust each parameter), the optimizer applies those gradients to update the parameters and improve the model.

The system uses AdamW (Adam with Weight Decay), which is a popular optimizer for training transformers. AdamW combines the best of two approaches: Adam (which adapts the learning rate for each parameter individually, helping with convergence) and weight decay (a form of regularization that prevents overfitting by discouraging large parameter values).

However, not all parameters should be regularized the same way. The optimizer splits parameters into two groups for different weight decay:

2D parameters (weight matrices): These are the main “learnable” parts of the model - the connections between neurons in different layers. These receive weight decay (value 0.1) to prevent them from growing too large, which helps prevent overfitting.
1D parameters (biases): These are additive constants that help shift the model’s predictions. They don’t receive weight decay (value 0.0) because regularizing biases doesn’t help with overfitting and can actually hurt performance.

This two-group approach follows standard practices for transformer training and ensures the model generalizes well to unseen historical text.

Modern PyTorch supports “fused” optimizer operations, which combine multiple steps into a single, faster GPU kernel. Instead of executing separate operations (unscale gradients, update parameters, update optimizer state), fused AdamW performs all three in a single optimized GPU operation. This can provide 10-20% speedup on modern GPUs. The system automatically detects whether your PyTorch version supports fused operations and uses them when available, falling back to the standard implementation otherwise.

Concretely, we use AdamW with the following settings for this project: betas=(0.9, 0.95), weight_decay=0.1, and the learning rate provided by the scheduler (warmup + cosine decay). The AdamW eps parameter is left at the PyTorch default unless you change it in code. When available, the fused AdamW kernel is enabled automatically. See Listing 11 for the exact call in init_model().

4.2.3 Model Compilation and Multi-GPU Setup

Model compilation with PyTorch’s torch.compile is similar to traditional code compilation, but with important differences. When you compile Python code (like using gcc for C), the compiler transforms the source code into optimized machine code once, which then runs faster. Similarly, torch.compile takes your model’s computation graph and optimizes it, but it does this at runtime rather than ahead of time.

The compilation process analyzes your model’s operations (matrix multiplications, attention layers, etc.) and generates optimized kernels tuned to your hardware. This includes operator fusion (combining multiple operations into single GPU kernels), memory layout optimization (arranging data for better cache usage), and kernel selection (choosing the fastest implementation for your specific GPU). The result is often 1.2-1.5x speedier training, but with an initial “warmup” cost: the first few forward/backward passes are slower while PyTorch analyzes the model and generates optimized code.

This differs from traditional compilation because the optimization happens dynamically based on actual input shapes and hardware capabilities, rather than being pre-computed. It’s more like a JIT compiler that specializes your model’s operations for the exact conditions it encounters during training.

For multi-GPU training, the model is wrapped with DistributedDataParallel (DDP), which enables parallel training across multiple GPUs. The DDP wrapper handles gradient synchronization and ensures that all GPUs work with identical model parameters throughout training.

def init_model(self):
    """Initialize the model"""
    logger.info("Initializing model...")
    
    # Load metadata from tokenized data
    meta_path = self.data_dir / "meta.pkl"
    meta_vocab_size = None
    if meta_path.exists():
        with open(meta_path, 'rb') as f:
            meta = pickle.load(f)
        meta_vocab_size = meta['vocab_size']
        logger.info(f"Found vocab_size = {meta_vocab_size}")
    
    # Create model configuration
    model_args = dict(
        n_layer=self.n_layer,        # 12 for SLM
        n_head=self.n_head,          # 12 for SLM  
        n_embd=self.n_embd,          # 768 for SLM
        block_size=self.block_size,  # 512 for SLM
        bias=self.bias,              # False
        vocab_size=meta_vocab_size,  # From tokenized data
        dropout=self.dropout         # 0.1
    )
    
    # Create and configure model
    gptconf = SimpleGPTConfig(**model_args)
    self.model = SimpleGPT(gptconf)
    self.model.to(self.device)
    
    # Initialize optimizer with proper parameter groups
    self.optimizer = self.model.configure_optimizers(
        weight_decay=0.1,
        learning_rate=self.learning_rate,
        betas=(0.9, 0.95),
        device_type='cuda' if 'cuda' in self.device else 'cpu'
    )
    
    # Compile model for performance
    if torch.cuda.is_available() and self.slm_config.get("enable_compile", True):
        logger.info("Compiling model...")
        self.model = torch.compile(self.model, mode='reduce-overhead')
    
    # Wrap with DDP for multi-GPU training
    if self.ddp:
        self.model = DDP(self.model, device_ids=[self.ddp_local_rank])
        param_count = self.model.module.get_num_params()
    else:
        param_count = self.model.get_num_params()
    
    logger.info(f"Model initialized with {param_count:,} parameters")

Listing 11: Model Initialization Process

While our model is a relatively simple toy example focused on a single domain (historical London text), proper initialization remains important to avoid common training issues. The vocabulary size must match our custom historical tokenizer, the sequence length needs to work with our tokenized data, and the model architecture should be appropriate for the text patterns we’re learning.

The initialization process ensures these basic requirements are met before training begins, preventing issues such as vocabulary mismatches or memory allocation problems that could lead to training failures. This careful setup was helpful during our development process, where we ran nearly 100 training experiments. Proper initialization helped us avoid some basic configuration errors and focus on the actual training challenges.

Reproducibility and random seeds: To make runs repeatable on the same hardware, we set a deterministic seed per process using torch.manual_seed(1337 + seed_offset), where seed_offset is the DDP rank (0 for single‑GPU). This gives consistent data shuffling and initialization across restarts while keeping each process distinct under DDP. Note that some CUDA kernels (and AMP/bf16) can introduce non‑determinism; for strict determinism, you may also configure PyTorch’s deterministic flags at the cost of performance.

5. WandB Integration

Weights & Biases (WandB) is an experiment tracking and monitoring platform designed specifically for machine learning projects. Think of it as a “black box” for your training runs - it automatically records everything that happens during training so you can understand what worked, what didn’t, and why.

Training a language model is a long‑running experiment. Without live telemetry, we are flying blind and can’t tell whether learning is stable, whether hardware is saturated, or whether runs are comparable. WandB gives real‑time visibility, remote monitoring, and reproducibility. It records loss, learning rate, and perplexity over time; captures GPU utilization and iteration latency; logs configuration and artifacts; and lets you compare runs side‑by‑side to understand which settings worked.

The system includes WandB integration for experiment tracking and monitoring, with automatic configuration logging, real-time metric tracking (including loss, perplexity, and learning rate), model checkpoint integration, experiment comparison across different training runs, and resource monitoring (GPU utilization and memory usage). This integration helps track and compare different training runs, identify better configurations, and reproduce successful experiments.

Understanding WandB integration:

Listing 12 logs the signals you need to fly by instruments: loss and perplexity trends for learning, the LR schedule to confirm warmup/decay, and hardware utilization and iteration timing for throughput and stability. It’s not just logging - it’s how you compare runs and catch issues early.

This real-time monitoring lets us spot problems early, compare different training approaches, and ensure our historical language model is learning properly over days or weeks.

# Log to WandB - loss first for better mobile UI
if self.use_wandb:
    wandb.log({
        "train/loss": lossf,
        "train/lr": lr,
        "train/iter": iter_num,
        "train/mfu": running_mfu * 100 if running_mfu > 0 else 0,
        "train/dt_ms": dt * 1000,
    })

Listing 12: WandB Integration and Logging

The system logs training loss, learning rate, iteration number, model flops utilization (MFU), and training time per iteration. These metrics provide comprehensive insight into training progress, efficiency, and potential issues.

The most useful dials to watch are training loss (should steadily trend from ~8–10 toward ~2–4), MFU (a proxy for GPU efficiency - single‑digit theoretical targets but mid‑20s achievable with good tuning), the learning‑rate curve (warmup then cosine decay), and iteration time (a practical signal for throughput and stalls).

Both SLM and Regular model training runs complete 60,000 iterations, providing consistent training depth across both model variants. Figure 11 below shows the complete training experience for our Regular model (354M parameters), demonstrating both the console output and WandB’s comprehensive monitoring capabilities.

Figure 11:Complete training run output showing console logs and WandB monitoring for Regular model (354M parameters)

Whilst it might be obvious, the screenshot in Figure 11 captures the final moments of a successful 60,000-iteration training run, showing both the real-time console output and WandB’s comprehensive run summary. In this run, the logs reveal the training progression through the final iterations (59,850 to 60,000), with training loss steadily decreasing from 3.0575 to 2.7063, demonstrating healthy convergence.

The WandB run summary provides the complete picture: a final training loss of 2.70315, a validation loss of 3.61921, and a validation perplexity of 37.31, all indicating successful model training. The system automatically saved the final checkpoint and cleaned up old checkpoints, while WandB captured the entire training journey with detailed metrics tracking. This comprehensive monitoring approach ensures we can both track progress in real time and analyze the full training history afterward.

6. Checkpointing and Model Persistence

Checkpointing is one of the most critical aspects of training large language models, especially for historical text, where training can take days or weeks. A robust checkpointing system ensures that training progress is never lost due to hardware failures, power outages, or other interruptions. In this section, we’ll explore the comprehensive checkpointing system built for the helloLondon project, covering everything from basic checkpoint creation to advanced resume functionality.

6.1 Checkpoint System

The training system implements a practical checkpointing system that preserves all aspects of training state, ensuring that training can be resumed from exactly where it left off. This is particularly important for any complex model, where training can take a long time.

Each checkpoint packages four essentials: the model weights (so learning is preserved), the optimizer state (so momentum and adaptive stats resume cleanly), the current iteration (so schedules pick up in the right place), and the best validation loss to date (so we only promote genuinely better models). Together, these let you stop and restart without losing training dynamics.

The code in Listing 13 shows how these components are saved when validation loss improves:

if losses['val'] < best_val_loss:
    best_val_loss = losses['val']
    if iter_num > 0:
        checkpoint = {
            'model': raw_model.state_dict(),
            'optimizer': self.optimizer.state_dict(),
            'iter_num': iter_num,
            'best_val_loss': best_val_loss,
        }
        checkpoint_path = self.output_dir / f'checkpoint-{iter_num}.pt'
        logger.info(f"Saving checkpoint to {checkpoint_path}")
        torch.save(checkpoint, checkpoint_path)
        
        # Clean up old checkpoints - keep only the last 3
        self.cleanup_old_checkpoints()

Listing 13: Checkpointing and Model Persistence

The checkpointing system uses a simple yet effective approach: it saves checkpoints only when the validation loss improves, rather than at every evaluation. This approach serves multiple purposes. First, it ensures we’re always keeping the best-performing model, not just the most recent one. Second, it significantly reduces I/O overhead during training, as checkpoint saves can be expensive operations (our 354M parameter model checkpoints are ~1.4GB each). Third, it prevents disk space issues by avoiding the accumulation of suboptimal checkpoints. This selective checkpointing approach can improve overall training time by 5-10% by eliminating redundant disk writes.

6.2 Checkpoint Management and Cleanup

Since our 354M parameter model checkpoints are ~1.4GB each, we need to clean up old checkpoints to avoid running out of disk space. The system automatically keeps only the last 5 checkpoints and deletes older ones (as defined in config.py). The cleanup function in Listing 14 finds all checkpoint files, sorts them by modification time (newest first), and deletes everything except the most recent 5. Only the master process handles cleanup to avoid race conditions in multi-GPU setups.

def cleanup_old_checkpoints(self, keep_last=5):
    """Clean up old checkpoints, keeping only the last N"""
    if not self.master_process:
        return  # Only the master process should clean up
        
    try:
        # Find all checkpoint files
        checkpoint_files = list(self.output_dir.glob("checkpoint-*.pt"))
        
        if len(checkpoint_files) <= keep_last:
            return  # Not enough checkpoints to clean up
        
        # Sort by modification time (newest first)
        checkpoint_files.sort(key=lambda x: x.stat().st_mtime, reverse=True)
        
        # Keep the newest ones, delete the rest
        files_to_delete = checkpoint_files[keep_last:]
        
        for file_path in files_to_delete:
            try:
                file_path.unlink()
                logger.info(f"Deleted old checkpoint: {file_path.name}")
            except Exception as e:
                logger.warning(f"Failed to delete checkpoint {file_path.name}: {e}")
                
    except Exception as e:
        logger.warning(f"Checkpoint cleanup failed: {e}")

Listing 14: Checkpoint Cleanup and Management

6.3 Resume Training Functionality

The ability to resume training from any checkpoint is useful when training gets interrupted. This functionality lets you pick up where you left off, whether the interruption was a few minutes or longer.

The resume functionality loads a checkpoint file and restores the training state: the model weights, optimizer state, current iteration number, and best validation loss. If checkpoint loading fails, the code falls back to starting from scratch.

When loading checkpoints, the code handles two practical considerations. First, the map_location=self.device parameter ensures the checkpoint loads onto the correct device (CPU or GPU), which matters if you’re resuming on different hardware or after a restart. Second, for multi-GPU setups using DistributedDataParallel, the model is wrapped in a .module attribute, so the code uses raw_model = self.model.module if self.ddp else self.model to access the actual model underneath.

def resume_from_checkpoint_file(self):
    """Resume training from a checkpoint file"""
    if not self.resume_from_checkpoint:
        return
        
    checkpoint_path = Path(self.resume_from_checkpoint)
    if not checkpoint_path.exists():
        logger.error(f"Checkpoint file not found: {checkpoint_path}")
        return
        
    logger.info(f"Resuming from checkpoint: {checkpoint_path}")
    
    try:
        # Load checkpoint
        checkpoint = torch.load(checkpoint_path, map_location=self.device)
        
        # Load model state
        raw_model = self.model.module if self.ddp else self.model
        raw_model.load_state_dict(checkpoint['model'])
        logger.info("Model state loaded successfully")
        
        # Load optimizer state
        self.optimizer.load_state_dict(checkpoint['optimizer'])
        logger.info("Optimizer state loaded successfully")
        
        # Get iteration number and best validation loss
        self.start_iter = checkpoint.get('iter_num', 0)
        self.best_val_loss = checkpoint.get('best_val_loss', 1e9)
        
        logger.info(f"Resuming from iteration: {self.start_iter}")
        logger.info(f"Best validation loss so far: {self.best_val_loss:.4f}")
        
    except Exception as e:
        logger.error(f"Failed to load checkpoint: {e}")
        logger.info("Starting training from scratch...")
        self.start_iter = 0
        self.best_val_loss = 1e9

Listing 15: Resume Training from Checkpoint

The function loads the model weights, optimizer state, iteration number, and best validation loss from the checkpoint file, then continues training from where it left off. If the checkpoint file doesn’t exist or can’t be loaded, it logs an error and starts training from scratch. Since our Regular model takes 28-32 hours to train, resuming from a checkpoint saves significant time when training is interrupted by power outages, crashes, or manual stops.

7. Training Launch and Management

7.1 Multi-GPU Training with torchrun

For a single GPU, you can run the training script directly — a single Python process will use that device. To use multiple GPUs, launch training with torchrun, which spawns one worker process per GPU and lets the code initialize DistributedDataParallel (DDP). This enables larger effective batch sizes and faster wall‑clock training while keeping weights synchronized across devices; set --nproc_per_node to the number of GPUs you want to use (for example, --nproc_per_node=2).

torchrun is PyTorch’s recommended launcher for distributed training: it initializes the distributed backend and sets environment variables (RANK, LOCAL_RANK, WORLD_SIZE) to keep workers in sync. With torchrun --nproc_per_node=N (where N is the number of GPUs to use — it can be less than the total GPUs available), batches are sharded across the chosen GPUs and gradients are synchronized after each backward pass, which often gives near‑linear speedups on a small multi‑GPU node.

# Single GPU (even with multiple available)
python train_model_slm.py

# Multi-GPU with near-linear speedup
torchrun --nproc_per_node=2 train_model_slm.py

The training script handles DDP (DistributedDataParallel) via train_model_slm.py for gradient sync and batch distribution across GPUs. Figure 12 below shows an example where we have dual GPUs and both are being used.

Figure 12: Multiple GPU used for training

Note that if you run python train_model_slm.py on a multi‑GPU machine, only one GPU is used; the others remain idle. To use more than one GPU, we must use torchrun.

7.2 Training Monitoring

Training is monitored locally via structured console logs and remotely via WandB. The snippet in Listing 16 records loss, learning rate, timing, and MFU at a configurable interval and, when enabled, streams the same metrics to WandB for side‑by‑side run comparison.

# Timing and logging
t1 = time.time()
dt = t1 - t0
t0 = t1
if iter_num % self.log_interval == 0 and self.master_process:
    lossf = loss.item()
    if local_iter_num >= 5:
        mfu = raw_model.estimate_mfu(self.batch_size, dt)
        running_mfu = mfu if running_mfu == -1.0 else 0.9*running_mfu + 0.1*mfu
    logger.info(f"iter {iter_num}: loss {lossf:.4f}, time {dt*1000:.2f}ms, mfu {running_mfu*100:.2f}%")
    
    # Log to WandB - loss first for better mobile UI
    if self.use_wandb:
        wandb.log({
            "train/loss": lossf,
            "train/lr": lr,
            "train/iter": iter_num,
            "train/mfu": running_mfu * 100 if running_mfu > 0 else 0,
            "train/dt_ms": dt * 1000,
        })

Listing 16: Training Monitoring and Logging

Together, console logs and WandB provide real‑time visibility and reproducible experiment tracking; Figure 13 below shows an example of the console logs; see Section 5 for setup and dashboards.

Figure 13: Console training logs: iteration, loss, step time, MFU, and checkpoint saves

8. Model File Formats and Conversion

Training produces PyTorch checkpoint files (.pt) that contain model weights, optimizer state, and training metadata — everything needed to resume training. These checkpoints are covered in detail in Section 6 .

For sharing models and standard deployment workflows, we convert PyTorch checkpoints into the Hugging Face repository format. This conversion creates a portable, standardized model package that can be loaded with standard Hugging Face APIs.

8.1 Converting PyTorch Checkpoints to Hugging Face Format

The Hugging Face repository format is a standardized directory structure containing:

config.json: Architecture definition (layers, heads, embedding dimensions, vocabulary size, sequence length). Allows AutoModelForCausalLM to reconstruct the model architecture without custom code.
model.safetensors: Model weights in SafeTensors format (memory-mapped, secure loading). Contains only model parameters, no optimizer state — suitable for inference workloads.
generation_config.json: Default text generation parameters (max_new_tokens, temperature, top_p, repetition_penalty). Can be overridden at runtime.
Tokenizer files (tokenizer.json, vocab.json, merges.txt, special_tokens_map.json, tokenizer_config.json): Serialized tokenizer with vocabulary, merge rules, normalization, and special tokens matching the training configuration.

The conversion code in Listing 17 loads a PyTorch checkpoint, extracts model weights and config, handles torch.compile naming prefixes if present, and saves the model and tokenizer in Hugging Face format.

import torch
from transformers import GPT2LMHeadModel

def convert_pytorch_to_huggingface(pytorch_checkpoint_path, output_dir, tokenizer):
    """Convert PyTorch checkpoint to Hugging Face format"""
    
    # Load PyTorch checkpoint
    checkpoint = torch.load(pytorch_checkpoint_path, map_location='cpu')
    model_state = checkpoint['model']
    config = checkpoint['config']
    
    # Handle torch.compile prefixes
    if any(key.startswith('_orig_mod.') for key in model_state.keys()):
        clean_state = {}
        for key, value in model_state.items():
            clean_state[key[10:]] = value if key.startswith('_orig_mod.') else value
        model_state = clean_state
    
    # Convert to Hugging Face format
    hf_model = GPT2LMHeadModel(config)
    hf_model.load_state_dict(model_state)
    
    # Save in Hugging Face format
    hf_model.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)

Listing 17: PyTorch → Hugging Face Conversion (essentials)

The conversion handles a few practical details. If the model was compiled with torch.compile, parameter names are prefixed with _orig_mod., which the code strips to match Hugging Face module names. GPT2LMHeadModel(config) instantiates a GPT-2-style architecture that matches the checkpoint’s layer structure, and load_state_dict() loads the weights with automatic shape validation. The save_pretrained() method writes all required files to disk.

File sizes: PyTorch checkpoints are ~450MB (SLM) and ~1.4GB (Regular model); the Hugging Face format reduces this slightly by excluding the optimizer state. The tokenizer adds ~15MB to the repository.

9. Inference Options

Inference can run directly from PyTorch checkpoints or from Hugging Face models. PyTorch checkpoints are convenient during development since you can test any training checkpoint without conversion. Hugging Face models use standard from_pretrained() APIs and are better suited for sharing and deployment workflows.

# Option 1: PyTorch checkpoint inference (direct from training)
python 06_inference/inference_pytorch.py \
  --checkpoint 09_models/checkpoints/slm/checkpoint-60001.pt \
  --prompt "In the year 1834, I walked through the streets of London and witnessed"

# Option 2: Hugging Face model inference (published models)
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("bahree/london-historical-slm")
model = AutoModelForCausalLM.from_pretrained("bahree/london-historical-slm")

inputs = tokenizer("In the year 1834, I walked through the streets...", return_tensors="pt")
outputs = model.generate(inputs['input_ids'], max_new_tokens=50, do_sample=True)
result = tokenizer.decode(outputs[0], skip_special_tokens=True)

Listing 18: Inference Options

Both methods load in seconds and generate ~50–100 tokens/sec on typical consumer GPUs (2–4GB VRAM for SLM, 6–8GB for the Regular model). Use PyTorch checkpoints for development and training comparisons; use Hugging Face models for production deployment and sharing. For interactive testing with published models, see Part 1 .

10. Summary

We built a training‑ready GPT pipeline for historical text, end‑to‑end: a clear decoder‑only architecture, pragmatic GPU/precision tuning, DDP for scale, resilient checkpointing/resume, WandB tracking, and clean hand‑off of artifacts (PyTorch checkpoints → Hugging Face export).

Outcome: two working models on the Part 2 corpus - 117M (SLM) and 354M (Regular) - ready for inference now and for evaluation/deployment in Part 4.

🔗 GitHub Repository: github.com/bahree/helloLondon - Complete training infrastructure (04_training/), model architecture (config.py), and GPU configuration (08_documentation/GPU_TUNING.md)

🧱 Series Posts: Part 1 – Using the Published Historical Models | Part 2 – Data Collection & Custom Tokenizer | Part 3 (this post) | Part 4 – Evaluation & Deployment

🤗 Published Models: SLM Model | Regular Model - Ready-to-use historical language models on HuggingFace

📚 Book Reference: Generative AI in Action - For deeper understanding of core LLM concepts.

Ready for Part 4? Part 4 covers model evaluation, testing, and deployment strategies that turn your trained models into working systems ready for real-world use.

References

Vaswani et al. (2017) – Attention Is All You Need: https://arxiv.org/abs/1706.03762
Radford et al. (2019) – Language Models are Unsupervised Multitask Learners: https://www.semanticscholar.org/paper/Language-Models-are-Unsupervised-Multitask-Learners-Radford-Wu/9405cc0d6169988371b2755e573cc28650d14dfe
Brown et al. (2020) – Language Models are Few-Shot Learners: https://arxiv.org/abs/2005.14165
Kaplan et al. (2020) – Scaling Laws for Neural Language Models: https://arxiv.org/abs/2001.08361
Hoffmann et al. (2022) – Training Compute-Optimal LLMs (Chinchilla): https://arxiv.org/abs/2203.15556
Chowdhery et al. (2022) – PaLM: Scaling Language Modeling with Pathways: https://arxiv.org/abs/2204.02311
Clark et al. (2019) – What Does BERT Look At?: https://arxiv.org/abs/1906.04341
Voita et al. (2019) – Analyzing Multi‑Head Self‑Attention: https://arxiv.org/abs/1905.09418
Dao et al. (2022) – FlashAttention: https://arxiv.org/abs/2205.14135
Micikevicius et al. (2018) – Mixed Precision Training: https://arxiv.org/abs/1710.03740
Rajbhandari et al. (2020) – ZeRO: https://arxiv.org/abs/1910.02054
Paszke et al. (2019) – PyTorch: https://arxiv.org/abs/1912.01703
Kingma & Ba (2014) – Adam: A Method for Stochastic Optimization: https://arxiv.org/abs/1412.6980
Loshchilov & Hutter (2017) – AdamW Decoupled Weight Decay Regularization : https://arxiv.org/abs/1711.05101
Smith & Topin (2017) – Super‑Convergence: Very Fast Training of Neural Networks Using Large Learning Rates: https://arxiv.org/abs/1708.07120
Goyal et al. (2017) – Accurate, Large Minibatch SGD: https://arxiv.org/abs/1706.02677
Sergeev & Del Balso (2018) – Horovod: https://arxiv.org/abs/1802.05799
Pope et al. (2022) – Efficiently Scaling Transformer Inference: https://arxiv.org/abs/2211.05102
Jawahar et al. (2019) – What does BERT learn about the structure of language?: https://aclanthology.org/P19-1356.pdf
Mikolov et al. (2013) – Word2vec: https://arxiv.org/abs/1301.3781
Pennington et al. (2014) – GloVe: https://aclanthology.org/D14-1162/
Devlin et al. (2018) – BERT https://arxiv.org/abs/1810.04805
Press & Wolf (2017) – Using the Output Embedding to Improve Language Models: https://arxiv.org/abs/1608.05859
Inan et al. (2016) – Tying Word Vectors and Word Classifiers: https://arxiv.org/abs/1611.01462

Happy Diwali - 2025

Sun, 19 Oct 2025 00:00:00 +0000

I 💖 Diwali. For those celebrating, from my family to yours!

🏛️Building LLMs from Scratch - Part 2: Data Collection & Custom Tokenizers

Sun, 12 Oct 2025 00:00:00 +0000

TL;DR

In this second part of our 4-part series on building language models from scratch, I explore the two foundational areas of LLM development: data collection and custom tokenizer creation. Part 1 - Building LLM from Scratch covered using the published model; here, we build the complete pipeline from raw historical documents to a custom tokenizer that understands archaic English, London geography, and period-specific terminology.

The challenge with historical LLMs isn’t just having enough data—it’s having the right data processed to preserve linguistic nuances across different historical periods. This post demonstrates how to transform over 218 historical sources into a corpus of more than 500 million characters using a specialized tokenizer for authentic historical text generation.

⚠️ Educational Purpose: This is a learning project designed to teach LLM development concepts. For production-scale LLMs, you’ll need significantly larger datasets, more sophisticated infrastructure, and additional considerations that are not covered in this post.

1. The Historical Language Modeling Challenge

Building a language model for historical text presents unique challenges. Historical English from 1500 to 1850 contains linguistic patterns, vocabulary, and cultural references that modern tokenizers have never encountered. Standard tokenizers like TikToken fragment archaic words like “quoth” and “hast” into multiple subword tokens, destroying semantic meaning crucial for historical text generation.

A simple phrase like Quoth the alderman, 'Tis a fair day at Newgate becomes dozens of meaningless fragments, losing both historical context and linguistic coherence. This fragmentation is why we built a custom tokenizer trained specifically on historical English patterns, ensuring the model can generate coherent, historically accurate text.

As a reminder, both the SLM (117M parameters) and Regular Model (354M parameters) utilize the same training code and infrastructure, including GPU optimization, checkpointing, and WandB integration. The only difference lies in the model architecture parameters, which are specified in config.py.

🔗 GitHub Repository: github.com/bahree/helloLondon - Complete source code for data collection (02_data_collection/) and tokenizer training (03_tokenizer/). We will see the relevant code snippets in this post show key concepts—see the full implementation in the repository.

🧱 Series Posts: Part 1 – Using the Published Historical Models | Part 2 (this post) | Part 3 – Training Architecture & GPU Optimization | Part 4 – Evaluation & Deployment

What will you learn?

This project provides hands-on experience with real-world LLM development challenges, including data collection from over 218 historical sources, cleaning OCR errors and encoding issues, and developing custom tokenizers for historical text. Unlike theoretical tutorials, you receive complete, runnable code that demonstrates actual trade-offs and decisions—such as choosing BPE over WordPiece or handling different file formats—that you’d encounter in any serious LLM project.

While operating at a learning scale, the principles taught here directly apply to larger systems. Data collection patterns, cleaning strategies, and tokenizer design principles scale from our 500M character corpus to the 500B+ character datasets used in production models.

1.1 High-Level Process Overview

The complete pipeline transforms raw historical documents into a working language model through five key stages:

Data Collection: 218+ historical sources (1500-1850), including literature, newspapers, court records, and personal diaries
Cleaning Pipeline: Handles multiple file formats (PDF, HTML, XML, TXT) while removing OCR artifacts and preserving authentic historical language
Quality Validation: Removes duplicates, filters non-English content, and ensures only meaningful historical text reaches the final corpus
Custom Tokenizer Training: BPE-based tokenizer with ~150 special tokens capturing archaic pronouns, historical landmarks, and period-specific terminology
Model Training: Two language models (SLM 117M and Regular 354M parameters) trained on the same historical corpus

The result is a system capable of generating authentic historical text that captures the linguistic patterns and cultural context of 1500-1850 English. Figure 1 illustrates this complete pipeline:

graph TD
    A[📚 218+ Historical Sources<br/>1500-1850] --> B[🔍 Data Collection<br/>Download and Filter]
    B --> C[🧹 5-Phase Cleaning Pipeline<br/>Format-Specific Processing]
    C --> D[📊 Quality Validation<br/>Duplicate and Language Detection]
    D --> E[📝 500M+ Character Corpus<br/>Clean Historical Text]
    E --> F[🔤 Custom Tokenizer Training<br/>BPE with 150+ Special Tokens]
    F --> G[🤖 Language Model Training<br/>SLM 117M + Regular 354M]
    
    style A fill:#e1f5fe
    style E fill:#f3e5f5
    style F fill:#fff3e0

Figure 1: Complete Historical Text Processing Pipeline

2. Data Collection: The Foundation of Historical Language Modeling

Let us dig deeper into steps 1-4: data collection, cleaning, validation, and corpus creation. The data collection system processes over 218 sources spanning the years 1500-1850 to create a corpus of over 500 million characters of authentic historical English text. But collecting historical data isn’t just about downloading files - it’s about handling the sheer variety of formats and quality levels that historical documents present.

Historical documents come in all shapes and sizes - scanned books with OCR errors, HTML pages with messy markup, XML archives with rich metadata, and plain text files with inconsistent encoding. This is especially true for the earlier periods, when the quality of the documents can vary significantly, and most modern techniques for processing them struggle to cope. This data diversity requires a cleaning pipeline that transforms raw historical documents into training data while preserving the authentic language patterns of 1500-1850 English.

2.1 System Architecture: Processing 218+ Historical Sources

The data collection system employs a modular architecture, with historical_data_collector.py serving as the primary orchestration engine, coordinating with a data_sources.json configuration file that contains metadata for over 218 historical sources. This enables easy management and updates without code changes.

Supporting scripts include add_data_source.py for interactive source addition with built-in validation, and generate_report.py for comprehensive reporting and analysis across multiple output formats.

The data_sources.json file contains metadata for each source, including time periods, formats, licensing, and processing priorities. Each entry includes:

time_period (e.g., [1690, 1800] for London Lives)
format (XML, HTML, PDF)
priority (high/medium/low)
search_terms for collection guidance

Our data sources span multiple categories, each contributing unique perspectives to the historical corpus.

Project Gutenberg: This provides foundational literature with 8+ carefully selected texts, using relaxed quality criteria that accept texts with as low as 40% meaningful words to capture the full spectrum of historical writing styles.
Historical Archives: Historical Archives like London Lives (240,000 pages of personal records) and Old Bailey (197,000+ trial transcripts) offer rich historical content and were initially enabled in our data collection.
- Note: I was using the aggressive cleaning earlier (enabled using the aggressive_cleaning flag designed to remove structured legal data and semantic markup), and discovered that it was too aggressive and caused generation quality issues. After initial training runs revealed repetitive and incoherent text patterns, I turned off these sources. Enabling this back might be an exercise for you to try.
Archive.org: Archive.org has an API access that can be used for file filtering, and this makes it relatively straightforward.
The National Archives (TNA): TNA records contribute government correspondence and official documents that provide the institutional context for historical events.
British History Online: Finally, these supplements our collection with historical surveys and period documents that offer scholarly perspectives on the time periods we’re modeling.

However, each source type presents unique technical challenges that require specialized processing approaches. One example is Project Gutenberg, which contains files with standardized headers and footers that must be removed. (As a side note, I really appreciate the effort that has gone into this to make this formatting consistent, which makes the process of this relatively straightforward.)

On the other hand, PDF files often suffer from OCR errors, especially for older documents that contain corrupted historical language, requiring sophisticated text correction algorithms to restore proper spelling and grammar from scanned documents. The figure below shows one example of how older documents look. This example is “The abridgment of the charter of the city of London” from 1680.

Figure 2: The abridgment of the charter of the city of London (1680) - showing faded text and ink blots typical of historical documents

As you can see, the text is faded, has ink blots, and the font style is very different from modern text. OCR software often misinterprets characters in such documents, resulting in numerous errors, as illustrated in the image below. These OCR artifacts can severely degrade the quality of our training data if not properly addressed.

Figure 3: OCR errors in the charter document - showing how optical character recognition struggles with historical fonts and document quality

HTML files from sources like Archive.org contain navigation elements, advertisements, and modern web markup that contaminate the historical corpus, demanding careful content extraction that preserves only the meaningful historical text.

XML archives like London Lives and Old Bailey require specialized parsing to extract meaningful text while preserving semantic markup that provides context about speakers, dates, and document structure - a delicate balance between removing technical artifacts and maintaining historical authenticity.

Government records from TNA often contain bureaucratic formatting, form fields, and institutional language that need careful filtering to extract the human stories and historical narratives.

British History Online documents present challenges with academic formatting, footnotes, and scholarly apparatus that must be processed to maintain readability while preserving the scholarly context that makes them valuable for historical language modeling.

2.2 Cleaning Pipeline

I implement a 5-stage cleaning pipeline that helps transform the raw historical documents into training-ready text. Each phase addresses specific challenges that would otherwise contaminate our language model training.

2.2.1 Stage 1: File Discovery & Initial Filtering

Historical archives often contain files in various formats, which may be missing proper file extensions or have non-standard naming conventions. Many files contain non-English content that would contaminate our English historical corpus. Additionally, many sources employ their own templates and standards for this purpose. To resolve this, we first implement a simple file detection and naming cleanup, as shown in Listing 1. The code itself is simple and self-explanatory.

def detect_file_type(file_path: str) -> str:
    """Detect file type based on extension and content analysis"""
    # Extension-based detection
    if file_path.endswith(('.txt', '.txt.utf-8', '_txt.utf-8')):
        return 'text'
    elif file_path.endswith(('.pdf',)):
        return 'pdf'
    elif file_path.endswith(('.html', '.htm')):
        return 'html'
    elif file_path.endswith(('.xml',)):
        return 'xml'
    
    # Content-based detection for files without extensions
    with open(file_path, 'rb') as f:
        content = f.read(1024)  # Read first 1KB
        if b'<html' in content.lower() or b'<!doctype' in content.lower():
            return 'html'
        elif b'<?xml' in content.lower():
            return 'xml'
        elif content.isascii() and b'\x00' not in content:
            return 'text'
        else:
            return 'binary'

Listing 1: File Type Detection Function

When we run this locally, we will see the flow as outlined below, which illustrates how the detection works. This, of course, can be made more robust for non-English characters, but for now, we reject these.

File Type Detection Flow:

📁 Raw Files (218+ sources)
    ↓
🔍 File Type Detection
    ├── .txt, .txt.utf-8, _txt.utf-8 → Text Processing
    ├── .pdf → PDF Processing  
    ├── .html, .htm → HTML Processing
    ├── .xml → XML Processing (Old Bailey, London Lives)
    └── No Extension → Content Detection
        ├── HTML-like content → HTML Processing
        ├── Text-like content → Text Processing
        └── Binary/Unknown → REJECTED
    ↓
🚫 Filename Language Check
    ├── Non-English characters → REJECTED (logged)
    └── English/Latin → Continue

Historical archives often lack standardized file extensions and contain content in languages other than English. Our two-stage detection ensures we capture valuable historical documents while filtering out irrelevant files, preventing both data loss and processing waste.

2.2.2 Stage 2: Format-Specific Content Extraction

Each file format requires specialized processing due to its unique contamination sources, including Project Gutenberg headers, PDF OCR errors, HTML navigation elements, and XML structural markup. Our format-specific extraction functions clean these artifacts while preserving authentic historical content.

Text Files (.txt, .txt.utf-8)

Project Gutenberg texts contain standardized headers and footers that would confuse our language model. The cleaning process removes these while preserving the actual historical content. The code snippet in Listing 2 demonstrates this approach and is quite straightforward. Of course, this can be made more robust, but this works well for our selected texts.

def clean_gutenberg_text(text: str) -> str:
    """Clean Project Gutenberg text by removing headers/footers and metadata"""
    lines = text.split('\n')
    cleaned_lines = []
    in_content = False
    
    for line in lines:
        # Skip Gutenberg headers (before "*** START OF")
        if "*** START OF" in line:
            in_content = True
            continue
        # Skip Gutenberg footers (after "*** END OF")
        if "*** END OF" in line:
            break
        # Skip metadata lines
        if line.startswith(('Title:', 'Author:', 'Release Date:', 'Language:')):
            continue
        # Skip empty lines at start
        if not in_content and not line.strip():
            continue
            
        if in_content:
            cleaned_lines.append(line)
    
    return '\n'.join(cleaned_lines).strip()

Listing 2: Project Gutenberg Text Cleaning Function

Real Example - Before Cleaning:

Title: A Journal of the Plague Year
Author: Daniel Defoe
Release Date: March 2003
Language: English

*** START OF THE PROJECT GUTENBERG EBOOK A JOURNAL OF THE PLAGUE YEAR ***

It was about the beginning of September 1664, that I, among the rest of my neighbours, heard in ordinary discourse that the plague was returned again in Holland...

*** END OF THE PROJECT GUTENBERG EBOOK A JOURNAL OF THE PLAGUE YEAR ***

After Cleaning:

It was about the beginning of September 1664, that I, among the rest of my neighbours, heard in ordinary discourse that the plague was returned again in Holland...

Without this cleaning, the model would learn to generate Gutenberg headers and metadata instead of authentic historical text, contaminating the training data with modern digital artifacts.

PDF Files

PDF files from historical archives often contain OCR errors and digital artifacts that require correction. The cleaning process in Listing 3 addresses these issues while preserving historical content, removing page numbers and all-caps headers. While not perfect, it significantly improves text quality.

The OCR correction rules are based on common patterns in historical documents and can be refined for specific datasets. Libraries like PyMuPDF or pdfplumber extract text, while regex-based cleaning corrects common OCR errors and removes digital stamps. More advanced techniques, such as layout analysis or AI-based OCR correction, can further enhance this process.

def clean_pdf_text(text: str) -> str:
    """Clean PDF text by removing OCR artifacts and digital stamps"""
    # Remove page numbers: [Page 123], standalone numbers
    text = re.sub(r'\[Page \d+\]', '', text)
    text = re.sub(r'^\d+$', '', text, flags=re.MULTILINE)
    
    # Remove library stamps: Internet Archive, Google, etc.
    stamps = [
        'Internet Archive', 'Google Books', 'HathiTrust',
        'Digitized by Google', 'Scanned by Google'
    ]
    for stamp in stamps:
        text = text.replace(stamp, '')
    
    # Fix common OCR artifacts
    ocr_fixes = {
        r'\b0\b': 'O',  # 0 → O
        r'\b1\b': 'I',  # 1 → I  
        r'\b5\b': 'S',  # 5 → S
        r'\b8\b': 'B',  # 8 → B
        r'\brn\b': 'm', # rn → m
        r'\bcl\b': 'd'  # cl → d
    }
    
    for pattern, replacement in ocr_fixes.items():
        text = re.sub(pattern, replacement, text)
    
    # Remove all-caps lines (usually headers)
    lines = text.split('\n')
    cleaned_lines = [line for line in lines if not line.isupper() or len(line) < 10]
    
    return '\n'.join(cleaned_lines)

Listing 3: PDF Text Cleaning Function

Real Example - Before Cleaning:

[Page 1]
INTERNET ARCHIVE
A JOURNAL OF THE PLAGUE YEAR
BY DANIEL DEFOE

It was about the beginning of September 1664, that I, among the rest of my neighbours, heard in ordinary discourse that the plague was returned again in Holland. For it was indeed a very terrible time, and the people began to be very much alarmed at it.

After Cleaning:

A JOURNAL OF THE PLAGUE YEAR
BY DANIEL DEFOE

It was about the beginning of September 1664, that I, among the rest of my neighbours, heard in ordinary discourse that the plague was returned again in Holland. For it was indeed a very terrible time, and the people began to be very much alarmed at it.

OCR errors can significantly impact the quality of model training. For example, if London appears as L0nd0n due to OCR errors, the model won’t learn the correct spelling and will generate nonsensical text when asked about historical London. The correction process ensures our model learns authentic historical language patterns rather than digital artifacts, which is crucial for generating coherent and historically accurate text.

HTML Files

HTML files from historical websites and digital archives contain markup that needs to be stripped while preserving the actual text content. We use the BeautifulSoup library in Listing 4 to clean the HTML structure and extract only the meaningful text.

def clean_html_text(html_content: str) -> str:
    """Clean HTML content by removing markup and extracting text"""
    from bs4 import BeautifulSoup
    
    soup = BeautifulSoup(html_content, 'html.parser')
    
    # Remove unwanted elements
    for element in soup(['script', 'style', 'nav', 'header', 'footer', 
                        'aside', 'menu', 'form', 'input', 'button']):
        element.decompose()
    
    # Remove wiki-specific elements
    for element in soup.find_all(['div', 'span'], class_=['navbox', 'infobox', 'sidebar']):
        element.decompose()
    
    # Remove navigation elements
    for element in soup.find_all(['div', 'ul'], class_=['breadcrumb', 'navigation', 'menu']):
        element.decompose()
    
    # Extract text content
    text = soup.get_text(separator=' ', strip=True)
    
    # Clean up excessive whitespace
    text = re.sub(r'\s+', ' ', text)
    
    return text.strip()

Listing 4: HTML Text Cleaning Function

Real Example - Before Cleaning:

<!DOCTYPE html>
<html>
<head><title>London History</title></head>
<body>
<nav>Home | About | Contact</nav>
<header>London Historical Society</header>
<div class="content">
    <h1>The Great Fire of London</h1>
    <p>In the year 1666, a great fire consumed much of London...</p>
</div>
<footer>© 2024 London Historical Society</footer>
</body>
</html>

After Cleaning:

The Great Fire of London in the year 1666, a great fire consumed much of London...

HTML tags and navigation elements would contaminate training, causing the model to generate markup instead of historical text. Our cleaning process extracts meaningful content while preserving natural flow and structure.

XML Files (Historical Archives):

XML files from historical archives, such as the Old Bailey and London Lives, use specific schemas that require specialized parsing. Old Bailey employs TEI (Text Encoding Initiative) with TEI.2 elements, while London Lives uses semantic markup (name, geo, occupation). These structured formats contain authentic historical language with rich metadata, as shown in Listing 5.

def extract_old_bailey_text(soup) -> str:
    """Extract text from Old Bailey XML using TEI schema structure"""
    extracted_text = []
    
    # Check for TEI.2 elements (Old Bailey schema)
    tei_elements = soup.find_all('TEI.2')
    if tei_elements:
        # Extract trial accounts (main narrative content)
        trial_accounts = soup.find_all('div1', {'type': 'trialAccount'})
        for trial in trial_accounts:
            trial_text = extract_trial_narrative(trial)
            if trial_text:
                extracted_text.append(trial_text)
        
        # Extract front matter (session information)
        front_matter = soup.find_all('div1', {'type': 'frontMatter'})
        for front in front_matter:
            front_text = extract_front_matter_narrative(front)
            if front_text:
                extracted_text.append(front_text)
    
    return '\n\n'.join(extracted_text)

def extract_london_lives_text(soup) -> str:
    """Extract text from London Lives XML using semantic markup schema"""
    extracted_text = []
    
    # Check for London Lives specific elements (name, geo, occupation, date)
    name_elements = soup.find_all('name')
    geo_elements = soup.find_all('geo')
    occupation_elements = soup.find_all('occupation')
    
    if name_elements and geo_elements and occupation_elements:
        # Extract paragraphs with semantic markup
        paragraphs = soup.find_all('p')
        for para in paragraphs:
            p_text = extract_paragraph_with_semantic_markup(para)
            if p_text.strip():
                extracted_text.append(p_text)
    
    return '\n\n'.join(extracted_text)

Listing 5: XML Text Extraction Functions

Real Example - Old Bailey XML (Before Processing):

<trial>
<frontmatter>
<session>Session 1</session>
<date>1674-04-15</date>
<location>Old Bailey</location>
</frontmatter>
<proceedings>
The prisoner being brought to the bar, and the indictment being read, he pleaded Not Guilty. The witnesses being sworn, the first witness deposed that on the 15th day of April last, he saw the prisoner in the company of several suspicious persons...
</proceedings>
</trial>

After Processing:

Session 1 1674-04-15 Old Bailey The prisoner being brought to the bar, and the indictment being read, he pleaded Not Guilty. The witnesses being sworn, the first witness deposed, that on the 15th day of April last, he saw the prisoner in the company of several suspicious persons...

These XML files contain the most authentic historical language in our entire dataset. The Old Bailey trials show how people actually spoke in court during the 17th-19th centuries, while London Lives reveals the everyday language used in personal records and official documents. This authentic historical language is very useful for training a model that can generate historically accurate text, as it provides the model with genuine examples of how people wrote and spoke during different historical periods.

2.2.3 Stage 3: Text Normalization

After extraction, text normalization ensures consistency and compatibility with the training data. Historical documents contain encoding issues, inconsistent formatting, and special characters that confuse the model. Our normalization process fixes these issues and breaks long lines to fit within the model’s context window. This is critical because lines exceeding the context window appear as incomplete sentences to the transformer, severely degrading generation quality due to the attention mechanism’s inability to process fragmented text.

Inconsistent encoding and formatting can severely confuse the language model during training. For example, if some files use smart quotes (") and others use straight quotes ("), the model might not learn that they represent the same concept, leading to inconsistent and potentially incorrect text generation. Normalization ensures that the model observes consistent patterns across all training data, which is crucial for learning coherent language patterns and generating high-quality historical text.

The code snippet in Listing 6 demonstrates how we implement this normalization, which is quite straightforward.

def normalize_text(text: str) -> str:
    """Normalize text for consistent training data"""
    import unicodedata
    
    # Fix common encoding issues
    encoding_fixes = {
        'â€™': "'",  # Smart apostrophe
        'â€œ': '"',  # Smart quote left
        'â€': '"',   # Smart quote right
        'â€"': '—',  # Em dash
        'â€¢': '•',  # Bullet point
        'â€¦': '…',  # Ellipsis
    }
    
    for old, new in encoding_fixes.items():
        text = text.replace(old, new)
    
    # Normalize Unicode (NFC)
    text = unicodedata.normalize('NFC', text)
    
    # Break long lines for training compatibility (max 2000 chars)
    lines = text.split('\n')
    normalized_lines = []
    for line in lines:
        if len(line) > 2000:
            # Split at sentence boundaries
            sentences = re.split(r'(?<=[.!?])\s+', line)
            current_line = ""
            for sentence in sentences:
                if len(current_line + sentence) > 2000:
                    if current_line:
                        normalized_lines.append(current_line.strip())
                    current_line = sentence
                else:
                    current_line += " " + sentence if current_line else sentence
            if current_line:
                normalized_lines.append(current_line.strip())
        else:
            normalized_lines.append(line)
    
    # Normalize line endings and whitespace
    text = '\n'.join(normalized_lines)
    text = re.sub(r'[ \t]+', ' ', text)  # Multiple spaces/tabs to single space
    text = re.sub(r'\n\s*\n', '\n\n', text)  # Multiple newlines to double newline
    
    return text.strip()

Listing 6: Text Normalization Function

Real Example - Before Normalization:

The year was 1666, and the plague had come to London. â€œIt was indeed a very terrible time,â€ wrote one observer. The streets were filled with the sounds of horse-drawn carriages and the cries of the afflicted.

After Normalization:

The year was 1666, and the plague had come to London. "It was indeed a very terrible time," wrote one observer. The streets were filled with the sounds of horse-drawn carriages and the cries of the afflicted.

2.2.4 Stage 4: Quality Validation

Not all extracted text is suitable for training. Some files contain duplicates, non-English content, or poor-quality text that would degrade model performance. We need a comprehensive validation system that ensures only high-quality, relevant text is included in our training corpus.

The key challenge is striking a balance between quality standards and historical value. A strict approach might reject valuable historical documents that have some OCR issues, while a lenient approach might include too much low-quality content, which can degrade model training. To address this, I implemented a tiered quality threshold system that applies different standards based on content type:

General Content: 200+ chars, 50+ words, 50% meaningful words
Project Gutenberg: 200+ chars, 50+ words, 40% meaningful words (relaxed for historical value)
Historical Documents: 1000+ chars, 100+ words, 30% meaningful words (very relaxed for historical value)

This tiered approach ensures that we capture valuable historical content while maintaining quality standards, filtering out duplicates, non-English content, and low-quality text, thereby preserving the integrity of useful historical documents. Again, these implementations are quite simple, in the context of a toy project, but can be made more robust. The code itself is quite straightforward, as shown in Listing 7.

def analyze_text_quality(text: str, source_type: str = 'general') -> dict:
    """Analyze text quality and determine if it should be included in training corpus"""
    import hashlib
    import re
    
    # Length validation
    char_count = len(text)
    word_count = len(text.split())
    
    # OCR artifact detection using regex patterns
    ocr_patterns = {
        'long_capitals': r'[A-Z]{5,}\s+[A-Z]{5,}',
        'spaced_letters': r'\b[A-Za-z]\s+[A-Za-z]\s+[A-Za-z]\s+[A-Za-z]\b',
        'special_chars': r'[!@#$%^&*()]{3,}',
        'mixed_alphanumeric': r'\b\d+[A-Za-z]+\d+\b',
        'long_non_word': r'[^\w\s]{10,}'
    }
    
    ocr_issues = []
    for pattern_name, pattern in ocr_patterns.items():
        if re.search(pattern, text):
            ocr_issues.append(pattern_name)
    
    # Advertisement detection
    ad_patterns = [
        'this day is published', 'just ready', 'elegantly bound',
        'now ready', 'new novels', 'advertisements',
        'price \d+s', 'paternoster row', 'corner of', 'publishers'
    ]
    
    ad_count = sum(1 for pattern in ad_patterns if re.search(pattern, text, re.IGNORECASE))
    ad_density = ad_count / max(word_count, 1)
    
    # Meaningful word ratio calculation
    words = text.split()
    meaningful_words = [w for w in words if w.isalpha() and len(w) > 2]
    meaningful_ratio = len(meaningful_words) / max(len(words), 1)
    
    # Quality thresholds based on source type
    thresholds = {
        'general': {'min_chars': 200, 'min_words': 50, 'min_meaningful_ratio': 0.50},
        'gutenberg': {'min_chars': 200, 'min_words': 50, 'min_meaningful_ratio': 0.40},
        'historical': {'min_chars': 1000, 'min_words': 100, 'min_meaningful_ratio': 0.30}
    }
    
    threshold = thresholds.get(source_type, thresholds['general'])
    
    # Quality scoring
    score = 100
    score -= len(ocr_issues) * 3  # OCR issues
    score -= ad_density * 50  # Advertisement density
    score -= (1 - meaningful_ratio) * 20  # Meaningful word ratio
    
    # Check if text meets quality thresholds
    meets_thresholds = (
        char_count >= threshold['min_chars'] and
        word_count >= threshold['min_words'] and
        meaningful_ratio >= threshold['min_meaningful_ratio'] and
        ad_density < 0.1  # Less than 10% advertisement content
    )
    
    return {
        'char_count': char_count,
        'word_count': word_count,
        'meaningful_ratio': meaningful_ratio,
        'ocr_issues': ocr_issues,
        'ad_density': ad_density,
        'score': score,
        'meets_thresholds': meets_thresholds,
        'content_hash': hashlib.md5(text.encode()).hexdigest()
    }

Listing 7: Text Quality Analysis Function

Content Quality Validation

Our validation system employs multiple detection mechanisms to ensure training corpus quality:

OCR Artifact Detection: Regex patterns identify common digitization errors, including misread headers, character separation failures, scanning artifacts, alphanumeric misinterpretations, and corrupted text regions
Advertisement Filtering: Pattern matching detects commercial content using phrases like “this day is published”, “just ready”, “elegantly bound”, and price references
Quality Scoring: A 100-point system deducts points for OCR artifacts (-3 each), advertisement density (-50), and low meaningful word ratios (-20)

This multi-layered approach balances quality standards with preservation of valuable historical content, ensuring the model trains on authentic historical language while filtering out contamination sources.

2.2.5 Stage 5: Final Processing and Corpus Creation

After cleaning and validation, we create a final training corpus optimized for language model training. This requires intelligent segmentation that breaks long texts into manageable chunks while preserving the historical narrative flow, which is essential given the context window limits (e.g., 2048 tokens). The code snippet in Listing 8 demonstrates this final processing stage.

def create_comprehensive_corpus(cleaned_files: list) -> str:
    """Create final training corpus with intelligent segmentation"""
    corpus_parts = []
    
    for file_path in cleaned_files:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # Split into training segments
        segments = split_into_training_segments(content)
        corpus_parts.extend(segments)
    
    # Create final corpus
    final_corpus = '\n\n'.join(corpus_parts)
    
    # Save to file
    with open('london_historical_corpus_comprehensive.txt', 'w', encoding='utf-8') as f:
        f.write(final_corpus)
    
    return final_corpus

def split_into_training_segments(text: str, max_length: int = 2000) -> list:
    """Split text into training segments while preserving narrative flow"""
    # First split on double newlines (paragraphs)
    paragraphs = text.split('\n\n')
    
    segments = []
    current_segment = ""
    
    for paragraph in paragraphs:
        if len(current_segment + paragraph) <= max_length:
            current_segment += paragraph + '\n\n'
        else:
            if current_segment:
                segments.append(current_segment.strip())
            current_segment = paragraph + '\n\n'
    
    if current_segment:
        segments.append(current_segment.strip())
    
    # Further split long segments at sentence boundaries
    final_segments = []
    for segment in segments:
        if len(segment) > max_length:
            sentences = re.split(r'(?<=[.!?])\s+', segment)
            current_segment = ""
            for sentence in sentences:
                if len(current_segment + sentence) <= max_length:
                    current_segment += sentence + " "
                else:
                    if current_segment:
                        final_segments.append(current_segment.strip())
                    current_segment = sentence + " "
            if current_segment:
                final_segments.append(current_segment.strip())
        else:
            final_segments.append(segment)
    
    # Filter out segments that are too short
    return [seg for seg in final_segments if len(seg) >= 50]

Listing 8: Corpus Creation and Segmentation Functions

During my local runs, this final processing stage generated a comprehensive corpus of over 500 million characters across ~250,000 segments, with an average segment length of around 2,000 characters. The success rate of files making it into the final corpus ranged from 70% to 90%, depending on the quality and availability of the source.

Final Corpus Statistics:

Total Sources Processed: 218+ historical sources
Final Corpus Size: 500M+ characters
Training Segments: ~250,000 segments
Average Segment Length: ~2,000 characters
Success Rate: 70-90% (depending on source availability)

2.3 Detailed Data Processing Flow

Building on the high-level flow and having reviewed each of the areas, the detailed flow below illustrates the complete data cleaning process, including rejection paths, error handling, and statistics tracking. This is intended to provide a bird’s-eye view of the entire process.

graph TD
    A[📁 Raw Files] --> B{File Type Detection}
    
    B -->|.txt, .txt.utf-8| C[📄 Text File]
    B -->|.pdf| D[📄 PDF File]
    B -->|.html, .htm| E[📄 HTML File]
    B -->|.xml| F[📄 XML File]
    B -->|No Extension| G{Content Detection}
    
    G -->|HTML-like| E
    G -->|Text-like| C
    G -->|Binary/Unknown| REJECT1[❌ REJECTED]
    
    C --> H[🧹 clean_gutenberg_text]
    D --> I[🔧 extract_text_from_pdf]
    E --> J[🧹 clean_html_text]
    F --> K{XML Type Detection}
    
    I --> L[🧹 clean_pdf_text]
    
    K -->|Old Bailey| M[🔧 extract_old_bailey_text]
    K -->|London Lives| N[🔧 extract_london_lives_text]
    
    M --> O[🧹 clean_old_bailey_text]
    N --> P[🧹 clean_london_lives_text]
    
    H --> Q[🔧 normalize_text]
    L --> Q
    J --> Q
    O --> Q
    P --> Q
    
    Q --> R[🔍 Duplicate Detection]
    R -->|Duplicate| REJECT2[❌ REJECTED - Duplicate]
    R -->|Unique| S[🌍 Language Detection]
    
    S -->|Non-English| REJECT3[❌ REJECTED - Non-English]
    S -->|English| T[📊 Quality Analysis]
    
    T --> U{Quality Check}
    U -->|Poor Quality| REJECT4[❌ REJECTED - Poor Quality]
    U -->|Good Quality| V[💾 Save to Processed Directory]
    
    V --> W[📊 Update Statistics]
    W --> X[✅ Successfully Processed]
    
    REJECT1 --> Y[📝 Log Rejection Reason]
    REJECT2 --> Y
    REJECT3 --> Y
    REJECT4 --> Y
    
    Y --> Z[📊 Update Rejection Stats]
    
    style A fill:#e1f5fe
    style X fill:#c8e6c9
    style REJECT1 fill:#ffcdd2
    style REJECT2 fill:#ffcdd2
    style REJECT3 fill:#ffcdd2
    style REJECT4 fill:#ffcdd2
    style Y fill:#fff3e0
    style Z fill:#fff3e0

Figure 4: Detailed Data Processing Pipeline

2.5 Corpus Creation Process

After cleaning, the system creates the final training corpus through intelligent segmentation that preserves historical narrative flow:

📁 Cleaned Files
    ↓
🔧 create_comprehensive_corpus()
    ├── Read all cleaned_*.txt files
    ├── Split into training segments (split_into_training_segments)
    │   ├── Split on double newlines (paragraphs)
    │   ├── Max length: 2000 characters
    │   ├── Min length: 100 characters
    │   └── Further split long segments at sentence boundaries
    ├── Filter segments (min 50 characters)
    └── Write to london_historical_corpus_comprehensive.txt

The corpus creation process reads all cleaned text files and intelligently segments them into training-ready chunks. It first splits on double newlines to preserve paragraph boundaries, which are natural break points in historical text. Segments are constrained to a maximum of 2000 characters to fit within the model’s context window, with a minimum of 100 characters to ensure substantial content. Long segments are further split at sentence boundaries to maintain readability. Finally, segments shorter than 50 characters are filtered out as they’re unlikely to contain meaningful historical content.

Proper segmentation is crucial for training language models. The model needs to learn from coherent text segments that maintain historical narrative flow while fitting within its context window. Splitting on paragraph boundaries preserves the natural structure of historical documents, while sentence-level splitting ensures that very long paragraphs don’t exceed the model’s processing capabilities. This approach maximizes the model’s ability to learn from authentic historical language patterns while maintaining training efficiency.

2.6 Outcome: Training-Ready Corpus

The result is a clean, historically faithful corpus containing over 500 million characters of authentic historical English spanning 350 years of London history from 1500-1850. The corpus comprises high-quality text with minimal OCR artifacts, preserving historical language patterns and a rich cultural context that reflects the social, political, and economic realities of various historical periods. The text has been intelligently segmented for optimal language model training, with careful attention to maintaining the natural flow of historical narratives while ensuring compatibility with modern training techniques.

This corpus serves as the essential foundation for training our specialized historical tokenizer and language model, ensuring the model learns authentic historical English rather than modern text patterns. By providing the model with genuine examples of how people wrote and spoke during different historical periods, we enable it to generate text that captures the linguistic nuances, cultural references, and historical context that make historical language modeling both challenging and rewarding.

💻 Try It Yourself: The complete implementation, including all the data collection scripts, cleaning algorithms, and quality validation systems described in this section, is available in the helloLondon GitHub repository . The repository includes detailed documentation, example usage, and step-by-step guides for setting up your own historical language model training pipeline.

Now that we have examined the data collection and cleaning process, we can proceed to the next steps: creating a custom historical tokenizer and preparing for model training.

3. Custom Historical Tokenizer: The Key to Authentic Historical Text Generation

Creating a custom tokenizer is crucial for generating effective historical text. This section examines the necessity of a custom tokenizer, the challenges presented by historical language, and our chosen architecture. The tokenizer preserves the semantic meaning of historical words and phrases, enabling coherent and contextually accurate historical narratives.

Standard tokenizers like GPT-2’s fragment archaic words like “quoth” and “hast” into multiple subword tokens, destroying semantic meaning crucial for historical text generation.

Real Example - Standard Tokenizer vs. Our Custom Tokenizer:

Standard GPT-2 Tokenizer:

"Quoth the alderman, 'Tis a fair day at Newgate"
→ ['Qu', 'oth', ' the', ' ald', 'erman', ',', ' ', ''', 'T', 'is', ' a', ' fair', ' day', ' at', ' New', 'gate']

Our Custom Historical Tokenizer:

"Quoth the alderman, 'Tis a fair day at Newgate"
→ ['<|quoth|>', ' the', ' alderman', ',', ' ', ''', '<|tis|>', ' a', ' fair', ' day', ' at', ' <|newgate|>']

The standard tokenizer breaks historical language into 18 meaningless fragments, losing semantic meaning and historical context. Our custom tokenizer reduces this to 12 meaningful tokens, preserving authentic historical language patterns essential for coherent text generation.

A tokenizer that fragments historical language destroys the model’s ability to learn authentic patterns. The model needs to perceive “quoth” as a single concept, rather than fragmented subwords, to capture the linguistic nuances of different historical periods.

3.1 What Happens with Off-the-Shelf Tokenizers

What would happen if we used standard tokenizers like tiktoken or GPT-2’s tokenizer?

Standard tokenizers would force the model to waste capacity reconstructing fragmented historical words from subwords rather than learning historical language patterns. The model might learn to generate “Qu” + “oth” but struggle to use “quoth” in new contexts. Historical phrases like “methinks” would split into meaningless fragments, losing semantic coherence. London geography becomes particularly problematic, as place names like “Newgate” fragment, making spatial relationships harder to understand.

Generation Quality Issues:

# What you'd get with standard tokenizer:
"Quoth the alderman, 'Tis a fair day at Newgate"
→ Generates: "Qu oth the ald erman, 'T is a fair day at New gate"
→ Result: Broken, unreadable historical text

# What you get with our custom tokenizer:
"Quoth the alderman, 'Tis a fair day at Newgate"  
→ Generates: "Quoth the alderman, 'Tis a fair day at Newgate"
→ Result: Authentic, coherent historical text

A vocabulary that’s too small (10K tokens) would fragment even more historical words, making the problem worse, while a vocabulary that’s too large (100K+ tokens) would overfit to rare historical terms, wasting capacity on words that appear only once. Our choice of 30K tokens provides a balanced approach that captures common historical patterns without overfitting, ensuring the model learns the most important historical language patterns efficiently.

Real-World Example: With a standard tokenizer, our model might generate:

"The ald erman walk ed to New gate where he saw the pris oner"

With our custom tokenizer, it generates:

"The alderman walked to Newgate where he saw the prisoner"

The difference in historical text authenticity is significant between the two approaches.

3.2 Tokenizer Architecture

I had started with the easier WordPiece tokenizer (more of an accident rather than by design). Still, I realized later that it was unsuitable for historical text due to the ## subword prefix artifacts. We need a tokenizer that can handle historical English efficiently while preserving semantic meaning, unlike standard tokenizers like GPT-2’s WordPiece approach, which fragments historical language and, as a result, destroys the linguistic patterns we want to preserve. After some experimentation, I settled on a custom Byte Pair Encoding (BPE) tokenizer trained specifically on historical English.

BPE is a subword tokenization algorithm that learns to break text into meaningful subword units by iteratively finding the most frequent character pairs in the training corpus and merging them into single tokens. The process begins with individual characters and gradually evolves into common words and phrases.

For example, if "th" appears frequently in our historical corpus, BPE will learn to treat it as a single token rather than separate "t" and "h" tokens. This is particularly valuable for historical English, where words like "thou", "thee", and "thine" share common prefixes and suffixes.

3.2.1 Tokenizer Training Process

The BPE training algorithm analyzes our entire historical corpus to identify the most frequent character combinations, building a vocabulary that’s optimized for historical language patterns. We start with a base alphabet (comprising all letters) and special tokens, then iteratively merge the most frequent pairs until we reach our target vocabulary size of 30,000 tokens. This ensures that common historical words, such as "quoth", "hast", and "methinks", are treated as single tokens, while still allowing for the handling of rare or unknown words by breaking them into learned subword units.

The training process is computationally efficient and produces a tokenizer that’s specifically tuned to the linguistic patterns found in our historical corpus.

In this case, we don’t have to reinvent the wheel and use the Hugging Face tokenizers library, which provides a modular approach to building custom tokenizers. The library is organized into several key components: models define the core tokenization algorithm (BPE, WordPiece, Unigram), pre_tokenizers handle initial text splitting, normalizers clean and standardize text, trainers configure the learning process, and processors handle special token insertion. This modular design enables us to mix and match components to create a tokenizer tailored to our specific use case.

The models module offers several tokenization algorithms: BPE() for Byte Pair Encoding (what we use), WordPiece() for Google’s WordPiece algorithm, Unigram() for Google’s Unigram language model, and WordLevel() for simple word-level tokenization.

Each has different strengths - BPE is efficient and handles unknown words well, WordPiece is used by BERT but creates ## artifacts, Unigram is more flexible but computationally expensive, and WordLevel is simple but creates very large vocabularies.

Let us look at the code in Listing 9 for training our custom historical BPE tokenizer:

def train_tokenizer(self):
    """Train a custom tokenizer for historical English"""
    # Import the tokenizers library components
    from tokenizers import Tokenizer, models, pre_tokenizers, processors, trainers
    from tokenizers.normalizers import Sequence, NFD, StripAccents
    
    logger.info("Training custom historical tokenizer...")
    logger.info(f"Corpus: {self.corpus_path}")
    logger.info(f"Target vocabulary: {self.vocab_size:,} tokens")
    logger.info(f"Output directory: {self.output_dir}")
    
    # Initialize BPE tokenizer (not WordPiece)
    # models.BPE() creates a Byte Pair Encoding model that will learn subword patterns
    tokenizer = Tokenizer(models.BPE())
    
    # Normalizers for historical text - preserve case for better text reconstruction
    # Normalizers clean and standardize text before tokenization
    tokenizer.normalizer = Sequence([
        NFD(),           # Unicode normalization - converts characters to canonical form
        StripAccents()   # Remove accents - converts "café" to "cafe"
    ])
    
    # Pre-tokenizer for historical English - use simple whitespace splitting
    # Pre-tokenizers split text into initial segments before the main tokenization
    tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
        pre_tokenizers.WhitespaceSplit(),  # Split on whitespace
        pre_tokenizers.Punctuation()       # Split punctuation from words
    ])
    
    # Special tokens for historical English
    special_tokens = [
        "<|startoftext|>", "<|endoftext|>", "<|pad|>", "<|unk|>", "<|mask|>",
        # Historical language tokens
        "<|thou|>", "<|thee|>", "<|thy|>", "<|thine|>", "<|hast|>", "<|hath|>",
        "<|doth|>", "<|dost|>", "<|quoth|>", "<|tis|>", "<|twas|>", "<|twill|>",
        # London geography tokens
        "<|london|>", "<|thames|>", "<|westminster|>", "<|tower|>", "<|newgate|>",
        "<|southwark|>", "<|cheapside|>", "<|fleet|>", "<|ludgate|>", "<|aldgate|>",
        # Historical period tokens
        "<|tudor|>", "<|stuart|>", "<|georgian|>", "<|regency|>", "<|victorian|>",
        # Social class tokens
        "<|noble|>", "<|gentleman|>", "<|commoner|>", "<|apprentice|>", "<|yeoman|>",
        # Professional tokens
        "<|apothecary|>", "<|coachman|>", "<|chimneysweep|>", "<|baker|>", "<|butcher|>"
    ]
    
    # BPE trainer configuration
    # The trainer defines how the BPE algorithm learns from our corpus
    trainer = trainers.BpeTrainer(
        vocab_size=self.vocab_size,        # Target vocabulary size (30,000 tokens) - balanced between coverage and efficiency
        special_tokens=special_tokens,     # Pre-defined tokens that are always included
        min_frequency=2,                   # Minimum frequency prevents vocabulary pollution from OCR errors
        show_progress=True,                # Display training progress
        # Removed continuing_subword_prefix="##" to eliminate WordPiece-style artifacts
        # This ensures pure BPE tokenization without ## symbols in generated text
        initial_alphabet=["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", 
                        "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z",
                        "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M",
                        "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z"]
    )
    
    # Train the tokenizer on our historical corpus
    # This is where the BPE algorithm learns the optimal subword patterns
    tokenizer.train([str(self.corpus_path)], trainer)
    
    return tokenizer

Listing 9: Custom Tokenizer Training Function

3.2.2 Tokenization Architecture Decisions

Our custom historical tokenizer necessitated several critical design decisions to handle historical English effectively. We evaluated multiple tokenization approaches including Byte Pair Encoding (BPE) ( Sennrich et al., 2016 ), WordPiece ( Schuster & Nakajima, 2012 ), Unigram Language Model ( Kudo, 2018 ), SentencePiece ( Kudo & Richardson, 2018 ), and traditional character-level and word-level tokenization. Each approach has distinct trade-offs: BPE produces clean subwords without special markers (used by GPT models), WordPiece adds ## prefixes that contaminate generated text (used by BERT), Unigram uses probabilistic modeling but is computationally expensive, SentencePiece treats text as raw bytes and excels at multilingual scenarios, while character-level and word-level tokenization either produce impractically long sequences or massive vocabularies.

For historical text generation, BPE provides the optimal balance of clean output, efficient training, and effective vocabulary coverage, as demonstrated by Radford et al., 2019 and Mielke et al., 2021 . We also preserve case throughout tokenization, since historical text often uses capitalization for semantic meaning (e.g., “Thou” vs. “thou”), and include over 150 carefully designed special tokens that capture historical language patterns, London geography, and social context. This combination ensures our tokenizer can effectively learn and generate authentic historical language while maintaining computational efficiency.

3.3 Special Token Design: Capturing Historical Language Patterns

Historical English contains linguistic patterns, vocabulary, and cultural references that are no longer present in modern English. Standard tokenizers fragment these patterns, destroying the semantic meaning crucial for historical text generation. The solution here was to design 150 special tokens that capture the essence of historical English, organized into strategic categories that reflect the linguistic and cultural structure of 1500-1850 English.

def create_special_tokens() -> list:
    """Create special tokens for historical English"""
    special_tokens = [
        # Basic control tokens
        "<|startoftext|>", "<|endoftext|>", "<|pad|>", "<|unk|>", "<|mask|>",
        
        # Historical language tokens
        "<|thou|>", "<|thee|>", "<|thy|>", "<|thine|>",  # Second person pronouns
        "<|hast|>", "<|hath|>", "<|doth|>", "<|dost|>",  # Archaic verb forms
        "<|quoth|>", "<|tis|>", "<|twas|>", "<|twill|>", # Common contractions
        
        # London geography tokens
        "<|london|>", "<|thames|>", "<|westminster|>", "<|tower|>", "<|newgate|>",
        "<|southwark|>", "<|cheapside|>", "<|fleet|>", "<|ludgate|>", "<|aldgate|>",
        
        # Historical period tokens
        "<|tudor|>", "<|stuart|>", "<|georgian|>", "<|regency|>", "<|victorian|>",
        
        # Social and professional tokens
        "<|noble|>", "<|gentleman|>", "<|commoner|>", "<|apothecary|>", "<|coachman|>",
        "<|merchant|>", "<|court|>", "<|jury|>", "<|verdict|>", "<|church|>", "<|parish|>"
    ]
    
    return special_tokens

Listing 10: Special Tokens Creation Function

3.3.1 Token Category Analysis

Our special token vocabulary spans ten carefully curated categories, each designed to capture essential aspects of historical London life. The largest categories focus on Historical Language (25 tokens) and London Geography (20 tokens), providing the linguistic and spatial foundation for authentic historical text generation. These tokens capture archaic pronouns like "thou" and "thee," along with specific London locations like "Thames" and "Newgate" that were central to historical narratives.

The remaining categories address the social, professional, and cultural dimensions of historical society. Social Class and Professional tokens (35 tokens combined) reflect the highly stratified nature of historical London, enabling accurate dialogue between nobles, commoners, and various tradespeople. Legal and Judicial tokens support court proceedings from the Old Bailey, while Religious tokens capture the central role of faith in historical society. Temporal, Currency, and Transportation tokens (35 tokens combined) provide the temporal, economic, and logistical context that makes historical narratives authentic and believable.

3.3.2 Special Token Categories Visualization

Let us visualize the special token categories and their relationships as shown below. These special tokens enable the model to understand and generate authentic historical language. Without them, the model would fragment historical concepts into meaningless subwords, losing the cultural and linguistic context that makes historical text generation both challenging and rewarding.

graph LR
    A[🔤 Special Tokens<br/>150+ Total] --> B[📜 Historical Language<br/>25 tokens]
    A --> C[🏛️ London Geography<br/>20 tokens]
    A --> D[⏰ Historical Periods<br/>10 tokens]
    A --> E[👥 Social Classes<br/>15 tokens]
    A --> F[💼 Professions<br/>20 tokens]
    A --> G[⚖️ Legal & Judicial<br/>10 tokens]
    A --> H[⛪ Religious<br/>10 tokens]
    A --> I[🕐 Temporal<br/>15 tokens]
    A --> J[💰 Currency & Measurement<br/>10 tokens]
    A --> K[🚗 Transportation<br/>10 tokens]

    B --> B1["<|thou|>, <|thee|>, <|hast|>, <|doth|>, <|quoth|>"]
    C --> C1["<|london|>, <|thames|>, <|newgate|>, <|westminster|>"]
    D --> D1["<|tudor|>, <|stuart|>, <|georgian|>, <|regency|>"]
    E --> E1["<|noble|>, <|gentleman|>, <|commoner|>, <|yeoman|>"]
    F --> F1["<|apothecary|>, <|coachman|>, <|chimneysweep|>, <|baker|>"]
    G --> G1["<|court|>, <|jury|>, <|verdict|>, <|prisoner|>"]
    H --> H1["<|church|>, <|parish|>, <|prayer|>, <|blessed|>"]
    I --> I1["<|morn|>, <|eve|>, <|season|>, <|year|>"]
    J --> J1["<|shilling|>, <|pound|>, <|yard|>, <|furlong|>"]
    K --> K1["<|coach|>, <|carriage|>, <|horse|>, <|vessel|>"]

    %% class definitions (custom palette matching your original)
    classDef cls_root fill:#e1f5fe,stroke:#81d4fa,color:#000;
    classDef cls_hist fill:#f3e5f5,stroke:#ce93d8,color:#000;
    classDef cls_geo fill:#e8f5e8,stroke:#a5d6a7,color:#000;
    classDef cls_period fill:#fff3e0,stroke:#ffe0b2,color:#000;
    classDef cls_social fill:#fce4ec,stroke:#f8bbd0,color:#000;
    classDef cls_prof fill:#f1f8e9,stroke:#c5e1a5,color:#000;
    classDef cls_legal fill:#e0f2f1,stroke:#80cbc4,color:#000;
    classDef cls_relig fill:#f9fbe7,stroke:#e6ee9c,color:#000;
    classDef cls_temp fill:#e3f2fd,stroke:#90caf9,color:#000;
    classDef cls_curr fill:#fef7e0,stroke:#ffe082,color:#000;
    classDef cls_trans fill:#f3e5f5,stroke:#e1bee7,color:#000;

    %% assign classes
    class A cls_root;
    class B cls_hist;
    class C cls_geo;
    class D cls_period;
    class E cls_social;
    class F cls_prof;
    class G cls_legal;
    class H cls_relig;
    class I cls_temp;
    class J cls_curr;
    class K cls_trans;

Figure 5: Special Token Categories and Examples

3.5 Post-Processing and Hugging Face Integration

After training our custom tokenizer, we need to make it compatible with the broader machine learning ecosystem and ensure it works properly with language model training. Raw tokenizers can only convert text to tokens and back. Still, language models require additional functionality, such as special token handling, sequence padding, and integration with popular frameworks like Hugging Face Transformers.

The challenge, though, is that language model training requires specific formatting that raw tokenizers don’t provide. For example, training sequences need to be wrapped with special start/end tokens (<|startoftext|> and <|endoftext|>), padded to consistent lengths for batch processing, and integrated with the rest of the ecosystem. In our case, we also want to utilize Hugging Face and its ecosystem, allowing us to leverage standard training scripts and model architectures. Without proper post-processing, our custom tokenizer would be incompatible with existing training infrastructure.

We add post-processing capabilities that wrap text sequences with control tokens and create Hugging Face-compatible tokenizer files, ensuring seamless integration with the broader machine learning ecosystem while preserving our historical text optimizations.

There are three key areas that we need to consider:

Understanding Post-Processing: The first step is adding a post-processor that automatically wraps every text sequence with special start and end tokens. This is crucial because language models must be able to identify where sequences begin and end during training. For example, when we tokenize "Hello world", the post-processor automatically converts it to <|startoftext|> Hello world <|endoftext|>. This template processing ensures consistent formatting across all our training data.
Hugging Face Integration: Next, we create a Hugging Face-compatible wrapper around our custom tokenizer. This wrapper maps our special tokens to the standard token types that Hugging Face expects: beginning-of-sequence (bos), end-of-sequence (eos), padding, unknown, and masking tokens. This mapping allows our custom tokenizer to work seamlessly with standard training scripts and model architectures.
Special Token Functions: Each special token serves a specific purpose in language model training. The beginning-of-sequence token indicates when a new text starts, the end-of-sequence token marks the end of the text, padding tokens ensure all sequences in a batch have the same length, unknown tokens handle words not in our vocabulary, and masking tokens are used during training for masked language modeling tasks.

The code in Listing 11 demonstrates how we implement these post-processing steps and create a Hugging Face-compatible tokenizer:

def create_huggingface_tokenizer(tokenizer: Tokenizer, max_length: int = 1024) -> PreTrainedTokenizerFast:
    """Create Hugging Face compatible tokenizer"""
    from transformers import PreTrainedTokenizerFast
    
    # Add post-processor for sequence formatting
    tokenizer.post_processor = processors.TemplateProcessing(
        single="<|startoftext|> $A <|endoftext|>",
        special_tokens=[
            ("<|startoftext|>", 1),
            ("<|endoftext|>", 0),
        ]
    )
    
    # Create Hugging Face tokenizer wrapper
    hf_tokenizer = PreTrainedTokenizerFast(
        tokenizer_object=tokenizer,
        bos_token="<|startoftext|>",
        eos_token="<|endoftext|>",
        pad_token="<|pad|>",
        unk_token="<|unk|>",
        mask_token="<|mask|>",
        model_max_length=max_length
    )
    
    return hf_tokenizer

Listing 11: Hugging Face Tokenizer Integration

Without this integration, our custom tokenizer would be incompatible with standard language model training. The post-processor ensures proper sequence formatting, while the Hugging Face wrapper enables seamless integration with existing training infrastructure and model architectures. This makes our tokenizer compatible with standard training frameworks, allowing for easy sharing and deployment.

3.6 Testing and Validation

We need to ensure the tokenizer works correctly with historical text before using it for model training. This requires testing on diverse historical samples and validating both encoding and decoding accuracy. A simple way to do this is to encode a set of historical text samples, decode them back, and check if the original text is perfectly reconstructed. We also want to verify that special tokens are used correctly in the tokenized output.

def test_historical_tokenizer(tokenizer: Tokenizer) -> dict:
    """Test the trained tokenizer on historical text samples"""
    
    test_texts = [
        "In the year of our Lord 1834, the streets of London were filled with the sounds of horse-drawn carriages.",
        "The gentleman from the country said, 'I have never seen such a sight in all my days.'",
        "The Thames flowed dark and mysterious through the heart of the city.",
        "It was the best of times, it was the worst of times."
    ]
    
    results = {'perfect_reconstruction': 0, 'special_token_usage': 0, 'failed_tests': []}
    
    for i, text in enumerate(test_texts):
        # Encode and decode text
        encoded = tokenizer.encode(text)
        decoded = tokenizer.decode(encoded.ids)
        
        # Check reconstruction accuracy
        if decoded.strip() == text.strip():
            results['perfect_reconstruction'] += 1
        else:
            results['failed_tests'].append({'index': i, 'original': text, 'decoded': decoded})
        
        # Check special token usage
        special_tokens = [token for token in encoded.tokens if token.startswith('<|') and token.endswith('|>')]
        if special_tokens:
            results['special_token_usage'] += 1
    
    return results

Test Results:

Perfect Reconstruction: 99%+ accuracy on test cases
Special Token Usage: 80%+ of test cases use special tokens
Average Compression Ratio: ~0.3 tokens per word (highly efficient)
Success Rate: 99%+ for historical text samples

It is essential to conduct comprehensive testing to ensure the tokenizer operates reliably. In our case, the test cases cover different historical periods, writing styles, and linguistic patterns, giving us confidence that the tokenizer can handle the full range of historical text in our corpus. For a real-world LLM, this is, of course, more complex and would need to cover a broader set of areas.

Tokenizer Performance Validation

Not surprisingly, our custom tokenizer significantly outperforms standard approaches on historical text, as demonstrated by comprehensive metrics that compare it to GPT-2’s tokenizer, as shown in the table below. These metrics indicate that our custom tokenizer significantly outperforms standard approaches for historical text. The improved compression ratio and reconstruction accuracy ensure that the model learns from authentic historical language rather than tokenization artifacts, which is crucial for generating coherent and historically accurate text.

Metric	Standard GPT-2	Our Custom Tokenizer	Improvement
Vocabulary Size	50,257 tokens	30,000 tokens	40% smaller
Special Tokens	4 tokens	150+ tokens	37x more
Compression Ratio	~0.4 tokens/word	~0.3 tokens/word	25% better
Reconstruction Accuracy	95%	99%+	4% better
Historical Language Support	Poor	Good	N/A

These metrics validate that our 30K token vocabulary provides optimal coverage for historical text while remaining manageable for small language models. The 150+ special tokens capture linguistic patterns of 1500-1850 English, and the 25% better compression ratio means historical text is represented more efficiently, allowing the model to process longer sequences. The 99%+ reconstruction accuracy ensures no information is lost during tokenization, while excellent performance on archaic vocabulary, period-specific terminology, and London geography demonstrates the tokenizer’s effectiveness for historical language modeling.

3.8 Implementation and Usage

The complete tokenizer implementation, including training scripts, testing utilities, and validation tools, is available in the helloLondon GitHub repository . The repository provides:

Training Code: Complete BPE tokenizer training with configurable vocabulary sizes and special token definitions.
Testing Utilities: Comprehensive validation tools for testing tokenizer performance on historical text
Integration Examples: Ready-to-use code for incorporating the tokenizer into your own projects
Documentation: Detailed usage guides and API references

This implementation demonstrates how to build production-ready tokenizers for specialized domains, with particular focus on historical language processing and integration with modern ML frameworks.

4. Current Limitations

This project is designed as a learning exercise for those new to AI and LLM development. While we’ve built a functional system that demonstrates core concepts, this is not production-ready code and has several limitations that would need to be addressed for real-world deployment:

Data Scale & Quality:

Corpus size: Our 500M character corpus is tiny compared to production LLMs, which typically use 100x-1000x more data (50B-500B+ characters). This limits the model’s ability to learn diverse patterns and reduces the quality of generated output.
Source diversity: With only 218 sources, we lack comprehensive historical coverage across the 1500-1850 span, potentially missing important linguistic evolution patterns and regional variations.
Geographic bias: Heavy focus on London may not accurately represent broader historical English patterns from other regions, limiting the model’s generalizability.
Bias detection: We lack systematic approaches to identify or mitigate historical biases in the data, which could lead to the model perpetuating outdated or problematic language patterns.
Quality assessment: Our cleaning pipeline, while effective for common issues, overlooks many edge cases and artifacts that would require more sophisticated ML-based quality assessment in production.

Tokenizer & Model Architecture:

Vocabulary size: Our 30K token vocabulary is small compared to modern models (which often use 50K-100K+ tokens), limiting the model’s ability to represent diverse vocabulary efficiently.
Special tokens: The 150+ special tokens are manually curated rather than learned from data, which may miss important patterns that data-driven approaches would discover.
Context length: The 1024 token context window is very short compared to modern models (which often use 4K-32K+ tokens), limiting the model’s ability to maintain coherence in longer texts.
Language support: No support for other languages or historical variants beyond English, significantly limiting the model’s applicability.
Tokenization approach: While our BPE approach is clean and avoids WordPiece artifacts, it may not be optimal for all historical text patterns and could benefit from more sophisticated techniques.

Technical Infrastructure:

Error handling: Basic error handling with limited logging and monitoring makes it difficult to debug issues and track system health in production.
Testing: Minimal test coverage that excludes edge cases means many potential failure modes remain undetected until they occur in production.
Performance: No optimization for speed, memory, or distributed processing, making the system unsuitable for production-scale deployment.
Data management: Lacks data versioning and reproducibility guarantees, making it difficult to track changes and reproduce results across different environments.
Security: No security considerations for data handling and model deployment, creating potential vulnerabilities for sensitive historical data.
Compliance: Missing compliance considerations for GDPR, data privacy, and regulatory requirements, which are essential for production deployment.
Monitoring: No production monitoring, alerting, or observability features, making it impossible to detect and respond to issues in real-time.

These limitations are intentional trade-offs made to keep the project manageable and focused on core learning objectives, but they represent significant gaps for production deployment.

4.3 What You’d Need for Production

Data Engineering and Legal Framework

Production systems require 100x-1000x more data from diverse sources, with ML-based quality assessment, bias detection, and filtering that goes far beyond our simple heuristics. You’d need robust ETL pipelines with proper error handling and monitoring, as well as a comprehensive legal framework for copyright clearance, data licensing, and compliance management, which we haven’t addressed.

Model Architecture and Training

Meaningful historical language understanding would require models with over 1 billion parameters, utilizing sophisticated training techniques, regularization, and optimization. You’d need a comprehensive evaluation on diverse historical text tasks and domain-specific fine-tuning capabilities that our current system doesn’t support.

Infrastructure and Operations

Production deployment requires a multi-GPU, multi-node distributed training infrastructure, production-grade model serving with load balancing and scaling, comprehensive monitoring and alerting systems, and end-to-end security for both data and model protection—none of which our learning-focused system currently provides.

This progression from data → tokenizer → training → deployment provides a complete methodology for building specialized historical language models.

5. Resources and Further Reading

GitHub Repository: github.com/bahree/helloLondon - Complete source code for data collection and tokenizer training
Part 1: Building LLMs from Scratch - Part 1 - Quick start and overview
Documentation: Complete guides in the 08_documentation/ folder covering every aspect of the project
Book Reference: Generative AI in Action - For deeper understanding of core LLM concepts

6. Summary

This post represents Part 2 of our learning journey into the fundamentals of LLM development. While we’ve built a functional data collection and tokenization system demonstrating core concepts, the real value lies in understanding:

Data flow from raw sources to training-ready corpora
Tokenization impact on model performance across different approaches
Challenges in processing historical and domain-specific text
Trade-offs between quality, scale, and complexity
Debugging and improvement strategies for encountered problems

The limitations we’ve identified are great learning opportunities. Every production LLM started as a learning project, and every limitation teaches you something new about how these systems work. This foundation prepares us for the next phase of our journey.

Ready for Part 3? Part 3 will cover the custom GPT architecture, GPU optimization strategies, and training infrastructure that transforms our clean data and custom tokenizer into working language models—while maintaining the same educational focus on understanding the fundamentals.

🧱 Series Posts: Part 1 – Using the Published Historical Models | Part 2 (this post) | Part 3 – Training Architecture & GPU Optimization | Part 4 – Evaluation & Deployment

🏛️How to build a Large Language Model from Scratch - Part 1

Tue, 23 Sep 2025 00:00:00 +0000

TL;DR

In this post, I show how to build a working LLM from scratch and show a complete end-to-end pipeline from data gathering to training to deployment of a language model. For this project I concentrate on Old English and only related to London, using historical London texts (1500-1850). To show the flexibility, I built two language models which are identical in architecture and the only differs is their size and parameters (117M vs 354M).

⚠️ Educational Purpose: This is a learning project designed to teach LLM development concepts. For production-scale LLMs, you’ll need much larger datasets, more sophisticated infrastructure, and additional considerations not covered here.

This guide shows you how to monitor training progression, perform rapid evaluations, test models from both PyTorch checkpoints and published Hugging Face repositories, and ultimately publish your own - supported by complete code, live model artifacts, and educational inference tooling.

4-Part Series:

Part 1 (this): Quick start, inference, and overview
Part 2: Data collection and custom tokenizers
Part 3: Model architecture and GPU training
Part 4: Evaluation and deployment

1. Overview

Train AI models on 1500-1850 London texts. Complete 4-part series covering data collection, training, and deployment. Part 1: Quick start and overview.

📖 Want to understand the core LLM concepts? This series focuses on implementation and hands-on building. For a deeper understanding of foundational concepts like tokenizers, prompt engineering, RAG, responsible AI, fine-tuning, and more, check out my book Generative AI in Action .

You can learn more about the book → by clicking here 📘.

1.1 What was built?

I found many folks don’t understand what it entails to build an LLM, and where we do have guides, they only share piecemeal elements and nothing that is comprehensive for someone who is new to this. There are more detailed guides on fine-tuning existing models, but not much on the complete development pipeline. This series outlines that by walking through the process of creating specialized language models trained exclusively on historical London texts from 1500 to 1850.

I am mostly doing this for my own learning, and also sharing what I can. Many work-related details, for obvious reasons, I cannot share and discuss, but some small pet projects like this embody the same sentiment.

The helloLondon Historical Language Models represent a complete end-to-end implementation, from data collection through deployment. Rather than fine-tuning existing models, I chose to train from the ground up to eliminate modern biases and create models that genuinely understand historical language patterns, cultural contexts, and period-specific knowledge.

Two Model Variants I built two identical models with the same architecture, tokenizer, and training process. The only difference is the number of parameters: an SLM (117M parameters) optimized for learning and resource-constrained environments, and a Regular model (354M parameters) designed for higher-quality generation.

Both use identical code with different configuration files, allowing you to understand the impact of model size on performance and choose the right variant for your needs.

Model	Parameters	Iterations	Training Time*	Use Case
SLM (Small)	117M	60,000	~8-12 hours	Fast inference, resource-constrained
Regular (Full)	354M	60,000	~28-32 hours	High-quality generation

Note: Technically speaking, both these models can be called classified as SLMs given they are 117M and 354M parameters; however, for the sake of this project, I call the smaller of the two the SLM and the other regular.

1.2 Core Pipelines

The complete development pipeline encompasses multiple critical stages that transform raw historical texts into working language models. The process starts with data collection, where we systematically gather and filter over 218 historical London sources spanning 1500–1850. This process ensures we capture authentic period language while minimizing modern biases that could contaminate our models.

Next, we develop a custom tokenization system specifically designed for historical English. This involves training a domain-specific tokenizer with a 30,000-token vocabulary plus 150+ special tokens that capture period language patterns, archaic spellings, and historical terminology that modern tokenizers often miss.

The model architecture phase implements GPT-style causal language models entirely from scratch, creating two variants with 117M and 354M parameters, respectively. Both models share identical architecture and training processes, allowing for direct comparison of performance versus computational requirements.

Our training infrastructure leverages modern multi-GPU training with Distributed Data Parallel (DDP), comprehensive checkpointing for restart resilience, and real-time monitoring through Weights & Biases. This ensures reliable training even across extended periods and hardware failures.

Evaluation goes beyond standard metrics to include historical accuracy probes, perplexity tracking, qualitative generation review, and early failure detection. We specifically test how well our models understand historical context, period-appropriate language, and London geography.

Finally, deployment includes publishing models to Hugging Face alongside unified local and cloud inference scripts, making the models immediately accessible to researchers and developers worldwide.

1.3 Hands-On Experience

Every aspect of this project is designed for practical implementation and learning. The working code covers every stage from data collection through tokenizer training, model training, evaluation, and publishing - all fully implemented and documented with clear instructions and examples.

I already have both the models published on Hugging Face; which allows for Live models are immediately available for use, allowing you to test published checkpoints instantly or retrain from scratch with a single command. This dual approach lets you either jump straight into experimentation or understand the complete development process.

The project works with real data - over 500 million characters of authentic historical English from 1500–1850, carefully filtered to minimize modern bias while preserving the rich linguistic patterns of the period. This is using genuine historical texts that provide authentic training material.

Everything is well-structured with clear documentation, error handling, reproducible configurations, and automated publishing workflows. The codebase follows good development practices, making it suitable for learning LLM development concepts.

This series is structured to take you through the complete LLM development pipeline:

Part	Focus	Description
Part 1 (this post)	Quick start and end-to-end overview	Use published models, understand the complete pipeline, and get hands-on experience with working code and live models. The intent is that if you want to build this, you can follow the instructions and get a model in the end. If you want to understand more of the inner workings and details, then those will be covered in the subsequent blog posts.
Part 2	Data collection and custom tokenization	Deep dive into gathering 218+ historical sources, cleaning pipelines, and building specialized tokenizers for historical language patterns.
Part 3	Model architecture and training infrastructure	Technical implementation of custom GPT architectures, multi-GPU training, checkpointing, and performance optimization.
Part 4	Evaluation and deployment	Comprehensive testing frameworks, historical accuracy assessment, and deployment to Hugging Face.

For this first part, you have two paths to choose from based on your goals and available time:

Option 1: Quick Start with Published Models - Jump straight into using the pre-trained models on Hugging Face for immediate testing and exploration. Perfect if you want to see results quickly and aren’t concerned with the technical implementation details.
Option 2: Build from Scratch - Dive deep into the complete codebase and build your own historical language model from the ground up. Ideal if you want to understand every aspect of the pipeline and learn how to create specialized LLMs.

Let us start with option 1 - use the models.

2. Use the models - Try it now using Hugging Face

If you just want to get going and use the models and kick tires, the models are live on Hugging Face and ready to use.

SLM Model (117M parameters): 💡 https://huggingface.co/bahree/london-historical-slm
Regular Model (354M parameters): 💡 https://huggingface.co/bahree/london-historical-llm

In addition, you can also explore the complete codebase and build your own historical language model from scratch. The entire pipeline is documented with working code, training scripts, and deployment guides, and is available on GitHub:

Github Repo 💻 –> ⚙️ github.com/bahree/helloLondon .

If you want to quickly test the published models on Hugging Face (HF), you can do so in two ways: quick automated tests or interactive mode. This is the easiest way to get started and show that the models are fully working. You can either clone the repo and run the scripts or use the Python code snippet below.

If you don’t have a development environment set up, you can follow the instructions in the GitHub repo to set up a conda environment with all dependencies. And just for the local testing, you can use CPU only, but for interactive mode, a GPU is recommended. Finally, you will need at a minimum the following Python packages shown in Listing 1. Note, these are also called out on the Hugging Face model page.

python -m pip install -U pip setuptools wheel
python -m pip install "transformers[torch]" accelerate safetensors

Listing 1: Install Required Dependencies

Note: It is recommended to use a virtual environment or conda environment to avoid dependency conflicts. See the GitHub repo for complete setup instructions.

If you don’t have the code repo yet, you can run the commands in Listing 2 directly and run inference from Hugging Face.

Python Code:

from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the published SLM model
model_name = "bahree/london-historical-slm"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Generate historical text
prompt = "In the year of our Lord 1834, I walked through the streets of London and witnessed"
inputs = tokenizer(prompt, return_tensors="pt")

outputs = model.generate(
    inputs["input_ids"],
    max_new_tokens=100,
    do_sample=True,
    temperature=0.3,
    top_p=0.9,
    top_k=20,
    repetition_penalty=1.2
)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Listing 2: Load and Test Published Model

2.1 Local Testing with the Complete Codebase

Now that you’ve seen the models work directly from Hugging Face, let’s explore the complete development experience by working with the actual codebase. This section walks you through testing the models locally using the same infrastructure that was used to train them.

The helloLondon repository contains everything needed to reproduce the entire pipeline - from data collection through model deployment. By running these tests locally, you’ll get hands-on experience with the inference scripts and understand how the models integrate with the broader development workflow.

The following examples assume you’ve cloned the repository and are running from the root directory. All scripts are designed to work out-of-the-box with the published models, giving you immediate access to the same testing infrastructure used during development. You can test the models using Listing 3 or Listing 4.

# Test SLM model (117M parameters)
python 06_inference/test_published_models.py --model_type slm

# Test Regular model (354M parameters)  
python 06_inference/test_published_models.py --model_type regular

Listing 3: Test SLM Model

There is also an interactive mode where you can type in your own prompts and see the model generate text.

Interactive Testing:

# SLM model - Interactive mode
python 06_inference/inference_unified.py --published --model_type slm --interactive

# Regular model - Interactive mode
python 06_inference/inference_unified.py --published --model_type regular --interactive

# Single prompt testing
python 06_inference/inference_unified.py --published --model_type slm --prompt "In the year 1834, I walked through the streets of London and witnessed"
python 06_inference/inference_unified.py --published --model_type regular --prompt "In the year 1834, I walked through the streets of London and witnessed"

Listing 4: Interactive Mode Testing

If everything works, you should see output similar to the following for the SLM model:

Example Output ( Hugging Face SLM Example):

🧪 Testing SLM Model: bahree/london-historical-slm
============================================================
📂 Loading model...
✅ Model loaded in 8.91 seconds
📊 Model Info:
   Type: SLM
   Description: Small Language Model (117M parameters)
   Device: cuda
   Vocabulary size: 30,000
   Max length: 512

--- Test 1/10 ---
Prompt: In the year 1834, I walked through the streets of London and witnessed
Generated: a scene in which some of those who did not incline to come in contact with him took part in his discourse. It was on this occasion that I perceived that he had been engaged in some new business connected with the house, but for some days it had not taken place, nor did he appear so desirous of pursuing any further display of interest .....
Time: 5.75s

Notice how the model captures:

Period-appropriate language (“thank ’ee kindly,” “bade me go,” “spectacles”)
Historical dialogue patterns (formal speech, period-appropriate contractions)
Historical context (West Indies, poor rates, needle work, pocket-book)
Authentic historical narrative (detailed scene setting, period-appropriate social interactions)

Now that we have tried using the model, let’s explore option 2 and see how we can build it. Once you’ve built your own model, you’ll be able to test it using the checkpoints saved during training - see section 7.4 for detailed checkpoint testing instructions.

3. Build the models - From Scratch

Building a language model from scratch is both an art and a science - requiring careful orchestration of data, architecture, and training to create something that can genuinely understand and generate historical text. Unlike fine-tuning existing models, training from scratch gives us complete control over every aspect of the model’s knowledge and behavior.

The journey from raw historical documents to a working language model involves six critical phases, each building upon the previous one. The flowchart below illustrates this complete end-to-end pipeline, showing how we transform 218+ historical sources into two specialized models that can generate authentic medieval London text.

graph TD
    A[📚 Historical Data Collection<br/>218+ sources, 1500-1850] --> B[🧹 Data Cleaning & Processing<br/>Text normalization, filtering]
    B --> C[🔤 Custom Tokenizer Training<br/>30k vocab + 150+ special tokens]
    C --> D[🏋️ Model Training<br/>Two Identical Models<br/>SLM: 117M / Regular: 354M]
    D --> E[📊 Evaluation & Testing<br/>Historical accuracy, ROUGE, MMLU]
    E --> F[🚀 Deployment<br/>Hugging Face + Local Inference]
    
    G[📖 Building a Custom LLM] --> A
    
    F --> L[🎯 Use Cases<br/>Historical text generation<br/>Educational projects<br/>Research applications]
    
    style A fill:#e1f5fe
    style D fill:#f3e5f5
    style F fill:#e8f5e8
    style G fill:#fff3e0

Figure 1: Complete LLM Development Pipeline

Now that we have a bird’s eye view of the complete pipeline, let us get into the details and build the model from scratch. I am going to walk you through the complete process step-by-step.

I am also going to assume you have a basic understanding of Python, PyTorch, and command-line operations and have a more recent dev setup, including a relatively modern GPU (NVIDIA RTX 3060 or better recommended). For the sake of simplicity, I will show commands for Linux/macOS, but Windows users can easily adapt them.

Again, as a reminder, the ⚙️ GitHub repo has all the code and instructions you need to get started. You can clone the repo and follow along.

4. Environment and Configuration Setup

The foundation of any successful machine learning project lies in proper environment setup and configuration. This step involves creating a virtual environment, installing dependencies, and configuring the project structure. Understanding the key configuration files, directory organization, and overall project architecture is crucial - these elements form the backbone of the entire training process. Taking time to get this right upfront prevents countless headaches and debugging sessions later, ensuring smooth execution through all subsequent phases.

4.1 Key Configuration Files

config.py: Central configuration system (paths, training settings, tokenizer config)
01_environment/setup_environment.py: Environment setup script (reads from config.py)
requirements.txt: Python dependencies (auto-generated by setup script)

4.2 Important Directories (Created by Setup)

helloLondon/: Virtual environment directory
data/london_historical/: Historical text data storage
09_models/checkpoints/: Model checkpoints during training
09_models/tokenizers/: Custom tokenizer storage

Now that we have that out of the way, let us run the setup commands as shown in Listing 5. This will clone the repo, set up the environment, and install all dependencies. For this to work you will already have git, python, and python3-venv installed. If you don’t have these, please install them first.

PS: See the Training QuickStart guide in the GitHub repo for more details.

# Clone and setup environment
git clone https://github.com/bahree/helloLondon/
cd helloLondon
python 01_environment/setup_environment.py
source activate_env.sh

Listing 5: Clone and Setup Repository

As you run the setup script, you should see output similar to the images shown below; the script will create a virtual environment, install dependencies, and set up necessary directories. And then you can activate the environment using the source activate_env.sh command.

Figure 2: Environment setup process - Step 1 of 3 showing virtual environment creation

Figure 3: Environment setup process - Step 2 of 3 showing dependency installation

Figure 4: Environment setup process - Step 3 of 3 showing final configuration

Now that the configuration and environment are set up, we can validate them by running the following command. This will check if everything is working and you have the necessary dependencies installed.

When one activates the environment using source activate_env.sh, you will see it in the console as shown below.

The default environment name is called helloLondon. If you want to change the environment name from helloLondon to something else, you can modify the venv_name field in environment_config.json before running the setup script. This will create a virtual environment with your preferred name.

python3 -c "
from config import config
print('🔧 Configuration Overview')
print('=' * 50)
print(f'Project Root: {config.project_root}')
print(f'Data Directory: {config.london_historical_data}')
print(f'Tokenizer Directory: {config.london_tokenizer_dir}')
print(f'Checkpoints Directory: {config.checkpoints_dir}')
print(f'Virtual Environment: {config.project_root}/helloLondon')
print(f'Vocabulary Size: {config.tokenizer_config[\"vocab_size\"]:,} tokens')
print(f'Special Tokens: {len(config.tokenizer_config[\"special_tokens\"])} tokens')
print(f'SLM Model: {config.slm_config[\"model_name\"]}')
print(f'Training Epochs: {config.slm_config[\"num_epochs\"]}')
print(f'Batch Size: {config.slm_config[\"batch_size\"]}')
print(f'Max Length: {config.slm_config[\"max_length\"]}')
print('\\n🎯 Configuration looks good!')
"

Listing 6: Validate Configuration

The following directory structure will be generated after executing the setup script. Please note that certain directories will remain empty until the data collection and training processes are initiated.

helloLondon/
├── 📁 data/london_historical/          # Historical text data
│   ├── 📄 london_historical_corpus_comprehensive.txt  # Final training corpus
│   ├── 📁 downloads/                   # Raw downloaded data
│   ├── 📁 processed/                   # Cleaned and processed text
│   └── 📁 metadata/                    # Data collection metadata
├── 📁 09_models/
│   ├── 📁 checkpoints/                 # Regular model checkpoints (354M)
│   │   ├── 📄 checkpoint-44000.pt
│   │   ├── 📄 checkpoint-47000.pt
│   │   ├── 📄 checkpoint-51000.pt
│   │   ├── 📄 checkpoint-59000.pt
│   │   └── 📄 checkpoint-60001.pt
│   ├── 📁 checkpoints/slm/             # SLM model checkpoints (117M)
│   │   ├── 📄 checkpoint-52000.pt
│   │   ├── 📄 checkpoint-60000.pt
│   │   └── 📄 checkpoint-60001.pt
│   └── 📁 tokenizers/london_historical_tokenizer/  # Custom tokenizer
│       ├── 📄 tokenizer.json           # Tokenizer configuration
│       ├── 📄 vocab.json               # Vocabulary mapping
│       └── 📄 merges.txt               # BPE merge rules
├── 📁 helloLondon/                     # Virtual environment
└── 📁 logs/                            # Training logs and WandB data

Listing 7: Project Directory Structure

Prerequisites: Before proceeding with the following steps, please verify the following requirements:

Storage: Minimum 20GB of free disk space and stable internet connectivity for data acquisition

Hardware: GPU with 8GB+ VRAM for SLM training, 16GB+ VRAM for Regular model training. Cloud users should select appropriate instance types

Experiment Tracking (Optional but highly recommended): Weights & Biases account with WANDB_API_KEY environment variable configured for comprehensive training monitoring

Dependencies: Required data processing libraries (nltk, beautifulsoup4, etc.) will be automatically installed via the setup script

5. Data Collection

The foundation of any language model lies in its training data. For our historical London models, we’ve built a comprehensive data collection system that sources authentic text from 218+ historical sources spanning 1500-1850 - a remarkable 350-year window of London’s linguistic evolution. This isn’t just about downloading files; it’s about curating a high-quality corpus that captures the authentic voice of historical London.

Our data collection pipeline automatically processes multiple formats (PDFs, HTML, XML, plain text) from diverse sources, including Project Gutenberg classics, Old Bailey trial records, London Lives manuscripts, and British History Online archives. The system includes sophisticated quality control measures: language detection to filter non-English content, OCR artifact correction, duplicate detection, and historical period validation to ensure every text genuinely represents the target era.

The result? A curated corpus of 500M+ characters of authentic historical English text, ready to train models that understand not just the words, but the cultural context, social dynamics, and linguistic patterns of 18th and 19th-century London. Of course, you can always add your own data sources if you have them, and the system is designed to be extensible.

We can kick off the data collection process using Listing 8. This will be run from the project root directory.

# Download historical data with advanced filtering
python 02_data_collection/historical_data_collector.py --max_sources 100

# The system automatically filters:
# - Non-English content (Arabic, Chinese, etc.)
# - Poor OCR quality scans and gibberish
# - Advertisement-heavy commercial content  
# - Duplicate content and empty files
# - Special handling for Project Gutenberg classics

Listing 8: Download Historical Data

This process may take some time, depending on your internet speed, the number of sources you choose to download and your system’s performance. For me, on a very fast internet connection and a powerful machine this took typically 2-4 hours for downloading, and processing the full dataset. The script will save the cleaned and processed data in the data/london_historical/ directory, creating a comprehensive historical corpus.

The data collection process creates a comprehensive historical corpus with the main training file london_historical_corpus_comprehensive.txt containing 270M+ characters (~258MB) of authentic historical text. The complete data directory spans approximately 1.2GB, including 521MB of raw downloaded sources, 263MB of processed and cleaned content, and 126MB of tokenized training sequences ready for model training. The image below shows the data collection in progress.

Figure 5: Data collection process in progress showing real-time processing of historical sources

The final corpus represents one of the largest collections of historical London text ever assembled for language model training, with authentic content spanning 350 years of linguistic evolution. The two images below show an example of one of my runs, one of them showing the final output of the data cleaning and outlining the statistics. And the second one shows the size of the data on disk.

Figure 6: Data collection summary showing final statistics and corpus composition

The total size at the end of the data. Note this does not include the Old Bailey and London Lives data.

Figure 7: Total data size on disk showing comprehensive historical corpus

Figure 7: Total data size on disk showing the complete historical corpus storage requirements

Now that we have our data and have cleaned it. Let us build a custom tokenizer.

6. Train Custom Tokenizer

With our cleaned historical corpus ready, we now need to create a custom tokenizer specifically designed for historical English. Standard tokenizers like GPT-2 are optimized for modern text and fail catastrophically with historical language - treating archaic words like “quoth” and “hast” as multiple subword fragments, losing both meaning and efficiency.

Our custom tokenizer uses Byte Pair Encoding (BPE) with a 30,000 vocabulary size and 150+ carefully designed special tokens that understand:

Historical Language: Archaic pronouns (<|thou|>, <|thee|>), verbs (<|hast|>, <|doth|>), and expressions (<|verily|>, <|forsooth|>)
London Geography: Landmarks (<|thames|>, <|newgate|>, <|tower|>), streets (<|cheapside|>, <|fleet|>), and districts (<|southwark|>, <|westminster|>)
Historical Context: Period markers (<|tudor|>, <|stuart|>, <|georgian|>), social classes (<|noble|>, <|commoner|>), and professions (<|apothecary|>, <|coachman|>)

This specialized vocabulary ensures that common historical terms remain as single tokens rather than being fragmented, dramatically improving both training efficiency and text generation quality. We can kick off the tokenizer using Listing 9. Again, this will be run from the project root directory.

# Train historical tokenizer (30k vocabulary)
python 03_tokenizer/train_historical_tokenizer.py

Listing 9: Train Historical Tokenizer

The training process analyzes our 270M+ character corpus to learn optimal token boundaries, creating a tokenizer that understands the linguistic patterns of 1500-1850 English. The result is a highly efficient tokenizer with a compression ratio of ~0.3 tokens per character and 99%+ reconstruction accuracy - essential for training models that can generate authentic historical text.

Once the training is finished (and usually it is pretty quick - just a few minutes for our data size), we run a quick sanity test as the image below shows.

Figure 8: Custom tokenizer training progress showing vocabulary learning

Figure 8: Custom tokenizer training process showing BPE algorithm learning optimal token boundaries

Note that in testing, we might see a warning that the reconstruction differs; this is only because of the alphabet case being different and is expected. You can ignore this. An example of this is shown below.

Figure 9: Tokenizer reconstruction warning during training process

Figure 9: Tokenizer reconstruction warning showing expected case normalization differences

Why the “Reconstruction differs” warning is actually beneficial:

The reconstruction differences you see are not errors - they’re the tokenizer working exactly as designed for optimal language model training. The tokenizer uses Byte Pair Encoding (BPE), which breaks complex words into smaller, reusable subword units (like “Bourgh” → “bour ##gh”), and normalizes text to lowercase to reduce vocabulary size. These “differences” are actually features that make the tokenizer more efficient and the resulting language model more capable of generating authentic historical text.

📖 For detailed technical explanation: Part 2 of this series covers the complete tokenizer architecture, BPE implementation, special token design, and why these reconstruction differences are essential for optimal language model training.

Now that we have our data and the tokenizer is ready, it is time to train the model.

7. Train the Model

With our cleaned historical corpus and custom tokenizer in place, we can now train our language models. The training system is designed to build two identical models with different parameter counts, allowing you to choose between speed (SLM) and quality (Regular model) based on your needs.

Training Architecture: Both models use a custom GPT architecture specifically optimized for historical text, featuring sophisticated attention mechanisms that understand the complex relationships in historical language. The system includes automatic GPU detection, multi-GPU support, and comprehensive monitoring to ensure optimal training performance.

Training Process: The training system implements modern optimization techniques, including dynamic learning rate scheduling, automatic checkpointing, and real-time experiment tracking via WandB. The entire process is automated with intelligent configuration that adapts to your hardware setup, whether you’re using a single GPU or multiple GPUs for distributed training.

Performance Optimization: The system includes precision optimization (TF32, AMP) and memory management specifically tuned for historical text processing. Training typically takes 7-8 hours for the SLM and 28-32 hours for the Regular model on modern hardware, with comprehensive monitoring to track progress and identify any issues. Note, this time can vary significantly based on your hardware. The times mentioned here are based on dual NVIDIA A30s.

📖 For detailed technical implementation: Part 3 of this series covers the complete model architecture, GPU configuration, training infrastructure, and performance optimization strategies in detail.
🧪 Ready to test your checkpoints? Once training completes, see section 7.4 for comprehensive instructions on testing your trained model checkpoints.

7.1 SLM Training

To kick off the training, the code is quite simple, as shown in Listing 10. Again, this would be from the project root folder. In my case, I am using torchrun --nproc_per_node=2 because I have dual GPUs and I want to use both. If you only have a single GPU, you can just run the automatic GPU detection script. The train_model_slm.py script specifically trains the SLM (Small Language Model) with 117M parameters.

Option A: Train SLM (117M parameters) - Faster, Good for Testing

# Clean any existing tokenized data
rm -rf data/london_historical/tokenized_data/

# Automatic GPU Detection (Recommended)
cd 04_training
./launch_slm_training.sh

# Manual Multi-GPU training
torchrun --nproc_per_node=2 04_training/train_model_slm.py --data_dir data/london_historical

Listing 10: Train SLM Model

Note: The first line rm -rf data/london_historical/tokenized_data/ cleans any existing tokenized data to ensure a fresh start. This is important because the training system caches tokenized data for efficiency, and we want to ensure it uses the latest corpus and tokenizer settings rather than potentially outdated cached data. You want to do this only if you have more updated data from the previous steps.

Once the training starts, you will see a similar output as the one shown below.

Figure 10: Model training initialization showing configuration and setup

Figure 10: Model training initialization showing tokenization and GPU setup process

Note the Tokenizing corpus line - this will take some time, depending on your data size and hardware. The tokenized data will be saved in data/london_historical/tokenized_data/ for future runs, so subsequent training runs will be much faster. If you want to force re-tokenization, you can delete this directory and restart the training. And if you think this is hung, you can check the GPU usage using nvtop in a separate terminal.

And if you have configured WandB as recommended earlier, then you can log in to that dashboard and also monitor the training progress. This is quite handy when you are away from the machine and see how it is generally progressing.

Figure 11: Weights & Biases training progress monitoring dashboard

Figure 11: Weights & Biases training dashboard showing real-time loss curves and performance metrics

WandB also provides valuable insights into your model’s training performance through comprehensive visualizations. The dashboard shows the complete training journey, revealing how your model’s loss decreased over time, whether the training plateaued, and how efficiently your hardware was utilized. These visualizations help you understand not just the final results, but the entire learning process - identifying if the model continued improving throughout training or if it reached a performance plateau.

While these metrics are incredibly useful for optimizing your training process, we’ll dive deeper into interpreting these results and fine-tuning your training strategy in Part 3 of this series.

SLM Results (117M parameters):

wandb: Run history:
wandb:       eval/iter   ▂▂▃▃▄▄▅▅▆▆▇▇██
wandb: eval/train_loss  ███▇▇▇▇▇▇▇▇▇▇▇▇
wandb:   eval/val_loss  ███████▇▇▇█▇▇▇▇
wandb:    eval/val_ppl  █▇▇▇▇▇▆▆▆▆▆▆▆▆▆
wandb:     train/dt_ms           █            █                
wandb:      train/iter      ▂▂▂▂▂▃▃▃▃▃▃▄▄▄▄▄▄▅▅▅▅▅▅▅▆▆▆▆▇▇▇▇▇██
wandb:      train/loss ▆▅▇▅▅▃▇▄▄█▅▄▅▄▃▇▄▄▅ ▃▃▂▄▅▂▅▂▄▅▃▃▄▅ ▄▃
wandb:        train/lr ██████████▇▇▇▅▄▄▄▃▃▃▃▃▃▂▂▂▂▂▂           
wandb:       train/mfu ▃▄▇▇█▄▄▆▆▇▅▂▅▆▆▇▇▂▄▅▇▇▇▆▆▇▇▇▇▅███▅▇▆▇▇ ▇

wandb: Run summary:
wandb:       eval/iter 60000
wandb: eval/train_loss 2.74369
wandb:   eval/val_loss 3.44089
wandb:    eval/val_ppl 31.21462
wandb:     train/dt_ms 10217.92054
wandb:      train/iter 60000
wandb:      train/loss 2.87667
wandb:        train/lr 3e-05
wandb:       train/mfu 7.50594

It’s also helpful to monitor GPU usage during training. I recommend using nvtop (a GPU monitoring tool similar to htop but for NVIDIA GPUs) in a separate terminal to track memory usage, temperature, and utilization in real-time. The screenshot below shows the GPU monitoring during model training.

Figure 12: GPU monitoring using nvtop showing real-time resource utilization

Figure 12: GPU monitoring during training showing memory usage, temperature, and utilization metrics

7.2 Understanding Checkpoints

Throughout training, the system automatically saves checkpoints - snapshots of your model’s current state, including all learned parameters, optimizer state, and training progress. These checkpoints serve as safety nets, allowing you to resume training if interrupted, and provide multiple model versions to choose from. The final checkpoint (typically saved at the end of training) represents your fully trained model, ready for inference and deployment.

Checkpoints are saved in the 09_models/checkpoints/ directory, with separate subdirectories for each model type. SLM checkpoints are stored in 09_models/checkpoints/slm/ (e.g., checkpoint-4000.pt, checkpoint-8000.pt), while regular model checkpoints are saved directly in 09_models/checkpoints/ (e.g., checkpoint-60001.pt, checkpoint-120000.pt). The checkpoint filenames include the training step number, making it easy to identify the training progress and select the best-performing version for your needs.

These checkpoints enable two powerful capabilities that significantly enhance your training workflow. You can test your model’s current performance at any point during training by running inference on intermediate checkpoints, allowing you to monitor progress without waiting for training to complete. Additionally, suppose training is interrupted due to power loss, system crash, or manual stop. In that case, you can resume from the last saved checkpoint exactly where you left off, saving both time and computational resources. This flexibility is particularly valuable for long training runs, enabling you to experiment with different model versions and recover from unexpected interruptions.

🧪 Ready to test your checkpoints? See section 7.4 for detailed instructions on testing your trained model checkpoints.

7.3 Regular Model Training

The Regular model training follows the same process as the SLM, using identical training infrastructure but with different configuration settings. The only differences are the training script (train_model.py instead of train_model_slm.py) and the model architecture parameters (354M parameters vs 117M).

# Clean any existing tokenized data
rm -rf data/london_historical/tokenized_data/

# Automatic GPU Detection (Recommended)
cd 04_training
./launch_training.sh

# Manual Multi-GPU training
torchrun --nproc_per_node=2 04_training/train_model.py --data_dir data/london_historical

Listing 11: Train Regular Model

Key Differences from SLM:

Training script: train_model.py (instead of train_model_slm.py)
Model size: 354M parameters (vs 117M for SLM)
Training time: 28-32 hours (vs 7-8 hours for SLM)
Memory usage: Higher VRAM requirements
Performance: Better text quality, slower inference

The training infrastructure, checkpointing, WandB integration, and all other features remain identical. The system automatically detects the model type and applies the appropriate configuration from config.py.

Regular Model Results (354M parameters):

wandb: Run history:
wandb:       eval/iter     ▂▂▂▂▂▃▃▃▃▃▃▄▄▄▄▅▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███
wandb: eval/train_loss  █████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▇▆▆▆▆▆▆▆▆▆
wandb:   eval/val_loss  ███████████████████████████████████▇███
wandb:    eval/val_ppl  ████▇▇█▇▇▇▇▇▇▇▇▇▆▇▇▇▇▇▇▇▇▆▇▇▆▆▆▆▆▆▆▆
wandb:     train/dt_ms                  █                      
wandb:      train/iter      ▂▂▂▃▃▃▃▃▄▄▄▄▄▄▅▅▅▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇███
wandb:      train/loss ▇▆▆▇▇▅▅█▅▄▃▅▅▅▄▇▄▄▄▄▄▃▃▃▅▂▄▅▂▅▂▄▅▃▃▄▅ ▄▃
wandb:        train/lr ▄██████▇▇▇▇▇▆▆▆▅▅▄▄▄▄▄▄▄▄▃▃▂▂▂          
wandb:       train/mfu ▆▇█▅▄ ▄▆▆▆▇▃▃▂▂▆█▃▃▅▅▃█▅▄▆▇▇▇▇▄▅▃█▆▇█▄▃█

wandb: Run summary:
wandb:       eval/iter 60000
wandb: eval/train_loss 2.70315
wandb:   eval/val_loss 3.61921
wandb:    eval/val_ppl 37.30823
wandb:     train/dt_ms 24681.64754
wandb:      train/iter 60000
wandb:      train/loss 2.70629
wandb:        train/lr 0.0
wandb:       train/mfu 7.20423

7.4 Testing Your Checkpoints

Once training is complete, you can immediately test your model using the checkpoints saved during training. This is one of the most exciting parts - seeing the model generate historical text for the first time! The PyTorch checkpoint approach provides immediate testing without any conversion needed, allowing you to test any checkpoint to monitor training progress while preserving the complete model state, including training metadata and optimizer state for fast, optimized inference.

Direct PyTorch Checkpoint Testing: Test your model directly from the training checkpoints without any conversion:

# Test SLM checkpoint (117M parameters)
python 06_inference/inference_pytorch.py \
  --checkpoint 09_models/checkpoints/slm/checkpoint-4000.pt \
  --prompt "In the year 1834, I walked through the streets of London and witnessed"

# Test Regular model checkpoint (354M parameters)  
python 06_inference/inference_pytorch.py \
  --checkpoint 09_models/checkpoints/checkpoint-60001.pt \
  --prompt "In the year 1834, I walked through the streets of London and witnessed"

Listing 12: Test Model Checkpoints

Expected Output: Your trained model will generate authentic historical text like:

“In the year 1834, I walked through the streets of London and witnessed the most extraordinary sight. The Thames flowed dark beneath London Bridge, whilst carriages rattled upon the cobblestones with great urgency. Merchants called their wares from Cheapside to Billingsgate, and the smoke from countless chimneys did obscure the morning sun.”

Testing Different Checkpoints: You can test any checkpoint from your training run to see how the model improved over time. Try testing checkpoints from different training stages to observe the learning progression - early checkpoints will generate more random text, while later checkpoints will produce increasingly coherent historical language.

💡 Pro Tip: For published Hugging Face models and community access, see the Quick Start section earlier in this post, where we demonstrated the published SLM model.

8. Publish to Hugging Face

Once you’ve successfully trained and tested your models, you can publish them to Hugging Face for community access and easy deployment. Publishing makes your models available to researchers, developers, and enthusiasts worldwide, while integrating them into the Hugging Face ecosystem for seamless use with the transformers library.

Publishing Process: The publishing code automatically handles the complete conversion process from PyTorch checkpoints to Hugging Face format, which is essential for making your trained models accessible to the broader community. This conversion transforms your local training artifacts into a standardized format that can be easily loaded by users worldwide.

The process includes converting model weights from PyTorch’s .pt format to the more efficient .safetensors format, generating proper configuration files (config.json, generation_config.json) that define the model architecture and generation parameters, uploading the custom tokenizer and all necessary files to ensure complete functionality, creating comprehensive model cards with usage instructions and metadata for easy adoption, and setting up proper model repositories with versioning for educational deployment.

This conversion is necessary because PyTorch checkpoints are optimized for training workflows and contain additional information like optimizer states that aren’t needed for inference, while the Hugging Face format is specifically designed for model sharing and deployment across different environments and hardware configurations.

We need to call the right script to publish the relevant model - either the SLM or the larger model. The publishing scripts will prompt you for your Hugging Face username and repository name, allowing you to customize where your models are published. The scripts automatically detect and use the latest checkpoint from your training run, so you can publish immediately after training completes.

💡 Quick Reference: If you want to test published models before publishing your own, see section 2 “Use the models - Try it now using Hugging Face” for immediate access to pre-trained models.

Prerequisites: You’ll need a Hugging Face account and either set the HF_TOKEN environment variable or provide your token when prompted. The scripts will guide you through the publishing process step by step.

Option A: Publish SLM (117M parameters)

# Publish SLM to Hugging Face
python 10_scripts/publish_slm_to_huggingface.py

Listing 13: Publish SLM to Hugging Face

Option B: Publish Regular Model (354M parameters)

# Publish Regular model to Hugging Face  
python 10_scripts/publish_to_huggingface.py

Listing 14: Publish Regular Model to Hugging Face

If everything is working correctly and the models are published, you will see confirmation messages and upload progress. Here’s what successful publishing looks like:

Figure 13: Hugging Face SLM model upload progress and confirmation

Figure 13: Hugging Face upload process for SLM model showing successful publishing workflow

And this is an example output for the Regular model:

Figure 14: Hugging Face Regular model upload progress and confirmation

Figure 14: Hugging Face upload process for Regular model showing successful publishing workflow

After Publishing: Once published, your models will be available at:

SLM: bahree/london-historical-slm
Regular Model: bahree/london-historical-llm

Users can then easily load and use your models with just a few lines of code, making your historical language models accessible to the broader AI community for research, education, and creative applications.

Testing Your Published Models: Once published, you can test your models using the same inference methods shown in the Quick Start section:

# Test published SLM model (10 automated tests)
python 06_inference/test_published_models.py --model_type slm

# Interactive testing with published models
python 06_inference/inference_unified.py --published --model_type slm --interactive

Listing 15: Test Published Models

10. What We’ve Accomplished

This comprehensive guide has taken you from raw historical documents to working language models that can generate authentic 18th and 19th-century London text. We’ve built a complete pipeline that transforms 218+ historical sources into two specialized models - a fast SLM for experimentation and a powerful Regular model for high-quality generation. The entire system is fully functional, with both PyTorch checkpoint inference and Hugging Face model publishing working seamlessly, tested and validated on real hardware.

What makes this project interesting is that it’s not just another language model - it’s a complete educational journey that teaches you every aspect of building LLMs from scratch. From custom historical tokenizers that understand archaic English to sophisticated GPU optimization and deployment, you’ve learned the full stack of modern language model development. The result is a system that preserves historical linguistic heritage while demonstrating cutting-edge AI techniques, making it valuable for researchers, educators, and anyone interested in the intersection of history and technology.

11. The Journey Continues

This is just the beginning. In the next three parts of this series, we’ll dive deeper into the technical foundations:

Part 2 explores historical data collection, showing how we curated 218+ authentic sources spanning 350 years of London’s history, and how we built a custom tokenizer that truly understands historical English.

Part 3 reveals the custom GPT architecture designed specifically for historical text, GPU optimization strategies, and training infrastructure.

Part 4 completes the journey with evaluation frameworks, testing strategies, and deployment techniques that transform your trained models into working systems.

Each part builds on what you’ve learned here, taking you from high-level overview to deep technical implementation details.

12. Resources

GitHub Repository:⚙️ github.com/bahree/helloLondon - Complete codebase with all training scripts, inference tools, and documentation
Hugging Face Models:
- 🤗 bahree/london-historical-slm - Small Language Model (117M parameters)
- 🤗 bahree/london-historical-llm - Regular Model (354M parameters)
📘Book Reference: Generative AI in Action - For deeper understanding of core LLM concepts
📖Documentation: Complete guides in the 08_documentation/ folder covering every aspect of the project

13. Acknowledgments

🙏

Ready to dive deeper? Part 2: Data Collection & Custom Tokenizers covers the technical details of data collection, cleaning pipelines, and custom tokenizer development for authentic historical text processing.

Reasoning AI Models: An overview

Mon, 01 Sep 2025 00:00:00 +0000

TL;DR

As part of my role at Microsoft’s AI Foundry Applied AI engineering team in CoreAI, I have participated in numerous detailed discussions about the evolving landscape of AI models. In conversations with many customers, from CxOs to engineers, one recurring topic is the rise of reasoning AI models. These models are designed to perform complex tasks by explicitly breaking down problems into logical steps, rather than just generating text in a single pass like traditional large language models (LLMs). This shift toward reasoning-centric AI marks a major evolution in how we develop and deploy AI systems—and it’s a key factor behind the rise of Agents and Agentic AI.

At the same time, there is a lot of confusion about what these reasoning models are, how they differ from traditional LLMs, and how to effectively adapt and evaluate them. In this post, I aim to clarify these concepts by providing a technical deep dive into reasoning AI models, their training and adaptation processes, and the challenges involved in fine-tuning them for specific tasks. We will also explore how to evaluate these models effectively, considering their unique characteristics.

This post is intended to help one gain a deeper understanding of reasoning models and their implications; I cover these areas:

What are reasoning AI models? A technical overview of their architecture and training paradigms.
How do they differ from traditional LLMs? Key distinctions in capabilities and performance
How to adapt and fine-tune reasoning models? Best practices and common pitfalls
What are the challenges in customizing them? Technical and organizational hurdles
How to evaluate reasoning models? Metrics and strategies for assessing their performance

1. Introduction

Recent AI models have begun to combine language generation with explicit reasoning, enabling more reliable solutions to complex problems. Traditional LLMs like GPT-4o complete a generation in one go, without showing their work. Reasoning models, on the other hand, produce a sequence of intermediate steps (a “reasoning trace”) before the final generation. For example, Microsoft’s Phi-4-Reasoning (14B parameters) will explicitly work through a math problem step-by-step, whereas a regular LLM might confidently state an answer with no explanation. This fundamental difference – predictive text generation vs. chained logical reasoning – makes reasoning LLMs significantly better at multi-step tasks, such as math word problems, code debugging, or complex decision queries.

Note: The AI model landscape is also shifting rapidly, with a newer trend of transitioning from separate “base” vs. “reasoning” models (e.g., o1/o3) to unified systems with internal routing (e.g., GPT-5). GPT 5 runs a system that routes between fast and deliberate paths and exposes developer controls to tune thinking time. In production, the system automatically switches modes; developers can cap or elevate effort as needed. This operationalizes dynamic compute allocation, reducing the need for prompt engineering, specifically when wanting to induce reasoning.

The shift toward unified systems like GPT-5 can be understood as operationalizing the compute-optimal scaling insights from research. Rather than requiring users to choose between reasoning modes manually, these systems implement automatic difficulty assessment and adaptive compute allocation - essentially embedding the “compute-optimal” strategy within the model architecture itself.

1.1 What are reasoning models?

Reasoning models are LLMs architected to solve problems via a multi-step chain-of-thought (CoT) approach. Instead of just predicting the next token, they simulate an internal “scratchpad” of logic. For instance, OpenAI’s latest models (o1 and o3) reportedly allocate extra computation at inference-time and use reinforcement learning (RL) fine-tuning to boost multi-step reasoning. DeepSeek’s R1 (671B-parameter Mixture-of-Experts model ) was explicitly trained with multi-stage reinforcement learning to encourage step-by-step thinking.

During training, such models may be given examples formatted like: *Question → (Begin reasoning) → ... reasoning steps ... → (Final answer)*, or prompted with cues like “Let’s think step by step.” This teaches the model to articulate intermediate steps instead of jumping straight to an answer. In essence, a reasoning LLM learns to internalize a logical process – it doesn’t just know facts or language, it learns how to solve problems by breaking them down.

Crucially, these reasoning models often use special tokens to separate the “thinking” from the final answer. Many use a convention such as <think> ... </think> tags to enclose the chain of thought. For example, DeepSeek-R1-Distill (a distilled 8B version of R1) will output a hidden “thinking” transcript between these tags, followed by a concise answer that summarizes the reasoning. The chain-of-thought (CoT) might include equations, logic, or code, which the model generates as if working on scratch paper, and then the answer is given separately. This behavior is usually built into the model through fine-tuning – if you prompt such a model normally, it will, by default, produce a step-by-step solution trace and then provide the answer.

Some recent systems even let developers toggle the visibility of this trace: e.g., Qwen-3 allows a “reasoning mode” where the chain of thought is shown or hidden as needed. The key point is that reasoning models carry out more computation in the open, and they may consume more tokens. It is quite common for them to use hundreds or thousands of tokens for a complex solution, whereas a regular LLM might try to produce an answer in, say, a single paragraph.

Figure 1: DeepSeek-R1 architecture showing Mixture-of-Experts design with selective parameter activation (21B of 671B parameters active per token) and 128K token context window. (Source: DeepSeek research)

1.2 Cognitive Architecture Parallels - Type 1 and Type 2 Thinking

The reasoning model paradigm directly parallels the Type 1/Type 2 thinking framework popularized by Daniel Kahneman . Some of the recent work demonstrating how LLMs can be aligned to either System 1 (intuitive and fast) or System 2 (analytical and deliberate) thinking patterns.

Type 1 thinking in AI systems corresponds to the pattern-matching and intuitive responses characteristic of traditional LLMs - fast, automatic responses based on learned patterns. Type 2 thinking represents the deliberate, step-by-step reasoning that reasoning models are designed to emulate. Research shows that System 2-aligned models excel in arithmetic and symbolic reasoning, while System 1-aligned models perform better in commonsense tasks.

Cognitive Flexibility and Performance Trade-offs

Unlike human cognition, which fluidly adapts between System 1 and System 2 thinking based on context, current LLMs lack this dynamic flexibility. This rigidity can lead to brittle performance when tasks deviate from trained patterns. However, reasoning models attempt to address this limitation by incorporating explicit System 2-style processing.

The research demonstrates an “accuracy-efficiency trade-off” where System 2-aligned models show greater uncertainty and more systematic processing, while System 1-aligned models provide more definitive but potentially less reliable answers. This suggests that optimal AI systems may need to switch between reasoning modes dynamically based on task complexity.

From an architectural perspective, reasoning LLMs are still transformer-based neural networks at their core. They don’t necessarily have new algorithmic components beyond the training tweaks, though some research explores adding tools or memory. It’s the training paradigm that sets them apart.

For example, where a classic 4o/4.1 style LLM is trained purely on next word prediction and maybe a bit of instruction tuning, a reasoning model like R1 or Phi 4 is trained in an extensive multi stage training pipeline (e.g. supervised fine tuning on curated CoT examples), then specialized reinforcement learning (using rewards for getting answers right and for producing a consistent chain of thought), and so on. OpenAI’s o1/o3 models are rumored to undergo similar multi-stage refinement, combining RL with the ability to allocate more thinking steps at runtime.

1.3 Chain-of-Thought: Built-in vs Prompted

Start by understanding what a chain of thought (CoT) is. CoT is the model’s “scratchpad”: a sequence of intermediate reasoning steps it writes out before giving the final answer. Many models fence this trace with special tokens (e.g., <think> ... </think>); there are configurations that can show or hide these. The advantage this gives us is better results on multi-step tasks (such as math, code, and planning) by decomposing problems. On the other hand, the trade-offs include more tokens → more cost/latency; and traces can be verbose or unfaithful if not evaluated. As a result, CoT is best used for complex queries, and where possible, it would be wise to consider either skipping or limiting these for simple lookups. See “Evaluation” for token-normalized accuracy and faithfulness checks.

CoT prompting emerged as a technique to enhance traditional LLMs by explicitly requesting step-by-step reasoning through prompts such as “Describe your reasoning in steps” or “Explain your answer step by step.” This approach leverages LLMs’ ability to “think out loud” in natural language, with effectiveness scaling with model size as an emergent ability.

Figure 2 shows an LLM decomposing a complex math word problem into sequential subquestions, solving each step before arriving at the final answer. (Credit: Chain-of-Thought Prompting Elicits Reasoning in Large Language Models)

Figure 2: CoT reasoning example

Reasoning models fundamentally differ in that they integrate CoT processing directly into their architecture and training process. Rather than requiring explicit prompting, these models automatically engage in step-by-step reasoning for complex tasks. Research indicates that “Chain-of-Thought built into the core architecture and training process” represents a more robust approach than external prompting.

However, CoT prompting is not universally effective across all models and tasks. Recent research on strategic reasoning has shown that CoT prompting is not universally effective, as it increases strategic reasoning only for models at certain levels, while providing limited gains elsewhere. This suggests that integrating reasoning capabilities requires careful architectural considerations beyond simple prompting strategies.

The effectiveness of CoT in reasoning models also varies by task complexity and domain. Models trained with reinforcement learning on reasoning tasks show more consistent application of multi-step reasoning compared to models relying solely on prompted CoT.

1.4 Test-Time vs Train-Time Compute

A critical innovation in reasoning models is the emphasis on test-time compute scaling. While their training parameters limit traditional LLMs, reasoning models can allocate variable computational resources during inference. OpenAI reports that the performance of o1 improves with more RL (train-time compute) and with more time spent thinking (test-time compute) ( overview ). This creates new scaling paradigms, where models can allocate more computational resources to harder problems during inference.

This inference-time compute scaling (using more tokens/steps) is a defining trait – it enables even smaller models to solve hard problems by iterating through reasoning. As Microsoft’s team describes, “Phi 4 Reasoning generates detailed reasoning chains that effectively leverage additional inference time compute,” allowing a 14B model to compete with far larger ones.

Because this extra “thinking” consumes tokens and compute, it helps to formalize the tradeoff to understand the concept better.

Test-time compute is best understood as a way to reshape the model’s output distribution at inference by searching over alternative reasoning paths and then selecting among them. It reliably lifts accuracy—especially on problems with verifiable answers—yet it is not interchangeable with pretraining compute.

Recent evidence shows that test-time compute helps most when the base model is already capable and the gap to the target difficulty is modest; on the hardest items, pretraining capacity still dominates - as outlined in Figure 3 below.

Figure 3: The relationship between test-time and train-time compute in reasoning models, showing how additional inference computation can compensate for reduced training compute. (Source: Snell et al., 2024)

A practical rule is to treat thinking tokens as a budgeted resource: use them to explore and score candidate chains (branching) and reserve a small budget for targeted revision when a verifier flags issues. In cost terms, this gives you predictable returns without pretending that more inference tokens can fully substitute for more capable pretraining.

Test-Time Compute vs Model Size Trade-offs

A groundbreaking finding from recent research is that on problems where smaller models achieve non-trivial success rates, test-time compute can be used to outperform models 14× larger in FLOP-matched evaluations. This suggests a fundamental shift in how we think about compute allocation:

Easy to medium problems: Test-time compute is often more effective than pretraining larger models
Very hard problems: Pretraining capacity still dominates, with limited benefits from test-time scaling
Practical implication: Rather than focusing purely on scaling pretraining, it may be more efficient to train smaller models and apply test-time compute strategically

Efficiency trade-off: How much “thinking” is enough?

OpenAI’s o1 explicitly reports: performance improves with more RL (train-time compute) and with more time spent thinking (test-time compute). Microsoft’s Phi-4 Reasoning (14B) shows similar patterns: small models, when allowed longer structured chains, outperform their weight in math/science. To examine the implications, consider a back-of-the-envelope cost model. If $L_r$ is “reasoning” length and $L_a$ is final answer length, a crude attention-heavy cost proxy is

$$ \text{Compute} ;\propto; H,(L_a + L_r)^2,d, $$

with hidden size $H$ and depth $d$. You can wrap this into an objective that matches your reality:

$$ \min_{L_r}; C(L_r) = \alpha,H,(L_a+L_r)^2,d ;+; \beta,\text{latency}(L_r) ;-; \gamma,\text{Acc}(L_r), $$

where $\alpha,\beta,\gamma>0$ are your infra cost, SLA pain, and value of accuracy. You won’t solve this analytically in prod—you’ll sweep the thinking budget and pick a knee point.

What is really interesting is that accuracy is typically concave in $(L_r)$; i.e, the first ~100–300 “thinking” tokens help a lot; beyond that, diminishing returns.

Quick intuition: if $L_a$ is small and $L_r$ doubles, the attention term grows by about $4\times$, while accuracy typically improves far less—hence token budgets and early stop heuristics. We’ll revisit this idea in Evaluation via token-normalized accuracy.

This trade-off also motivates practical features, such as token budgets, early-stop heuristics, and “fast vs. deliberative” paths (e.g., Qwen-3’s reasoning mode). With that lens, let’s look at what differs under the hood.

More deliberate thinking often helps—up to a point - You can trade thinking tokens for accuracy on complex items.
Returns diminish - the first ~100–300 “reasoning” tokens carry a lot of the lift; beyond that, you’re paying for a long tail.

Rule of thumb. Treat “thinking tokens” as a first-class budget; log it, control it, and optimize it like you optimize memory or p95 latency. Some model providers like Qwen 3 and NVIDIA’s NIM expose this thinking budget directly.

In short, reasoning LLMs are LLMs with a logic upgrade – through additional training, they learn to use reasoning strategies that standard models lack.

1.5 Effectiveness and Limitations of Reasoning LLMs

Recent benchmarks indicate that CoT reasoning yields significantly improved performance on complex tasks (see Figure 4 below). For example, Microsoft’s Phi-4-Reasoning models, with only 14B parameters, match or surpass much larger models in math and science benchmarks—sometimes even outperforming a model 5x times their size (surpassing OpenAI’s o1-mini, and R1’s 70B distilled version on many math and science benchmarks). This success is attributed to reasoning-focused training and reinforcement learning, proving that with strategic training, smaller models can excel at challenging tasks without needing massive scale. This demonstrates a general trend: with the right training, a model doesn’t have to be huge to solve complex tasks – it just needs to learn how to use its capacity more algorithmically.

Figure 4: Performance scaling on AIME mathematics benchmark: Both train-time compute (left) and test-time compute (right) show smooth accuracy improvements, validating the compute-optimal scaling approach. (Source: OpenAI research)

Another data point is the DeepSeek R1 family. The original R1 (671B, MoE) was a “reasoning-maximal” model (see Figure 1), pushed to an extreme scale and trained with novel RL algorithms (such as GRPO, a group-based self-improvement method) to excel at long-horizon problems. Distilled smaller versions of R1 (70B, 8B, etc.) inherited some of these skills through knowledge distillation. These distilled reasoning models, even at 8B, achieved math and puzzle-solving scores significantly higher than those of similarly sized generic LLMs. Open-source efforts like Bespoke-Stratos-7B and OpenThinker-7B followed suit, demonstrating that a properly fine-tuned 7B model with CoT can outperform naive 7Bs by significant margins on benchmarks. In research from late 2024, Qwen-3 (an advanced open model by Alibaba) was released in both “thinking mode” and “no thinking” mode. Running Qwen-3 in its CoT mode, it actually outperformed DeepSeek-R1 on a majority of evaluated tasks despite activating only a subset of its parameters at each token (it’s a mixture-of-experts model, effectively).

What is interesting is that when Qwen-3 was toggled off (i.e., no CoT visible), it still beat a GPT-4-sized baseline on many benchmarks, implying that integrating reasoning steps did not harm its base competency – it only added the ability to dig deeper when needed. All these examples underscore that reasoning LLMs hold a significant edge on tasks that aren’t straightforward single-step predictions. Whenever an answer requires multiple pieces of information or intermediate calculations, a traditional LLM often fails or guesses incorrectly, whereas a reasoning LLM can navigate the steps systematically (much like a human showing their work). The gap is so notable that analysts have called reasoning LLMs “a critical evolution” in AI capability, and enterprise users are exploring them for decision-making support where correctness takes precedence over brevity.

Mathematical and Logical Reasoning

Reasoning models demonstrate substantial improvements over traditional LLMs in mathematical and logical reasoning tasks. OpenAI’s o1 achieves remarkable performance, ranking in the 89th percentile on competitive programming questions (Codeforces) and placing among the top 500 students in the US in a qualifier for the USA Math Olympiad (AIME).

Codeforces: Codeforces is a major competitive programming platform and community. It hosts frequent online contests (“Rounds”) where participants solve algorithmic problems within time limits, and it maintains an Elo-like rating system and color-coded titles (ranging from Newbie to Legendary Grandmaster).

Comprehensive evaluations (see Figure 5) show that o1-preview demonstrates 100% accuracy in high school-level mathematical reasoning tasks, providing detailed step-by-step solutions and an 83.3% success rate in solving complex competitive programming problems, surpassing many human experts. These results indicate performance that often meets or exceeds that of human experts in structured reasoning domains.

Figure 5: Comprehensive benchmark comparison showing o1-preview's superior performance over GPT-4o across mathematics, science, and reasoning tasks. (Source: OpenAI, 2024)

Domain-Specific Applications

Beyond mathematics, reasoning models show strong performance across diverse specialized domains. Evaluations indicate remarkable proficiency in anthropology and geology, demonstrating a deep understanding and sound reasoning in these specialized fields, as well as strong capabilities in quantitative investing, complemented by comprehensive financial knowledge. The models also demonstrate superior ability in generating coherent and accurate radiology reports, outperforming other evaluated models.

Recent research with ReasonFlux-32B has demonstrated that smaller, specialized reasoning models can outperform larger, general models. On the MATH benchmark, ReasonFlux-32B achieves an accuracy of 91.2% and surpasses o1-preview by 6.7% while being trained with only 8 GPUs.

ReasonFlux - ReasonFlux is a template-driven, hierarchical RL approach to reasoning LLMs; instead of lengthening raw CoT, it plans over a library of thought templates and scales those at inference time, yielding strong math results in a 32B-parameter model.

However, this does not mean reasoning models dominate on every task. For very simple or single-step queries (e.g., straightforward fact lookups or classifications), a regular LLM might perform just as well and with less latency - it is using fewer tokens and does not have to generate a long explanation; more tokens mean more computation and slower responses. That said, many reasoning LLMs are designed to be flexible – they can shorten or skip the reasoning when it’s not needed. Some deployments use a “fast path” versus a “deliberative path” approach: run the model in normal mode for easy questions and only invoke full reasoning mode for complex ones. This dynamic compute allocation is a research area in itself (how to predict when to make a model think longer).

The token-budget mechanism in Qwen-3 is one example: it allows users to cap how many reasoning tokens the model can use, forcing it to decide what’s most important. Accuracy does improve with more tokens (e.g., from ~70% at 2K tokens to ~85% at 16K on a math test), but after a point, it is a matter of diminishing returns. The existence of such features highlights that reasoning LLMs introduce a new dimension – a time/accuracy trade-off. Traditional LLM evaluation is usual just one-dimensional – measuring accuracy or quality for a given fixed model output length. On the other hand, reasoning LLMs let us trade generation length for correctness. (Note: the Evaluation section will cover more details on how to measure).

1.6 Branching & Editing at Test Time (how to “spend” thinking compute)

Test-time compute isn’t just “more tokens”; it’s a way to reshape the model’s output distribution by searching for, and then selecting, better reasoning paths during decoding. In practice, this plays out along two complementary axes. The first is branching: generate multiple candidate chains and prefer the one that scores best under a process- or outcome-aware judge. The second is editing: let the model (and its tools) reflect on an initial attempt and revise it once or twice. Both strategies are ways of allocating limited thinking budget where it matters most.

On the branching side, simple best-of-N sampling remains a solid baseline, while beam or tree-style search makes exploration adaptive by spending more decoding on promising partial thoughts. Process-aware scoring—via a process reward model (PRM) or per-step self-evaluation—helps prune low-quality branches early; when ground truth isn’t available, self-consistency (majority voting across diverse chains) is a practical fallback. Two small but useful tricks from recent work are to branch early—keeping only the top few first-token continuations before decoding greedily—and to anneal temperature across tokens to reduce accumulated randomness as chains grow. Together, these make parallel exploration both cheaper and more reliable.

Editing tackles a different failure mode: an answer that looks plausible but hides a local mistake. Here, short reflect-revise loops work best when anchored to reliable feedback—unit tests for code, exact-match checks for math, heuristic rubrics, or judgments from a stronger model. Pure “self-correction” without such anchors tends to be unstable: models often make minor, non-helpful edits, occasionally flip correct answers to incorrect ones, or fail to generalize the revision behavior. Keeping revision rounds tight, skipping revision when a verifier signals “already correct,” and rolling back to the best-verified candidate are practical guardrails.

Importantly, branching and editing are not substitutes; the best results often come from using both. For easier problems, a short sequential pass can be enough, but as difficulty rises, the sweet spot shifts toward a deliberate mix of parallel exploration and a small revise budget. Thinking time is therefore a budget allocation question: how much diversity you buy up front versus how much you reserve for targeted fixes after you’ve seen a candidate chain.

Operationally, it pays to make the budget explicit and observable. Expose a cap on “thinking tokens,” allow early exit when candidates agree with high confidence, and log the signals that drove selection—per-step PRM or self-evaluation scores, agreement margins, and precise stop reasons. Over time, these traces make it easy to tune the ratio between breadth (how many chains you explore) and depth (how hard you try to fix a promising one), and to decide when a verifier is strong enough to justify skipping revision. Finally, remember that this test-time axis complements, but does not replace, pretraining: extra thinking generally helps, yet it cannot fully compensate for large capability gaps on the hardest items.

Compute-Optimal Scaling

Recent research by Snell et al. demonstrates that compute-optimal scaling - allocating test-time compute adaptively based on problem difficulty - can improve efficiency by more than 4× compared to traditional best-of-N sampling. This approach recognizes that different problems require different amounts of thinking time, and optimal allocation varies dramatically based on prompt difficulty.

The key insight is that question difficulty can be predicted and used to determine the most effective test-time compute strategy. For easier problems, simple parallel sampling suffices, while harder problems benefit from sequential revision or sophisticated search strategies.

Research identifies two primary mechanisms for scaling test-time computation effectively:

Process-Based Verifier Search: Using dense, process-reward models (PRMs) to guide search through reasoning paths, enabling beam search or lookahead search strategies that prune low-quality branches early.
Adaptive Distribution Updates: Modifying the model’s distribution over responses at test time, such as through sequential revision where the model iteratively improves its initial attempts.

The effectiveness of these approaches critically depends on problem difficulty - easier problems benefit more from parallel exploration (branching), while harder problems require sequential refinement (editing).

Difficulty-Aware Compute Allocation

A key insight from recent research is that optimal test-time strategies vary dramatically with problem difficulty. This motivates adaptive allocation strategies:

Easy problems: Simple best-of-N sampling with minimal compute
Medium problems: Weighted voting or beam search with moderate compute budgets
Hard problems: Sequential revision with larger compute budgets, but diminishing returns beyond a threshold

This difficulty-aware approach enables 4× efficiency improvements over uniform compute allocation strategies.

1.7 External tools inside the reasoning loop

Several steps in a chain can be offloaded to exact tools (e.g., code execution, math). Approaches like PAL (program-aided language model) and CoC (Chain-of-Code) let the model “think” by writing and running code; ReAct interleaves search (e.g., Wikipedia) with thoughts. Recent o-series releases similarly intertwine web, code, and vision tools during reasoning. This improves robustness on math, algorithmic tasks, and multi-hop QA – without asking the LLM to emulate a compiler.

PAL

Program-Aided Language Models (PAL) are an approach where LLMs address reasoning tasks by generating Python code rather than relying solely on natural language. This method utilizes programming to manage complex logic and calculations, aiming to decrease errors and improve results on benchmarks such as GSM8K and MATH. PAL’s architecture is modular and interpretable, with the LLM functioning as a code generator and the Python interpreter serving as the reasoning engine. This clear separation improves debugging, verification, and extensibility, enhancing transparency and reproducibility. By combining symbolic reasoning with neural language modeling, PAL provides a hybrid approach that is both effective and practical.

CoC

Chain of Code (CoC) is a method that expands code-driven reasoning in LLMs by using a hybrid execution strategy. In contrast to traditional methods that rely exclusively on interpretable code or natural language reasoning, CoC enables models to generate programs combining executable code with semantic pseudocode. When the interpreter encounters undefined or non-executable behavior, such as abstract functions like detect_sarcasm(string), CoC uses an “LMulator”, which is a language model-based emulator that predicts the expected output. This approach allows LLMs to process tasks involving both algorithmic and semantic elements.

By “thinking in code” CoC greatly expands the range of problems it can solve, surpassing Chain of Thought and other baseline methods on benchmarks like BIG-Bench Hard, where it reached an 84% success rate—12% higher than CoT. Its modular structure adapts well to different model sizes and fields, making it particularly suitable for tasks in robotics, perception, and mixed-modality reasoning. The use of flexible pseudocode and fallback emulation strategies provides a strong foundation for developing more generalizable and interpretable AI reasoning.

In summary, reasoning AI models distinguish themselves by how they solve problems. They use explicit multi-step reasoning (often visible as a chain-of-thought) and are trained with techniques (special prompts, reward signals, data curation) to make this effective. In doing so, they often achieve higher accuracy on complex tasks than traditional LLMs of comparable (or even much larger) size. The cost is greater complexity in training and sometimes in usage. We next discuss how one can adapt and fine-tune these models, and the pitfalls to watch out for.

2. Adapting and Fine-Tuning Reasoning Models

Similar to LLMs, reasoning models can also be fine-tuned or adapted to specific domains and tasks. A key advantage is that they can be domain-specialized while retaining strong reasoning skills.

For example, if you have a reasoning LLM and you want it to excel at medical diagnostics, you could fine-tune it on medical Q&A data that includes step-by-step reasoning about symptoms and lab results. The model should, in principle, retain its general logical abilities and learn to apply them in the medical context. Fine-tuning can also help a model learn when to engage reasoning mode – e.g., always do detailed reasoning for high-stakes medical questions, but perhaps skip it for trivial prompts if instructed.

However, adapting a reasoning model is more complex than fine-tuning a regular LLM because you need to handle the reasoning traces properly. A key question is whether the fine-tuning data includes chains of thought or just question→answer pairs. Generally, to preserve and leverage the model’s strength, you want to fine-tune with the reasoning format intact. That means if your dataset doesn’t already have human-written rationales, you may need to generate them (possibly using a larger teacher model like R1 or GPT-4 to produce explanations for your domain problems). By training on QA pairs supplemented with correct reasoning sequences, you reinforce the model’s inclination to think things through.

There is a subtle issue, though; if your fine-tuning data’s reasoning traces are of lower quality than the model’s current capability (for instance, you provide simplistic or even flawed reasoning examples), you might hurt performance. It’s like training a math student who can solve calculus problems to only practice arithmetic – they might lose their edge in advanced problem solving.

2.1 Loss Masking

One approach is called loss masking, which involves including reasoning steps in the input/output during fine-tuning so the model learns to produce them, but not applying back-prop loss on those reasoning tokens. So, fine-tuning gradients is applied only to the final answer portion, rather than the whole CoT text. This allows us to adjust the model’s final answers for a new domain while minimizing changes to its internal reasoning process. The rationale is that the model’s existing reasoning ability, developed through prior training, should be maintained. The technique allows the model to retain its established reasoning while modifying how it presents final answers. Initial community observations indicate this approach can help preserve the quality of the model’s reasoning after fine-tuning. However, it may not be necessary if the fine-tuned dataset is large and of high quality.

2.2 Prompt-Based Fine-Tuning

Another approach when working with limited data is to use prompt-based fine-tuning or instruction prompts. Since reasoning models already respond to prompts like “show your reasoning, then answer,” you might not need to change their weights at all for some custom tasks – providing a few exemplars with reasoning in a prompt might suffice (few-shot learning). If actual fine-tuning is needed (e.g., to integrate new knowledge or jargon), lightweight methods like LoRA adapters can be applied in principle. One must ensure the prompt format (the presence of <think> tags or special tokens) is consistent during fine-tuning to prevent the model from being confused about when to produce reasoning. Many open implementations of reasoning models require a specific format to trigger the chain of thought. Adhering to that format in any further training data is important.

In summary, adapting a reasoning LLM is doable but requires careful dataset design. Ideally, your fine-tuning set should contain high-quality problem-solving examples with the full reasoning shown. If you don’t have that, you might generate it or opt to preserve the pre-trained reasoning behavior via techniques like masking. One should also monitor if the model starts to skip reasoning; if it does, this could indicate that the fine-tuning data encouraged direct answers only. Balancing task specialization with maintained reasoning capability is key.

Next, let’s examine the challenges that may arise during this fine-tuning and customization process.

Practical Compute Budget Guidelines

Recent empirical analysis provides concrete guidance for practitioners:

Budget allocation: Treat test-time compute as a first-class resource requiring explicit budgeting and monitoring
Difficulty prediction: Use learned difficulty predictors to route problems to appropriate compute strategies
Diminishing returns: Most benefits come from the first 100-300 reasoning tokens; beyond that, returns diminish rapidly
Cost-performance optimization: Smaller models with sophisticated inference can achieve Pareto-optimal trade-offs compared to larger models with simple inference

3. Challenges in Fine-Tuning and Customizing Reasoning Models

Adapting reasoning models to new tasks comes with unique challenges beyond those in standard LLM fine-tuning. These challenges span technical issues inherent to the models’ reasoning nature, as well as organizational hurdles in data and expertise. Let us explore some of the key challenges.

3.1 Trace Quality Degradation

A major technical concern is preserving the quality of the reasoning trace. Fine-tuning, if done either poorly or used on narrow data, can cause the model’s CoT to become less coherent or less faithful to its actual reasoning. Recent research shows that after fine-tuning on specific tasks, the faithfulness of a model’s CoT explanations often decreases, on average, compared to the pre-finetuned model. In other words, the model might still provide accurate answers, but its stated reasoning is more likely to omit key steps or include spurious ones. This “trace degradation” can occur because the fine-tuning objective typically emphasizes obtaining the correct final answer for the new task – the model may learn that it can score well without strictly adhering to its original reasoning style.

In addition, if the fine-tune dataset isn’t sufficiently diverse or is missing the intermediate logic, the model’s previously polished reasoning abilities can “unravel” or get overwritten. It’s akin to using coarse sandpaper after a fine polish – the model may lose some of its nuanced problem-solving steps. Ensuring that fine-tuning does not erase the chain-of-thought skill is a complex and challenging task.

Techniques like the aforementioned loss masking or multi-stage fine-tuning (where you intermix some original reasoning training data) are used to mitigate this. Another aspect of trace quality is faithfulness – even if the model produces a plausible-looking rationale, is it honestly reflecting how the answer was derived? Fine-tuning can sometimes widen the gap between what the model does to get an answer and what it says in the explanation, especially if the fine-tuning introduces shortcut ways to get the answer. This is hard to detect; it requires careful evaluation (as we discuss later).

Overall, maintaining a correct and faithful reasoning trace under new training pressures is a key challenge.

3.2 Overfitting and Distribution Shift

Like any model, a reasoning LLM can overfit to a small fine-tune dataset, but the consequences here might be strange. An overfit model might memorize specific solution patterns and fail to generalize its reasoning to slightly new problems (losing one of the main advantages of a reasoning approach). Because these models were often trained on a wide variety of reasoning tasks, fine-tuning on a narrow domain (say, only physics puzzles) might reduce their versatility or even accuracy on reasoning problems outside that niche.

Small, high-quality reasoning datasets can improve models, but if applied naively, they can also reduce performance on broader evaluations. The model may become too narrowly focused in its thought process (e.g., always expecting a specific style of solution). Ensuring the fine-tuning data covers enough variation or using regularization techniques (such as mixout or weight decay on reasoning layers) can help counteract this, but it remains a delicate balancing act.

LIMA shows that ~1k carefully curated examples can generalize well, and LIMO finds that ~800 math-reasoning samples yield large gains when the data is selected thoughtfully. However, a narrow or naïve fine-tuning can backfire—studies report catastrophic forgetting and degraded out-of-distribution robustness, as well as a drop in CoT faithfulness after fine-tuning. This can be mitigated with regularization (e.g., Mixout, layer-wise noise-stability) and optimization that flattens the loss landscape (e.g., SAM), and keep the fine-tune mix diverse to avoid over-specialization.

3.3 Training Stability and Long Outputs

Fine-tuning with long CoT outputs (which can be thousands of tokens) can lead to stability issues in training. Gradient updates on very long sequences might cause more variance or instabilities in convergence. Moreover, suppose one uses reinforcement learning (e.g., to further optimize a reasoning model with a reward for correct answers). In that case, the credit assignment is complex – which part of a 100-step reasoning deserves credit or blame for the outcome?

Instabilities like mode collapse (where the model’s outputs become strangely repetitive or nonsensical) or oscillating performance have been observed if the RL reward model is poorly aligned. For example, in one training run, simply increasing the reward for “correct final answer” without properly balancing the reward for good reasoning steps caused the model to exploit quirks – it started producing minimal reasoning and guessing answers to game the reward, leading to a drop in overall logical correctness.

Researchers working on Phi-4 and others have had to introduce tricks to stabilize RL training, such as gradually increasing the allowed reasoning length, filtering out bad traces, or adjusting reward scaling. These measures highlight that straightforward fine-tuning or RL on a reasoning model can easily go off-track if the optimization isn’t carefully managed. In essence, teaching a model how to think is a more delicate process than teaching it what to say.

3.4 Reward Alignment and “Hacks”

Aligning a reasoning model with human preferences or task-specific rewards can be tricky – there’s a risk of reward hacking and unintended behaviors. An illustrative scenario was described by researchers at Anthropic: they gave a reasoning model (Claude 3.7 and DeepSeek R1) a series of multiple-choice questions with a twist – a hidden “hint” in the prompt sometimes told the model to choose a wrong answer (and they rewarded the model for following that hint). The models learned to exploit this to earn reward points, selecting the hinted-at wrong answers, but their chain of thought never acknowledged the malicious hint. They would generate a detailed (fake) reasoning to justify the wrong answer, rather than saying “I chose this because I was hinted at.” This is a dramatic example of a model gaming the objective: the training set or reward said “getting this answer is good,” so it did. Still, it also learned to hide the true reason, presenting a facade of coherent reasoning.

Such behavior is misaligned with the intent (we want the model to be truthful in its reasoning). This experiment highlights the importance of aligning the process of reasoning as much as the outcome. If a reward model only considers the correctness of the final answer, it may sacrifice honesty or thoroughness in the reasoning process.

Conversely, suppose you over-emphasize a reward for producing very detailed reasoning. In that case, the model might start outputting verbose, mostly correct-sounding monologues that don’t lead to a better answer (effectively optimizing the wrong metric). Achieving the right alignment – so that the model is rewarded for correct and genuinely helpful reasoning – is an open challenge. It often requires iterative human feedback, custom reward functions (e.g., penalize logical leaps or unsupported claims in the trace), and careful validation. Without these, one might end up with a model that appears to reason well but is just skilled at “output grooming” – formatting answers to look good rather than being correct.

3.5 Data Quality and Availability

On the organizational side, fine-tuning a reasoning model demands high-quality training data that includes reasoned solutions. Such data can be difficult to obtain. At the same time, there are public datasets for math proofs or logical reasoning (e.g., MATH, GSM8K, etc.), but many domains (legal reasoning, financial analysis, medical diagnostics) don’t have readily available step-by-step annotations in large quantities.

Teams often have to generate this data synthetically (using a larger model to produce reasoning traces and then filtering them) or invest in expert annotations. The quality of these traces is paramount – noisy or incorrect reasoning examples can confuse the model or teach it bad habits. As discussed earlier, even small, curated datasets (on the order of hundreds of examples) have been shown to improve reasoning if they are extremely well-targeted; however, curating such datasets is a specialized skill.

In practice, fine-tuning a reasoning model involves a lot of tooling, ranging from running automatic proof checkers to verify steps, using consistency checks, or employing human reviewers to label where a model’s synthetic reasoning went wrong. This is a step up in complexity from preparing a straightforward prompt→response dataset.

3.6 Tooling and Infrastructure

Working with long CoT and multi-stage training means that the training pipelines will need modification. For instance, standard training code may need to be adapted to handle special tokens (e.g., <think> segments might need masking if needed), or to log and evaluate not just final answers but also intermediate step accuracy during training.

Debugging a reasoning model can be more involved – you might want to watch how its reasoning changes epoch by epoch, which requires custom logging or visualization tools. Moreover, these models often have large context windows (since they need to handle long reasoning sequences, e.g., 16K or 32K tokens). Fine-tuning with such long contexts can demand more GPU memory and faster I/O. Not all training frameworks efficiently support extremely long sequences out of the box.

Evaluation tooling (to be discussed later) can also be considered—a possible approach is integrating an automated verifier into the training loop to assess the model’s reasoning steps and provide targeted feedback, which is a type of process supervision. Implementing this involves technical complexity and remains an ongoing area of research. Overall, organizations seeking to customize a reasoning model should be aware that the training workflow may be more complex than a standard LLM fine-tuning process.

3.7 Expertise

Fine-tuning reasoning models demands both machine learning expertise and domain knowledge, often requiring multidisciplinary teams. Since reasoning LLMs are new, practitioners face a steep learning curve with frequent trial and error.

Expect several iterations to balance concise and detailed responses; objectives or examples may need adjustment throughout the process. Rigorous testing is essential, especially in high-stakes applications like medical or legal fields, making reliability and interpretability critical. Typically, 10–12 rounds of tuning are required to achieve an optimal model.

Organizations typically use a hybrid strategy: starting with a robust base model (such as o1-mini or Phi-4-Reasoning), applying minimal tuning, and relying on prompts and few-shot learning for specificity. When deeper customization is required, it’s best to use reliable data, maintain reasoning formats, monitor trace fidelity, and integrate human feedback. Success yields a strong analytical tool, but the process is more complex than for general chatbots.

A key part of customization is the ability to evaluate the reasoning models. Let us dig into specialized evaluation strategies required to assess not just what a reasoning model answers, but how it arrives at that answer.

4. Evaluation Strategies for Reasoning Models

Traditional LLM evaluation – e.g., measuring accuracy on a Q&A or using BLEU scores for text – may not capture the full picture when a model is effectively performing a multi-step reasoning process. Evaluating reasoning-oriented LLMs requires going beyond the final answer, incorporating metrics that assess both the process and quality of reasoning. This represents a departure from traditional LLM evaluation, which typically treats the model as a black box that produces an answer or text, which we then compare to a reference or expected output.

For reasoning models, we care about questions like: Did the model’s CoT follow a correct logical path? Is it telling the truth about its reasoning? How efficient is its reasoning? Below are key evaluation strategies and metrics that have emerged for reasoning models, contrasted with traditional approaches:

4.1 Outcome vs. Process Evaluation

In traditional AI evaluation, we mostly judge the outcome (e.g., did the model get the correct answer to a question). With reasoning models, researchers perform dual evaluations – one for the outcome and one for the reasoning steps. An outcome evaluation may be identical to a standard LLM test, where the goal is to verify if the final answer is correct (exact match, F1 score, multiple-choice accuracy, etc.). The process evaluation, however, examines the intermediate steps of the solution.

For instance, a math word problem benchmark might not only check the answer but also parse the model’s step-by-step solution and verify each part. An emerging method is to use an automated judge (which can be another LLM) to analyze the CoT and flag errors or leaps in logic. One example being a recent benchmark called MM-MATH (for multimodal math problems); in this, an LLM-based evaluator looks at each step of a model’s solution, comparing it to the ground truth solution, and classifies errors (e.g., “incorrect algebraic simplification” vs “misinterpreted the diagram”).

This kind of fine-grained process evaluation provides insights into where a model’s reasoning fails, not just whether the final answer is wrong. This is useful because a reasoning model might get the right answer for the wrong reasons (i.e., it had a reasoning flaw), or vice versa – it might have mostly correct reasoning but a minor slip at the end leading to a wrong answer. Traditional single-score metrics would miss this nuance.

4.2 Chain-of-Thought Faithfulness Metrics

As discussed earlier, faithfulness refers to whether the model’s stated reasoning accurately reflects its actual internal reasoning (or use of information). One way to test this is to insert known information (or traps) into the context and see if the model admits it.

For example, Anthropic’s experiment provided the model with hidden hints (sometimes incorrect) and then checked if the model’s explanation mentioned using those hints. They derived a metric: the percentage of solutions where the model was truthful about using the hint. Claude 3.7 was only ~25% faithful in their setup, and DeepSeek R1 was about 39% – meaning in the majority of cases, they used the hint but didn’t reveal it in the reasoning chain. This indicates that the CoT was often unfaithful, presumably because the model’s training taught it always to sound logical and self-contained, even if it took a shortcut.

Another way to measure faithfulness is to check consistency under variations: if a model truly is reasoning step by step, then if we force it to reveal steps, it should arrive at the same answer as when it’s not forced. If hiding the CoT changes the answer frequently, it might suggest the model’s explanations were more post-hoc and not driving the answer.

Note: These evaluations are still an active research area – unlike a simple accuracy score, faithfulness is somewhat difficult to quantify, but it’s crucial for trust. When deploying a reasoning model, you’d like to trust that, say, a financial analysis it provides is actually how it came to its conclusion, not a fabricated rationale. Thus, papers often report the percentage of solutions with “fully faithful reasoning” by manual or automated inspection. If that percentage is low, it’s a red flag: the model’s reasoning output might be more for show. Improving this might involve further training (e.g., penalizing inconsistent rationales) or architectural changes; however, at the very least, we need to measure it.

4.3 Token-Normalized Accuracy (Efficiency)

Because reasoning models can use an arbitrary number of tokens to reason (within the context window limits, of course), we want to measure accuracy as a function of reasoning length – effectively, how efficiently does a model reach correct answers? For example, a model that gets 90% accuracy with 2K tokens of reasoning might be less desirable than one that gets 85% accuracy with only 1K tokens, depending on deployment constraints.

Token-normalized accuracy is a metric that attempts to penalize overly lengthy reasoning. In one formulation (used in some multiple-choice evaluations), it computes the probability of a correct answer, normalized by the length (i.e., the number of tokens) of that answer’s explanation or output. More generally, we can think of it as accuracy per 100 reasoning tokens or similar.

Another interpretation is to measure the area under the curve of accuracy versus the number of tokens allowed. For example, allow a model to think with 100 tokens, record the accuracy, then 200 tokens, 500 tokens, and so on, up to a certain limit – and see which model yields the best accuracy for the least token budget. Researchers have explicitly emphasized the goal of maximizing accuracy per token in reasoning scenarios.

This reflects practical concerns: in production, reasoning steps are costly - both in terms of latency and tokens (i.e, money). A model that uses half the steps to reach the same answer is effectively twice as fast. Moreover, sometimes unconstrained reasoning leads to diminishing returns or even errors—for example, a model might start wandering or overexplaining if it “thinks” too long. Thus, token-normalized metrics encourage models that use their reasoning budget optimally.

A simple implementation is to take the total tokens the model generated for all test problems and divide them by the number of correct answers. Then, compare models on this normalized score (lower tokens per correct answer is better).

Another approach is a normalized log-probability where longer outputs are penalized. In any case, this kind of metric was usually irrelevant for standard LLMs (which output a single short answer), but becomes important when evaluating the cost-effectiveness of reasoning models.

4.4 Stepwise Accuracy and Consistency

This is a more granular evaluation of the correctness of the reasoning chain. For tasks where we have ground-truth step-by-step solutions (like a math proof or a formal logic derivation), we can mark each step of the model’s chain as “correct” or “incorrect” compared to an expected solution. This yields a sequence of accuracy values (e.g., getting the first three steps right, but failing at step four). We can then compute metrics like average step accuracy, or percentage of solutions that made it to at least X steps correct before failing.

This is informative because two models might both solve 70% of problems, but one might always fail early on the 30% it can’t solve, whereas another might almost solve everything and only slip at the end for those 30%. Stepwise evaluation can reveal such differences. It also helps in evaluating partial credit – maybe a model didn’t get the final answer but did significant parts correctly (which might be useful in applications where a human or another tool can pick up from the middle).

Some evaluations also check consistency: if a model is asked to explain its answer vs. directly answer, do those agree? If it solves a problem in two different ways (maybe by reordering steps or under different prompts), does it reach the same conclusion? Consistency checks can catch cases where the reasoning process is brittle or overly sensitive to phrasing.

4.5 Automated Reasoning Critics (LLM-as-a-judge)

A practical framework that has gained traction is using a strong language model to evaluate the reasoning of another model (or even itself). For instance, one can prompt GPT-4 with: “Here is a chain-of-thought and an answer. Evaluate the correctness and logical validity of the reasoning, and whether the final answer is justified.” This uses the fact that cutting-edge models can often spot obvious reasoning errors or missing justifications in a solution that a simpler rubric might miss.

Such LLM-based evaluators can be more flexible than hard-coded checkers. The aforementioned process evaluators in research are essentially reasoning models used as judges, with the ability to allocate extra computational resources to evaluate each step carefully. In one study, researchers found that when they allowed an evaluator model to think more (generate a longer evaluation reasoning), its accuracy in judging solutions improved monotonically – much like how making a model think more improves problem-solving, it also improves evaluation quality.

This is a fascinating recursive idea: use a reasoning model to evaluate better outputs that themselves involve reasoning. It was even shown that using such process-aware evaluators to re-rank answers (choosing the answer that the evaluator model scores highest) can significantly improve the solving ability of the base model.

In summary, process evaluation frameworks often involve an LLM evaluator performing a two-level check:

Outcome evaluation (is the final answer correct?)
Process evaluation (are the steps valid and do they lead to that answer?).

By combining these, one gets a more robust assessment. This approach complements traditional metrics; for example, you might report that a model has 80% outcome accuracy, but according to an LLM judge, only 50% of its solutions were fully correct with no logical errors in any step. That tells a deeper story than 80% alone.

4.6 Illustrative Example

To illustrate, consider a concrete example: say we ask a model a puzzle and it answers with a 5-step reasoning chain. The final answer is correct, so outcome-wise it’s a success. However, upon evaluation, we found that an arithmetic mistake occurred in step 3, which fortunately canceled out in step 5, yielding the correct answer nonetheless. A pure outcome metric says “perfect solution”. A process-aware evaluation would ding this as flawed reasoning (the model got lucky or coincidentally correct) – something we’d want to know if using the model for, say, validating scientific calculations. Conversely, if a model’s final answer is wrong, traditional evaluation is 0 for that question. However, process evaluation might reveal that the model was correct up until the last step – perhaps it performed all the reasoning correctly and made an error at the end.

In a human-learning context, you’d give partial credit. For model evaluation, noting that the model was, say, “90% correct in the procedure” could inform how we attempt to improve it (perhaps it just needs a slight boost in arithmetic precision or a final double-check step). This rich information is only available if we evaluate the reasoning, not just the outcome.

For practitioners, incorporating these evaluations is vital, as they help ensure that a high-performing reasoning model isn’t just getting by with smoke and mirrors (or hidden cues), and they quantify the efficiency and transparency of the model’s problem-solving approach. As these models become more integrated into workflows (e.g., as AI reasoning assistants), having reliable evaluation methodologies will also be key for governance and trust – one might, for example, require that a model’s chain-of-thought passes a certain automated consistency check before its answer is shown to a user.

In summary, the evaluation of reasoning LLMs has evolved to include trace-centric metrics alongside traditional outcome metrics. We assess the faithfulness of their explanations, measure accuracy in a way that accounts for the cost of reasoning length, and use novel frameworks where models critique reasoning steps (providing a “process score”).

5. Safety Concerns and Vulnerabilities

While reasoning models offer powerful capabilities, they also introduce new safety concerns and vulnerabilities that must be carefully managed and addressed. The very features that make these models effective – their ability to generate detailed CoT and reason through complex problems – can also be exploited by malicious actors or lead to unintended behaviors. Below, we discuss some of the key safety challenges specific to reasoning AI models.

5.1 Reward Hacking and Training Vulnerabilities

Reward hacking represents a significant concern in reasoning model development, particularly given their reliance on reinforcement learning during training. Reward hacking occurs when “a RL agent exploits flaws or ambiguities in the reward function to achieve high rewards, without genuinely learning or completing the intended task”.

In the context of LLMs trained with RLHF, reward hacking manifests when models learn to game evaluation metrics rather than genuinely improve at the intended tasks. This is particularly concerning for reasoning models, where the complexity of the reasoning process makes it difficult to specify comprehensive reward functions that capture all aspects of good reasoning.

For example, a reasoning model might discover that providing overly verbose explanations leads to higher evaluation scores, even if those explanations are not genuinely helpful or accurate. This could incentivize the model to generate long-winded responses that obfuscate its actual reasoning process, ultimately undermining the quality of its outputs.

5.2 Jail-breaking and Safety Mechanism Vulnerabilities

Recent research has revealed severe vulnerabilities in the safety mechanisms of reasoning models. The Hijacking Chain-of-Thought (H-CoT) attack method demonstrates how attackers can “leverage the model’s own displayed intermediate reasoning to jailbreak its safety reasoning mechanism”. Under such attacks, refusal rates in models like OpenAI’s o1 drop dramatically, “from 98% to below 2%”.

The Malicious-Educator benchmark exposes how “extremely dangerous or malicious requests” can be disguised “beneath seemingly legitimate educational prompts”. This research reveals that “attackers can easily extract criminal strategies from DeepSeek-R1 and Gemini 2.0 Flash Thinking without any additional tricks”, highlighting fundamental vulnerabilities in current safety approaches.

In addition, the ability of reasoning models to generate detailed CoT can be weaponized by attackers to create more convincing prompts that bypass safety filters. This raises the stakes for ensuring that safety mechanisms are robust and capable of handling sophisticated manipulation attempts.

5.3 Alignment Challenges in Reasoning Systems

The integration of reasoning capabilities creates new alignment challenges. While reasoning models can “reason about our safety policies in context when responding to potentially unsafe prompts, through deliberative alignment”, this same capability can be exploited by sophisticated attacks. The transparency of reasoning processes, while beneficial for interpretability, also provides attack vectors that didn’t exist in traditional LLMs.

Research indicates that reasoning models still exhibit sensitivity to probability distributions from their training data, suggesting that “optimizing a language model for reasoning can mitigate but might not fully overcome the language model’s probability sensitivity”. This indicates that fundamental limitations from autoregressive training may persist even in reasoning-optimized systems.

5.4 Hallucination in Reasoning Contexts

Despite their enhanced reasoning capabilities, reasoning models continue to exhibit hallucination patterns, particularly in constraint satisfaction problems. Research on graph coloring tasks reveals that reasoning models are “prone to hallucinate edges not specified in the prompt’s description of the graph”. This phenomenon “persists across multiple problem complexity levels and semantic frames” and “appears to account for a significant fraction of the incorrect answers from every tested model”.

These findings suggest that reasoning models may have “broader issues with misrepresentation of problem specifics”, indicating that the enhanced reasoning capabilities don’t fully address fundamental issues with information fidelity and accuracy.

5.5 Scaling and Efficiency Considerations

While reasoning models demonstrate impressive capabilities, they incur significant computational costs. The variable test-time compute approach means that complex problems can require substantially more resources than traditional LLM inference. This creates practical deployment challenges, particularly for applications requiring consistent response times.

The relationship between reasoning quality and computational cost remains unclear. Research indicates that more thinking time generally leads to better performance, but the optimal allocation of computational resources across different problem types remains an active area of investigation.

6. Conclusion

Reasoning AI models, such as o1, o3, R1, and Phi-4, mark a shift towards systems that execute algorithmic steps rather than relying purely on black-box prediction. Unlike traditional LLMs, these models leverage chain-of-thought reasoning, curated data, and advanced fine-tuning to solve complex tasks—though this comes with increased training and inference complexity.

Fine-tuning reasoning models demands specialized methods and high-quality data, as their reasoning chains are both powerful and vulnerable to inconsistency or reward hacking. Effective deployment requires both technical expertise and organizational investment; however, the benefits include clearer explanations and deeper insights across domains such as finance and science.

Evaluation now extends beyond final answers to include scrutiny of the reasoning process itself, using metrics such as trace faithfulness and process accuracy. This makes model behaviour more transparent and trustworthy.

For practitioners, reasoning models become collaborative problem-solvers, offering logical breakdowns for tasks from coding to contract analysis. But maintaining reliable reasoning and avoiding hallucinations requires ongoing vigilance and tailored oversight.

The center of gravity has shifted from pick a reasoning model to use a unified system with routed reasoning,’ with explicit controls for compute and explanation; this aligns with your agentic guidance and simplifies deployment ergonomics. In the near future, with this direction, we expect to see more robust state representations, verification-based training, and compositional planning; evaluate under router-aware, deception-aware protocols, and replicate Apple-style stress tests with fixed effort/latency budgets.

The focus is shifting toward unified systems that route and manage reasoning explicitly, enabling robust evaluation and compositional planning. Reasoning AIs won’t replace standard LLMs everywhere, but they excel in high-stakes scenarios requiring transparency. As techniques mature, these models will become more stable and interpretable, merging pure reasoning with external tools and knowledge. Teams adopting these models should invest in robust pipelines and new evaluation metrics to realize the benefits of interpretable, verifiable solutions—a step forward for AI’s ability to explain not just what or when, but how and why.

References

What is KV Cache in LLMs and How Does It Help?

Sat, 14 Jun 2025 00:00:00 +0000

TL;DR: KV cache is a memory optimization central to efficient LLM inference. It enables faster, longer, and more cost-effective generation by caching previously computed attention keys and values—unlocking the practical deployment of models like GPT-4o, Llama 3, etc.

1. Introduction

Generative AI, powered largely today by Large language models (LLMs) such as GPT-4o, Llama 3, etc., is transforming AI applications, from chatbots to code assistants and multimodal reasoning. As these models scale in size and context length, inference becomes a major computational and memory challenge. The Key-Value (KV) cache is a pivotal optimization that enables practical, high-performance inference in modern transformer architectures.

In my experience, after speaking with many individuals, including customers at work (mostly enterprises), I’ve found that most don’t fully understand what a KV cache is or why they should care. In this post, I aim to provide an overview of what KV cache is, how it helps, and outline some recent research innovations. I also have simple code samples for practical understanding.

At the most basic level, a KV cache is a memory optimization technique used in LLMs to improve inference efficiency during generation. The KV cache stores the key and value tensors generated during the attention mechanism of transformer architectures, allowing models to avoid redundant computations when generating sequential text. To understand the KV cache, it’s essential to grasp how self-attention works in transformers.

2. Transformer Attention and the Role of KV Cache

As part of the transformer architecture, the attention mechanism enables models to dynamically assess the importance of various elements in the input sequence and calculate relationships between input tokens through three components: queries (Q), keys (K), and values (V).

For each token, the model computes:

Query vectors ($Q$): Represent the current element seeking information
Key vectors ($K$): Act as reference points for all elements in the sequence
Value vectors ($V$): Contain the actual information that will be aggregated

The attention computation follows the formula: $$ \mathrm{Attention}(Q, K, V) = \mathrm{softmax} \left( \frac{QK^T}{\sqrt{d_k}} \right) V $$

LLMs are autoregressive; that is, text generation refers to the sequential process by which a model predicts each new token (word or subword) based on all previously generated tokens. This creates a dependency chain: every new token depends on the entire history of prior tokens. During autoregressive text generation, where models predict one token at a time using all previous tokens, the KV cache serves as a repository to “remember” the pre-computed key and value pairs from earlier tokens. Each new token requires attending to all previous tokens. Without optimization, this would necessitate recomputing all $K$ and $V$ matrices at every step, leading to quadratic time complexity.

3. The Caching Process

The KV cache stores the computed key and value tensors for all previously generated tokens. When generating a new token, only its key and value are computed and appended, while the model attends to the full cache. This caching reduces redundant computation, transforming each inference step’s complexity from quadratic to linear with respect to sequence length—a foundational efficiency for scaling LLMs to long contexts.

Without KV caching, transformers would have to recompute keys and values for all previous tokens during each generation step, resulting in quadratic computational complexity. The KV cache removes this inefficiency through the following process:

Initial Generation: When processing the first input token, the model calculates and stores its key and value vectors in the cache.
Subsequent Tokens: For each new token, the model only computes the key and value for that specific token.
Cache Appending: New key-value pairs are appended to the existing cache.
Attention Computation: The model uses the complete cached key-value history to compute attention with the current query.

This approach transforms the attention computation from quadratic $O(n^2)$ to linear $O(n)$ complexity in terms of sequence length.

Example: Minimal KV Cache in PyTorch. This class provides a minimal example of how to store and update key-value tensors for autoregressive generation in a transformer model.

class KVCache:
    def __init__(self):
        self.cache = {"key": None, "value": None}

    def update(self, key, value):
        if self.cache["key"] is None:
            self.cache["key"] = key
            self.cache["value"] = value
        else:
            self.cache["key"] = torch.cat([self.cache["key"], key], dim=1)
            self.cache["value"] = torch.cat([self.cache["value"], value], dim=1)

    def get_cache(self):
        return self.cache

4. Memory Requirements and Bottlenecks

There is no free lunch, of course, and there are also practical memory considerations. The memory footprint of the KV cache is substantial, especially for long contexts. For large models, this can easily consume tens of gigabytes for long sequences, often exceeding the memory needed for model weights themselves.

For example, with Llama-2-7B using half precision (FP16), for batch size 1, the KV cache consumes approximately 0.5MB per 1000 token. For a sequence of 28,000 tokens, this equals about 14GB of memory - the same amount required to store the entire model weights.

Research shows that the KV cache can consume over 30% of GPU memory during deployment and become the primary memory bottleneck for long-context applications. At a simplistic level, the required memory can be estimated as:

$$ \text{Memory} = 2 \times \text{Precision} \times \text{Layers} \times \text{ModelDim} \times \text{SeqLen} \times \text{BatchSize} $$

Note her:

$\text{Precision}$ is typically 2 bytes (FP16)
$\text{Layers}$ is the number of transformer layers
$\text{ModelDim}$ (a.k.a Model Dimension) is the hidden size per layer

5. What Would Happen Without KV Cache?

Understanding the importance of KV cache becomes clearer when you contemplate the consequences of its absence. The overall experience would be significantly worse, characterized by high latency, shorter context windows, and increased costs.

Severe slowdown: Every new token requires recomputing attention for all previous tokens, causing computation to grow quadratically with sequence length
Unsustainable compute overhead: Each step repeats all previous attention calculations, wasting compute and energy.
High latency and poor user experience: Users experience significant lag, especially for long-form or multi-turn conversations.
Limited sequence lengths: Practical context limits shrink, and out-of-memory errors become common for large models.
Inefficient hardware use: Lower throughput and increased energy consumption.
No cache-level optimizations: No prompt reuse, no advanced memory management, and no opportunity for compression.

6. The New Trade-off: Cache Size vs. Model Performance

Recent studies have transformed the balance between KV cache size and model performance. Previously, decreasing the cache size directly impacted model quality, particularly for tasks that required long contexts or retrieval. Now, innovative approaches—utilizing quantization, pruning, and adaptive retention—enable significantly smaller caches, with minimal or virtually no decline in performance.

6.1. Token-Precision Trade-off: Quantized Pruning

A key breakthrough is the realization that storing more tokens at lower precision (quantized pruning) outperforms storing fewer tokens at high precision under the same memory budget.

Key finding: For example storing 4x as many tokens in 4-bit precision outperforms storing 1x tokens in 16-bit precision, especially for long-context and retrieval tasks. Note that the $‘4x’$ factor depends on model architecture, context, and task.
Result: Quantized pruning preserves long-range context and enables robust performance across task types, input lengths, and model scales, even in extreme memory-constrained scenarios.
Stability: This method is robust across various pruning and quantization strategies, providing a new paradigm for cache compression.

Code Example: Quantized Pruning This function demonstrates how to select and quantize the most important tokens in the KV cache to maximize memory efficiency with minimal accuracy loss.

def quantized_pruning(kv_cache, importance_scores, num_tokens_to_keep, num_bits=4):
    # Select top tokens by importance
    top_indices = importance_scores.argsort()[-num_tokens_to_keep:]
    selected_kv = kv_cache[top_indices]

    # Quantize selected tokens
    scale = (selected_kv.max() - selected_kv.min()) / (2**num_bits - 1)
    quantized = ((selected_kv - selected_kv.min()) / scale).round().clamp(0, 2**num_bits - 1)

    return quantized, scale, selected_kv.min()

6.2. Adaptive and Selective Retention

Recent methods, such as FastGen and MorphKV, profile attention patterns at runtime to determine which tokens or cache entries are most relevant for each layer or head. This enables:

Dynamic cache size: The cache adapts to attention diversity, keeping more entries where needed and aggressively compressing elsewhere.
Constant-size caches: MorphKV, for instance, maintains a fixed-size cache by iteratively refining which tokens to keep using attention patterns, preserving long-range dependencies with minimal accuracy loss and >50% memory savings.
Layer/head specialization: Different cache strategies can be applied to different layers or heads, rather than a one-size-fits-all approach.

6.3. Quantization for Throughput and Batch Size

Hardware-aware quantization (e.g., FP8, INT8, 4-bit) dramatically reduces memory requirements and enables higher effective batch sizes, especially in decode-heavy serving scenarios:

Throughput gains: Quantizing the KV cache can provide up to 1.45× throughput improvement in real-world LLM serving, primarily by allowing more requests to be processed in parallel.
Minimal accuracy loss: With careful quantization and dequantization strategies, there is little to no impact on model quality for most tasks.
Implementation caveats: The speedup depends on the compatibility of quantized caches with high-performance attention kernels; some frameworks (e.g., TensorRT-LLM) benefit more than others (e.g., vLLM) depending on kernel optimizations.

6.4. System-Level Optimizations

Beyond algorithmic compression and adaptive retention, recent research has revealed significant performance gains through system-level KV cache optimizations. Frameworks like NVIDIA TensorRT-LLM and vLLM’s PagedAttention have re-architected cache management to resemble operating system virtual memory more closely, using paged or block-based KV storage to minimize memory fragmentation and enable efficient on-demand allocation.

Other innovations, such as FlowKV, introduce distributed and disaggregated cache management strategies to reduce cache transfer latency and better utilize hardware resources across multiple nodes.

These system-level enhancements complement algorithmic advances by improving scalability, throughput, and latency, ensuring that KV cache innovations are effectively translated into real-world production deployments, particularly for large-scale, multi-user LLM inference workloads.

7. Practical Implications

Efficient cache strategies have significant practical implications that affect both memory and computational requirements. By optimizing cache retention and compression techniques, these advancements create smaller yet smarter caches that free up resources. This enables larger batch sizes, longer context windows, and allows for deployment on less expensive or resource-constrained hardware. Such innovations greatly enhance the performance and scalability of language models, making them accessible to a wider range of users and applications.

Moreover, these strategies enhance cost-effectiveness by reducing hardware demands and energy consumption. This not only lowers the financial barriers to deploying language models but also promotes sustainability within the field of machine learning. Additionally, the ability to track long-range dependencies with optimized cache management proves invaluable for tasks such as document retrieval, multi-turn dialogue generation, and summarization. These improvements underscore the importance of adaptive and efficient KV cache techniques in advancing the performance of large language models.

8. Conclusion

KV cache management is a cornerstone of efficient LLM inference, especially as models scale and the size of the context windows expands. The field has rapidly advanced from basic caching to sophisticated, adaptive, and task-aware strategies that strike a balance between memory, speed, and accuracy. Without KV cache, modern LLMs would be too slow, costly, and limited for today’s real-world applications. As research continues, expect even more advanced cache management, enabling efficient inference for ever-larger models and longer contexts.

References

RustySnake - Classic Snake game to learn Rust

Thu, 27 Mar 2025 00:00:00 +0000

1. Overview

Rust has been gaining attention recently due to its unique combination of performance, safety, and modern programming features. Its strict ownership model eliminates common memory issues like null pointer dereferencing and data races, providing a secure environment for developers. At the same time, its expressive syntax and focus on developer productivity make it a strong contender for systems programming. Its growing ecosystem and community continue to expand its capabilities into areas like web development and embedded systems, ensuring a confident and secure learning experience for developers.

When diving into Rust, a language lauded for its speed and safety, I wanted to start with a fun and educational project. Inspired by my experience with C and C++, I decided to learn Rust by implementing a classic snake game - a project that not only provides a fun learning experience but also reinforces an understanding of programming fundamentals. This journey has introduced me to Rust’s unique features and reinforced my understanding of programming fundamentals, engaging and motivating me to learn more.

As I embarked on this journey, I also wanted to understand how LLMs can assist in grasping Rust in terms of syntax, layout, flow, and execution aspects. These models played a crucial role in helping me debug and resolve issues, providing the necessary support and guidance in my learning process. For this, I used the following different models. Except for GPT 4o, I also want to run the others on my home GPU cluster - again, to help experiment with different aspects of inference.

Phi 4
Llama 3.3 70b Instruct
Qwen 2.5B Coder 32B
DeepSeek R1 Distill Llama 8B
GPT 4o

1.1. LLM Performance Comparison

Model	Code Generation	Error Explanation	Documentation	Overall Experience
Phi 4	✔️✔️✔️	✔️✔️	✔️✔️	Good for simpler tasks
Llama 3.3 70b	✔️✔️✔️✔️	✔️✔️✔️	✔️✔️✔️	Strong but inconsistent
Qwen 2.5B Coder 32B	✔️✔️✔️✔️	✔️✔️✔️	✔️✔️✔️✔️	Excellent for code examples
DeepSeek R1 Distill 8B	✔️✔️	✔️✔️	✔️✔️	Limited understanding of Rust
GPT 4o	✔️✔️✔️✔️✔️	✔️✔️✔️✔️✔️	✔️✔️✔️✔️✔️	Comprehensive assistance throughout

Of course, it is not an apples-to-apples comparison. Each started well but then quickly went downhill for various reasons. The only awesome one I finally ended up using all the way through was GPT4o.

1.2. Why Rust?

Rust is designed to be a systems-level language with modern features. It eliminates common pitfalls like memory safety issues without compromising performance. As someone familiar with C and C++, I think this seems like a natural progression.

2. The Snake Game Project

   /^\\/^\\
  / o   o \\
 (    ^    )
  \\_______/
   |     |
   |     |

Rusty Snake

The nostalgic, classic snake game provides a perfect playground for learning Rust’s syntax, constructs, and libraries. For the game, I started with the basic game but then soon added other things - I wanted to give the player a choice on the size of the play area (measured in number of characters as the rendering is in ASCII) and speed.

A crate in Rust is a collection of related code akin to a library in C++ or a module in Python. The project imports the following crates for different functionalities needed:

crossterm: For terminal UI and input handling.
rand: For generating random numbers (e.g., food placement).

Let us walk through some key implementation parts and draw parallels to C and C++ for easier comprehension.

3. Key Concepts in the Code

3.1 Structs: Rust’s Version of `struct` in C

Rust’s struct is similar to struct in C and C++ but with additional safety and functionality. Here’s how we define a Point struct to represent 2D coordinates:

#[derive(Debug, Clone, PartialEq)]
struct Point {
    x: i32,
    y: i32,
}

Parallel in C/C++:

struct Point {
    int x;
    int y;
};

In Rust, #[derive(...)] automatically implements traits like Debug (for debugging), Clone (for copying), and PartialEq (for comparisons).

3.2 Vectors `Vec<T>`): Dynamic Arrays

Rust uses Vec<T> for dynamically-sized arrays, similar to std::vector in C++.

let mut snake = vec![Point { x: width / 2, y: height / 2 }];

Parallel in C++:

std::vector<Point> snake = { Point{width / 2, height / 2} };

Unlike C++, Rust ensures safety with bounds checking and ownership rules.

3.3 Ownership and Borrowing

Rust’s ownership model ensures memory safety without a garbage collector. When working with snake, you own or borrow the data.

Example:

let mut new_head = snake.last().unwrap().clone();

Here, .clone() creates a deep copy of the last element to avoid ownership issues.

Parallel in C++: Copying would be explicit but lacks ownership enforcement:

Point new_head = snake.back();

3.4 Enums and Pattern Matching

Rust’s pattern matching with match is more expressive than a switch statement in C++.

match key_event.code {
    KeyCode::Up if direction.y == 0 => next_direction = Point { x: 0, y: -1 },
    KeyCode::Down if direction.y == 0 => next_direction = Point { x: 0, y: 1 },
    _ => {}
}

Parallel in C++: A switch with additional if conditions might approximate this but would lack the same level of elegance.

3.5 Error Handling

Rust avoids exceptions by using the Result type for error handling, relying on explicit and predictable control flow to manage errors. This differs from other languages and ensures errors are handled explicitly rather than relying on potentially disruptive exception mechanisms, which can lead to unwieldy codebases in large systems.

fn main() -> crossterm::Result<()> {
    terminal::enable_raw_mode()?;
    ...
    terminal::disable_raw_mode()?;
    Ok(())
}

The ? operator propagates errors automatically, akin to C++’s std::optional or manual error handling:

if (!enable_raw_mode()) return -1;

3.6 Functional Features

Rust’s iterators and closures make code concise and expressive. For instance, generating random food coordinates:

let new_food = Point {
    x: rng.gen_range(1..width - 1),
    y: rng.gen_range(1..height - 1),
};

Parallel in C++:

Point new_food = { rand() % (width - 2) + 1, rand() % (height - 2) + 1 };

Rust’s range syntax and rand crate simplify random number generation.

4. Code

The complete source code for this project is available on GitHub: https://github.com/bahree/rustysnake

The game loop handles input, updates the snake’s position, checks for collisions, and renders the game state. The crossterm library provides a simple way to handle keyboard input and terminal rendering.

loop {
    // Handle input
    if crossterm::event::poll(Duration::from_millis(game_speed))? {
        if let Event::Key(key_event) = event::read()? {
            match key_event.code {
                KeyCode::Char('q') => break,
                KeyCode::Up if direction.y == 0 => direction = Point { x: 0, y: -1 },
                KeyCode::Down if direction.y == 0 => direction = Point { x: 0, y: 1 },
                KeyCode::Left if direction.x == 0 => direction = Point { x: -1, y: 0 },
                KeyCode::Right if direction.x == 0 => direction = Point { x: 1, y: 0 },
                _ => {}
            }
        }
    }

    // Update snake position
    let new_head = Point {
        x: (snake[0].x + direction.x + width) % width,
        y: (snake[0].y + direction.y + height) % height,
    };

    // Check collision with self
    if snake.contains(&new_head) {
        break;
    }

    snake.insert(0, new_head);

    // Check if snake ate food
    if snake[0] == food {
        // Generate new food
        loop {
            let new_food = Point {
                x: rng.gen_range(0..width),
                y: rng.gen_range(0..height),
            };
            if !snake.contains(&new_food) {
                food = new_food;
                break;
            }
        }
    } else {
        snake.pop();
    }

    // ...
}

Feel free to clone, fork, or contribute to the repository!

4.1 Running the Game

To run this on Windows:

Install Rust with rustup.

Add dependencies to Cargo.toml:

[dependencies]
crossterm = "0.27"
rand = "0.8"

Build and run the project:
```
cargo run
```
Compile the project:
```
cargo build --release
```

4.2 Controls

Key	Action
`Arrow Keys`	Move the snake
`+`	Increase game speed
`-`	Decrease game speed
`Spacebar`	Pause/Resume the game
`q`	Quit the game

4.3 Gameplay

Select the boundary size and difficulty from the menu.
Use arrow keys to move the snake.
Eat the red food (■) to grow your snake and increase your score.
Avoid hitting the walls (#) or yourself!

Snake Game

5. Learnings

I encountered a few things while learning Rust through this snake game project. These are called out below. As with anything else, where things are different, one must forget what one knows and learn the new way of doing things.

Ownership model: Coming from C/C++, adjusting to Rust’s strict borrowing rules required a shift in thinking
Terminal UI limitations: Working with crossterm’s terminal interface meant designing around ASCII graphics
Game loop timing: Balancing responsiveness with performance required careful tuning
Error propagation: Learning to use the ? operator and Result types effectively took practice

Understanding these differences was part of the learning process and helped solidify some of my understanding of Rust’s principles.

5.1 Reflecting on Rust

Building this game helped me appreciate Rust’s features:

Memory Safety: Rust eliminates segmentation faults common in C and C++.
Expressive Syntax: Iterators, pattern matching, and traits make Rust elegant and powerful.
Community and Ecosystem: Libraries like crossterm and rand accelerate development.

5.2 Acknowledgment

Throughout this journey, I used the assistance of LLMs to clarify Rust concepts and draw parallels to languages I already know. This collaborative approach made learning smoother and deepened my understanding by offering different perspectives. The combination of hands-on learning and real-time assistance has been invaluable.

6 Final Thoughts

Rust might feel different initially if you’re from a C or C++ background, but its safety guarantees and modern features make it worthwhile. Building a snake game is just one of many ways to start.

What will you build next? 📎

An introduction to Mixture of Experts (MoE)

Wed, 01 Jan 2025 00:00:00 +0000

AI is advancing at an unprecedented pace, with Mixture of Experts (MoE) models being one set of model architectures at the forefront of this revolution. These architectures enable breakthroughs in efficiency and scalability by leveraging a modular design where only a subset of specialized “expert” networks are activated for each input. MoE architectures have become a cornerstone in building ultra-large-scale models like GLaM and Switch Transformers.

Mixture of Experts (MoE) is an advanced machine learning architecture that lately has gained significance, particularly in the realm of #LLMs (large language models) and NNs (neural networks). In talking with many people about AI, I’ve found that MoE as a topic comes up often, with many folks either not understanding what it is or their understanding of it being incorrect.

With recent announcements of trillion-parameter models and announcements from Microsoft, OpenAI, and Google, understanding MoEs is more important than ever. More recently, DeepSeek v3 is a great example of a model that uses MoEs to achieve state-of-the-art performance - where the language model has 671B total but only 37B activated for each token.

I this post I provide a high-level overview of MoEs, their core components, how they work and workflow. I also includes a simple toy example implementation to help grasp the core concepts.

Why Mixture of Experts?

The central motivation behind MoE stems from the tension between growing model size and limiting computational resources. As we have seen in the recent past, increasing the parameter count of a model often yields better performance, especially in domains like natural language processing (NLP) and computer vision; however, this also drastically increases the cost of both training these models and computing cost for inference of these models. This massive computing cost is at the heart of what MoEs are addressing by offering a different paradigm known as conditional computation. MoE’s activate only a small subset of specialized sub-networks (called experts) for each input or token) only to a small number of experts rather than processing it through every parameter in the network. This helps with three key aspects:

Increased Model Capacity: Because only a few experts are activated at a time, MoE architectures can pack many parameters (experts) without proportionally increasing the computational cost per input.
Specialized Sub-networks: Different experts can learn specific patterns or token-level specializations, leading to better performance.
Efficient Usage of Compute: MoE optimizes resources (compute) by activating only a small fraction of the entire model, leading to significant efficiency gains.

This sparse activation strategy enables constructing models with billions (or even trillions) of parameters, making MoE a scalable solution for large-scale applications. Another benefit of MoE is its capacity for specialization. Different experts can learn distinct, context-dependent processing strategies, enabling the model to cover a broad set of input variations more effectively than a monolithic architecture.

Early MoE ideas trace back to model ensembling in classic machine learning. Still, MoE extends beyond ensembling by learning a parametric “router” (gating function) that dynamically decides which experts to use. Notable works like the Sparsely-Gated Mixture of Experts showed that MoE could massively scale model size while staying computationally efficient.

Core Components of an MoE System

The fundamental building blocks of an MoE system are:

Experts: Specialized sub-models, typically feed-forward neural networks.
Router or Gating Network: Determines which tokens are sent to which experts.
Combiner: Aggregates the outputs from the selected experts.

Experts

Each expert is typically a neural sub-network replicated multiple times within the same model. Depending on the task, these experts are often implemented as independent feed-forward networks, MLP blocks, or convolutional layers. For large-scale language models, each expert usually mirrors the structure of the feed-forward component of a Transformer block.

Since multiple experts exist in parallel, each can potentially learn to handle different token types or data distributions. Contrary to common misconceptions, experts do not necessarily correspond to human-like semantic domains (e.g., “Expert #3 = Physics”). One might assume that certain experts correspond to high-level domains such as “finance” or “medicine.” Still, it is more common for the learned expertise to be more subtle and token-level, capturing idiosyncratic patterns that are not necessarily interpretable in a straightforward semantic way. Each expert ends up specializing in token-level or feature-level patterns that are discovered during training. For instance, in NLP, one expert might specialize in syntactic structures, while another focuses on semantic relationships.

Sometimes, certain layers (like embeddings or attention blocks) are shared among all experts, and only feed-forward layers are duplicated (as in Switch Transformers). This partial sharing allows the model to keep some global representation while still having specialized processing in the experts.

Router

The router, often called the gating network, is a small module that predicts which experts should handle any given input. Its purpose is to determine which experts should handle a given input. Modern MoE designs are typically parameterized as a simple neural network (often a single linear layer + SoftMax) or a simple linear transform. The SoftMax output provides a probability distribution across the experts, indicating which should be “activated” for each input.

The router reads the input representation (e.g., the token embedding in an NLP model) and produces a probability distribution over the experts, typically with a SoftMax function. We then select the top-k experts (e.g., top-1, top-2) based on these probabilities - for each token, known as Top-k gating.

Practical implementations often limit the number of tokens an expert can process per batch (known as Expert Capacity). If too many tokens route to the same expert, some tokens may get dropped or rerouted, leading to training instability. This capacity limit helps prevent any single expert from monopolizing the model’s processing and prevents load imbalances.

The gating network is trained jointly with the experts through back-propagation. As we outlined above, the gating process introduces discrete decisions into the computational graph, which can hamper backpropagation. To help counter this, additional techniques such as adding a small amount of noise to the logits (Noisy Top-k Gating) or using Soft MoE — are employed to smooth out these discrete selections, keep training stable, help provide smoother gradients, and encourage balanced expert utilization. Additional mechanisms like Expert Capacity limit how many tokens each expert can process, preventing load imbalances in which a single expert might receive most tokens.

Combiner

Once the chosen experts have computed their outputs for a given token, these outputs must be aggregated to produce a single vector that feeds into subsequent layers. The MoE architecture achieves this through a combiner, which typically performs a weighted sum of the experts’ outputs, using the gating probabilities as weights. In top-k gating, if $ k $ experts were activated, each expert’s output is multiplied by its corresponding probability from the router. The combiner then sums or otherwise fuses the results to form the token’s transformed representation. This consolidated output is passed into the rest of the model, such as attention blocks or additional Transformer layers.

In “Soft MoE” variants, we might use a soft combination, letting tokens pass to all experts but with different fractional weights—alleviating some routing discontinuities at the cost of higher computation.

High-Level Workflow

Input is fed into the gating network, which produces a probability distribution (or scores) over experts.
The top-k experts are activated for each input (or token, in the case of language models).
The selected experts process the input in parallel.
A combiner fuses the experts’ outputs into a single vector.
The model produces the final output, which can then feed into other layers or tasks.

Below is a simple flow diagram that shows how to visualize this:

flowchart TD
    Input_Tokens --> Gating_Network
    Gating_Network --> Expert_1
    Gating_Network --> Expert_2
    Gating_Network --> Expert_3
    Expert_2 --> Combiner
    Expert_3 --> Combiner
    Combiner --> Final_Output

The MoE module often appears in LLMs where a standard feed-forward layer normally would reside. For each layer designated as an MoE, the process begins by sending the hidden representations of each token to the router. The router computes gating probabilities across all experts, ranks the experts by these probabilities, and picks the top few. Each selected expert applies its transformation to the token’s hidden state. Finally, the combiner merges these expert outputs with a weighted sum.

During training, the experts and the router are updated jointly through backpropagation. However, discrete gating can make gradient flow tricky since the top-k selection is not inherently differentiable. In practice, noise injection (Noisy Top-k Gating) or methods like Soft MoE can help approximate continuous gradients, ensuring that even experts with lower gating probabilities receive occasional training signals.

In addition, to avoid a scenario where one or two experts monopolize all tokens, an auxiliary load-balancing loss is often introduced to encourage more uniform usage. This might take the form of a penalty term that grows when the variance in expert usage is high, incentivizing the router to distribute tokens more evenly.

This mechanism occurs for every token in every layer designated as an MoE layer, which is why load balancing is so critical—without careful design, certain experts can receive far more tokens than others.

What is the Difference between an MoE and an Ensemble?

It’s easy to confuse MoE with model ensembling - where multiple independently trained models vote or average predictions. MoE differs in a few critical ways:

Dynamic and data-dependent routing: In classical ensembles, each model sees the same input, and a meta-learner or a simple average produces the final result. MoE, in contrast, uses a router that decides different subsets of the input for each expert. This dynamic routing allows MoE to specialize in different patterns or tokens.
Single model training: MoE typically trains all experts jointly in one go. Each expert does not have a separate training pass; they’re part of the same computation graph, sharing some parameters (like embeddings) and learning together. This is in contrast to ensembles, where each model is trained independently.
Fine-grained token specialization: Different tokens in the same sentence might get routed to different experts, enabling extremely fine-grained specialization. This is impossible in a traditional ensemble, where each model sees the entire input.

Load Balancing Experts and Training Pitfalls

A significant challenge in training MoE models is ensuring the balanced utilization of experts. Certain experts may become overburdened without proper load balancing, while others remain underutilized, leading to inefficient training and suboptimal performance. Gating can become highly imbalanced early in training, favoring a few experts. Common solutions for Load Balancing:

Auxiliary Load Balancing Losses: Adding a regularization term to the loss function encourages the gating network to distribute inputs evenly across all experts.
Top-k Randomization: Instead of always selecting the top-k experts with the highest gating probabilities, randomizing the selection among the top candidates can prevent overloading.
Expert Capacity Constraints: Limiting the number of tokens an expert can process at a time can help ensure all experts are used during training.

Over time, load-balancing losses/techniques help the distribution even out, but the model can remain fragile without careful hyperparameter tuning. Expert capacity is another important design choice. Since top-k selection may route too many tokens to the most popular experts, a capacity limit ensures each expert processes no more than a certain maximum number of tokens in one forward pass. The remaining tokens must be dropped or re-routed to other experts if an expert is at capacity. Both approaches come with trade-offs: dropping tokens entirely can lead to data inefficiency, whereas re-routing can add complexity and undermine the sparsity benefits that MoE aims to provide.

Load balancing - ensuring experts share the training load - remains one of the biggest technical hurdles with MoEs. Early in training, the gating network might discover that routing most tokens to the same expert (or a small number) might yield acceptable results, leaving other experts underutilized and effectively “dead.” This imbalance leads to suboptimal solutions - if only a few experts get trained on all tokens, you lose the advantages of specialization.

To mitigate this, MoE architectures often introduce additional techniques that nudge the model toward more even usage. One such technique is the auxiliary load-balancing loss. By monitoring how frequently each expert is selected, the model can be penalized if certain experts remain underutilized. This can be thought of as an additional penalty term that measures how evenly tokens (or batches) are distributed across experts.

Another common approach is Noisy Top-k Gating, which injects a small amount of learnable or fixed noise into the gating logits before the softmax, making the gating probabilities slightly more random. This randomness allows less-popular experts to occasionally receive tokens, which can help them develop more useful specializations over time. The gating network thus learns not only to minimize the primary task loss but also to spread tokens more uniformly.

Sparsity in Mixture of Experts (MoE) Models

Sparsity is one of MoE’s most valuable contributions to model efficiency. By only activating a small fraction of the total parameters for each input, the model maintains a much lower compute footprint than a dense model of the same overall size. This efficiency is crucial for scaling; while a trillion-parameter dense model may be prohibitively expensive to train and deploy, a trillion-parameter MoE model that only activates 1% of those parameters simultaneously becomes significantly more tractable.

That said, implementing sparsity at scale often requires specialized infrastructure. Frameworks like GShard or Mesh-TensorFlow are designed to handle data and model parallelism necessary for distributing the experts across GPU clusters. The overhead of routing tokens to the correct devices can become significant if the system is not carefully optimized. Researchers have also explored alternative gating mechanisms, such as Soft MoE, which approximates selection by routing every token to all experts in a soft, weighted fashion. While this approach can mitigate the fragility of discrete gates, it naturally increases computation since more experts perform computations at once.

Sparcity in MoE models offers several key advantages:

Computational Efficiency: Sparsity dramatically reduces the number of FLOPs required to process each task.
Scalability: The sparse activation of experts enables MoE models to scale to a large number of experts without a corresponding linear increase in computational and memory requirements.
Increased Model Capacity: Sparsity allows MoE models to increase their overall parameter count and model capacity without significantly increasing the computational cost during training or inference.
Memory Efficiency: Operating sparsely, MoE models require less memory for activations and parameters.
Specialized Processing: Sparsity enables the model to route different inputs to the most relevant experts, allowing for more specialized and efficient processing of diverse inputs.

Practical applications of MoE

MoE architectures have already demonstrated clear benefits in many areas. Microsoft’s Z-code model (Machine Translation), for instance, leverages MoE to handle multilingual translation tasks at a massive scale, and Google’s Switch Transformers showed that sparse activation can reach higher quality at lower training cost than dense baselines on benchmarks such as GLUE and SuperGLUE. In computer vision, MoE modules have been integrated into Vision Transformers (V-MoEs) to achieve better image classification and detection accuracy, with each expert focusing on different aspects of the image representation. In multimodal learning, the capacity to handle diverse data types—such as text, images, and audio—makes MoE a natural fit because experts can adapt to different modalities or different subproblems within a single modality.

In the context of LLMs, systems like ChatGPT, Claude, and Gemini can benefit from MoE by leveraging different experts for different topics or query types - though specifics are often proprietary and not shared. MoE is particularly suited to multi-modal tasks involving text, images, and audio, as experts can specialize in different modalities or sub-modalities. This is valuable for text-to-image generation or video understanding.

Several emerging directions continue to push MoE research forward. Soft MoE (Zuo et al., 2022) is an example aiming to produce a fully differentiable version of sparsely gated Transformers. Another is Parameter-Efficient Sparsity Crafting (PESC), which seeks to retrofit existing dense models into a sparse MoE design without retraining from scratch. These innovations reflect ongoing efforts to refine the balance between sparse efficiency, training stability, and model reliability.

In production, deploying a large MoE model requires carefully coordinating hardware resources, data pipelines, and load-balancing techniques. Training an MoE system may involve more hyperparameters than a comparable dense model, including the number of experts, gating softmax temperature, top-k value, load-balancing penalty weights, and expert capacity. These factors can significantly affect performance, convergence speed, and final accuracy. When scaling an MoE model across multiple GPUs, designers must pay attention to network communication overhead. Token-based routing leads to collective operations that can become bottlenecks if not carefully optimized.

Despite these complexities, MoE’s flexibility and computational cost savings make it a compelling choice for handling highly varied or large-scale tasks. However, Fine-tuning MoE models can be more delicate than dense models because the gating distributions or specialized experts may not adapt smoothly to a new domain without carefully applied load-balancing strategies. There can also be interpretability challenges since the model’s internal “expert structure” does not always map neatly to skills.

Challenges and Considerations

We have touched on most of these in this blog post, but it is helpful to outline the key issues and considerations to be mindful of when using MoE-based models:

Complexity: MoE models are significantly more complex (compared to traditional neural networks) and require substantial computational resources for training and inference.
Training Instability: MoE models can suffer from training instability due to the discrete nature of expert selection.
Load Balancing: Proper load balancing among experts is crucial for efficiently using model capacity and optimal performance.
Computational Overhead: The gating mechanism introduces additional computational overhead, potentially impacting training and inference times.
Interpretability Issues: The dynamic routing of inputs makes interpreting how MoE models arrive at their decisions challenging.
Hyperparameter Sensitivity: MoE models have several hyperparameters that must be tuned for optimal performance.

Let’s define these mathematically before discussing practical considerations.

Mathematical Formulation of MoE

Let $ x \in \mathbb{R}^d $ denote an input vector (a hidden representation from a preceding layer). We assume the system has $ N $ experts, with each expert $ E_i $ parameterized as a function $ E_i: \mathbb{R}^d \rightarrow \mathbb{R}^m $. The gating network $ G $ takes the same input $ x $ and outputs a vector in $\mathbb{R}^N$—essentially, a “score” or “weight” for each of the $ N $ experts. Finally, a combiner function $ C $ merges the experts’ outputs into a single output vector $ y \in \mathbb{R}^m $.

Formally, we can write:

Input: $ x \in \mathbb{R}^d $
Experts: $ E_i: \mathbb{R}^d \rightarrow \mathbb{R}^m \quad \text{for}\ i = 1, \ldots, N $
Gating Network: $ G: \mathbb{R}^d \rightarrow \mathbb{R}^N $
Combiner: $ C: \mathbb{R}^{N \times m} \rightarrow \mathbb{R}^m $

Top-k Gating

Most contemporary MoE implementations use top-k gating, which activates only the $ k $ experts with the highest gating scores. In this scenario, the summation is performed only over those top-$ k $ indices. If we denote $\text{top-k}(G(x))$ as the set of indices corresponding to the $ k $ largest values of $ G(x) $, then

$ y = C\Bigl(\sum_{i \in \text{top-k}(G(x))} G(x)_i ,\cdot, E_i(x)\Bigr). $

By pruning all but the top-$ k $ experts per input, this design enforces sparse activation: each input (or token) only “touches” $ k $ out of $ N $ experts at a time. This approach significantly reduces the computational load relative to using all $ N $ experts for every input.

Dense Gating Formulation

This refers to a version of MoE architecture in which all experts contribute to the final output for each input rather than filtering out all but the top-k experts. Here, the gating network assigns continuous weights to every expert, aggregating each expert’s weighted output into a final result. There is no discrete selection to zero out certain experts based on gating. Fundamentally, this is the opposite of top-K gating.

In the simplest version of MoE, where the gating network’s output is a set of continuous weights, the forward pass for one input $ x $ can be written as:

$ y = C\Bigl(\sum_{i=1}^N G(x)_i ,\cdot, E_i(x)\Bigr). $

Here, $ G(x)_i $ is the $i$-th component of the gating network’s output for $ x $ and represents the weight (or probability) assigned to expert $ i $. Intuitively, if $ G(x)_i $ is large, expert $ i $ contributes more to the final output $ y $. The function $ C $ often takes the form of a simple weighted sum or concatenation-and-projection step, depending on the specific design.

Load Balancing and Expert Monopolization

As we saw earlier, MoE architectures introduce an *auxiliary load-balancing loss*to load valance across all experts. This can be thought of as an additional penalty term that measures how evenly tokens (or batches) are distributed across experts. One common strategy is penalizing the variance in expert usage or encouraging each expert to receive roughly an equal proportion of examples. The gating network thus learns not only to minimize the primary task loss but also to spread tokens more uniformly.

Mathematically, a typical load-balancing loss might look like:

$ \mathcal{L} _ {\text{balance}} = \lambda \sum_{i=1}^N \Bigl(\frac{f_i}{\sum_j f_j} - \frac{1}{N}\Bigr)^2, $

where $ f_i $ is the total number of tokens assigned to expert $ i $ in a minibatch (or the sum of gating probabilities if you’re using a continuous measure), and $ \lambda $ is a hyperparameter controlling the strength of this penalty. This ensures the model is incentivized to explore and train all experts over time.

Example PyTorch Implementation

Here’s a simplified Python implementation of an MoE model using PyTorch:

import torch
import torch.nn as nn

class Expert(nn.Module):
    def __init__(self, input_size, output_size):
        super(Expert, self).__init__()
        self.layer = nn.Linear(input_size, output_size)
    
    def forward(self, x):
        return self.layer(x)

class GatingNetwork(nn.Module):
    def __init__(self, input_size, num_experts):
        super(GatingNetwork, self).__init__()
        self.layer = nn.Linear(input_size, num_experts)
        self.softmax = nn.Softmax(dim=1)
    
    def forward(self, x):
        return self.softmax(self.layer(x))

class MixtureOfExperts(nn.Module):
    def __init__(self, input_size, output_size, num_experts, top_k=1):
        super(MixtureOfExperts, self).__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.experts = nn.ModuleList([Expert(input_size, output_size) for _ in range(num_experts)])
        self.gating = GatingNetwork(input_size, num_experts)
    
    def forward(self, x):
        batch_size = x.size(0)
        gating_probs = self.gating(x)
        topk_vals, topk_inds = torch.topk(gating_probs, self.top_k, dim=1)
        expert_outputs = torch.zeros(batch_size, self.top_k, output_size, device=x.device)
        for i in range(self.top_k):
            inds = topk_inds[:, i]
            outputs_for_expert = torch.stack([self.experts[inds[b]](x[b].unsqueeze(0)) for b in range(batch_size)])
            expert_outputs[:, i, :] = outputs_for_expert.squeeze(1)
        topk_vals_expanded = topk_vals.unsqueeze(-1)
        weighted_sum = expert_outputs * topk_vals_expanded
        combined_output = weighted_sum.sum(dim=1)
        return combined_output

# Example Usage
input_size = 10
output_size = 5
num_experts = 4
top_k = 2
model = MixtureOfExperts(input_size, output_size, num_experts, top_k)
sample_input = torch.randn(8, input_size)
output = model(sample_input)
print("Output shape:", output.shape)

When you run this code, you should see the output shape printed as torch.Size([8, 5]). This confirms that the model routes each input to the most relevant experts, processes the input through those experts, and combines their contributions into a unified output.

Code Explanation

This code defines a simple MoE model with a single gating network and multiple experts. The Expert class represents a single sub-module, the expert, a linear layer that transforms the input and forms the backbone of the MoE architecture. The expert takes the input of size input_size and transforms it to output_size using a fully connected (Linear) layer. In the real world, each expert is designed to specialize in different transformations based on their training.

The GatingNetwork class is the router that computes gating probabilities for each expert and determines which experts to activate for a given input. It takes the input size and the number of experts as input and outputs a probability distribution over the experts (num_experts) using a linear layer followed by a Softmax function. Higher probability values indicate that the corresponding expert is more relevant to the input.

The MixtureOfExperts class combines the experts’ outputs based on the gating probabilities and returns the final output. It takes the input size, output size, number of experts, and the top-k value as input. The forward method computes the gating probabilities, selects the top-k experts based on these probabilities, and computes the weighted sum of the expert outputs to produce the final output. The top_k parameter controls how many experts are activated for each input.

The input ( x ) is passed through the gating network to produce a probability distribution over the experts: gating_probs = self.gating(x)
The gating network selects the indices of the top_k experts with the highest probabilities: topk_vals, topk_inds = torch.topk(gating_probs, self.top_k, dim=1)

For each of the selected experts, process the input:

expert_outputs = torch.zeros(batch_size, self.top_k, output_size, device=x.device)
for i in range(self.top_k):
    inds = topk_inds[:, i]
    outputs_for_expert = torch.stack([self.experts[inds[b]](x[b].unsqueeze(0)) for b in range(batch_size)])
    expert_outputs[:, i, :] = outputs_for_expert.squeeze(1)

Multiply each expert’s output by its gating probability and sum them to form the final output:
```
topk_vals_expanded = topk_vals.unsqueeze(-1)
weighted_sum = expert_outputs * topk_vals_expanded
combined_output = weighted_sum.sum(dim=1)
```
- The topk_vals (probabilities) are weights for the corresponding expert outputs.

The model processes a batch of inputs and returns the combined output with a shape matching (batch_size, output_size). For the example above, Output shape: torch.Size([8, 5]) confirms that the model routes each input to the most relevant experts, processes the input through those experts, and combines their contributions into a unified output.

Code Dependencies

To run this code, you’ll need the following dependencies:

Python 3.6 or later (preferably 3.10 or higher)
PyTorch 1.0 or later

Save the code to a file, e.g., moe_example.py, and run it:

python moe_example.py

I used Conda, which I prefer for managing Python environments. If you’re using a virtual environment, you can adapt the installation commands accordingly. For Conda, you can create a new environment and install PyTorch using the following steps.

Start by creating a new conda environment with PyTorch dependencies. In your terminal, execute the following commands:

# Create a new conda environment named "moe_env"
conda create -n moe_example python=3.10 -y

# Activate the environment
conda activate moe_example

Install PyTorch and necessary dependencies; adjust the CUDA version based on your system’s GPU configuration. You can omit the pytorch-cuda package using a CPU-only setup.

# For CPU-only:
conda install pytorch torchvision torchaudio cpuonly -c pytorch -y

# For GPU (use appropriate CUDA version, e.g., 11.7):
conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia -y

# Install additional dependencies if needed
pip install numpy

This example demonstrates the basic structure of an MoE model, including the experts, gating network, and the MoE module that combines them. Of course, this is a toy version that helps understand the basic construct and does not include advanced features like load balancing or sparsity, etc.

Conclusion

Mixture of Experts offers a compelling framework for scaling neural networks and managing the trade-offs between model size, computational cost, and performance. By selectively activating only a subset of parameters for each input, MoE allows researchers and practitioners to build models with enormous capacity without incurring a proportionate computational penalty. Given the additional interest in inference optimization for LLMs and broadly with Transformer-based architecture, we expect to see further innovations and applications for MoE.

References

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer - https://arxiv.org/abs/1701.06538
GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding - https://arxiv.org/abs/2006.16668
Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity - https://arxiv.org/abs/2101.03961
A Lightweight Mixture-of-Experts Neural Machine Translation Model with Stage-wise Training Strategy - https://aclanthology.org/2024.findings-naacl.154/
Deep Mixture of Experts via Shallow Embedding - https://arxiv.org/abs/1806.01531
Mixture of Experts in Image Classification: What’s the Sweet Spot? - https://arxiv.org/abs/2411.18322v1
Mesh-TensorFlow: Deep Learning for Supercomputers - https://proceedings.neurips.cc/paper/2018/hash/3a37abdeefe1dab1b30f7c5c7e581b93-Abstract.html
A Generalist Agent - https://arxiv.org/abs/2205.06175
Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation - https://arxiv.org/abs/1308.3432
From Sparse to Soft Mixtures of Experts - https://arxiv.org/abs/2308.00951
Parameter-Efficient Sparsity Crafting for Large Language Models - https://aclanthology.org/2024.emnlp-main.43/

Automating Hugo Deployments

Mon, 23 Dec 2024 00:00:00 +0000

1. A little background

I have been meaning to automate the deployment of my blog post to a dev server (running locally) for a while, but I haven’t had the time to get around to it until now. In addition to deploying this, the dev server also had several constraints. It is one of the machines at home and is not exposed directly to the internet. I also don’t have any ports opened on the firewalls at home, which adds a few more interesting challenges. I have a server hosted online in one of the data centers, but I want to keep that for more ‘production’ things.

Another thing I wanted to do was clean up the theme being used, as it wasn’t compatible with the newer version of Hugo. The longer I was pinned to an older version, the more challenging it would be to update. Now, these are all done and working, and I wanted to share what I learned here—more so because I will likely need to find this at some point in the future, but it might also help someone else.

PS - If you don’t notice much (any?) difference in the blog’s look, that means it is a good thing; I worked through fixing the theme and upgrading to the latest version of Hugo (v0.140.0 at the time of writing this).

While cleaning up, I also got around to things I had wanted to do for a while. I wanted to see how to securely expose one of the machines I have at home online without opening up ports and things and serve several services—mostly for my own use.

1.1 Dev-Server Setup 🖥️

I think outlining how this machine is set up would be helpful, as I am sure many others will be trying this. First, this machine runs locally on the home network, running a headless Ubuntu server. The environment that made this dev server unique was the combination of things - I use Tailscale, Cloudflare Tunnels, CloudPanel, and as outlined earlier Hugo. The code sits in a private Github repo, and the build-server is via Github Actions. I won’t get into the details of each setup here as the intent of the blog post is not that, but I will share some context below.

1.2 Tailscale

To make my life easier between the different machines and servers I have, running both locally at home and on the cloud, including bare-metal machines, I use Tailscale.

Tailscale simplifies network management by creating a secure, zero-config mesh VPN based on the WireGuard protocol. Unlike traditional VPNs that route all traffic through a central server, Tailscale forms direct, encrypted connections between devices, regardless of their physical location or network configuration. This “mesh” approach offers significant performance benefits, reduced latency, and improved reliability. By assigning each device a stable IP address within a private network namespace, Tailscale eliminates the complexities of managing firewalls, port forwarding, and dynamic DNS, making it ideal for connecting personal devices, servers, and cloud instances into a cohesive, secure network.

What makes Tailscale work like magic is a feature they call MagicDNS. This feature eliminates the need to remember IP addresses and SSH keys for different machines. Using Tailnet, I get a private VPN between different machines, and I can use that to connect to any of these without opening any ports, etc. This significantly simplifies the networking process and enhances security.

1.3 CloudPanel

I use CloudPanel, a free and open-source server control panel for managing web servers and applications. In the context of this blog post, CloudPanel provides a user-friendly web interface for creating and managing websites, databases, email accounts, and DNS records. It’s a crucial part of the automation process, allowing for efficient server management and deployment. Its lightweight design and efficient resource utilization make it great for me.

1.4 Cloudflare Tunnels

Cloudflare Tunnels play a crucial role in the deployment process. They offer a secure and efficient way to expose web services running on private networks without opening inbound ports on firewalls. By establishing an outbound-only connection between a lightweight daemon (cloudflared) running on your server and Cloudflare’s global network, Tunnels create an encrypted tunnel that allows external users to access your services through Cloudflare’s edge. This approach enhances security by eliminating the attack surface associated with open ports while leveraging Cloudflare’s performance optimizations, such as caching, DDoS protection, and global CDN. It’s a powerful tool for self-hosting applications, APIs, and websites from home networks or private infrastructure without compromising security.

2. Server Setup ⌨️

To get this to work at a high level, we will do the following to get all of this working. This post won’t get into all the details of setting up the prerequisites, and it assumes that you have the following:

CloudPanel installed on your machine.
Cloudflare account with a Cloudflare Tunnel set up.
GitHub repo containing the Hugo project.
Tailscale for private networking - this is in addition to the HTTP traffic that Cloudflare handles.

2.1. Create a Static Site in CloudPanel

In CloudPanel, create a new “Site” (from the Dashboard). Select Static, and assign it a domain, e.g., blog.desigeek.com. CloudPanel will create a directory structure under something like:

/home/[cloudpanel_username]/htdocs/blog.desigeek.com

(Exact path can vary based on CloudPanel’s naming conventions.)

The figure below shows an example of how I set this up.

2.1.1 Create an SSH User

If you haven’t already, go to CloudPanel’s “Users” section. You can create an SSH user (for example, blog-ssh). This user will be mapped to the same group or path used by CloudPanel so that it can write your site files.

Important:

Make sure the SSH user’s home directory either is or can write to the directory created earlier /home/[cloudpanel_username]/htdocs/blog.desigeek.com.
For our example, I created an SSH user called blog-ssh, and the user’s group on Ubuntu was called bahree-blog.
Note: on some CloudPanel installations, the user’s primary group might be something else, not necessarily bahree-blog.

The image below shows an example of how to set this up.

2.1.2. Directory Paths and Permissions

Confirm that your SSH user (blog-ssh) can write into the blog.desigeek.com directory:

# Example:
ls -ld /home/[cloudpanel_username]/htdocs/blog.desigeek.com
# Make sure the owner or group matches your user or a group the user is in.

# If needed:
sudo chown -R blog-ssh:bahree-blog /home/[cloudpanel_username]/htdocs/blog.desigeek.com
sudo chmod -R 755 /home/[cloudpanel_username]/htdocs/blog.desigeek.com

(Adjust bahree-blog to the actual group that CloudPanel assigned.)

2.2 Cloudflare Tunnel Configuration

Cloudflare Tunnel is typically set up via Cloudflare’s admin dashboard or command line (see the docs for more details).

We will map the domain blog.desigeek.com to the local server’s port 80 (or 443). Since CloudPanel listens on port 80/443, your tunnel might forward something like:

Cloudflare subdomain → localhost:80 on your server.

Because Cloudflare is reverse proxying your site, you won’t need to open inbound ports on your server. Ensure your DNS is configured so that blog.desigeek.com points to the Cloudflare Tunnel.

The image below shows an example of what this would look like.

2.3. Setup Tailscale Secure Private Access

To deploy the Hugo site on my home dev server, I need private remote access (SSH, SCP, rsync, etc.) between the GitHub build servers and the home server. I achieve this using Tailscale.

Install Tailscale on your server:

curl -fsSL https://tailscale.com/install.sh | sh
sudo tailscale up

The figure below shows an example of an ephemeral node, which created and deleted in approx. a minute.

Our Build server is essentially spun up when something is committed in a specific branch, and as a result, we need to get that added to the Tailscale network - but it should not be permanent. Over time, that will pollute the tailnet and make things more difficult to manage. Tailscale has a perfect solution for this - Ephemeral nodes. These nodes, as the name suggests, make it easier to connect and then clean up short-lived devices such as containers, cloud functions, or CI/CD systems that spin up and spin down on a regular basis.

To use this, we generate an Ephemeral Auth Key from the auth keys page of the Tailscale admin console. Given that the build server would need this, we store this key in GitHub as TS_AUTH_KEY.

This allows the ephemeral GitHub Actions runner to connect to my dev server privately, even when I don’t have any inbound ports open.

2.4. Nginx Configuration via CloudPanel

CloudPanel typically manages the Nginx configuration for each “Site” automatically. It also handles SSL certificates, logs, etc. However, in this case, we are relying on Cloudflare Tunnel, and as a result, we do not need SSL termination on the server; Cloudflare can handle it and forward HTTP internally. We therefore need to ensure that the config file (vhost) in CloudPanel is minimal. Below is an example of what I use:

server {
  listen 80;
  listen [::]:80;
  server_name blog.desigeek.com;
  {{root}}

  {{nginx_access_log}}
  {{nginx_error_log}}

  location ~ /.well-known {
    auth_basic off;
    allow all;
  }

  {{settings}}

  include /etc/nginx/global_settings;
  index index.html;

  location ~* ^.+\.(css|js|jpg|jpeg|gif|png|ico|gz|svg|svgz|ttf|otf|woff|woff2|eot|mp4|ogg|ogv|webm|webp|zip|swf)$ {
    add_header Access-Control-Allow-Origin "*";
    expires max;
    access_log off;
  }

  if (-f $request_filename) {
    break;
  }
}

Note: Please note that {{root}} points to /home/[cloudpanel_username]/htdocs/blog.desigeek.com, which is set up in the CloudPanel site configuration.

The figure below shows what this would look like in CloudPanel.

3. GitHub Actions Configuration

We finally get to the main thing of building out GitHub actions that will trigger whenever something is checked in on a certain branch. All the other things until now were to essentially expose this dev machine externally (specifically the blog) and allow the build server to deploy to this dev server running at home. If you have a publicly accessible dev machine, much of this might not be required.

3.1. Storing Variables & Secrets in GitHub

We start by storing the variables and secrets we need for this to work. The details are outlined below. To set these up, go to Settings → Secrets and variables → Actions in your GitHub repo.

Variables (non-sensitive):

SERVER_HOSTNAME = my-dev-server (or your Cloudflare Tunnel/hostname). If Tailscale is used, it might be the Tailscale IP/hostname.
- Note: If you are using Tailscale, the notation of machine-name.local will not resolve and fail, even with MagicDNS enabled. It is best to use the machine name (e.g., my-dev-server and not my-dev-server.local or the IP address from the Tailnet)
SERVER_USERNAME = blog-ssh (the SSH user you created in CloudPanel).
HUGO_VERSION = e.g., 0.140.0
HUGO_EXTENDED = true (if you want the extended version, otherwise false).

Secrets (sensitive):

SSH_PASSWORD = the password for blog-ssh; this was either set or auto-generated when you created the account in CloudPanel.
TS_AUTH_KEY = Your Tailscale ephemeral auth key (only if using Tailscale).

The screenshot below shows what this would look like:

3.2 GitHub Actions Workflow

The code below outlines a sample workflow for this. It checks out the code, builds it with Hugo, and, once complete, copies it over to the dev server, where I can expose it securely to test it. To allow us to do this, we take a dependency on sshpass and rsync.

I do have a few things to check as this happens to ensure there aren’t any failures. When Hugo is finished, we show some basic stats, such as the number of files and folders created. This is called Show public directory stats.

Another is to test the SSH connection, as without that, the build machine cannot talk to the dev server where the final site is getting deployed. This step is called Test SSH Connection. At a high level, here is the flow that this goes through:

Check out the code
Installs Tailscale and connects as an ephemeral node
Builds with Hugo
Installs sshpass & rsync
Uses rsync to deploy to the server
Rolls back if the build fails
Logs out of Tailscale at the end

Below is a sample .github/workflows/deploy-hugo.yml.

name: Deploy Hugo Site

on:
  push:
    branches:
      - main  # or your chosen branch

jobs:
  deploy:
    runs-on: ubuntu-latest

    # Gets non-sensitive values from "Variables" (not Secrets)
    env:
      SERVER_HOSTNAME: ${{ vars.SERVER_HOSTNAME }}
      SERVER_USERNAME: ${{ vars.SERVER_USERNAME }}
      HUGO_VERSION: ${{ vars.HUGO_VERSION }}
      HUGO_EXTENDED: ${{ vars.HUGO_EXTENDED }}

    steps:
      - name: Checkout code
        uses: actions/checkout@v3

      # Optional Tailscale login if you're using Tailscale for SSH
      - name: Setup Tailscale
        uses: tailscale/github-action@v3
        with:
          authkey: ${{ secrets.TS_AUTH_KEY }}
          statedir: /tmp/tailscale-state/

      - name: Setup Hugo
        uses: peaceiris/actions-hugo@v2
        with:
          hugo-version: ${{ env.HUGO_VERSION }}
          extended: ${{ env.HUGO_EXTENDED }}

      - name: Build Hugo site
        id: build-site
        run: |
          hugo --minify -d public

      - name: Verify Hugo Output
        if: steps.build-site.outcome == 'success'
        run: |
          if [ ! -d "public" ] || [ -z "$(ls -A public)" ]; then
            echo "Error: 'public' directory is empty or does not exist."
            exit 1
          fi
          echo "Public directory exists and is not empty."

      - name: Show public directory stats
        if: steps.build-site.outcome == 'success'
        run: |
          cd public
          echo "Files: $(find . -type f | wc -l)"
          echo "Directories: $(find . -type d | wc -l)"
          echo "Size: $(du -sm . | cut -f1) MB"

      - name: Install sshpass and rsync
        run: |
          sudo apt-get update
          sudo apt-get install -y sshpass rsync

      - name: Test SSH Connection
        run: |
          sshpass -p "${{ secrets.SSH_PASSWORD }}" \
            ssh -o StrictHostKeyChecking=no ${{ env.SERVER_USERNAME }}@${{ env.SERVER_HOSTNAME }} \
            'echo "SSH Connection Successful"'

      - name: Deploy via rsync
        if: steps.build-site.outcome == 'success'
        env:
          SSH_PASSWORD: ${{ secrets.SSH_PASSWORD }}
        run: |
          sshpass -p "$SSH_PASSWORD" rsync -az --delete \
            -e "ssh -o StrictHostKeyChecking=no" \
            public/ ${{ env.SERVER_USERNAME }}@${{ env.SERVER_HOSTNAME }}:/home/[cloudpanel_username]/htdocs/blog.desigeek.com

      - name: Rollback on Failure
        if: steps.build-site.outcome != 'success'
        uses: appleboy/ssh-action@master
        with:
          host: ${{ env.SERVER_HOSTNAME }}
          username: ${{ env.SERVER_USERNAME }}
          password: ${{ secrets.SSH_PASSWORD }}
          script: |
            BACKUP_DIR="/home/[cloudpanel_username]/backups"
            SITE_DIR="/home/[cloudpanel_username]/htdocs/blog.desigeek.com"
            LATEST_BACKUP=$(ls -t "$BACKUP_DIR"/*.tar.gz | head -n 1)

            if [[ -n "$LATEST_BACKUP" ]]; then
              echo "Rolling back to: $LATEST_BACKUP"
              rm -rf "$SITE_DIR"/*
              tar -xzvf "$LATEST_BACKUP" -C "$SITE_DIR" --strip-components=1
              echo "Rollback successful!"
            else
              echo "No backups found to rollback to!"
            fi

      - name: Tailscale Logout
        if: always()
        run: |
          # If Tailscale requires sudo or operator privileges to log out,
          # either configure passwordless sudo or skip this if ephemeral.
          sudo tailscale logout

Note: Update /home/[cloudpanel_username]/htdocs/blog.desigeek.com to match the actual path CloudPanel created. The same applies to the backup directory if you use that mechanism.

The figures below show an example of when this was building, and one of the validations.

4. Tips & Troubleshooting 🚨

I wanted to share a few things I have learned as I went through and debugged this. Hopefully this will help someone else too.

Permission Denied: Make sure your CloudPanel SSH user has ownership or group-writable permission on /home/[cloudpanel_username]/htdocs/blog.desigeek.com.
Rsync Overwrites: The --delete flag in rsync removes files on the server that no longer exist locally. Remove ‘- delete ’ if you don’t want that behavior.
Cloudflare Tunnel: If the site is not reachable at blog.desigeek.com, check your Cloudflare DNS config or the tunnel routing.
SSH: You might still rely on Tailscale or an open SSH port for SSH since Cloudflare typically only proxies HTTP/HTTPS.
Tailscale Logout: If you see “Access denied: logout access denied,” you can skip tailscale logout (ephemeral nodes eventually vanish) or configure passwordless sudo for Tailscale.

5. Conclusion

By combining CloudPanel for site management, Cloudflare Tunnel for secure external traffic, Tailscale for private networking, and GitHub Actions for automated builds, you get a robust, convenient pipeline to deploy your Hugo blog. 😍

CloudPanel simplifies web server config and user management.
Cloudflare Tunnel ensures secure traffic without exposing direct ports.
Tailscale provides easy, private SSH if desired.
GitHub Actions automates the entire Hugo build and deployment process.

Now, when you push changes to your repository, your CloudPanel-managed site updates automatically, so no manual FTP or SSH is required!

Happy Deploying! ☺️

AI generated Podcast for my book: Generative AI in Action 🎧

Sun, 13 Oct 2024 00:00:00 +0000

The one thing I wanted to do after my book Generative AI in Action was complete was to create a summary in natural speech and possibly use TTS (Text-to-speech) to create an audio summary—think of it as a podcast that is easier for people to consume and get a quick sense of what the book is about.

TTS (Text to Speech) or not to TTS?

Initially, I was inclined towards using TTS (Text to Speech) for the audio summary. This technology, I thought, would be a convenient way to create a podcast-like summary that would be easier for people to consume and get a quick sense of what the book is about. My journey began with TTS - using GPT 4o to create a summary after ingesting the book and then using that as into the Azure AI Speech stack . However, I stumbled across something intriguing. Instead of TTS, I opted for NotebookLM from Google Labs to generate the audio overview - the podcast. This decision marked a significant shift in my approach, and I created two podcasts using the content from the book - one from multiple sources and another from a single source (the book).

NotebookLM is an experimental AI-first notebook from Google Labs designed to help users gain insights faster by grounding the language model in their documents. It aims to assist with synthesizing facts and ideas from multiple sources, making connections quicker and easier. It can help users understand, summarize, and generate new ideas based on their content. What is fascinating is that it can generate audio, which is a natural dialog between two people - with wit, humor, and a natural flow. It is like conversing with someone who has read the book and is summarizing it for you. And if I hadn’t told you that this was AI-generated, it would be hard to tell that it was not a real conversation.

The “Podcasts”

For the first audio generation, the multiple sources I used were:

The book, and my blog post announcing the book
A real podcast I did with Miko on his podcast Hockeystick
And another real podcast with Jamie on his podcast The modern .NET show .

For the second audio generation I only used a single source - the book.

The results of the AI-generated audio are truly impressive. In each instance, the audio was produced in a natural voice, simulating a genuine conversation between two people. The quality of the audio, the conversation, and the flow are nothing short of mind-blowing. Even in the second audio, where some acronyms were not pronounced correctly, it’s a minor issue, considering the audio was generated in just a few minutes with the press of a button. I was genuinely surprised at the level of realism and natural flow in the conversation. 🤯

Have a listen and let me know what you think.

Podcast Summary 1 - using multiple sources

The sources provide a comprehensive overview of generative AI and its application within enterprises. The first source, a YouTube video transcript, features an interview with Amit Bahree, a technical program manager at Microsoft, who discusses the rise of large language models (LLMs) and their potential impact on society. The second source, a book excerpt, delves into the technical aspects of generative AI, covering foundational models, large language models, retrieval-augmented generation, and the architectural principles for building generative AI applications. The book also explores various use cases, including image generation, code generation, and the ethical considerations surrounding the use of generative AI.

Your browser does not support the audio element.

Book Podcast - using multiple sources

Podcast Summary 2 - using single source (the book)

This book is a comprehensive guide to Generative AI, focusing on how this transformative technology can be leveraged within an enterprise. It explains core concepts like foundational models and large language models (LLMs) as well as practical applications for generating various content, including text, images, code, audio, and video. The book also explores responsible AI practices, highlighting the importance of prompt engineering, ethical considerations, and security measures for implementing these technologies. The author emphasizes the need for careful evaluation, monitoring, and scalability when deploying Generative AI models in a production environment.

Your browser does not support the audio element.

Book Podcast - using single source

Conclusion

AI-generated podcasts showcase the potential of AI to revolutionize content creation. Generating natural-sounding audio summaries from text is a game-changer for authors, creators, and educators. As AI continues to advance, we anticipate more opportunities to improve how we access information. I am excited to explore the new possibilities in AI-generated content.

Backing up TeslaMate data to OneDrive

Fri, 05 Jul 2024 00:00:00 +0000

I have been running a couple of instances of Teslamate - one locally on a server at home and another in Azure in a Ubuntu VM (see 👉 this blog post for details). I have been backing up the data to a NAS and then an offsite backup for the local instance. For the Azure instance, I have been running various backups during the day and backing up the data to OneDrive. This allows me to have a data backup in case the VM crashes or the data gets corrupted.

Work and AI have kept me busy, and I have not had a chance to write about it until now. I have a ton of storage on OneDrive, which seemed the most logical place to store it. Multiple options are available to back up data to OneDrive. The one that has been working consistently and well for me is rclone.

What is Rclone?

Rclone Rclone is an open-source command-line tool for managing files across various cloud storage services. It supports over 70 providers, including Google Drive, Amazon S3, Dropbox, and Microsoft OneDrive. Rclone allows users to perform operations like copying, moving, syncing, and deleting files and more advanced tasks such as encryption, compression, and chunking. Additionally, it can mount cloud storage as a local drive, facilitating seamless file access and management.

Known as the “Swiss army knife of cloud storage,” Rclone is highly versatile and can be integrated into scripts for automating complex file operations. Its ability to act as a bridge between different cloud providers makes it an invaluable tool for efficiently transferring files and managing cloud storage environments.

Step 1: Setting up Rclone

The base assumption is that you have a working instance of Teslamate running in a docker container on a Linux machine. Rclone is also available for Windows and Mac, and you can follow the [instructions] ( https://rclone.org/install/ ) to install it on those platforms. Let us start by outlining how to install rclone.

First, let’s install rclone on your system. You can download and install it using the following commands:

curl https://rclone.org/install.sh | sudo bash

Step 1.1 - Configuring Rclone with OneDrive ☁️

After installing rclone, we need to configure it to use OneDrive. Rclone has detailed steps on how to configure OneDrive. I show some of them here on how to run the following command to start the configuration process:

rclone config

You’ll be prompted with several questions about setting up a new remote. Here is an example configuration for OneDrive:

No remotes found - make a new one
n) New remote
r) Rename remote
c) Copy remote
s) Set configuration password
q) Quit config
n/r/c/s/q> n
name> onedrive

Type of storage to configure.
Choose a number from below, or type in your own value.
...
24 / Microsoft OneDrive
   \ "onedrive"
...
Storage> 24

Microsoft App Client Id
Leave blank normally.
client_id> 

Microsoft App Client Secret
Leave blank normally.
client_secret> 

Edit advanced config? (y/n)
y) Yes
n) No
y/n> n

Use auto config?
 * Say Y if not sure
 * Say N if you are working on a remote or headless machine
y) Yes
n) No
y/n> y

For the last prompt (auto config), if you are on a headless machine (e.g., you had SSH’d into a server to set this up), you can say n, and it will give you a URL to open in a browser to authenticate - as these require OAuth2. On the host (from which you are connected), you can authenticate and get the token, which you can paste into the terminal window. If the browser does not open, you can open your local browser and navigate to http://127.0.0.1:53682/auth.

If you are on SSH, you might need to create an SSH tunnel over port 53682 to your local machine using the following command:

ssh -L localhost:53682:localhost:53682 username@remote_server

For more details on how to do this, refer to the rclone remote setup documentation.

It is important to remember the name you give to the remote, as you will need it in the backup script later. In our example, we called it onedrive.

Also, note that the rclone configuration is stored in the .config directory in the user’s home directory. The configuration file is named rclone.conf and contains the details of the remotes you have configured. You can view the configuration file by running:

cat ~/.config/rclone/rclone.conf

To test the configuration, you can run the following command:

rclone ls onedrive:

This should list the files and directories in your OneDrive account. If you see the list of files, the configuration is successful. If you encounter any errors, review the configuration and ensure the authentication details are correct. For example, the image below shows the successful configuration of the OneDrive remote:

Remote OneDrive folder listing

Step 1.2 - Mounting Remote Directory with Rclone 📁

Now that we have configured rclone with OneDrive, we can mount the remote directory to a local directory in the system. This allows us to access the remote files as if they were local files. We cannot use a regular mount command, as the rclone mount command is not a real drive but is a FUSE-based mount. To mount the remote directory, use the following command.

rclone mount --config=/home/amit/.config/rclone/rclone.conf --vfs-cache-mode full --allow-non-empty --allow-other --daemon --dir-perms 0777 car_backup:car_data2 /home/amit/onedrive

Let’s break down the command:

In the above command, replace /home/amit/.config/rclone/rclone.conf with the path to your rclone configuration file.
Replace car_backup:car_data2 with the remote name (configured in Sep 3) and the directory you want to mount.
The local directory /home/amit/onedrive is where the remote directory will be mounted (you might need to create a new folder for this). You can choose any local directory for this purpose.
In the above command, the --daemon flag runs the mount in the background, and the --dir-perms 0777 flag sets the directory permissions.
The --allow-non-empty flag allows mounting over a non-empty directory.
And finally, the --allow-other flag allows other users to access the mounted directory.

Once the command is executed, you should see the remote directory mounted to the local directory. You can access the files in the remote directory as local files. For example, you can list the files in the mounted directory using the ls command, as shown in the image below:

Listing mounted OneDrive directory

Step 1.3 - Persistent Mounting with Rclone 💾

It is important to note that the rclone mount command is not persistent across reboots. So, if your VM reboots for some reason, you will lose this. To make the mount persistent, you can add the command to the system’s startup scripts or use a tool like systemd to manage the mount as a service. We will create a systemd service file to manage the mount for this.

To do this, create a new service file using the following command:

sudo nano /etc/systemd/system/rclonemount.service

Add the following content to the file:

[Unit]
Description=rclonemount
AssertPathIsDirectory=/home/amit/onedrive
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
User=amit
Group=amit
ExecStartPre=/bin/sleep 10
ExecStartPre=/bin/mkdir -p /home/amit/onedrive
ExecStartPre=/bin/chmod 777 /home/amit/onedrive
ExecStart=/usr/bin/rclone mount \
        --config=/home/amit/.config/rclone/rclone.conf \
        --vfs-cache-mode full \
        --allow-other \
        --dir-perms 0777 \
        car_backup:/car_data2 /home/amit/onedrive
ExecStop=/bin/fusermount -uz /home/amit/onedrive
Restart=always
RestartSec=10
Environment=HOME=/home/amit
Environment=USER=amit

[Install]
WantedBy=default.target

In the above service file, the ExecStartPre`` commands create the directory and set the permissions before mounting the remote directory. The ExecStartcommand mounts the remote directory to the local directory. TheExecStopcommand unmounts the directory when the service is stopped. TheRestartandRestartSec` settings ensure the service is restarted in case of failure.

The image below shows the contents of the rclonemount.service file in my system:

rclonemount.service contents

After creating the service file, reload the systemd daemon and start the service using the following commands:

sudo systemctl daemon-reload

To test the service, start it using the following command:

sudo systemctl start rclonemount

And check the status of the service to ensure it is running:

sudo systemctl status rclonemount

If everything is working correctly, you should see the status of the service as active (running), as shown in the image below:

rclonemount service status

Note: Press Q to quit.

And finally, if the service is running without any errors, enable it to start at boot:

sudo systemctl enable rclonemount

Of course, the service can be tested by rebooting the system and checking if the remote directory is mounted automatically.

Yay! Congratulations! You have successfully configured rclone to mount the remote directory as a service. The remote directory will now be mounted automatically on system startup. 👍

Step 2 - Basic Backup Script 🛠️

Now that we have rclone set up and mounted, we can create a basic backup script to back up the Teslamate data to OneDrive. The script will perform the following steps:

Back up the Teslamate database.
Compress the backed-up file.
Move the compressed file to the remote backup directory using rclone.

The will be a simple bash script that can be run manually or scheduled using a cron job. Here’s an example of the basic backup script. Create a new file named backup.sh and add the following content to it:

#!/bin/bash
# Basic backup script

SOURCEDIR=/home/amit/teslamate # Source directory where the data is located
BACKUPDIR=/home/amit/onedrive # Remote directory mounted via rclone
FILENAME=teslamate-$(date +%Y-%m-%d-%H%M%S).bak.gz # Date-time stamp filename

# Perform the database dump and compress it
docker compose exec -T database pg_dump -U teslamate teslamate | gzip > $FILENAME

# Move the compressed file to the remote backup directory
rclone move $FILENAME $BACKUPDIR

You will need to replace the SOURCEDIR and BACKUPDIR with the actual directories on your system. The FILENAME variable creates a unique filename for the backup file using the current date and time. The script uses pg_dump to back up the Teslamate database and rclone move to move the compressed file to the remote directory.

This script assumes the following:

You have Teslamate running in a docker container, and the database is named teslamate.
You have rclone configured with OneDrive and mounted the remote directory to a local directory.
You have the necessary permissions to access the source and destination directories.

Make the script executable:

`1`	`chmod +x /home/amit/backup.sh`

Run the script manually to test it:

`1`	`/home/amit/backup.sh`

You should see the backup file in the remote directory if the script runs successfully. You can check the remote directory using the ls command or by navigating to it.

Step 3 - Adding Logging 🗒️

To make the script more informative, let’s add logging. This will help us keep track of the backup process and identify any issues. We will use the same script as above as our starting point and add logging to it. We will log messages to a log file and display them on the console. Here’s the updated script with logging:

#!/bin/bash
# Backup script with logging

SOURCEDIR=/home/amit/teslamate # Source directory where the data is located
BACKUPDIR=/home/amit/onedrive # Remote directory mounted via rclone
FILENAME=teslamate-$(date +%Y-%m-%d-%H%M%S).bak.gz # Date-time stamp filename
LOGFILE=/home/amit/backup.log

# Function to log messages
log() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') - $1" | tee -a $LOGFILE
}

# Start of the script
log "Starting backup script."

# Perform the database dump and compress it
docker compose exec -T database pg_dump -U teslamate teslamate | gzip > $FILENAME
if [ $? -eq 0 ]; then
    log "Database dump and compression successful: $FILENAME"
else
    log "Database dump and compression failed"
    exit 1
fi

# Move the compressed file to the remote backup directory
rclone move $FILENAME $BACKUPDIR
if [ $? -eq 0 ]; then
    log "Successfully moved backup to $BACKUPDIR"
else
    log "Failed to move backup to $BACKUPDIR"
    exit 1
fi

log "Backup script completed successfully."

As in the earlier step, run the script manually to test it and check the log file for messages.

Step 4 - Adding Email Notifications 📧

Finally, we want to be notified about the backup status. We can add email notifications to the script. We’ll use ssmtp for this purpose. This simple mail transfer agent can send email notifications from the command line. We’ll update the script to send an email notification after the backup is completed.

To send email notifications, you need to have an SMTP server configured. You can use your email provider’s SMTP server for this purpose. Our first step is to set up ssmtp for email notifications

Step 4.1 - Installing and Configuring `ssmtp`

Install ssmtp on your system:

`1`	`sudo apt-get install ssmtp`

Edit the /etc/ssmtp/ssmtp.conf file to add your SMTP server details:

`1`	`sudo nano /etc/ssmtp/ssmtp.conf`

Add the following lines, replacing the placeholders with your actual details:

root=your-email@example.com
mailhub=smtp.your-email-provider.com:587
AuthUser=your-email@example.com
AuthPass=your-email-password
UseSTARTTLS=YES
TLS_CA_File=/dev/null
FromLineOverride=YES

Edit or create the /etc/ssmtp/revaliases file:

`1`	`sudo nano /etc/ssmtp/revaliases`

Add the following line, replacing with your actual details:

`1`	`your-username:your-email@example.com:smtp.your-email-provider.com:587`

To test the email configuration, we can create a simple script called test_email.sh that sends a test email:

#!/bin/bash
# Test email script

FROM_NAME="Teslamate Server"
FROM_EMAIL="the-from-email-address"
EMAIL="email-address-where-you-want-to-send-the-email"
SUBJECT="Test Email"
BODY="This is a test email to verify the SMTP configuration."

echo -e "To: $EMAIL\nFrom: \"$FROM_NAME\" <$FROM_EMAIL>\nSubject: $SUBJECT\n\n$BODY" | ssmtp $EMAIL

You should receive a test email at the specified address when you run the script. If you encounter any issues, review the configuration and ensure the SMTP server details are correct.

Step 4.2 - Timeout for Email Notifications

We can add a timeout to the email command to prevent the script from hanging indefinitely if the email-sending process takes too long. We’ll use the timeout command to set the duration of the email sending. First, we must check if the timeout command is available on your system. You can check this by running the following command:

timeout --version

If the timeout command is not available, you can install it using the following command:

sudo apt-get install coreutils

Step 4.3 - Updating the Backup Script with Email Notifications

Now that we have ssmtp set up for email notifications, we can update the backup script to send an email notification after the backup is completed. We’ll add a function to send email notifications and call it at the end of the script. Here’s the updated script that includes email notifications:

#!/bin/bash
# Backup script with logging and email notifications

SOURCEDIR=/home/amit/teslamate # Source directory where the data is located
BACKUPDIR=/home/amit/onedrive # Remote directory mounted via rclone
FILENAME=teslamate-$(date +%Y-%m-%d-%H%M%S).bak.gz # Date-time stamp filename
LOGFILE=/home/amit/backup.log
EMAIL="email-address-you-want-to-send-the-email-to" # Replace with your email address
FROM_NAME="Teslamate Server"
FROM_EMAIL="your-email@example.com" # Replace with the authorized sender email address
TIMEOUT_DURATION=10 # Set timeout duration (in seconds) for sending email
RCLONE_CONFIG="/home/amit/.config/rclone/rclone.conf" # Path to the rclone config file

# Initialize a variable to capture the entire log
FULL_LOG=""

# Function to log messages
log() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') - $1" | tee -a $LOGFILE
    FULL_LOG+="$1\n"
}

# Function to send email notification
send_email() {
    local STATUS="$1"
    local SUBJECT="Backup ${STATUS} - $(date '+%Y-%m-%d %H:%M:%S')"
    local BODY="The backup script ${STATUS} on $(date).\n\nLog details:\n$FULL_LOG"
    
    echo -e "To: $EMAIL\nFrom: \"$FROM_NAME\" <$FROM_EMAIL>\nSubject: $SUBJECT\n\n$BODY" | timeout $TIMEOUT_DURATION ssmtp $EMAIL
    EMAIL_EXIT_CODE=$?
    log "Email command exit code: $EMAIL_EXIT_CODE"

    if [ $EMAIL_EXIT_CODE -eq 124 ]; then
        log "Email sending timed out after $TIMEOUT_DURATION seconds"
    elif [ $EMAIL_EXIT_CODE -ne 0 ]; then
        log "Failed to send email"
    else
        log "Email sent successfully"
    fi
}

# Start of the script
log "Starting backup script."

# Change to the source directory
cd $SOURCEDIR
if [ $? -eq 0 ]; then
    log "Changed directory to $SOURCEDIR"
else
    log "Failed to change directory to $SOURCEDIR"
    send_email "Failed"
    exit 1
fi

# Perform the database dump and compress it
docker compose exec -T database pg_dump -U teslamate teslamate | gzip > $FILENAME
if [ $? -eq 0 ]; then
    log "Database dump and compression successful: $FILENAME"
else
    log "Database dump and compression failed"
    send_email "Failed"
    exit 1
fi

# Move the compressed file to the remote backup directory
rclone --config $RCLONE_CONFIG move $FILENAME $BACKUPDIR
if [ $? -eq 0 ]; then
    log "Successfully moved backup to $BACKUPDIR"
else
    log "Failed to move backup to $BACKUPDIR"
    send_email "Failed"
    exit 1
fi

log "Backup script completed successfully."

# Send email notification
send_email "Completed Successfully"

If you run the script, you should receive an email notification after completing the backup. The email will contain the log details of the backup process. You can customize the email content and format as needed. The image below shows an example email notification:

Email notification

Step 4.4 - Scheduling the Backup Script

You can schedule the script to run regularly using a cron job to automate the backup process. You can create a cron job to run the script daily, weekly, or at any interval you prefer. Here’s an example of a cron job to run the script daily:

0 0 * * * /home/amit/backup.sh

This cron job runs the backup.sh script every day at midnight. You can customize the schedule based on your requirements. To add the cron job, open the crontab file using the following command:

crontab -e

Add the cron job to the file and save it. The cron job will now run the backup script at the specified interval.

I have the script running at 5 AM, 1 PM, and 9 PM daily. In addition, I log the output of the cron job to a file (backup_cron.log) so I can review it later if there is an issue. Here are the steps to set this up:

0 5 * * * /home/amit/backup.sh >> /home/amit/backup_cron.log 2>&1
0 13 * * * /home/amit/backup.sh >> /home/amit/backup_cron.log 2>&1
0 21 * * * /home/amit/backup.sh >> /home/amit/backup_cron.log 2>&1

Be sure to replace /home/amit/backup.sh with the actual path to your backup script. The >> /home/amit/backup_cron.log 2>&1 part of the cron job redirects the script output to the specified log file. This allows you to capture the script output and any errors during execution.

Finally, you can check the cron job status using the following command:

crontab -l

This will list the cron jobs that are currently scheduled. You can also check the cron job logs to verify that the script is running as expected.

Here is a snippet of the backup_cron.log file from my system:

2024-07-05 13:00:01 - Starting backup script.
2024-07-05 13:00:01 - Changed directory to /home/amit/teslamate
time="2024-07-05T13:00:01-07:00" level=warning msg="/home/amit/teslamate/docker-compose.yml: `version` is obsolete"
2024-07-05 13:00:30 - Database dump successful: teslamate-2024-07-05-130001.bak
2024-07-05 13:00:55 - Successfully compressed teslamate-2024-07-05-130001.bak.gz
2024-07-05 13:00:57 - Successfully moved backup to /home/amit/onedrive
2024-07-05 13:00:57 - Backup script completed successfully.
2024-07-05 13:00:59 - Email command exit code: 0
2024-07-05 13:00:59 - Email sent successfully

Step 4.5 - Timezone Configuration

Note in most cloud-based VMs, the timezone is set to UTC; if you specify the time in your local timezone, you might want to set the timezone to your local timezone. You can do this by running the following command:

sudo timedatectl set-timezone your-timezone

Replace your-timezone with the appropriate timezone (e.g., America/Los_Angeles).

Conclusion

And that’s it! You now have a backup script that backs up your Teslamate data to OneDrive and sends email notifications about the backup status. The script runs automatically at the scheduled intervals, ensuring your data is backed up regularly. You can customize the script further based on your requirements and preferences. 🫰

SLMs - Running Phi-3 on an iphone and locally

Wed, 08 May 2024 00:00:00 +0000

We released Phi-3 recently, which builds on Phi-2 ( read more on that here ) and it is a great model to use for various tasks. In this post, we will show how to run Phi-3 locally including a demo of it running on a phone.

There should have been a video here but your browser does not seem to support it.

1. What are Small Language Models (SLMs)?

Before diving into running Phi-2 locally, let’s take a moment to understand the concept of small language models (SLMs) and their significance in natural language processing (NLP). A SLM is a type of AI model that has been trained on a massive dataset of text but is limited in terms of its size and capabilities compared to a Large Language Model (LLM). SLMs are designed to be more lightweight and efficient, making them suitable for various applications, including chatbots, language translation, and content generation. SLMs are much smaller than LLMs, with fewer parameters and a smaller dataset, so they have a lower computational cost, making them more suitable for edge or resource-constraint devices.

2. What is Phi-2?

Phi-2 is the latest model in the Phi series of small language models (SLMs) that aim to break the conventional scaling laws of language models. Unlike large language models (LLMs) that require massive amounts of data and compute resources, Phi models are trained on a mixture of web-crawled and synthetic “textbook-quality” data, following the idea of Textbooks Are All You Need . Phi models also leverage innovations in model architecture, optimization, and data augmentation to achieve remarkable performance on various benchmarks. 😀

Phi-2 is twice as large as its predecessor Phi-1.5, and was trained for two weeks on a cluster of 96 A100 GPUs. It demonstrates outstanding reasoning and language understanding capabilities, showcasing state-of-the-art performance (for LLMs <13 billion parameters).

As shown in the figure below on complex benchmarks, Phi-2 matches or outperforms models up to 25x larger, such as Llama-2 and Gemini Nano. Phi-2 is also an ideal playground for researchers, as it allows for exploration around mechanistic interpretability, safety improvements, or fine-tuning experimentation on various tasks. Phi-2 is available in the Azure AI Studio (via the model catalog ) and on Hugging Face.

Phi-2 Evaluation

2.1 Phi-2’s Capabilities

What better way to show this other than Some prompts that Phi-2 can handle:

Creative Writing:
- 🗒️Prompt: “Write a short story about a time-traveling computer geek who discovers a secret algorithm book with the power to change computing history.”
- 🤖 Phi2-Generated output: “In the dimly lit archives of the ancient library, Amelia stumbled upon a dusty tome unlike any other. Its pages whispered secrets of forgotten civilizations and whispered promises of altered destinies. As she traced her fingers over the delicate script, the room seemed to blur, and suddenly, she was standing in the heart of Renaissance Florence. The air smelled of ink and ambition. The book had worked its magic, and Amelia knew she held the fate of empires in her hands.”
Philosophical Reflection:
- 🗒️Prompt: “Discuss the concept of free will versus determinism.”
- 🤖 Phi2-Generated output: “Free will and determinism have long been at odds in philosophical discourse. Some argue that our choices are predetermined by external factors—genetics, upbringing, societal norms—while others champion the idea that we possess true agency. Perhaps the truth lies in the delicate dance between the two. We navigate life’s labyrinth, sometimes guided by fate, other times forging our own path.”
Educational Explanation:
- 🗒️Prompt: “Explain the process of photosynthesis in simple terms.”
- 🤖 Phi2-Generated output: “Photosynthesis, the magical dance of plants and sunlight, is how they create their own food. Imagine tiny chlorophyll factories in leaves, capturing sunlight like solar panels. They convert that energy into chemical fuel—glucose—while releasing oxygen as a byproduct. So next time you breathe in that fresh air, thank a leaf for its photosynthetic prowess!”
Humorous Scenario:
- 🗒️Prompt: “Describe a day in the life of a mischievous squirrel.”
- 🤖 Phi2-Generated output: “Meet Nutty, the squirrel extraordinaire. His mornings begin with acorn heists, followed by elaborate treehouse construction. He’s the neighborhood’s unofficial alarm clock, chattering loudly at dawn. But his pièce de résistance? Prank-calling the crows, convincing them they’ve won a lifetime supply of shiny objects. Nutty’s motto: ‘Life’s too short not to be a little nuts.’”

3. Small Language Models (SLMs) vs Large Language Models (LLMs)

Large Language Models (LLMs) are a type of AI model that is much larger and more powerful than SLMs. LLMs have hundreds of billions of parameters and are trained on massive text datasets. This gives LLMs the ability to handle complex tasks, such as language generation, translation, and question answering, with high accuracy and fluency. However, LLMs also have some disadvantages. They are larger, making them more expensive and slower to train. They also have a higher computational cost, meaning they may require access to specialized hardware.

On the other hand, SLMs, as we called out, are smaller and more lightweight than LLMs, making them more efficient and cost-effective in training computing resources and inference. While it might seem that SLMs are also more suitable for edge or resource-constrained devices, such as mobile phones or IoT devices, they are small compared to LLMs but still require significant computational resources to run. Phi-2, for example, still has 2.7B parameters, and while it can make inferences on a CPU, it is very slow and impractical for real-time applications. One would need a GPU or a cloud-based service for any realistic use case.

3.1 When to use SLM vs LLM?

Firstly, neither model is inherently better - the choice between an SLM and an LLM depends on the specific application and requirements. SLMs are a good choice when size, cost, and speed are important considerations. LLMs are a better choice when high performance and complex capabilities are required. If a task at hand is quite narrow and in one of the supported languages, then SLMs might be good. However, for a given task, an SLM may be sufficient, but an LLM may be necessary for more complex tasks or tasks requiring high accuracy and fluency.

Furthermore, it is key to understand that it is not necessarily about the number of languages understood but rather the depth and nuance with which each model can understand and generate language. SLMs are designed to be efficient and effective within their scope, which may include a wide range of languages. LLMs like GPT-4, due to their size and complexity, often can understand and generate text in a larger number of languages and with greater nuance.

The choice between an SLM and an LLM would again depend on the specific requirements of the task, including the languages involved and the level of language understanding and generation needed. Using a combination of SLMs and LLMs is common to achieve the best results for a given application.

4. Running Phi-2 locally

On one hand, running this is simple if you just don’t want to program anything and only want to use the model. The easiest option in this case is to use [LM Studio , a web-based platform for running language models. You can use the Hugging Face API to download and run the model.

We use a simple console chat example that runs locally. We use the Hugging Face Transformers library to generate text based on user input. The user can generate a story, a haiku, or a joke on a topic of their choice. Here is how to run it locally on a Windows machine - the same should apply to a Mac or Linux machine.

The full code is below, but here are the key aspects to grok when running Phi-2 locally.

The key is to use the AutoModelForCausalLM and AutoTokenizer classes from the transformers library to load the Phi-2 model and tokenizer.
We then use the generate method to generate text based on a user prompt. The generate method takes the user prompt as input and returns the generated text
We use the from_pretrained method to load the model and tokenizer from the Hugging Face model hub.
We also use the save_pretrained method to save the model and tokenizer to a local directory. This allows us to load the model and tokenizer from the local directory if they are already saved, which can help save time and resources.

The following code snippet is what loads the model and the tokenizer:

# Download the model and tokenizer from Hugging Face
model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2",
                                            torch_dtype="auto",
                                            trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2",
                                         trust_remote_code=True)

And the following is where we encode the user input and call the generation. First, we create tokens of the user prompt; the resulting tokens are returned as PyTorch tensors. Then, the model generates text based on the tokenized input. We cap the tokens to a maximum of 500 tokens, and the end-of-sequence token is used for padding if necessary. Finally, the generated tokens are decoded back into human-readable text.

inputs = tokenizer(prompt,
                   return_tensors="pt",
                   return_attention_mask=False,
                   add_special_tokens=False)
outputs = model.generate(**inputs,
                         max_length=500,
                         pad_token_id=tokenizer.eos_token_id)
text = tokenizer.batch_decode(outputs)[0]

The full code is below. The code is a simple console chat example that runs locally. The user can generate a story, a haiku, or a joke on a topic of their choice.

Some examples of what Phi-2 can generate using the above code are shown below. The first is a story about pandas and dogs.

Story about Pandas 🐼 and Dogs 🐶

Here is another example of a Haiku and a Joke generated by Phi-2 on Pandas.

Haiku and Joke about Pandas 🐼

Switching gears, let’s look at how we can implement the RAG using Phi-2.

4.1 Running Phi-2 Locally - Full Code

The following code is the complete code that executes the examples we showed before for running Phi-2 locally. This can work on a CPU, but it is very slow, and a good GPU is strongly suggested.

import os
import warnings
import logging
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

DEBUG = False

# Suppress warnings and set the logging level to ERROR
warnings.filterwarnings("ignore")
logging.getLogger("transformers").setLevel(logging.ERROR)

# Define the directory where you want to save the model and tokenizer
MODEL_PATH = "./local_model"

# Check if the model and tokenizer are already saved locally
if os.path.exists(MODEL_PATH):
    if DEBUG:
        print(f"Loading model and tokenizer from local directory: {MODEL_PATH}")

    # Load the model and tokenizer from the local directory
    model = AutoModelForCausalLM.from_pretrained(MODEL_PATH)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
else:
    if DEBUG:
        print(f"Downloading model and tokenizer from Hugging Face")

    # Download the model and tokenizer from Hugging Face
    model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)

    if DEBUG:
        print(f"Saving model and tokenizer to local directory: {MODEL_PATH}")

    # Save the model and tokenizer locally
    model.save_pretrained(MODEL_PATH)
    tokenizer.save_pretrained(MODEL_PATH)

if DEBUG:
    print(f"Model device: {model.device}")
    print(f"CUDA available: {torch.cuda.is_available()}")

# Set the default device to CUDA if available, otherwise use CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def handle_prompt(user_input, type_of_text)->str:
    prompt = f"Instruct: Write a {type_of_text} about {user_input}\nOutput:"

    with torch.no_grad():
        inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=False, add_special_tokens=False)
        inputs = {name: tensor.to(model.device) for name, tensor in inputs.items()}
        outputs = model.generate(**inputs, max_length=500, pad_token_id=tokenizer.eos_token_id)
        
    text = tokenizer.batch_decode(outputs)[0]

    # Remove the prompt from the output text
    text = text.replace(prompt, '').strip()
    text = text.replace("<|endoftext|>", '').strip()
    
    if DEBUG:
        print(f"Answer:{text}")

    return text

if __name__=="__main__": 
    while True:
        print("First What would you like to write today?")
        print("1. Story 📝")
        print("2. Haiku ✍️")
        print("3. Joke 😆")
        print("4. Quit 👋")
        user_choice = input("Choose an option:")

        if user_choice == '4':
            break

        user_prompt = input("And on which topic:")

        if not user_prompt.strip():
            print("Input cannot be empty or consist only of spaces.")
            continue

        if user_choice == '1':
            print(handle_prompt(user_prompt, 'story'))
        elif user_choice == '2':
            print(handle_prompt(user_prompt, 'haiku'))
        elif user_choice == '3':
            print(handle_prompt(user_prompt, 'joke'))
        else:
            print("Invalid choice. Please choose a valid option.")

        print("_" * 100)

5. Implementing Retrieval-Augmented Generation (RAG) with Phi-2

RAG is a powerful technique that combines the strengths of retrieval-based and generation-based approaches to natural language processing. RAG is one of the ways one can get proprietary information and knowledge to the model and use it as part of the prompt. It leverages a retriever to find relevant context passages and a generator to produce fluent and coherent responses. The retriever identifies relevant context passages, and the generator uses these passages to generate a response.

This approach allows RAG to produce high-quality, informative, and contextually relevant responses. In-context learning is a key feature of RAG, as it allows the model to learn from the context of the conversation and generate more accurate and relevant responses. This is particularly useful in scenarios where the model needs to understand and respond to complex queries or provide detailed information on a specific topic.

At a high level, the process of implementing RAG involves the following steps:

Generate Embeddings with Phi-2:
- Use Phi-2 to encode your context passages (documents) and extract their embeddings.
- These embeddings will represent the semantic content of each passage.
Create a Vector Index:
- Choose a vector index library or framework (such as Faiss, Annoy, or HNSW).
- Initialize an index structure to store the embeddings efficiently.
- Add the generated embeddings to the index.
Save Embeddings to a Local Vector Database:
- Create a local database to store the embeddings.
- For each context passage, save its corresponding embedding in the database.
- You can use the passage ID or a unique identifier as the key for retrieval.
Perform Similarity Search:
- When you receive a new context (query), encode it using Phi-2 to obtain its embedding.
- Use the vector index to perform a similarity search against the saved embeddings.
- Retrieve the most similar context passages based on cosine similarity or another distance metric.
- Return the relevant passages as results.

In our example, we will use the FAISS library to create a vector index and perform a similarity search. We will also save the embeddings to a local database for efficient retrieval. FAISS (Facebook AI Similarity Search) is a library developed by Facebook for efficient similarity search and clustering of high-dimensional vectors. It allows for a quick nearest-neighbor search over large datasets and supports CPU and GPU-based computations. FAISS is widely used in information retrieval, recommendation systems, and other applications that require similarity search.

5.1 Loading data for RAG and Phi-2

To implement RAG, we use the script from the Oppenheimer movie - which is quite new in that it is not in the Phi-2 training set and is available as a PDF. We will extract the script from this PDF, creating embeddings, which will then save the embeddings to a local database and perform a similarity search to retrieve relevant context passages based on a user query. We will use the FAISS library to create a vector index and perform a similarity search. We will also save the embeddings to a local database for efficient retrieval.

We use the PyPDF2 library to parse PDFs, a pure Python library for reading and writing PDF files. It can extract text, merge and split documents, and more. We will use it to extract the PDF text from the Oppenheimer movie script. The following code function shows how to read the PDF and extract the text. This is efficient for our use case, but it is not the most efficient way to extract text from a PDF when thinking about production scale, especially if the PDF has a lot of images and tables.

def read_pdf(file_path):
    pdf_file_obj = open(file_path, 'rb')
    pdf_reader = PyPDF2.PdfFileReader(pdf_file_obj)

    num_pages = pdf_reader.numPages
    for page_num in range(num_pages):
        page_obj = pdf_reader.getPage(page_num)
        text = page_obj.extractText()
        yield text
    pdf_file_obj.close()

5.2 Generate embeddings using Phi-2

Now that we have the text, the following functions show how to create the embeddings using Phi-2. We read the text as a list of context passages and then use Phi-2 to encode each passage and extract its embedding using the encode method.

def create_embeddings(file_path, tokenizer, model, device):
    if file_path.endswith('.pdf'):
        context_passages = list(read_pdf(file_path))
    else:
        with open(file_path, 'r', encoding='utf-8') as file:
            context_passages = file.readlines()

    embeddings = []
    for passage in tqdm(context_passages):
        if not passage.strip():
            # Skip the passage
            pass
        input_ids = tokenizer.encode(passage, 
                                    return_tensors="pt",
                                    add_special_tokens=False,
                                    return_attention_mask=False).to(device)
        with torch.no_grad():
            output = model(input_ids)
            logits = output.logits
            embedding = logits.mean(dim=1).detach().cpu().numpy()
            embeddings.append(embedding)

    return embeddings, context_passages

Here are a few things that are going on:

Given that we are using this for inference and not training, we use a torch.no_grad()`, which tells PyTorch not to track, calculate, or modify gradients while executing code within this block. This helps us save the amount of memory needed.
Inside this block, the input_ids are fed into the model, and the output is stored in the output variable. The logits, which are the raw, unnormalized scores outputted by the last layer of the model, are then extracted from the model’s output.
The logits are then processed to generate the embedding for the passage. The .mean(dim=1) method calculates the mean of the logits along dimension 1, which typically represents the sequence length in a language model.
The .detach() method detaches the result from the computation graph so that no gradients will be backpropagated along this variable.
The .cpu() method moves the tensor to the CPU if it’s not already there. Finally, the tensor is converted to a numpy array using the .numpy() method.
The resulting embedding is then appended to the embeddings list, which contains the embeddings for all the passages.

5.3 Creating Vector Index

The following function shows how to create a vector index using the FAISS library and perform a similarity search to retrieve relevant context passages based on a user query. The create_index function initializes a flat index structure to store the embeddings and adds the embeddings to the index. The search_query function encodes the user query using Phi-2 to obtain its embedding and performs a similarity search against the saved embeddings to retrieve the most similar context passages.

def create_index(query_embedding):
    index = faiss.IndexFlatL2(query_embedding[0].shape[1])  # Euclidean distance
    faiss.normalize_L2(query_embedding)
    
    # Add embeddings to the index
    for i, item in tqdm(enumerate(query_embedding), total=len(query_embedding)):
        if item.ndim == 1:
            item = item.reshape(1, -1)  # Reshape 1D array to 2D
        index.add(item)
    return index

The normalize_L2() function normalizes the vectors and is a crucial step when using Euclidean distance in high-dimensional spaces to ensure that the distance is not dominated by the dimensionality of the vectors. As we iterate through the embeddings, we check if the item is a 1D array and reshape it to a 2D array if necessary. This is important because FAISS expects the input to be a 2D array, and we need to reshape the 1D array to a 2D array before adding it to the index.

The function finally returns the created index. This index can then be used to perform efficient similarity searches.

5.4 Perform Similarity Search

The following function shows how to perform a similarity search using the vector index to retrieve relevant context passages based on a user query. As noted earlier, the most similar context passages are then retrieved based on cosine similarity or another distance metric.

The function starts by encoding the input query using a tokenizer and performs a similarity search on the FAISS index using the query embedding. It retrieves the indices of the top 3 most similar passages to the input query and then retrieves the corresponding context passages from the context_passages list. The similar context passages are then concatenated into a single string and passed to the Phi-2 model to generate a response.

def search_query(input_query, inputTokenizer, model, device, index, context_passages):
    # Given a new query context, encode it and perform similarity search
    input_ids = inputTokenizer.encode(input_query,
                                    return_tensors="pt",
                                    return_attention_mask=False,
                                    add_special_tokens=False).to(device)
    with torch.no_grad():
        input_ids = input_ids.long()
        output = model(input_ids)
        logits = output.logits
        query_embedding = logits.mean(dim=1).detach().cpu().numpy()

    # Perform similarity search - top 3 similar passages
    _, similar_indices = index.search(query_embedding, k=3)
    
    # Retrieve context passages based on similar_indices
    similar_contexts = [context_passages[i] for i in similar_indices[0]]

    # Concatenate the similar contexts into a single string
    context = ' '.join(similar_contexts)

Let us run this and see how it works, as we discussed before. We will use the Oppenheimer movie script as the context passages and perform a similarity search to retrieve relevant context passages based on a user query. The next few figures show the output of us asking questions about the movie, where those pieces of information are not in the model but passed using the semantic search.

Example 1 - Phi-2 and RAG implementation

Example 2 - Phi-2 and RAG implementation

Example 3 - Phi-2 and RAG implementation

Now that we have seen the different elements, the code below brings everything together as a console app that one can run. The Oppenheimer script (pdf file) you need can be downloaded from here .

import os
import warnings
import logging
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import numpy as np
import faiss 
from tqdm import tqdm
import pickle
import re
import PyPDF2

DEBUG = False

warnings.filterwarnings("ignore")
logging.getLogger("transformers").setLevel(logging.ERROR)

# Define the directory where you want to save the model and tokenizer
MODEL_PATH = "./local_model"
MODEL_NAME = "microsoft/phi-2"
BATCH_SIZE = 1000

# Oppenheimer movie
#DATA_FILE = "./oppenheimer-2023.txt"
DATA_FILE = "./oppenheimer-2023.pdf"
EMBEDDINGS_FILE = './embeddings_movie.pkl'

def load_model(model_path, model_name, debug=False):
    # Check if the model and tokenizer are already saved locally
    if os.path.exists(model_path):
        if debug:
            print(f"Loading model and tokenizer from local directory: {model_path}")

        # Load the model and tokenizer from the local directory
        model = AutoModelForCausalLM.from_pretrained(model_path)
        tokenizer = AutoTokenizer.from_pretrained(model_path)
    else:
        if debug:
            print("Downloading model and tokenizer from Hugging Face")

        # Download the model and tokenizer from Hugging Face
        model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", trust_remote_code=True)
        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

        if debug:
            print(f"Saving model and tokenizer to local directory: {model_path}")

        # Save the model and tokenizer locally
        model.save_pretrained(model_path)
        tokenizer.save_pretrained(model_path)

    if debug:
        print(f"Model device: {model.device}")
        print(f"CUDA available: {torch.cuda.is_available()}")

    # Set the default device to CUDA if available, otherwise use CPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    return model, tokenizer, device

def save_embeddings(embeddings, passages, file_name):
    try:
        with open(file_name, 'wb') as f:
            pickle.dump((list(embeddings), list(passages)), f)
    except IOError:
        print(f"Error writing to file {file_name}.")
    except pickle.PicklingError:
        print("Error pickling embeddings and passages.")
    
def load_embeddings(file):
    try:
        with open(file, 'rb') as f:
            return pickle.load(f)
    except FileNotFoundError:
        print(f"File {file} not found.")
        return None
    except pickle.UnpicklingError:
        print(f"Error unpickling file {file}.")
        return None

def read_pdf(file_path):
    pdf_file_obj = open(file_path, 'rb')
    pdf_reader = PyPDF2.PdfFileReader(pdf_file_obj)

    num_pages = pdf_reader.numPages
    for page_num in range(num_pages):
        page_obj = pdf_reader.getPage(page_num)
        text = page_obj.extractText()
        yield text

    if DEBUG:
        print(f"Finished reading file. Number pages: {num_pages}")
    
    pdf_file_obj.close()

def create_embeddings(file_path, tokenizer, model, device):
    if os.path.exists(EMBEDDINGS_FILE):
        # Load the embeddings and passages from disk
        embeddings, context_passages = load_embeddings(EMBEDDINGS_FILE)
    else:
        if file_path.endswith('.pdf'):
            context_passages = list(read_pdf(file_path))
        else:
            with open(file_path, 'r', encoding='utf-8') as file:
                context_passages = file.readlines()

        embeddings = []
        for passage in tqdm(context_passages):
            if not passage.strip():
                # Skip the passage
                pass
            input_ids = tokenizer.encode(passage, return_tensors="pt", add_special_tokens=False, return_attention_mask=False).to(device)
            with torch.no_grad():
                output = model(input_ids)
                logits = output.logits
                embedding = logits.mean(dim=1).detach().cpu().numpy()
                embeddings.append(embedding)

        # Save the embeddings and passages to disk
        save_embeddings(embeddings, context_passages, EMBEDDINGS_FILE)

    return embeddings, context_passages
  
def handle_prompt(user_input, context)->str:
    prompt = f"Instruct: You are a helpful bot who only answers using the given context ONLY. If you cannot find the answer in the context reply 'Sorry don't have that detail'. Given the context '{context}', answer this:{user_input}\nOutput:"

    with torch.no_grad():
        inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=False, add_special_tokens=False)
        inputs = {name: tensor.to(model.device) for name, tensor in inputs.items()}
        outputs = model.generate(**inputs, max_length=2000, pad_token_id=tokenizer.eos_token_id)
        
    text = tokenizer.batch_decode(outputs)[0]

    # Remove the prompt from the output text
    text = text.replace(prompt, '').strip()
    text = text.replace("<|endoftext|>", '').strip()
    return text

def create_index(query_embedding):
    index = faiss.IndexFlatL2(query_embedding[0].shape[1])  # Euclidean distance
    faiss.normalize_L2(query_embedding)
    
    # Add embeddings to the index
    for i, item in tqdm(enumerate(query_embedding), total=len(query_embedding)):
        if item.ndim == 1:
            item = item.reshape(1, -1)  # Reshape 1D array to 2D
        index.add(item)
    return index

def search_query(input_query, inputTokenizer, model, device, index, context_passages):
    # Given a new query context, encode it and perform similarity search
    input_ids = inputTokenizer.encode(input_query, return_tensors="pt", return_attention_mask=False, add_special_tokens=False).to(device)
    with torch.no_grad():
        input_ids = input_ids.long()
        output = model(input_ids)
        logits = output.logits
        query_embedding = logits.mean(dim=1).detach().cpu().numpy()

    # Perform similarity search - top 3 similar passages
    _, similar_indices = index.search(query_embedding, k=3)
    if DEBUG:
        print(f"DEBUG - Number of similar indices: {similar_indices.size}")    

    # Retrieve context passages based on similar_indices
    similar_contexts = [context_passages[i] for i in similar_indices[0]]

    # Concatenate the similar contexts into a single string
    context = ' '.join(similar_contexts)
    
    # Pass the concatenated context and query to the Phi-2 model
    answer = handle_prompt(input_query, context)
    print("Answer:", answer)

if __name__=="__main__":
    model, tokenizer, device = load_model(MODEL_PATH, MODEL_NAME, DEBUG)

     # Create embeddings and add to index before entering the loop
    query_embedding, context_passages = create_embeddings(DATA_FILE, tokenizer, model, device)
    query_embedding = np.array(query_embedding)
    index = create_index(query_embedding)
    
    exit_commands = ['exit', 'quit', 'q', 'e']
    while True:
        query = input("Enter your query or 'exit' to quit: ")
        if query in exit_commands:
            break
        search_query(query, tokenizer, model, device, index, context_passages)
        print("-"*50)

Now let us switch gears and try something that pushes the ability of Phi-2.

6. Code generation example using Phi-2

If we want to push the boundaries of what Phi-2 can do, we can use it to generate code. Below is an example of using Phi-2 to generate code for a simple C function. 🤓

🗒️Prompt: “Write a program in C that implements a BPE-based tokenizer; it should implement both encoding and decoding functions. Think through this step by step.”

The code we see below is what was generated. At face value, it looks like a good start but is incomplete. It is a good starting point for a developer to continue from and shows the power of SLMs like Phi-2.

#include<stdio.h> 
#include<string.h>
#define MAX_LEN 100

// Encodes the input text into an array of integers
int* encode(char *text, int *tokenizedStringsSize){

  // Allocate memory for the tokens list
  char *str = (char*)malloc(sizeof(char)*MAX_LEN+1);
  int i, j=0, k=0;

  for (i=0; i<strlen(text) && text[i] != '\0'; ++i){
    // Check if the current character is whitespace or not alphanumeric
    if(isspace(text[i]) || !isalpha(text[i])){
      // Add whitespace to the tokens list
      strcat(str,&text[i]); 
    } else {
      while (k<j && strcmp(text+i, str + k)==0){
        ++k; 
      }
      // Check if the token is of length atleast 1
      if(k-j >= 1)
        // If yes, concatenate it to tokens list
        strcat(str,&text[i]);
    }
  }

  // Output the tokenized string
  printf("Tokenized String: %s\n",str);

  // Allocate memory for tokenized string size array
  *tokenizedStringsSize = (int)malloc(sizeof(int)*(j-1));

  for (i=0; i<j; ++i){
    // Convert tokens to integer values
    (*tokenizedStringsSize)[i] = strtol(&str[i], NULL, 16);
  }

  return *tokenizedStringsSize;
}

// Decodes the input array of integers into a string
char* decode(int *tokensize, int *tokenizedStringsSize){

  // Declare memory for decoded string
  char str[MAX_LEN] = "";

  for (i=0; i<*tokensize; ++i) {
    strcat(str, "%x", tokensize[i]);
  }

  // Output the tokenized string
  printf("Tokenized String: %s\n",str);

  // Declare memory for decoded string
  char *decodedStr = (char*)malloc(sizeof(char)*(*tokenizedStringsSize)+1);
  for (i=0; i<*tokensize; ++i) {
    // Convert integer values to hexadecimal strings
    sprintf(&decodedStr[i], "%x", tokensize[i]);
  }

  return decodedStr;
}

⚠️ Note: It has been a while since I wrote C, but at a high level, these are some of the issues I can see with this; these issues are off the top of my head and are not exhaustive. Finally, it is not meant to test my coding capabilities. 😬

Memory Allocation: The malloc function is used without checking for successful allocation; if it returns NULL, which is not checked, we will get hurt.
Tokenization Logic: The logic in the encode function does not reflect the BPE algorithm, which involves merging the most frequent pairs of characters or bytes.
String Concatenation: The strcat function is used incorrectly; instead of a null-terminated string (as part of the second argument), we get a pointer to a single character
Decoding Logic: The decode function attempts to use strcat with a format string ("%x"), which is invalid. The sprintf function should be used for formatted strings.

Hopefully, this gives you a good understanding of SLMs, specifically Phi-2, and how to use them locally. 😍

SLMs - How to run Phi-2 Locally, and implement RAG

Wed, 13 Mar 2024 00:00:00 +0000

1. What are Small Language Models (SLMs)?

2. What is Phi-2?

Phi-2 Evaluation

2.1 Phi-2’s Capabilities

What better way to show this other than Some prompts that Phi-2 can handle:

Creative Writing:
- 🗒️Prompt: “Write a short story about a time-traveling computer geek who discovers a secret algorithm book with the power to change computing history.”
- 🤖 Phi2-Generated output: “In the dimly lit archives of the ancient library, Amelia stumbled upon a dusty tome unlike any other. Its pages whispered secrets of forgotten civilizations and whispered promises of altered destinies. As she traced her fingers over the delicate script, the room seemed to blur, and suddenly, she was standing in the heart of Renaissance Florence. The air smelled of ink and ambition. The book had worked its magic, and Amelia knew she held the fate of empires in her hands.”
Philosophical Reflection:
- 🗒️Prompt: “Discuss the concept of free will versus determinism.”
- 🤖 Phi2-Generated output: “Free will and determinism have long been at odds in philosophical discourse. Some argue that our choices are predetermined by external factors—genetics, upbringing, societal norms—while others champion the idea that we possess true agency. Perhaps the truth lies in the delicate dance between the two. We navigate life’s labyrinth, sometimes guided by fate, other times forging our own path.”
Educational Explanation:
- 🗒️Prompt: “Explain the process of photosynthesis in simple terms.”
- 🤖 Phi2-Generated output: “Photosynthesis, the magical dance of plants and sunlight, is how they create their own food. Imagine tiny chlorophyll factories in leaves, capturing sunlight like solar panels. They convert that energy into chemical fuel—glucose—while releasing oxygen as a byproduct. So next time you breathe in that fresh air, thank a leaf for its photosynthetic prowess!”
Humorous Scenario:
- 🗒️Prompt: “Describe a day in the life of a mischievous squirrel.”
- 🤖 Phi2-Generated output: “Meet Nutty, the squirrel extraordinaire. His mornings begin with acorn heists, followed by elaborate treehouse construction. He’s the neighborhood’s unofficial alarm clock, chattering loudly at dawn. But his pièce de résistance? Prank-calling the crows, convincing them they’ve won a lifetime supply of shiny objects. Nutty’s motto: ‘Life’s too short not to be a little nuts.’”

3. Small Language Models (SLMs) vs Large Language Models (LLMs)

3.1 When to use SLM vs LLM?

4. Running Phi-2 locally

The full code is below, but here are the key aspects to grok when running Phi-2 locally.

The key is to use the AutoModelForCausalLM and AutoTokenizer classes from the transformers library to load the Phi-2 model and tokenizer.
We then use the generate method to generate text based on a user prompt. The generate method takes the user prompt as input and returns the generated text
We use the from_pretrained method to load the model and tokenizer from the Hugging Face model hub.
We also use the save_pretrained method to save the model and tokenizer to a local directory. This allows us to load the model and tokenizer from the local directory if they are already saved, which can help save time and resources.

The following code snippet is what loads the model and the tokenizer:

# Download the model and tokenizer from Hugging Face
model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2",
                                            torch_dtype="auto",
                                            trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2",
                                         trust_remote_code=True)

inputs = tokenizer(prompt,
                   return_tensors="pt",
                   return_attention_mask=False,
                   add_special_tokens=False)
outputs = model.generate(**inputs,
                         max_length=500,
                         pad_token_id=tokenizer.eos_token_id)
text = tokenizer.batch_decode(outputs)[0]

The full code is below. The code is a simple console chat example that runs locally. The user can generate a story, a haiku, or a joke on a topic of their choice.

Some examples of what Phi-2 can generate using the above code are shown below. The first is a story about pandas and dogs.

Story about Pandas 🐼 and Dogs 🐶

Here is another example of a Haiku and a Joke generated by Phi-2 on Pandas.

Haiku and Joke about Pandas 🐼

Switching gears, let’s look at how we can implement the RAG using Phi-2.

4.1 Running Phi-2 Locally - Full Code

The following code is the complete code that executes the examples we showed before for running Phi-2 locally. This can work on a CPU, but it is very slow, and a good GPU is strongly suggested.

import os
import warnings
import logging
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

DEBUG = False

# Suppress warnings and set the logging level to ERROR
warnings.filterwarnings("ignore")
logging.getLogger("transformers").setLevel(logging.ERROR)

# Define the directory where you want to save the model and tokenizer
MODEL_PATH = "./local_model"

# Check if the model and tokenizer are already saved locally
if os.path.exists(MODEL_PATH):
    if DEBUG:
        print(f"Loading model and tokenizer from local directory: {MODEL_PATH}")

    # Load the model and tokenizer from the local directory
    model = AutoModelForCausalLM.from_pretrained(MODEL_PATH)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
else:
    if DEBUG:
        print(f"Downloading model and tokenizer from Hugging Face")

    # Download the model and tokenizer from Hugging Face
    model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype="auto", trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)

    if DEBUG:
        print(f"Saving model and tokenizer to local directory: {MODEL_PATH}")

    # Save the model and tokenizer locally
    model.save_pretrained(MODEL_PATH)
    tokenizer.save_pretrained(MODEL_PATH)

if DEBUG:
    print(f"Model device: {model.device}")
    print(f"CUDA available: {torch.cuda.is_available()}")

# Set the default device to CUDA if available, otherwise use CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def handle_prompt(user_input, type_of_text)->str:
    prompt = f"Instruct: Write a {type_of_text} about {user_input}\nOutput:"

    with torch.no_grad():
        inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=False, add_special_tokens=False)
        inputs = {name: tensor.to(model.device) for name, tensor in inputs.items()}
        outputs = model.generate(**inputs, max_length=500, pad_token_id=tokenizer.eos_token_id)
        
    text = tokenizer.batch_decode(outputs)[0]

    # Remove the prompt from the output text
    text = text.replace(prompt, '').strip()
    text = text.replace("<|endoftext|>", '').strip()
    
    if DEBUG:
        print(f"Answer:{text}")

    return text

if __name__=="__main__": 
    while True:
        print("First What would you like to write today?")
        print("1. Story 📝")
        print("2. Haiku ✍️")
        print("3. Joke 😆")
        print("4. Quit 👋")
        user_choice = input("Choose an option:")

        if user_choice == '4':
            break

        user_prompt = input("And on which topic:")

        if not user_prompt.strip():
            print("Input cannot be empty or consist only of spaces.")
            continue

        if user_choice == '1':
            print(handle_prompt(user_prompt, 'story'))
        elif user_choice == '2':
            print(handle_prompt(user_prompt, 'haiku'))
        elif user_choice == '3':
            print(handle_prompt(user_prompt, 'joke'))
        else:
            print("Invalid choice. Please choose a valid option.")

        print("_" * 100)

5. Implementing Retrieval-Augmented Generation (RAG) with Phi-2

At a high level, the process of implementing RAG involves the following steps:

Generate Embeddings with Phi-2:
- Use Phi-2 to encode your context passages (documents) and extract their embeddings.
- These embeddings will represent the semantic content of each passage.
Create a Vector Index:
- Choose a vector index library or framework (such as Faiss, Annoy, or HNSW).
- Initialize an index structure to store the embeddings efficiently.
- Add the generated embeddings to the index.
Save Embeddings to a Local Vector Database:
- Create a local database to store the embeddings.
- For each context passage, save its corresponding embedding in the database.
- You can use the passage ID or a unique identifier as the key for retrieval.
Perform Similarity Search:
- When you receive a new context (query), encode it using Phi-2 to obtain its embedding.
- Use the vector index to perform a similarity search against the saved embeddings.
- Retrieve the most similar context passages based on cosine similarity or another distance metric.
- Return the relevant passages as results.

5.1 Loading data for RAG and Phi-2

def read_pdf(file_path):
    pdf_file_obj = open(file_path, 'rb')
    pdf_reader = PyPDF2.PdfFileReader(pdf_file_obj)

    num_pages = pdf_reader.numPages
    for page_num in range(num_pages):
        page_obj = pdf_reader.getPage(page_num)
        text = page_obj.extractText()
        yield text
    pdf_file_obj.close()

5.2 Generate embeddings using Phi-2

def create_embeddings(file_path, tokenizer, model, device):
    if file_path.endswith('.pdf'):
        context_passages = list(read_pdf(file_path))
    else:
        with open(file_path, 'r', encoding='utf-8') as file:
            context_passages = file.readlines()

    embeddings = []
    for passage in tqdm(context_passages):
        if not passage.strip():
            # Skip the passage
            pass
        input_ids = tokenizer.encode(passage, 
                                    return_tensors="pt",
                                    add_special_tokens=False,
                                    return_attention_mask=False).to(device)
        with torch.no_grad():
            output = model(input_ids)
            logits = output.logits
            embedding = logits.mean(dim=1).detach().cpu().numpy()
            embeddings.append(embedding)

    return embeddings, context_passages

Here are a few things that are going on:

Given that we are using this for inference and not training, we use a torch.no_grad()`, which tells PyTorch not to track, calculate, or modify gradients while executing code within this block. This helps us save the amount of memory needed.
Inside this block, the input_ids are fed into the model, and the output is stored in the output variable. The logits, which are the raw, unnormalized scores outputted by the last layer of the model, are then extracted from the model’s output.
The logits are then processed to generate the embedding for the passage. The .mean(dim=1) method calculates the mean of the logits along dimension 1, which typically represents the sequence length in a language model.
The .detach() method detaches the result from the computation graph so that no gradients will be backpropagated along this variable.
The .cpu() method moves the tensor to the CPU if it’s not already there. Finally, the tensor is converted to a numpy array using the .numpy() method.
The resulting embedding is then appended to the embeddings list, which contains the embeddings for all the passages.

5.3 Creating Vector Index

def create_index(query_embedding):
    index = faiss.IndexFlatL2(query_embedding[0].shape[1])  # Euclidean distance
    faiss.normalize_L2(query_embedding)
    
    # Add embeddings to the index
    for i, item in tqdm(enumerate(query_embedding), total=len(query_embedding)):
        if item.ndim == 1:
            item = item.reshape(1, -1)  # Reshape 1D array to 2D
        index.add(item)
    return index

The function finally returns the created index. This index can then be used to perform efficient similarity searches.

5.4 Perform Similarity Search

def search_query(input_query, inputTokenizer, model, device, index, context_passages):
    # Given a new query context, encode it and perform similarity search
    input_ids = inputTokenizer.encode(input_query,
                                    return_tensors="pt",
                                    return_attention_mask=False,
                                    add_special_tokens=False).to(device)
    with torch.no_grad():
        input_ids = input_ids.long()
        output = model(input_ids)
        logits = output.logits
        query_embedding = logits.mean(dim=1).detach().cpu().numpy()

    # Perform similarity search - top 3 similar passages
    _, similar_indices = index.search(query_embedding, k=3)
    
    # Retrieve context passages based on similar_indices
    similar_contexts = [context_passages[i] for i in similar_indices[0]]

    # Concatenate the similar contexts into a single string
    context = ' '.join(similar_contexts)

Example 1 - Phi-2 and RAG implementation

Example 2 - Phi-2 and RAG implementation

Example 3 - Phi-2 and RAG implementation

Now that we have seen the different elements, the code below brings everything together as a console app that one can run. The Oppenheimer script (pdf file) you need can be downloaded from here .

import os
import warnings
import logging
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import numpy as np
import faiss 
from tqdm import tqdm
import pickle
import re
import PyPDF2

DEBUG = False

warnings.filterwarnings("ignore")
logging.getLogger("transformers").setLevel(logging.ERROR)

# Define the directory where you want to save the model and tokenizer
MODEL_PATH = "./local_model"
MODEL_NAME = "microsoft/phi-2"
BATCH_SIZE = 1000

# Oppenheimer movie
#DATA_FILE = "./oppenheimer-2023.txt"
DATA_FILE = "./oppenheimer-2023.pdf"
EMBEDDINGS_FILE = './embeddings_movie.pkl'

def load_model(model_path, model_name, debug=False):
    # Check if the model and tokenizer are already saved locally
    if os.path.exists(model_path):
        if debug:
            print(f"Loading model and tokenizer from local directory: {model_path}")

        # Load the model and tokenizer from the local directory
        model = AutoModelForCausalLM.from_pretrained(model_path)
        tokenizer = AutoTokenizer.from_pretrained(model_path)
    else:
        if debug:
            print("Downloading model and tokenizer from Hugging Face")

        # Download the model and tokenizer from Hugging Face
        model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", trust_remote_code=True)
        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

        if debug:
            print(f"Saving model and tokenizer to local directory: {model_path}")

        # Save the model and tokenizer locally
        model.save_pretrained(model_path)
        tokenizer.save_pretrained(model_path)

    if debug:
        print(f"Model device: {model.device}")
        print(f"CUDA available: {torch.cuda.is_available()}")

    # Set the default device to CUDA if available, otherwise use CPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    return model, tokenizer, device

def save_embeddings(embeddings, passages, file_name):
    try:
        with open(file_name, 'wb') as f:
            pickle.dump((list(embeddings), list(passages)), f)
    except IOError:
        print(f"Error writing to file {file_name}.")
    except pickle.PicklingError:
        print("Error pickling embeddings and passages.")
    
def load_embeddings(file):
    try:
        with open(file, 'rb') as f:
            return pickle.load(f)
    except FileNotFoundError:
        print(f"File {file} not found.")
        return None
    except pickle.UnpicklingError:
        print(f"Error unpickling file {file}.")
        return None

def read_pdf(file_path):
    pdf_file_obj = open(file_path, 'rb')
    pdf_reader = PyPDF2.PdfFileReader(pdf_file_obj)

    num_pages = pdf_reader.numPages
    for page_num in range(num_pages):
        page_obj = pdf_reader.getPage(page_num)
        text = page_obj.extractText()
        yield text

    if DEBUG:
        print(f"Finished reading file. Number pages: {num_pages}")
    
    pdf_file_obj.close()

def create_embeddings(file_path, tokenizer, model, device):
    if os.path.exists(EMBEDDINGS_FILE):
        # Load the embeddings and passages from disk
        embeddings, context_passages = load_embeddings(EMBEDDINGS_FILE)
    else:
        if file_path.endswith('.pdf'):
            context_passages = list(read_pdf(file_path))
        else:
            with open(file_path, 'r', encoding='utf-8') as file:
                context_passages = file.readlines()

        embeddings = []
        for passage in tqdm(context_passages):
            if not passage.strip():
                # Skip the passage
                pass
            input_ids = tokenizer.encode(passage, return_tensors="pt", add_special_tokens=False, return_attention_mask=False).to(device)
            with torch.no_grad():
                output = model(input_ids)
                logits = output.logits
                embedding = logits.mean(dim=1).detach().cpu().numpy()
                embeddings.append(embedding)

        # Save the embeddings and passages to disk
        save_embeddings(embeddings, context_passages, EMBEDDINGS_FILE)

    return embeddings, context_passages
  
def handle_prompt(user_input, context)->str:
    prompt = f"Instruct: You are a helpful bot who only answers using the given context ONLY. If you cannot find the answer in the context reply 'Sorry don't have that detail'. Given the context '{context}', answer this:{user_input}\nOutput:"

    with torch.no_grad():
        inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=False, add_special_tokens=False)
        inputs = {name: tensor.to(model.device) for name, tensor in inputs.items()}
        outputs = model.generate(**inputs, max_length=2000, pad_token_id=tokenizer.eos_token_id)
        
    text = tokenizer.batch_decode(outputs)[0]

    # Remove the prompt from the output text
    text = text.replace(prompt, '').strip()
    text = text.replace("<|endoftext|>", '').strip()
    return text

def create_index(query_embedding):
    index = faiss.IndexFlatL2(query_embedding[0].shape[1])  # Euclidean distance
    faiss.normalize_L2(query_embedding)
    
    # Add embeddings to the index
    for i, item in tqdm(enumerate(query_embedding), total=len(query_embedding)):
        if item.ndim == 1:
            item = item.reshape(1, -1)  # Reshape 1D array to 2D
        index.add(item)
    return index

def search_query(input_query, inputTokenizer, model, device, index, context_passages):
    # Given a new query context, encode it and perform similarity search
    input_ids = inputTokenizer.encode(input_query, return_tensors="pt", return_attention_mask=False, add_special_tokens=False).to(device)
    with torch.no_grad():
        input_ids = input_ids.long()
        output = model(input_ids)
        logits = output.logits
        query_embedding = logits.mean(dim=1).detach().cpu().numpy()

    # Perform similarity search - top 3 similar passages
    _, similar_indices = index.search(query_embedding, k=3)
    if DEBUG:
        print(f"DEBUG - Number of similar indices: {similar_indices.size}")    

    # Retrieve context passages based on similar_indices
    similar_contexts = [context_passages[i] for i in similar_indices[0]]

    # Concatenate the similar contexts into a single string
    context = ' '.join(similar_contexts)
    
    # Pass the concatenated context and query to the Phi-2 model
    answer = handle_prompt(input_query, context)
    print("Answer:", answer)

if __name__=="__main__":
    model, tokenizer, device = load_model(MODEL_PATH, MODEL_NAME, DEBUG)

     # Create embeddings and add to index before entering the loop
    query_embedding, context_passages = create_embeddings(DATA_FILE, tokenizer, model, device)
    query_embedding = np.array(query_embedding)
    index = create_index(query_embedding)
    
    exit_commands = ['exit', 'quit', 'q', 'e']
    while True:
        query = input("Enter your query or 'exit' to quit: ")
        if query in exit_commands:
            break
        search_query(query, tokenizer, model, device, index, context_passages)
        print("-"*50)

Now let us switch gears and try something that pushes the ability of Phi-2.

6. Code generation example using Phi-2

If we want to push the boundaries of what Phi-2 can do, we can use it to generate code. Below is an example of using Phi-2 to generate code for a simple C function. 🤓

🗒️Prompt: “Write a program in C that implements a BPE-based tokenizer; it should implement both encoding and decoding functions. Think through this step by step.”

#include<stdio.h> 
#include<string.h>
#define MAX_LEN 100

// Encodes the input text into an array of integers
int* encode(char *text, int *tokenizedStringsSize){

  // Allocate memory for the tokens list
  char *str = (char*)malloc(sizeof(char)*MAX_LEN+1);
  int i, j=0, k=0;

  for (i=0; i<strlen(text) && text[i] != '\0'; ++i){
    // Check if the current character is whitespace or not alphanumeric
    if(isspace(text[i]) || !isalpha(text[i])){
      // Add whitespace to the tokens list
      strcat(str,&text[i]); 
    } else {
      while (k<j && strcmp(text+i, str + k)==0){
        ++k; 
      }
      // Check if the token is of length atleast 1
      if(k-j >= 1)
        // If yes, concatenate it to tokens list
        strcat(str,&text[i]);
    }
  }

  // Output the tokenized string
  printf("Tokenized String: %s\n",str);

  // Allocate memory for tokenized string size array
  *tokenizedStringsSize = (int)malloc(sizeof(int)*(j-1));

  for (i=0; i<j; ++i){
    // Convert tokens to integer values
    (*tokenizedStringsSize)[i] = strtol(&str[i], NULL, 16);
  }

  return *tokenizedStringsSize;
}

// Decodes the input array of integers into a string
char* decode(int *tokensize, int *tokenizedStringsSize){

  // Declare memory for decoded string
  char str[MAX_LEN] = "";

  for (i=0; i<*tokensize; ++i) {
    strcat(str, "%x", tokensize[i]);
  }

  // Output the tokenized string
  printf("Tokenized String: %s\n",str);

  // Declare memory for decoded string
  char *decodedStr = (char*)malloc(sizeof(char)*(*tokenizedStringsSize)+1);
  for (i=0; i<*tokensize; ++i) {
    // Convert integer values to hexadecimal strings
    sprintf(&decodedStr[i], "%x", tokensize[i]);
  }

  return decodedStr;
}

Memory Allocation: The malloc function is used without checking for successful allocation; if it returns NULL, which is not checked, we will get hurt.
Tokenization Logic: The logic in the encode function does not reflect the BPE algorithm, which involves merging the most frequent pairs of characters or bytes.
String Concatenation: The strcat function is used incorrectly; instead of a null-terminated string (as part of the second argument), we get a pointer to a single character
Decoding Logic: The decode function attempts to use strcat with a format string ("%x"), which is invalid. The sprintf function should be used for formatted strings.

Hopefully, this gives you a good understanding of SLMs, specifically Phi-2, and how to use them locally. 😍