Constructing Transformer Fashions from Scratch with PyTorch (10-day Mini-Course)

You’ve doubtless used ChatGPT, Gemini, or Grok, which exhibit how massive language fashions can exhibit human-like intelligence. Whereas making a clone of those massive language fashions at house is unrealistic and pointless, understanding how they work helps demystify their capabilities and acknowledge their limitations.

All these fashionable massive language fashions are decoder-only transformers. Surprisingly, their structure will not be overly advanced. When you might not have in depth computational energy and reminiscence, you’ll be able to nonetheless create a smaller language mannequin that mimics some capabilities of the bigger ones. By designing, constructing, and coaching such a scaled-down model, you’ll higher perceive what the mannequin is doing, moderately than merely viewing it as a black field labeled “AI.”

On this 10-part crash course, you’ll study by examples how you can construct and practice a transformer mannequin from scratch utilizing PyTorch. The mini-course focuses on mannequin structure, whereas superior optimization strategies, although necessary, are past our scope. We’ll information you from information assortment by to working your skilled mannequin. Every lesson covers a selected transformer element, explaining its function, design parameters, and PyTorch implementation. By the tip, you’ll have explored each facet of the mannequin and gained a complete understanding of how transformer fashions work.

Let’s get began.

 

Who Is This Mini-Course For?

Earlier than we start, let’s be sure you’re in the proper place. The record beneath gives common tips on whom this course is designed for. Don’t fear should you don’t match these factors precisely—you may simply have to brush up on sure areas to maintain up.

• Builders with some coding expertise. You need to be comfy writing Python code and organising your growth setting (a prerequisite). You don’t must be an knowledgeable coder, however you must be capable of set up packages and write scripts with out hesitation.
• Builders with fundamental machine studying data. You need to have a common understanding of machine studying fashions and really feel comfy utilizing them. You don’t must be an knowledgeable, however you shouldn’t be afraid to study extra about them.
• Builders conversant in PyTorch. This venture relies on PyTorch. To maintain it concise, we won’t cowl the fundamentals of PyTorch. You aren’t required to be a PyTorch knowledgeable, however you’re anticipated to have the ability to learn and perceive PyTorch code, and extra importantly, know how you can learn the documentation of PyTorch in case you encountered any features that you’re not conversant in.

This mini-course will not be a textbook on transformer or LLM. As a substitute, it serves as a project-based information that takes you step-by-step from a developer with minimal expertise to 1 who can confidently exhibit how a transformer mannequin is created.

Mini-Course Overview

This mini-course is split into 10 components.

Every lesson is designed to take about half-hour for the common developer. Whereas some classes could also be accomplished extra rapidly, others may require extra time should you select to discover them in depth.
You may progress at your personal tempo. We advocate following a cushty schedule of 1 lesson per day over ten days to permit for correct absorption of the fabric.

The subjects you’ll cowl over the following 10 classes are as follows:

• Lesson 1: Getting the Information
• Lesson 2: Practice a Tokenizer for Your Language Mannequin
• Lesson 3: Positional Encoding
• Lesson 4: Grouped Question Consideration
• Lesson 5: Causal Masks
• Lesson 6: Combination of Skilled Fashions
• Lesson 7: RMS Norm and Skip Connection
• Lesson 8: The Full Transformer Mannequin
• Lesson 9: Coaching the Mannequin
• Lesson 10: Utilizing the Mannequin

This journey shall be each difficult and rewarding.
Whereas it requires dedication by studying, analysis, and programming, the hands-on expertise you’ll achieve in constructing a transformer mannequin shall be invaluable.

Publish your leads to the feedback; I’ll cheer you on!

Cling in there; don’t surrender.

You may obtain the code of this submit here.

Lesson 01: Getting the Information

We’re constructing a language mannequin utilizing transformer structure. A language mannequin is a probabilistic illustration of human language that predicts the probability of phrases showing in a sequence. Fairly than being manually constructed, these chances are discovered from information. Due to this fact, step one in constructing a language mannequin is to gather a big corpus of textual content that captures the pure patterns of language use.

There are quite a few sources of textual content information accessible. Challenge Gutenberg is a superb supply of free textual content information, providing all kinds of books throughout totally different genres. Right here’s how one can obtain textual content information from Challenge Gutenberg to your native listing:

This code downloads every ebook as a separate textual content file. Since Challenge Gutenberg gives pre-cleaned textual content, we solely have to extract the ebook contents and retailer them as a listing of strings in Python:

The preprocess_gutenberg() perform removes the Challenge Gutenberg header and footer from every ebook and joins the traces right into a single string. The get_dataset_text() perform applies this preprocessing to all books and returns a listing of strings, the place every string represents an entire ebook.

Your Process

Attempt working the code above! Whereas this small assortment of books would sometimes be inadequate for coaching a production-ready language mannequin, it serves as a superb start line for studying. Discover that the books within the DATASOURCE dictionary span numerous genres. Can you concentrate on why having various genres is necessary when constructing a language mannequin?

Within the subsequent lesson, you’ll discover ways to convert the textual information into numbers.

Lesson 02: Practice a Tokenizer for Your Language Mannequin

Computer systems function on numbers, so textual content have to be transformed into numerical type for processing. In a language mannequin, we assign numbers to “tokens,” and these 1000’s of distinct tokens type the mannequin’s vocabulary.

A easy method could be to open a dictionary and assign a quantity to every phrase. Nevertheless, this naive technique can not deal with unseen phrases successfully. A greater method is to coach an algorithm that processes enter textual content and breaks it down into tokens. This algorithm, referred to as a tokenizer, splits textual content effectively and may deal with unseen phrases.

There are a number of approaches to coaching a tokenizer. Byte-pair encoding (BPE) is without doubt one of the hottest strategies utilized in fashionable LLMs. Let’s use the tokenizer library to coach a BPE tokenizer utilizing the textual content we collected within the earlier lesson:

This instance creates a small BPE tokenizer with a vocabulary dimension of 10,000. Manufacturing LLMs sometimes use vocabularies which are orders of magnitude bigger for higher language protection. Even for this toy venture, coaching a tokenizer takes time because it analyzes character collocations to type phrases. It’s really useful to avoid wasting the tokenizer as a JSON file, as proven above, so you’ll be able to simply reload it later:

Your Process

Apart from BPE, WordPiece is one other widespread tokenization algorithm. Attempt making a WordPiece model of the tokenizer above.

Why is a vocabulary dimension of 10,000 inadequate for a great language mannequin? Analysis the variety of phrases in a typical English dictionary and clarify the implications for language modeling.

Within the subsequent lesson, you’ll study positional encoding.

Lesson 03: Positional Encoding

In contrast to recurrent neural networks, transformer fashions course of total sequences concurrently. Nevertheless, this parallel processing means they lack inherent understanding of token order. Since token place is essential for understanding context, transformer fashions incorporate positional encodings into their enter processing to seize this sequential data.

Whereas a number of positional encoding strategies exist, Rotary Positional Encoding (RoPE) has emerged as probably the most broadly used method. RoPE operates by making use of rotational transformations to the embedded token vectors. Every token is represented as a vector, and the encoding course of includes multiplying pairs of vector parts by a $2times 2$ rotation matrix:

$$
mathbf{hat{x}}_m = mathbf{R}_mmathbf{x}_m = start{bmatrix}
cos(mtheta_i) & -sin(mtheta_i)
sin(mtheta_i) & cos(mtheta_i)
finish{bmatrix} mathbf{x}_m
$$

To implement RoPE, you should use the next PyTorch code:

The RotaryPositionalEncoding module implements the positional encoding mechanism for enter sequences. Its __init__ perform pre-computes sine and cosine values for all potential positions and dimensions, whereas the ahead perform applies the rotation matrix to remodel the enter.

An necessary implementation element is the usage of register_buffer within the __init__ perform to retailer sine and cosine values. This tells PyTorch to deal with these tensors as non-trainable mannequin parameters, making certain correct administration throughout totally different computing gadgets (e.g., GPU) and through mannequin serialization.

Your Process

Experiment with the code supplied above. Earlier, we discovered that RoPE applies to embedded token vectors in a sequence. Take a better have a look at the enter tensor sequence used to check the RotaryPositionalEncoding module: why is it a 4D tensor? Whereas the final dimension (128) represents the embedding dimension, are you able to establish what the primary three dimensions (1, 10, 4) signify within the context of transformer structure?

Within the subsequent lesson, you’ll study in regards to the consideration block.

Lesson 04: Grouped Question Consideration

The signature element of a transformer mannequin is its consideration mechanism. When processing a sequence of tokens, the eye mechanism builds connections between tokens to know their context.

The eye mechanism predates transformer fashions, and several other variants have advanced over time. On this lesson, you’ll study to implement Grouped Question Consideration (GQA).

A transformer mannequin begins with a sequence of embedded tokens, that are basically vectors. The fashionable consideration mechanism computes an output sequence primarily based on three enter sequences: question, key, and worth. These three sequences are derived from the enter sequence by totally different projections:

The projection is carried out by a fully-connected neural community layer that operates on the enter tensor’s final dimension. As proven above, the projection’s output is reshaped utilizing view() after which transposed. The enter tensor x is 3D, and the view() perform transforms it right into a 4D tensor by splitting the final dimension into two: the eye heads and the pinnacle dimension. The transpose() perform then swaps the sequence size dimension with the eye head dimension.

The ensuing 4D tensor has consideration operations that solely contain the final two dimensions. The precise consideration computation is carried out utilizing PyTorch’s built-in scaled_dot_product_attention() perform. The result’s then reshaped again right into a 3D tensor and projected to the unique dimension.

This structure known as grouped question consideration as a result of it makes use of totally different numbers of heads for queries versus keys and values. Sometimes, the variety of question heads is a a number of of the variety of key-value heads.

Since we’ll use such consideration mechanism so much, let’s create a category for it:

The ahead perform contains two non-compulsory arguments: masks and rope. The rope argument expects a module that applies rotary positional encoding, which was coated within the earlier lesson. The masks argument shall be defined within the subsequent lesson.

Your Process

Think about why this implementation known as grouped question consideration. The unique transformer structure makes use of multihead consideration. How would you modify this grouped question consideration implementation to create a multihead consideration mechanism?

Within the subsequent lesson, you’ll study masking in consideration operations.

Lesson 05: Causal Masks

A key attribute of decoder-only transformer fashions is the usage of causal masks of their consideration layers. A causal masks is a matrix utilized throughout consideration rating calculation to forestall the mannequin from attending to future tokens. Particularly, a question token $i$ can solely attend to key tokens $j$ the place $j leq i$.

With question and key sequences of size $N$, the causal masks is a sq. matrix of form $(N, N)$. The ingredient $(i,j)$ signifies whether or not question token $i$ can attend to the important thing token $j$.

In a boolean masks matrix, the ingredient $(i,j)$ is True for $i le j$, making all parts on and beneath the diagonal True. Nevertheless, we sometimes use a floating-point matrix as a result of we are able to merely add it to the eye rating matrix earlier than making use of softmax normalization. On this case, parts the place $i le j$ are set to 0, and all different parts are set to $-infty$.

Creating such a causal masks is easy in PyTorch:

This creates a matrix of form $(N, N)$ crammed with $-infty$, then makes use of the triu() perform to zero out all parts on and beneath the diagonal, creating an upper-triangular matrix.

Making use of the masks in consideration is easy:

In some circumstances, you may have to masks further parts, akin to padding tokens within the sequence. This may be finished by setting the corresponding parts to $-infty$ within the masks tensor. Whereas the instance above reveals a 2D tensor, when utilizing each causal and padding masks, you’ll have to create a 3D tensor. On this case, every ingredient within the batch has its personal masks, and the primary dimension of the masks tensor ought to match the batch dimension of the enter tensors q, ok, and v.

Your Process

Given the scaled_dot_product_attention() name above and a tensor q of form $(B, H, N, D)$ containing some padding tokens, how would you create a masks tensor of form $(B, N, N)$ that mixes each causal and padding masks to: (1) stop consideration to future tokens and (2) masks all consideration operations involving padding tokens?

Within the subsequent lesson, you’ll study MLP sublayer.

Lesson 06: Combination of Skilled Fashions

Transformer fashions include stacked transformer blocks, the place every block comprises an consideration sublayer and an MLP sublayer. The eye sublayer implements a multi-head consideration mechanism, whereas the MLP sublayer is a feed-forward community.

The MLP sublayer introduces non-linearity to the mannequin and is the place a lot of the mannequin’s “intelligence” resides. To boost the mannequin’s capabilities, you’ll be able to both enhance the dimensions of the feed-forward community or make use of a extra refined structure like Combination of Specialists (MoE).

MoE is a current innovation in transformer fashions. It consists of a number of parallel MLP sublayers with a router that selects a subset of them to course of the enter. The ultimate output is a weighted sum of the outputs from the chosen MLP sublayers. Many fashionable massive language fashions use SwiGLU as their MLP sublayer, which mixes three linear transformations with a SiLU activation perform. Right here’s how you can implement it:

For instance, in a system with 8 MLP sublayers, the router processes every enter token utilizing a linear layer to provide 8 scores. The highest 2 scoring sublayers are chosen to course of the enter, and their outputs are mixed utilizing weighted summation.

Since PyTorch doesn’t but present a built-in MoE layer, you might want to implement it your self. Right here’s an implementation:

The ahead() technique first makes use of the router to generate top_k_indices and top_k_probs. Based mostly on these indices, it selects and applies the corresponding consultants to course of the enter. The outcomes are mixed utilizing weighted summation with top_k_probs. The enter is a 3D tensor of form (batch_size, seq_len, hidden_dim), and since every token in a sequence may be processed by totally different consultants, the tactic makes use of masking to appropriately apply the weighted sum.

Your Process

Fashions like DeepSeek V2 incorporate a shared knowledgeable of their MoE structure. It’s an knowledgeable that processes each enter no matter routing. Are you able to modify the code above to incorporate a shared knowledgeable?

Within the subsequent lesson, you’ll study normalization layers.

Lesson 07: RMS Norm and Skip Connections

A Transformer is a typical deep studying mannequin that may simply stack a whole lot of transformer blocks, with every block containing a number of operations.
Such deep fashions are delicate to the vanishing gradient downside. Normalization layers are added to mitigate this situation and stabilize the coaching.

The 2 commonest normalization layers in transformer fashions are Layer Norm and RMS Norm. We’ll use RMS Norm as a result of it has fewer parameters. Utilizing the built-in RMS Norm layer in PyTorch is easy:

There are two methods to make use of RMS Norm in a transformer mannequin: pre-norm and post-norm. In pre-norm, you apply RMS Norm earlier than the eye and feed-forward sublayers, whereas in post-norm, you apply it after. This distinction turns into clear when contemplating the skip connections. Right here’s an instance of a decoder-only transformer block with pre-norm:

Every transformer block comprises an consideration sublayer (applied utilizing the GQA class from lesson 4) and a feed-forward sublayer (applied utilizing the MoE class from lesson 6), together with two RMS Norm layers.

Within the ahead() technique, we first normalize the enter earlier than making use of the eye sublayer. Then, for the skip connection, we add the unique unnormalized enter to the eye sublayer’s output. In a post-norm method, we might as an alternative apply consideration to the unnormalized enter after which normalize the tensor after the skip connection. Analysis has proven that the pre-norm method gives extra steady coaching.

Your Process

From the outline above, how would you modify the code to make it a post-norm transformer block?

Within the subsequent lesson, you’ll study to create the whole transformer mannequin.

Lesson 08: The Full Transformer Mannequin

Thus far, you’ve got created all of the constructing blocks of the transformer mannequin. You may construct an entire transformer mannequin by stacking these blocks collectively. Earlier than doing that, let’s record out the design parameters by making a dictionary for the mannequin configuration:

The variety of transformer blocks and the hidden dimension instantly decide the mannequin dimension. You may consider them because the “depth” and “width” of the mannequin respectively. For every transformer block, you might want to specify the variety of consideration heads (and in GQA, the variety of key-value heads). Since we’re utilizing the MoE mannequin, you additionally have to outline the whole variety of consultants and the top-k worth. Notice that the MLP sublayer (applied as SwiGLU) sometimes units the intermediate dimension to 4 occasions the hidden dimension, so that you don’t have to specify this individually.

The remaining hyperparameters don’t have an effect on the mannequin dimension: the utmost sequence size (which the rotary positional encoding is determined by), the vocabulary dimension (which determines the embedding matrix dimensions), and the dropout fee used throughout coaching.

With these, you’ll be able to create a transformer mannequin. Let’s name it TextGenerationModel:

On this mannequin, we create a single rotary place encoding module that’s reused throughout all transformer blocks. Because it’s a continuing module, we solely want one occasion. The mannequin begins with an embedding layer that converts token IDs into embedding vectors. These vectors are then processed by a sequence of transformer blocks. The output from the ultimate transformer block stays a sequence of embedding vectors, which we normalize and venture to vocabulary-sized logits utilizing a linear layer. These logits signify chance distributions for predicting the following token within the sequence.

Your Process

The mannequin is now full. Nevertheless, think about this query: Why does the ahead() technique settle for a masks as an non-compulsory argument? If we’re utilizing a causal masks, wouldn’t it make extra sense to generate it internally inside the mannequin?

Within the subsequent lesson, you’ll study to coach the mannequin.

Lesson 09: Coaching the Mannequin

Now that you just’ve constructed a mannequin, let’s discover ways to practice it. In lesson 1, you ready the dataset for coaching. The following step is to wrap the dataset as a PyTorch Dataset object:

This dataset is designed for mannequin pre-training, the place the duty is to foretell the following token in a sequence. The dataset object is a Python iterable that produces pairs of (x,y), the place x is a sequence of token IDs with fastened size, and y is the corresponding subsequent token. For the reason that coaching targets (y) are derived from the enter information itself, this method known as self-supervised studying.

Relying in your {hardware}, you’ll be able to optimize the coaching velocity and reminiscence utilization. When you have a GPU with restricted reminiscence, you’ll be able to load the mannequin onto the GPU and use half-precision (bfloat16) to scale back reminiscence consumption. Right here’s how:

For those who nonetheless encounter out of reminiscence error, you might need to scale back the mannequin dimension or batch dimension.

That you must write a coaching loop to coach the mannequin. In PyTorch, you might do as follows:

Whereas this coaching loop may differ from what you’ve used for different fashions, it follows greatest practices for coaching transformers. The code makes use of a cosine studying fee scheduler with a warm-up interval—the educational fee step by step will increase throughout warm-up after which decreases following a cosine curve.

To stop gradient explosion, we implement gradient clipping, which stabilizes coaching by limiting drastic adjustments in mannequin parameters.

The mannequin features as a next-token predictor, outputting a chance distribution over your entire vocabulary. Since that is basically a classification job (predicting which token comes subsequent), we use cross-entropy loss for coaching.

The coaching progress is monitored utilizing tqdm, which shows the loss for every epoch. The mannequin’s parameters are saved each time the loss improves, making certain we hold the very best performing model.

Your Process

The coaching loop above runs for under two epochs. Think about why this quantity is comparatively small, and what components may make further epochs pointless for this specific job.

Within the subsequent lesson, you’ll study to make use of the mannequin.

Lesson 10: Utilizing the Mannequin

After coaching the mannequin, you should use it to generate textual content. To optimize efficiency, disable gradient computation in PyTorch. Moreover, since some modules like dropout behave otherwise throughout coaching and inference, swap the mannequin to analysis mode earlier than use.

Let’s create a perform for textual content technology that may be referred to as a number of occasions to generate totally different samples:

The generate_text() perform implements probabilistic sampling for token technology. Though the mannequin outputs logits representing a chance distribution over the vocabulary, it doesn’t at all times choose probably the most possible token. As a substitute, it makes use of the softmax perform to transform logits to chances. The temperature parameter controls the sampling distribution: decrease values make the mannequin extra conservative by emphasizing doubtless tokens, whereas increased values make it extra artistic by decreasing the chance variations between tokens.

The perform takes a partial sentence as a immediate string and generates a sequence of tokens utilizing the mannequin. Though the mannequin is skilled with batches, this perform makes use of a batch dimension of 1 for simplicity. The ultimate output is returned as a decoded string.

Your Process

Have a look at the code above: Why does the perform want to find out the mannequin’s system in the beginning?

The present implementation makes use of a easy sampling method. A complicated method referred to as nucleus sampling (or top-p sampling) considers solely the most probably tokens whose cumulative chance exceeds a threshold $p$. How would you modify the code to implement nucleus sampling?

That is the final lesson.

The Finish! (Look How Far You Have Come)

You made it. Effectively finished!

Take a second and look again at how far you’ve got come.

• You found what are transformer fashions and their structure.
• You discovered how you can construct a transformer mannequin from scratch.
• You discovered how you can practice and use a transformer mannequin.

Don’t make gentle of this; you’ve got come a good distance in a short while. That is only the start of your transformer mannequin journey. Maintain practising and growing your abilities.

Abstract

How did you do with the mini-course?
Did you get pleasure from this crash course?

Do you’ve got any questions? Had been there any sticking factors?
Let me know. Depart a remark beneath.