Every time you say “Thank you” to a neural network, you trigger a pipeline that multiplies hundreds of matrices with billions of elements, burning as much electricity as an LED lamp in a few seconds.
This is the first article in a small series dedicated to networks for AI/ML clusters and HPC.
In this series, we will touch upon the principles of how models work and are trained, parallelization, DMA and RDMA technologies, network topologies, InfiniBand and RoCE, and we’ll also philosophize on the topic of general and special solutions.
Specifically in this article, we will figure out what a neural network is, how it works, how it is trained, and most importantly, why it needs hundreds of expensive GPU cards and some kind of special network.
The refrain of today’s story is: there is no magic in neural networks—it’s just a multitude of simple operations on numbers performed on computers with special chips. There is no magic in how they work, nor in the infrastructure on which they run.
Let’s dive in!
Table of Contents
How to read this article?
The thing is, I couldn’t resist—and instead of a superficial overview of neural networks, I dug deeper than necessary for the series. Therefore, everything excessive is hidden under a spoiler and can be skipped. Also, in this article, I describe everything several times with increasing levels of immersion. So if something is unclear or causes indignation on your first encounter, don’t rush—read on. However Within this article, I am not trying to provide a deep, comprehensive description of how neural networks work and how to build clusters for them. Everything you read below is given with one sole purpose—to outline the domain area and to emphasize what is important when choosing a network architecture (well, a computer network) and equipment. We will primarily examine the workings of neural networks using LLMs as an example. Although different types of neural networks are structured very differently, in the context of building infrastructure for their training and inference, their specific features can be disregarded.
Glossary
What you might need...
Neuron — a basic computational element that mimics the operation of a biological neuron. Essentially, something that activates if the input value is above a certain threshold. link
Layer — a group of neurons performing a specific function for the model. With each new layer, the model gets closer to the final result. link
Block — a “computational module” consisting of several layers, functions, and operations. link
Weights — the strength of the connections between neurons of different layers. Mathematically: the values in the weight matrices at each layer. Weights are just one type of parameter. link
Parameters — a set of numbers that transform input data into a result. These are weight matrices, bias vectors, and the embedding dictionary. Parameters are only the trainable values that do not depend on the input data during inference. And these are the very trillions of parameters of modern models that everyone talks about.
MLP — Multi-Layer Perceptron — a neural network architecture containing several layers and non-linear functions between them. link
Transformer — an architecture of modern neural networks (not just LLMs) where input data is processed in parallel, not sequentially as before. link
Inference — the process of using an already trained model. link
Training — the process of training a model, where parameters are adjusted for the best result. link
Hyperparameters — the external settings of a model: number of layers, embedding size, which activation function is used, gradient step, number of experts in MoE. link
Tokens — the smallest unit of text into which the model breaks down the input data. In this article, the words “token” and “word” are often used as synonyms, although this is not entirely accurate. link
Embeddings — numbers that describe the input data (text, images, sound, etc.). link
Activation — the result of a layer or block’s operation: a vector that activates the neurons of the next layer. link
Vector — a set of numbers, like [0, 1, 2, 3, 4]. Examples: input vectors, activation vectors, bias vectors, attention vectors, and others. It is also a tensor of rank 1. The dimensionality of a vector is the number of elements. The example above is a vector of dimension 5.
Matrix — a table of numbers
[00, 01, 02, 03, 04] [10, 11, 12, 13, 14] [20, 21, 22, 23, 24] [30, 31, 32, 33, 34] [40, 41, 42, 43, 44] [50, 51, 52, 53, 54]
Examples: model weight matrices, gradient matrices. It is also a tensor of rank 2. The dimensionality of a matrix is the number of elements in rows and columns. The example above is a 6x5 matrix (6 rows, 5 columns).
Tensors — arrays of numbers of arbitrary rank. A rank 0 tensor is just a number (scalar). The analogy in geometry is a point. A rank 1 tensor is a vector. The analogy is a line segment. A rank 2 tensor is a matrix. The analogy is a rectangle. A rank 3 tensor no longer has a name, but you can call it a 3D tensor. The analogy is a parallelepiped. And so on. The dimensionality of a tensor is the number of elements along each axis.
Backpropagation — a mechanism for training neural networks by changing parameters.
Gradients — tensors that allow adjusting the values of model parameters during training to obtain more accurate answers. link
Batch — a portion of data for one weight update.
Checkpoint — a saved state of the model at a specific point in time during the training process.
KV cache — a cache of previously calculated Key and Value values that is reused for optimization. link
Context — the entire history of interaction with the neural network within the current session. That is, when the neural network “remembers” what you “talked” about. link
PyTorch — the most common framework for machine learning, which implements an API for NCCL and other libraries. link
What does an AI/ML/HPC cluster consist of?
I’d like to touch on the real world before we dive headfirst into the structure of neural networks.
All we need to use and train neural networks is a lot of servers with a lot of GPUs, all connected by a network.
Graphics Processing Units, also known as GPUs or computational accelerators. They are the basic element of the cluster. To simplify as much as possible, training models involves operations with matrices—huge arrays of numbers that need to be added and multiplied. And GPUs are the best at this. Although there are also specialized ML accelerators like Google’s TPUs.
Servers, into which these GPUs are installed. Their function is ordinary: to ensure the operation of components, data storage, and a fast network between the cards.
Network, connecting everything together: switches, cables.
Cluster management infrastructure: schedulers, collective operations libraries, data storage, and management and monitoring systems.
There’s no question about why GPUs are needed, right? But why a special network?
Why are special networks needed for ML?
The time when it was possible to train relevant neural networks on a single machine, even with the most powerful video card (the common name for a GPU), is long gone.
You need many video cards. And you also need to be able to parallelize computations among them.
Consequently, we connect them in a network. There are dozens of different topologies with various optimizations: some are aimed at reducing the diameter and therefore network latency, some at increasing bandwidth, and some at reducing cost. We will look at this later. But what makes networks for ML and HPC so special that we need to think and write articles about them?
It’s all about how training happens and the parallelization of computations. If data moves too slowly somewhere, the cluster will be idle. Along with a burnt-out server, we will lose all the computations delegated to it, and they will need to be restarted from a checkpoint. If a block of data is lost during transmission, it will need to be re-requested, and during this time the cluster will again be idle.
For example, training GPT-3 on a cluster of several thousand GPUs took about two months. Two months of continuous computations on thousands of nodes, connected by hundreds of switches with thousands of cables. All of this breaks. And how we build networks and what parallelization schemes we use is crucial.
Therefore, ML requires a Low Latency Lossless network. Why—we’ll look at that today. How to make it so—in the following articles.
What is a neural network
In essence:
Neurons
The creation of neural networks was inspired by brain research in the 1940s, which found that a neuron can be represented by a simple function that takes signals from several neighbors as input, computes a result, and if it exceeds a certain threshold, the neuron activates and, creating a new signal, passes it on to its neighbors.
It is this principle of neuron activation through connections between them that is modeled in all modern artificial neural networks.
Each circle in the simple example below is a separate neuron in an artificial neural network:
But we are not yet very good at hardware implementation where electrical signals pass through millions of neurons (if we don’t count shady schemes with memristors), so almost any artificial neural network is hundreds of matrices and vectors, each containing several million numbers, stored in memory and processed by a processor.
The initial data is fed into the neural network, which is converted into vectors of numbers. Then these vectors, moving from layer to layer, are multiplied by matrices, added to other vectors, and various functions are applied to them. At the output of each layer, a new sequence of vectors is obtained, which moves to another layer, where the same thing happens to them. And so it continues until the final vectors reach the output layer.
And at the output layer, the neurons magically create for us what we asked for.
Some networks take text as input and create text as output—these are LLMs, like BERT or GPT. Others take an image as input and are supposed to output text, for example, convolutional neural networks (CNNs) that recognize images, including text on them. Still others take text as input and are supposed to generate an image (diffusion networks). A fourth type takes sound as input and is supposed to recognize text in it.
Fundamentally, they all work the same way—complex mathematical models with billions and trillions of parameters that can independently find patterns in data and “understand” them. And they are all not programmed, but trained based on a huge number of examples.
Layers and weights
Neural networks consist of layers:
input
hidden
output
The number of hidden layers can vary from a few to hundreds. Layers are connected to each other by connections, which have weights (w).
From layer to layer, the model recognizes increasingly complex patterns and relationships in the input data, which leads it to the result.
The weights of the connections (i.e., the values of the matrices), the values of the vectors—all of these are the model’s parameters. They are calculated during training and do not change thereafter. That is, all these billions of numbers are the model itself. When they say that GPT-3 contains 175 billion parameters, they mean that its matrices and vectors have a total of 175 billion trained values.
Activations
If each layer were just matrix multiplications, then all the layers could be collapsed into one by simply multiplying them once. And there would be no magic. The freedom, flexibility, and creativity of neural networks are due to the fact that each layer also has a non-linear activation function, which strengthens good connections and weakens parasitic ones. Functions like sigmoid, ReLU, or GELU can be used for this.
This is very similar to a real neuron—it gets excited and transmits a signal further only if the incoming signal has crossed the excitation threshold.
Let’s look at what happens inside, using a specific example of an MLP.
A specific example - Multi-Layer Perceptron (MLP)
This is based on the video But what is a neural network?, which discusses the MLP—Multi-Layer Perceptron—the simplest neural network that recognizes handwritten digits.
The model in question has one input layer, one output layer, and two hidden layers. Between each pair of layers, every neuron is connected to every other neuron.
The input is an image of size 28x28 pixels, which is a total of 784—this is the dimensionality of the input layer—784 neurons. Each neuron receives a value from 0-255, depending on the brightness of the pixel it is responsible for. The output is ten neurons, which correspond to the ten digits. The two hidden layers in the middle each have 16 neurons. This number was chosen more or less arbitrarily, as was the number of hidden layers. Each connection has its own weight.
And so it turns out that our model has 784 * 16 + 16 * 16 + 16 * 10 = 12,960 parameters—one for each connection between layers. And the MLP also has a bias vector—the size of each layer, so that’s another 16 + 16 + 10 = 42 parameters. In total, that’s 12,960 + 42 = 13,002 parameters that need to be trained in such a way that they can identify handwritten digits with high accuracy (>90%).
And on this same example, let’s take the first step towards understanding how neural networks work.
In neural networks, each hidden layer is responsible for extracting certain features. The further the layer is from the input, the more abstract and general features it identifies.
So, in an MLP, the first hidden layer might look for lines, corners, and connections, while the second might combine them into more complex shapes, like closed loops.
And the output layer already “understands” which digit is most likely formed from them:
Each neuron is responsible for a certain feature and activates depending on whether it was able to detect it in the input data or not. And whether it was able to detect it is determined by the level of the signal it received—if it’s powerful enough, the feature is present, and the neuron “lights up”; if it’s weak, its feature is not detected, and the neuron remains silent. And the signal strength is determined by the weights of the connections leading to it from each neuron of the previous layer.
So in the MLP example:
In the input layer, the neurons under the “white” pixels light up. Connections lead from each of them to the second layer. Each connection has its own weight.
In the first hidden layer, the neurons that identify different simple shapes light up. Very specific neurons light up because the connections with the highest weights lead to them.
In the second hidden layer, neurons that identify more complex shapes light up. Again, specific neurons are activated because the strongest connections lead to them from the active neurons of the previous layer.
The output layer can already determine the exact digit because one neuron lights up, to which connections with large weights lead from the active neurons of the previous layer.
As a distant analogy, you can think of a visit to the optometrist, when you are asked to look at an eye chart, and the letters are all blurry. You try new lenses, with increasing diopters—and with each new lens, the image becomes clearer and clearer.
In other words.
If strong connections with high weights lead from neurons 53 to 84 of the input layer to neuron 5 of the next layer, then this 5th neuron “understands” that a horizontal line has been detected, and it activates.
In turn, in the second layer, neurons 5-9 and 11 activated and sent a sufficiently powerful signal to the 7th neuron of the third layer for it to determine that a circle was formed, for the recognition of which it is responsible.
In this layer, neurons 1, 5, 7, and 8 activated, and they lit up the neuron at the output responsible for the digit 9.
All neural networks work on a similar principle. We have considered a relatively simple example with an MLP that recognizes handwritten characters in photographs.
It is now important to note that MLP is not specifically about image recognition. It is a neural network architecture that has several fully connected layers and a non-linear activation function between them:
Input → [Fully Connected Layer → Non-linearity] × N → Output Where N is the number of hidden layers
MLP is also used in transformers, on which all modern LLMs are built: GPT, Gemini, DeepSeek, Claude, QWEN. And now we’ll tackle them—to get a deeper understanding of the principles of how neural networks work.
Principles of how neural networks work
First, let’s clearly separate the processes of training a model and using it—Training and Inference.
Usage (inference)** is when you ask a model to tell you in which part of Moscow the bream are bigger, or to generate a picture of an anime girl for you. Or it writes code for a Telegram bot for you. Or it creates a strength training plan. That is, you are just using it.
Inference can also work on GPUs that are not assembled into a special cluster. In most cases, a trained model can fit on a single server with eight GPUs. A neural network should respond to user requests in seconds.
Training is a lengthy process of parameter selection, where millions of examples are run through the model, improving its performance with each cycle. That is, inference becomes possible after training.
You need dedicated clusters of thousands of cards, connected by a special network. Training can last continuously for several months.
Training and usage require fundamentally different approaches and use different infrastructure. And to understand why networks for ML are the way they are, we need to understand how they are structured.
Let’s start with inference as a simpler and more general operation for both processes.
Inference
Layers and blocks
All modern models are multi-layered. From a few layers (perceptrons) to hundreds (LLMs). Each layer performs a specific function. For example, when recognizing letters in images, the first layer finds corners, lines, and curves. The second combines them into sticks, circles, and loops. The third determines what the resulting symbols look like (“the probability that this is the digit 8 is 0.998472”).
LLMs now use blocks, each of which is a combination of layers. In GPT-3, for example, there are 96 blocks. Some blocks understand syntax, parts of speech, gender, tense, and names. Others understand semantics and context, tracking subject-object relationships. Still others recognize style and find historical context.
The further a block is from the input layer, the more complex and distant relationships it finds, and the more general meaning and context of the text it determines.
Each block, in turn, consists of layers. The layers are functional: one multiplies the input matrix by a weight matrix, another applies a bias, a third performs activation, and so on.
And here is an illustration from the classic paper Attention Is All You Need. It’s about how Attention adds a bit of magic (after all) to transformers.
We will be разбираться with this daunting picture next.
What’s important to understand about LLMs is that they DO NOT follow a specific plan: analyze the meaning of the question → prepare a plan for the answer → formulate the answer. They just find the next probable word. … They just find the next probable word.
How?
Let’s ask the neural network: "Continue the phrase “The quick brown fox jumps over the lazy dog.”
The neural network breaks it down into components, “understanding” the precise meaning of both individual words and the entire phrase. It understands that this is a pangram. It also “understands” what is expected of it—to continue the phrase.
After the first pass through all the blocks, it generates the word “And”. This is simply the most probable first word, given the current query and all the examples it has seen. Yes, yes! After all the billions of multiplications, it just produced “And”.
Now, the phrase “Continue the phrase “The quick brown fox jumps over the lazy dog.” And” is fed into the neural network. It is run completely through all the blocks again, and at the output, another word is added: “Continue the phrase “The quick brown fox jumps over the lazy dog.” And then” This is simply the most probable next word, given all the examples it has seen, the query, and the last generated word.
This new phrase is again fed entirely as input and run completely through all the blocks one more time, and at the end, the next word is chosen: “Continue the phrase “The quick brown fox jumps over the lazy dog.” And then it” This is simply the most probable next word, given all the examples it has seen, the query, the response, and the last generated word.
This process continues until the neural network decides it is finished:
And then it went on to list all the letters of the alphabet.
The quick brown fox jumps over the lazy dog, and so does the alphabet from A to Z. The quick brown fox jumps over the lazy dog, a sentence famous for its complete set of letters.
How does it know it’s finished? Simply, the next most probable word will be the end-of-generation symbol. How so? It was just that way in the examples it was trained on. Magic? Yes, but no.
Three fundamental things about LLMs:
They do not understand in the human sense. They see the proximity of different vectors in a mathematical sense and thereby find connections between words and phrases.
They do not invent or prepare an answer. They simply, as a result of all calculations, get the probabilities of each word (token) in the available vocabulary.
The text is run through all the model’s blocks many times: after each new generated word, the entire text is fed back into the model’s input.
Sequence of events during inference
Let’s consider, on a very simplified but real example, what happens when we send a phrase to an already trained LLM—that is, when we use it.
Let’s take GPT-3 and the phrase chosen above.
In fact, no better visualization of how LLMs work has been created than this: https://bbycroft.net/llm. I highly recommend clicking on it before, during, and after reading this article to stabilize the chaos in your head. And I will be pulling illustrations from there for my paragraphs.
1. Tokenization
The first thing the model does is break the phrase into parts, each equal to one token. These can be whole words or individual letters, but usually neither, but rather parts of equal or approximately equal size. Breaking it down by letters results in too many tokens and makes it too difficult to connect everything into a single meaning (try it yourself: ["T", "h", "e", "q", "u", "i", "c", "k", "b", "r", "o", "w", "n", "f", "o", "x"]). By words, we grow the vocabulary—we would have to create a separate entity for all forms of a word. For example, “run”, “running”, “ran”, “runner”, and so on.
For example, our phrase might be broken down like this: ["The", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog", "."]
Thus, we have broken the text into individual tokens.
2. Embedding
Next comes the first stage of magic (yes, yes, magic): the neural network determines the coordinates of the received words in a general semantic space. No more, no less.
But these are all words. And we need actions numbers! We need numbers because we cannot apply the mathematical apparatus of matrix operations to words.
Therefore, each token is associated with a vector of numbers. In the case of GPT-3, the length of this vector is 12,288. This is called the model’s dimensionality. Roughly speaking, dimensionality means how accurately the model can determine all possible meanings and variations of a word in all possible contexts. At the same time, the closer words are in meaning, the more similar their vectors are. So the vectors for “near” and “close” will be very similar, while “near” and “far” will not, although they are still more similar than “near” and “dinosaur”. At the same time, the words “bat” (the animal) and “bat” (the sports equipment), despite their different meanings, will have identical vectors. Which will then change as they pass through the layers.
How does the model do this? It has an Embedding dictionary—a complete dictionary of all possible words. For GPT-3, this is 50,257 words (or tokens). Each is represented by a vector of 12,288 numbers. And this is the first term of all the model’s parameters: 50257 * 12288 = 617,558,016 = ~617M. Essentially, in this matrix, the model finds the row that matches the word passed as input.
These were vocabulary embeddings. In addition, for transformers, the position of the word in the text is important, so another vector is generated for each word, describing its position in the text—positional embeddings—they can be static, but are usually trained. These vectors… are simply added to each other.
As a result, for our phrase, we will get a matrix of size 8×12,288.
This was the data preparation stage—we determined the coordinates of the text.
Next, this matrix is sequentially passed through dozens of uniform transformer blocks. In the case of GPT-3, there are 96 blocks, each consisting of several functional layers.
In the illustration, this is the transformer i block.
As I mentioned above, each block reveals new knowledge about this phrase. For example, it identifies the subject (“fox”), the predicate (“jumps”), the object (“dog”), their gender, number, how these words relate to each other, and in what context they were likely used.
3. Self-attention mechanism
First, the attention mechanism works within the block—how words are related to each other. Attention heads are responsible for this. This is where the full power of transformers lies! No other architectures allow tracking connections between words that are far apart in the text.
The paper that became a classic on the attention mechanism and how it transforms a regular MLP is: Attention Is All You Need.
In our case, the attention heads determine that the main word is “fox”, that it is singular, that the fox is jumping over the dog, etc. Each block has several Attention Heads (in GPT-3, there are 96 in each block), each of which “looks” at the word from different angles:
syntactic relationships (subject—predicate);
semantic connections (synonyms, antonyms);
short- and long-range dependencies;
thematic and stylistic patterns,
grammar
etc.
Just as in the MLP example above, where each neuron recognized its own feature in the image, in transformers, each head recognizes its own type of connection. In each block.
Self-attention in detail
For each token, three vectors are calculated:
Q (Query) — What I want to know
K (Key) — What I know about
V (Value) — What I ultimately convey
Let’s first understand this on an intuitive level.
Let’s take the word “jumps” from the phrase “The quick brown fox jumps over the lazy dog.” Possible Queries for the word “jumps” (Q_jumps). That is, what this token wants to know from others.
“Looking for syntactic connections with other words”
“Who is my subject?”
“Are there any adverbs nearby?”
The token “jumps” is set to look for: subjects, objects, adverbs, prepositions.
Possible Keys for the word “jumps” (K_jumps). That is, what this token knows about itself and how others can find it.
“I am a verb in the present tense”
“I can be a predicate”
“I describe an action”
It tells others: “I am an action word, I can answer the question ‘what does it do?’”
Possible Values for the word “jumps” (V_jumps). That is, what the token will tell others.
“Action of leaping”
“Context: movement, activity”
“Grammar: present tense, singular”
“Can be combined with: fox (subject), over (preposition)”
Contains enriched information to be passed into the context.
Now let’s take the word “fox”. Possible Queries for the word “fox” (Q_fox):
“Looking for a verb that I am performing”
“What is my predicate?”
“Are there any adjectives describing me?”
The token “fox” is set to look for: verbs, adjectives, especially right after it. The noun “fox” knows: “I am a subject! I need a verb!”
Possible Keys for the word “fox” (K_fox).
“I am a noun, a subject”
“I can perform an action”
“I am an animal”
“Pay attention to me if you are a verb after me!”
It signals to others: “I am the grammatical subject, important for nearby verbs.” For “jumps”:
K_foxshouts “I am your subject!” For “dog”:K_foxsays “I’m not that important to you, ‘jumps’ is between us.”
Possible Values for the word “fox” (V_fox):
“Semantic meaning: a carnivorous mammal”
“Syntactic role: subject”
“Grammatical connection: noun + verb”
“Adds meaning: the agent of the action”
Contains “I am the one performing the action.”
These vectors Q, K, V are calculated for all tokens.
And after the attention heads, the vector that contained only the embedding now contains: “jump + it is performed by the subject ‘fox’ + there is a related object ‘dog’ + the general context of an action.” And the vector for “fox”: “the noun ‘fox’ + it is the subject of the verb ‘jumps’ + the specific construction ‘fox jumps’ + context: action, lazy dog nearby.”
And so on for all tokens.
What is this from a mathematical point of view? Matrix multiplication, of course. What did you expect?
The vector of each word is multiplied by special (pre-trained) weight matrices Wq, Wk, Wv to get the vectors Q, K, V. Then, for each word, its Query is compared with the Keys of others.
For example, the Query of the word “jumps” and the Keys of other words (including itself):
Q_jumps · K_The = 0.2 (low connection) Q_jumps · K_quick = 0.3 (low connection) Q_jumps · K_brown = 0.3 (low connection) Q_jumps · K_fox = 0.9 (high connection — subject!) Q_jumps · K_jumps = 1.0 (with itself) Q_jumps · K_over = 0.6 (medium — related preposition) Q_jumps · K_the = 0.2 (low connection) Q_jumps · K_lazy = 0.4 (low connection) Q_jumps · K_dog = 0.7 (medium — related object)
Q is a vector, K are vectors. The numbers in the examples above are the dot product of Q_jumps and the K of other words. And this number shows how much token j “answers” the query of the token “jumps,” i.e., the strength of the connection.
The dot products of all pairs Qᵢ and Kⱼ are calculated, where i, j are all the tokens in the phrase. Then the softmax function is applied, and these connection strengths are turned into probabilities whose sum is 1. That is, Q relates to the current word we want to know about, K, V relate to all other words in the text, including the word itself.
All this is put into a weight matrix, i.e., the connection strengths of all pairs of words. And then a weighted sum is calculated for all words: with what strength the vector Vⱼ influences the vector i.
In other words: the input vector of a word is changed by V with a strength of Q * K—and so on in each block.
I won’t even pretend that we are supposed to understand anything here:
But try watching this in motion seven times, read a couple more articles, reread this article, and watch it in motion again—and it will start to become clear.
And… this is an example for only one attention head. There are several dozen or hundreds of them in a block, each doing the same thing but tracking different relationships.
The parameters in the attention mechanism are the matrices
Wq,Wk,Wv, and they are determined during the training process. They are immutable for the model during inference. The vectorsQ,K,Vare specific to each particular example and change from example to example. They are not parameters of the model.
That is, Q, K, V are tools for calculating “who influences whom and how” in the text. And this is the ultimate “understanding” mechanism in transformers.
Here we get a vector enriched with connections to other words, that is, with the meaning of the phrase.
4. FFN and non-linear transformations
Language is a strictly non-linear thing. Therefore, to process language, you need to add non-linearity. This is exactly what the FFN—Feed-Forward Network—a microscopic two-layer neural network, does. It applies a non-linear function to the vector. And this allows the neural network to be flexible, recognize complex patterns, and not be just a series of linear transformations that could actually just collapse into a single linear transformation.
Roughly speaking, FFN preserves strong connections and nullifies weak ones.
This is actually an MLP—Multi-Layer Perceptron, which we dealt with above. That is, MLP is an embedded component of the transformer block, which has its own name here—FFN.
These are two matrix multiplications, between which a non-linear function operates.
The first multiplication increases the dimensionality of the token’s vector by 4 times (in GPT-3, from 12288 to 49152). Then a Bias vector is added.
The non-linear function modifies the vector’s values: it strengthens strong connections and weakens weak ones.
Non-linear functions
There are many non-linear functions, here are the main ones used in ML:
Sigmoidtransforms all numbers into values from 0 to 1. The result is a smooth line resembling the letter S.
Sigmoid ReLU: zeros out negative values. The result is a horizontal line that transitions into a sloped line.
ReLU GELU: smoothly reduces small values while preserving information. Used in GPT.
GELU
More on non-linear functions below.
The second multiplication returns the vector to its original dimensionality.
A more detailed illustration:
In the illustration, you can also see the MLP Bias—this is a bias vector necessary to be able to cut off values not at zero, but at a value less than a certain threshold during activation. That is, roughly speaking, to raise or lower our graph.
Non-linear activation function in detail
The video Why Deep Learning Works Unreasonably Well explains why it is more beneficial for a model to have more hidden layers than neurons, i.e., it is preferable to be deeper rather than wider (just like in life).
Each hidden layer contains an MLP with a non-linear activation function. But WHY?
Well, they only work together: the non-linear function and the multi-layered structure.
Let’s start by understanding why linearity is… bad. I’ll be pulling illustrations from a good video that explains the principle of vanishing gradients: Activation Functions - EXPLAINED!
We all know what a linear function is… Well… It’s a line. And it works well only when we have a clear and linear boundary between the objects that need to be separated:
But in this situation, it no longer works:
But! If we add non-linearity to the model, it can draw non-straight boundaries:
And the irony is that no matter how huge and complex a network with millions of layers and multiplications we make, they all collapse into a single linear function if there is no non-linear activation function. wat.jpg
We’ll have to add a few formulas.
X— is the input dataW— are the connection weightsb— is the bias vectorh— is the activation vector at the output of the layer.y— is the result
The activation vector of the first layer will be calculated as follows:
h₁ = W₁x + b₁
We multiplied the input by the weights and added the bias vector. The activation vectors of the second and third layers are calculated similarly:
h₂ = W₂h₁ + b₂ y = W₃h₂ + b₃
Now let’s substitute one into the other:
y = W₃(W₂(W₁x + b₁) + b₂) + b₃ y = W₃W₂W₁x + W₃W₂b₁ + W₃b₂ + b₃
W₃W₂W₁ is just a constant matrix—the product of static weight matrices—we can replace it with W. And W₃W₂b₁ + W₃b₂ + b₃ is a constant vector—we can replace it with b:
y = Wx + b
And a three-layer network with 46 parameters turns into an elegant single-layer network with eight parameters.
Everything changes as soon as we add a non-linear function. For example, ReLU:
h₁ = ReLU(W₁x + b₁) h₂ = ReLU(W₂h₁ + b₂) y = W₃h₂ + b₃ y = W₃ReLU(W₂ReLU(W₁x + b₁) + b₂) + b₃
And it can no longer collapse. Roughly speaking, with ReLU, we add kinks to our line and can get an arbitrarily complex shape.
Here we get a vector cleared of noise and with enhanced context.
5. Layer Norm and Residual connection
One could say that an LLM normalizes the vector whenever possible and adds the input vector to it. Why? Because empirically, the model is more stable and gives better results this way.
Normalization or Layer Norm
This is needed to make the mean of the vector 0 and to fit the range of values within a certain range. That is, to avoid outlier large or small values. This makes the model more stable.
For this, two vectors of the model’s dimensionality are needed (12,288 for GPT-3): γ (scale) and β (shift). The output vector is obtained by the formula output=γ × normalized + β.
Residual connections
It turns out that if you add the input vector that entered the block to the output vector, the model seems to better remember what came before (some historical roots) and also works more stably.
Why? That’s an interesting question. At first, this was discovered empirically. But later, an explanation was found: Residual connections combat vanishing gradients and significantly improve model training. We will talk about gradients and training below.
At the output of the last layer, we get a vector representation of the last token, enriched with meaning.
6. Output level
For text generation tasks, which is what we mostly expect from LLMs, we only need the vector of the LAST token (word). It is on the basis of this vector that the model searches for the next token. The output vectors of all previous tokens are not used for this, but they were critically important for the attention mechanism to work. That is, the vector of the last token has “accumulated” the entire mass of meaning of all the input data.
In the model, for the output layer, there is a dictionary, the same as at the beginning (where words were turned into numerical vectors)—this is also a matrix of 50,257 by 12,288 (for GPT-3). Usually, this is the same embedding matrix.
The output token vector is semantic. Now it needs to be transformed into probabilities of the next token.
For this, the logit is first calculated: the token vector is multiplied by a special weight matrix. The result is a vector of the vocabulary’s dimensionality (50,257 for our example). A logit (a term from probability theory) is, let’s say, an estimate of the model’s confidence in the next word. It’s not yet a probability because the values can be negative and are not necessarily in the range of 0 to 1.
In our example for the phrase “Continue the phrase “The quick brown fox jumps over the lazy dog.” And then it”, it might be like this:
“went” +4.3 “continued” +2.1 “ate” -12
Next, the softmax function is applied to the logit vector, and it is transformed into a vector of probabilities, the sum of which is 1. There were 50,257 real numbers, and now there are 50,257 probabilities corresponding to the words in the dictionary.
Then it might be like this:
“went” 0.541 “continued” 0.311 “ate” 0.00001
It would be logical to take the word with the highest probability. But in practice, sampling is done—a word is chosen randomly, according to its probabilities. In our example, in 54.1% of cases, “went” will be chosen, in 31.1%—“continued”, and in 0.001%—“ate”.
And so on, each time using the vector of the last token, the model iteratively selects the most probable next word.
The vector’s path through the model
It is very important to understand that we represented a word as a vector at the input, and then we run this vector through the meat grinder of transformer blocks, modifying it, enriching it with meanings, connections to other words, and normalizing it. The numbers in it change in one way or another depending on the context in which the word is used in this phrase and was used in all the training examples. And at the output, we get a vector of the same dimensionality, which has become a derivative of the initial vector and all the transformations in the pipeline.
This is how two initially identical vectors for the words “bat” (animal) and “bat” (sports) become completely different by the output layer because they have different meanings.
And even the word “pit” will produce two completely different vectors at the output in the sentences “oh, there’s a hill by the pit” and “Everyone helps dig a deep pit, of course, so that everyone can fit in it”.
If only the vector of the last word is used to generate the next word, why calculate all the others?
Well, first of all, the attention mechanism—those very
KandVvectors are necessary to understand the meaning—and for this, all words are needed, both the original and the generated ones. Secondly, the vector of the last token is already enriched by all the previous ones. Without them, there would be no context in which this word is used. Thirdly, LLMs have other tasks where all vectors are indeed important. And fourthly, there is the KV cache, which optimizes these calculations.
Model temperature
There is also the term “model temperature”—this is how “bold” it is in choosing the next word or how creative it is. The hotter it is, the higher the probability of choosing a less popular word. For example, at a low temperature, the model will continue our phrase “The quick brown fox jumps over the lazy dog” like this:
and then it went to sleep.
At an increased temperature:
and danced the fandango.
If the model is heated up:
and my WiFi works better there.
The term comes from thermodynamics, where temperature essentially determines the degree of chaos in a system. The higher the temperature, the more actively the particles move, and the system is less predictable.
KV cache
What can be astonishing when first deeply understanding how neural networks work is that all the words of the text are repeatedly and completely ground through all the layers. Imagine you ask the network to retell “War and Peace” in no less than 20 pages. As long as the token limit is sufficient, for each new word, the neural network will restart the calculations from scratch for all the already generated words.
It’s not clear why. And it generally seems like a waste of energy. But without this, the model simply has no context to pick the next word that fits in form (gender, case, number) and meaning. Each new word changes the meaning of the entire phrase.
Okaaaay. It’s not that necessary, because there is a pretty good optimization—to avoid recalculating embeddings and FFN for all previous tokens for new ones.
Each block saves a KV cache (that is, a cache of the K and V vectors) of all previous tokens, so essentially only the new token is run through all the blocks, and the values for the old ones are taken from this saved KV cache. For the new token, all vectors are calculated: Q, K, V. But the Query from the new token needs to be recalculated against the old ones as well, which is logical.
This is a classic trade-off of computation for memory—this cache takes up more than a decent amount of space.
Detailed figures for storing the KV cache
For one layer: KV_cache_size = 2 × n_layers × batch_size × seq_len × n_heads × head_dim × bytes_per_param
For GPT-3 (175B):
2 — for K and V (two caches)
n_layers = 96 — number of transformer layers
batch_size = 1 — how many requests are processed in parallel
seq_len = 2048 — sequence length (context)
n_heads = 96 — number of attention heads
head_dim = 12,288 / 96 = 128 — dimensionality of each head (
d_model / n_heads)bytes_per_param = 2 — size of one number (2 for fp16)
KV_cache = 2 × 96 × 1 × 2048 × 96 × 128 × 2 = 9,663,676,416 bytes ≈ 9.0 GB
9 GB for all layers for a single prompt! And we must consider the linear growth with context length (as a reminder, today it’s up to 1 million tokens)—so for Claude with 200k tokens of context, the cache size will be 781 GB.
And the real cost of inference is often determined by the size of the KV cache, not the model weights!
Why KV cache and not QKV? Well, because the Query is not needed for new generated tokens, and the old tokens have already “asked their questions” and received answers when they themselves were new tokens.
An important point: the KV cache exists within a single request (to be more precise, within a single forward pass through the model). So you’ve written a question in a chat with a neural network. While it’s processing it, generating the first token, the second—it’s filling the KV cache. And so on until the answer is ready and the model has finished processing this request. After that, the cache is cleared.
When you ask the next question, everything starts over from scratch. If after processing the request “The quick brown fox jumps over the lazy dog” you write “Peter Piper picked a peck of pickled peppers”, the model will start working with the phrase “Peter Piper picked a peck of pickled peppers”, knowing nothing about the previous one. The model remembers nothing of what you talked about with it. That’s why they say LLMs have no long-term memory—only short-term memory within a single request.
If you want to have a dialogue with it, you need context.
Context
For a dialogue, you need history. Therefore, interfaces to models (for example, chats) themselves intentionally supplement the user’s request with the conversation history: both the user’s questions and the LLM’s answers. That’s why after two tongue twisters, you can ask “What do they have in common?”—and it will answer.
You can see an illustration of the difference between the KV cache and context when you watch a response being generated. First, the model thinks for a relatively long time, and then it types out the answer quite quickly, word by word. First comes the Prefill phase: it runs the entire context through itself, fills the KV cache, which takes some time. And for each new word of the answer, the model uses the cache and calculates the new word—this is the Decode phase.
With a new request, everything repeats: think, calculate the long context, then quickly generate the answer.
Since all calculations for the full context are started anew each time, it’s an increasing number of tokens each time.
First request and response: 200 tokens
Second: 150. Total 150 + 200 = 350
Third: 400. Total 400 + 350 = 750 etc.
This means that requests will be “more expensive”—in terms of money, computations, and time.
How does it really work?
In reality, they optimize everything they can, including context. Many AI interface providers implement Prefix Caching (or Prompt Caching). In this case, the context is stored in something like Redis/Memcached for a few minutes after the request. When you send a request to the model, it first checks if there is a cache that needs to be loaded into the GPU.
To avoid excessive context growth, all models have a limit: from a few thousand to several hundred thousand tokens. For GPT-3, it was 2048, and for Claude-3: 200,000. For modern models at the time of writing (Gemini-3 Pro), the context is already 1M tokens. If the dialogue is longer than the limit, either the oldest part is cut off, or the history is optimized.
What’s the problem with context length? As soon as we overflow it, the model starts to get weird—it forgets what it was doing, it can’t hold all the data “in its head,” and it starts to break the logic of its answers. For example, you’re writing some program with the neural network, and suddenly it starts inserting random characters or breaking the application’s architecture or suggesting very bad solutions—most likely, it’s due to exceeding the context.
Summarizing inference
Attention Heads find relationships and patterns.
FFN first expands the dimensionality of the matrix, that is, its resolution. Then, after the non-linear function, it narrows it back down.
The non-linear function allows “cutting off” noise in the data, thereby strengthening useful weights.
The bias vector indicates exactly along which line to cut the noise.
Layer Norm reduces what is too large.
The Residual Connection adds the original input vector to the vector obtained in the block to strengthen the connection with it.
The KV cache can significantly reduce response time during inference by temporarily storing calculations for previous tokens.
So what is a neural network?
It’s billions to trillions of numbers in matrices and vectors, arranged in a correct, empirically determined order.
Let’s look at the example of GPT-3.
Embedding matrix 50,257 (vocabulary size) × 12,288 (model dimensionality) = ~617M
Positional embeddings 2048 × 12,288 = ~25M
96 sequential transformer blocks
Attention heads: Wq, Wk, Wv — 3 × (12,288 × 12,288) = ~453M. And the Output projection back to the token vector: 12,288 × 12,288 = ~151M. Total ~604M
FFN: first, we expand the dimensionality by 4 times: 12,288 × 49,152 = ~604M. Then we reduce it back: 49,152 × 12,288 = ~604M. Total ~1208M
Two normalization vectors with two parameters each: 2 (γ) × 2 (β) × 12,288 = 49,152
Let’s calculate: 642M (embeddings) + 96 (layers) × (604M (attention) + 1208M (FFN) + 0.049M (normalization)) = ~175 billion parameters.
All these 175 billion parameters describe the entire world available to the neural network (not so big, is it? Although modern models already have trillions of parameters). And all of them must be calculated during the training process.
Training
The process described above is so-called Inference, or the use of an already trained neural network. We give it data as input—and at the output, we very quickly (in seconds and even milliseconds) get an answer.
“Trained” means that the matrices and vectors already contain numbers that give a good result.
The task of the training stage is to find such numbers. Each of the billions, or even trillions, of numbers. The training process for modern networks, like GPT, takes months of computation on hundreds and thousands of GPUs.
Key stages of training:
Forward pass. We take a phrase, predict the next word (essentially what happens during inference).
Calculate the loss function. We find out how poorly the parameters are currently selected.
Backward pass. Layer by layer, correction matrices are calculated—that is, how each parameter of the model needs to be changed to reduce the error.
Update parameters. In a separate step, all changes are applied.
And all this is done millions of times on millions of training examples—until the predictions become satisfactory. That is, during training, we are looking for the minimum of the loss function. In addition, there is monitoring and logging, hyperparameter tuning, control of overfitting, and a whole bunch of other things. We will look at some of this below.
The most important thing is that as a result of the process, we get hundreds of correct weight matrices and vectors.
The vast majority of modern neural networks exist only thanks to the backpropagation mechanism. Without it, training is practically impossible.
Here’s what training looks like on an oscilloscope:
Hyperparameters
These are the global parameters of the model, such as the dimensionality of the embedding matrix, the number of layers, the number of neurons in the layers, and the size of the gradient step. They also define the structure of the blocks, the number of attention heads, and which function will perform the non-linear transformations.
They are set once before the start of training and are then its immutable laws. That is, they are not trained like regular parameters, but set the rules of the game—that’s why they are hyperparameters.
What the training process looks like
At step 0, we have several hundred huge matrices filled with random numbers. The model will not be able to predict even the completion of the phrase “I love”. More precisely, it will, but it will be gibberish, like “I love Sayan-Shushenskaya ultraviolet liquid to mess around in a million”.
Step 1: to teach it to compose text, you need to show it an example. Many examples. A lot. These are websites, books, periodicals, archives, Wikipedia. In total, GPT-3 had about 300 billion training tokens.
For each example, the model must predict the next word. And depending on how well it guessed, all its parameters are adjusted. We repeat this until the predictions become accurate enough.
Let’s be more specific. For example, we chose the phrase “The cat sat on the windowsill and watched the birds”. It is sliced into training and validation data. We will make the model finish the phrase (input → target)
The cat → sat The cat sat → on The cat sat on → the The cat sat on the → windowsill The cat sat on the windowsill → and The cat sat on the windowsill and → watched The cat sat on the windowsill and watched → the The cat sat on the windowsill and watched the → birds
And there are millions of such examples.
In the second step, the model does exactly the same thing as during inference: self-attention, normalize, FFN, residual connection, and so on through all the blocks. In the same way, the model tries to “understand” the text, to find connections between words. But for now, it’s very bad at it. However, an algorithm is an algorithm, and it will choose the next word anyway: “The cat sat on the windowsill and watched purple”
Step three. At the output, the new word is compared with what it should actually be. The loss (error) function is calculated—that is, how much it missed and how much the model needs to be changed.
“Purple” is not at all like “the”, so the weights need to be changed… All… 175 billion of them.
Fourth step: all the model’s weights are changed. This process proceeds sequentially, layer by layer, in the reverse direction: from end to beginning, which is why it’s called the backward pass or Backpropagation. You can imagine a wave that sweeps away everything in its path and changes the landscape. In the same way, the correction matrices change all the numbers they encounter on their way.
And then we start over: we try other examples with the new numbers. The error should become a little smaller. And again. And again. And so, millions of times.
These are all words. But we need numbers—the correction matrices for each layer are calculated using gradient descent.
Gradient descent and backpropagation in detail
These are two different mechanisms:
Gradient descent is used to calculate how to adjust the weights at each layer.
Backpropagation is used to transmit the error information for calculating gradients from layer to layer.
Let’s first understand what this is on a very simple example. I’ve based this on the video The Most Important Algorithm in Machine Learning.
We have some points on a plane. From calculus, we know that a polynomial of degree n-1 can be drawn through any n points on a plane. In practice, this is difficult to do, and in the case of neural networks, it’s also harmful due to possible overfitting. But we can find a function that passes well, okay, close enough to the points, that is, close enough to all of them, but not necessarily through them.
In the chosen example, there are 34 points, unevenly distributed on the plane. Let’s try to describe them all with a fifth-degree polynomial.
y(x) = k₀ + k₁x + k₂x² + k₃x³ + k₄x⁴ + k₅x⁵
Why the fifth degree? Just because we think this accuracy will be sufficient. At the same time, the number of coefficients (parameters) is quite small—and they are easy to find.
So, at the training stage, we need to find these very coefficients kᵢ so that the resulting function passes as close as possible to the points.
Below, on the left graph, we have chosen them quite well, and on the right, terribly—the error is very high (even the provocatively red lines say that everything is bad)—which means the function will be very bad at predicting the value of y at a certain point x.
Backpropagation
This very error is calculated: for each known point x, we can calculate the deviation of the predicted value y from the actually known one. We can add them all up and get a single number that characterizes our model (the chosen coefficients kᵢ)—whether it turned out bad or okay. This, by the way, is the only place where we can use the term “scalar”: the error is a scalar.
We can randomly change each of the kᵢ and see what the final error is.
And this will generally work as long as there are a few to a hundred coefficients. But the number of attempts to achieve success will grow exponentially with the increase in the number of coefficients.
And this was one of the strongest limiting factors in the development of neural networks until the mid-80s. When you have a model the size of several thousand parameters, and you don’t have a mathematically precise way to adjust each coefficient, you won’t be able to do anything good on a human timescale.
But, thanks to the trinity of scientists who in '86 showed the practical applicability of the backpropagation mechanism. They found a way to adjust all-all-all the coefficients based on the error, and each one separately, going from the end to the beginning of the input data processing pipeline—layer by layer.
That is, in our example, knowing the final error for a given set of coefficients, we can create a special vector with 5 values, each of which tells us how to change the corresponding coefficient to make the error a little smaller. We adjusted the coefficients, tried again, got a smaller error—which means we chose the vector values correctly in the previous step—let’s change the coefficients a little more in the same direction. And so on, many times.
Gradient Descent
We won’t dive into the math and formulas here. The essence is that we are at a certain error point in the multidimensional space of the model’s coefficients. And from this point, we perform a gradient descent to the minimum of the error function, calculating partial derivatives for each coefficient.
Okay, let’s start with an intuitive level.
Imagine you are a person… Oh, right. Imagine you are an error… No… That’s also bad.
Anyway, you are standing on the slope of some hill. You are cold. And you want to go down to the valley where it’s warmer. You see in which direction you need to take a step to go down. You take it. And you find yourself in a new point, where you again think about where to take a step. And so, step by step, all the way to the bottom.
For clarity, let’s take the linear regression function y(x) = k₀ + k₁x.
Here is an example of gradient descent for it:
Well, each such step you take is in the direction of the gradient—that is, the direction where the descent is steepest. And each of your steps is a change in the coordinates k₀ and k₁ on the ground plane—you are changing them. You are finding a combination of them that reduces the height of the point where you are standing—you are tuning the parameters, it turns out.
In our hill example, there are two parameters—k₀ and k₁. And the height, which depends on them. This is what we can represent as a 3D map. If there are 3 parameters, it will already be a four-dimensional space. There are problems with representing it, although some claim to visualize tesseracts at night.
And if we return to the original example with the 5th-degree polynomial, there are 6 parameters, which means the loss function will be in a 7-dimensional space. The essence is the same—we are a material point moving downwards in a multidimensional space.
The gradient, as we remember from calculus, is calculated as the derivative of a function. In the case of many variables, it will be a vector—a set of partial derivatives for each variable. That is, for each parameter. In our 3D example, the gradient is a vector of two numbers—roughly speaking, the slope along two axes. And when there are 6 axes (parameters), there will be 6 components in the vector.
Let’s try to “digitize” this a bit
So we have randomly chosen 6 parameters (coefficients). For example:
y(x) = 3 + 2x + 8.3x² + 4.1x³ + x⁴ + 0.23x⁵
We tried them for a certain value of x, for example 6.1. And it turns out that y = 4581.82, but it should be -0.22 (according to the graph). Let’s calculate the error using the simple Mean Squared Error formula: (4581.82 - (-0.22))²/34 = (4582.04)²/34 = 617,502.66. And at x=0, it should be -1.2, but we get 3. Error: 0.52. And at x=8, it should be -0.8, but we get 14282.04. Error: 5,999,985.83.
Going through all the known values like this, we calculate the sum: 617,502.66 + 0.52 + 5,999,985.83 + .... And so we have ended up somewhere very high on the error function graph (that is, we have given a numerical assessment of the quality of the current model). And now we need to reduce this error by choosing new parameter values. And it is precisely here that the partial derivatives of the resulting function are taken for each of the parameters.
It’s important to understand here that there is a function
yof the variablex(y = f(x, k₁, k₂, k₃, k₄, k₅)), which we are trying to obtain by selectingkᵢ—this is our neural network. And there is a pre-known error function of the variableskᵢ(Loss function, for exampleMSE = (1/n)×Σ(y_true - y_pred)², which we used above). It is from this function that we calculate the derivatives for each of thekᵢ.
And so we have found the derivatives, let’s say—for each parameter—we got some vector. And it is according to this vector that we change the parameters.
Another important point here is the step size of the descent. The larger it is, the faster we descend into the valley, that is, the faster we find the parameters and train the model. But with too large a step, you can overshoot the minimum—as if jumping to the other side of the hill.
After many iterations, we will find the optimal coefficients, and the model will be close to correctly predicting points both in known places and in unknown places within the studied range.
So, what does all this mean for our real model like GPT-3? It’s actually all the same, only the error function is defined in a 175-billion-dimensional space, that is, where there are 175 billion axes, not 3. So the whole complexity here is only in how to calculate such giant numbers, but the operations themselves are as simple as when we had two coefficients in linear regression.
Here is the sequence of actions that takes place:
The model performed a forward pass and predicted the probabilities of each word from the vocabulary (50,257 for GPT-3) appearing as the next one. That is, this is a probability vector with the dimensionality of the vocabulary, the sum of the numbers in which is 1. At the same time, the model knows which specific word is expected for this example (since we specially sliced and prepared the training data). And it prepares a vector of the same dimensionality, where there will be all 0s and one 1 at the position corresponding to the expected word—this is called a
one-hotvector. And then the difference between the predicted vector and the real one is calculated.Loss function in detail
Above, we used Mean Squared Error (MSE) to calculate the error because we were solving a regression problem—predicting a number. And it is excellent for continuous data: amount of precipitation, price per square meter, road congestion, etc. In transformers, we are solving a classification problem: which next word from the list to choose. And here MSE performs poorly, and Cross-Entropy Loss is used instead.
The one-hot vector is element-wise multiplied by the logarithm of the prediction vector, all its values are summed, and a single number is obtained.
L = -Σ(y_i × log(ŷ_i)) Where - y_i — is the true label (0 or 1) for class i (One-hot) - ŷ_i — is the predicted probability of class i - Σ — is the sum over all classes (tokens) from i=1 to CBut, since one-hot is a vector with one 1 and all other zeros, the formula simplifies and looks like this:
L = -log(ŷ_correct) Where - ŷ_correct — is the probability of the correct class.And since we are not calculating on a single example, but on a batch, we average the error over all examples. Like this:
L_batch = -(1/N) × Σ(log(ŷ_j_correct))This difference is the error. It tells us where we are. If we go back to the mountain example, we can say that the error tells us at what altitude. The higher we’ve climbed, the worse. And we need to start descending, into the valley. Into the valley of the smallest error.
To understand in which direction and of what length to take a step, we need the correction matrices (gradients), which will reduce the error. If the error tells us at what altitude we are, then the gradient tells us in which direction to move to descend. A gradient is a tensor of the same shape as the parameters, containing for each parameter of a given layer the number by which it needs to be changed.
Gradients are calculated sequentially through the entire model, layer by layer. This is the backward pass or backpropagation.
In the final step, the changes are applied—all the model’s weights are updated.
So, the model performed a forward pass, calculated the word probabilities, found the error, performed a backward pass, calculating the gradients, and updated all the weights—this is one training iteration. And we repeat.
After hundreds of thousands, or even millions, of such forward and backward passes for the entire dataset, the billions of weights stabilize, and the model is considered trained.
It is important that only the weights of the model’s matrices and vectors are saved as parameters: FFN, normalization, bias, and the attention matrices
Wq,Wk,Wv. The activation vectors of the layers,Q,K,V, are NOT saved because they are unique to each example.
Training data
There should be a sufficient number of training examples. For example, modern LLMs have seen almost all digitized and publicly available text information—this amounts to hundreds of billions of tokens.
Usually, not one example and not the entire dataset is processed, but a batch—a specific piece of training data. Processing one example at a time is, first, super slow, and second, the gradients turn out to be very noisy, and the model trains either too slowly or incorrectly. Processing the entire dataset is very expensive, and, rather, even impossible—it’s impossible to stuff all the training data into the cluster’s memory at once.
In other words, the model’s weights are updated once for the entire batch based on the averaged gradient over all examples. The batch size can be chosen simply by the length of the tokens. During the training of GPT-3, for example, the batch size was 3.2 million tokens. The larger the batch, the more stable the gradient. More stable means it’s not thrown around by each new example from the heat of cold numbers to the gift of divine visions.
In the case of LLMs, on average, each training example passes through the model only once. But high-quality data can be run through several times. At the same time, during the Fine-tuning stage, some examples from additional datasets can be run through several times.
Model size and quality, MoE
Obviously, there is a connection between the number of parameters and the quality of the model’s results. It seems that the more parameters, the better the answers. In general, the number of parameters does indeed correlate with the volume of training data. But this is only true up to a certain point, until the model becomes saturated.
If the number of parameters is significantly larger, then, firstly, “dead neurons” may appear, which do not affect anything but participate in calculations—we are just burning electricity, but not hitting the mark. And, secondly, what is scarier, the model can “overfit”—it starts not to look for patterns, but to “memorize examples,” it starts to remember noise, and the quality of the answer drops because of this.
Research and experiments show that sparse activation can help with this. It assumes that for specific examples/requests, not all neurons participate in the calculations—only a part.
One of the methods of sparse activation is MoE—Mixture of Experts. Inside the neural network, groups of neurons appear that specialize in a certain topic—experts. It could be like this:
programming,
science,
fiction,
history,
recipes
etc.
A router is placed before the experts, which determines which of them will best process the request.
But in reality, the specialization of experts is not so clear-cut, and, quite possibly, not even interpretable by humans, because it happens during the training process. That is (important point): unlike inference, during training, all experts are used, as well as the router, because it also needs to be trained.
And then, during inference, only the part of the parameters that is needed to process requests on that topic is used. Thus, we increase the number of parameters and the quality of the model without complicating the calculations.
Let’s repeat that the specialization of experts may not be understandable to a human—they do not determine it. It is not necessarily based on semantic topics. Sometimes experts specialize in positions in a sequence. Or on the frequency of tokens (rare vs. common words).
How is this implemented technically? In fact, it’s almost the same as for a regular model, with two differences:
Instead of one pair of FFN matrices in each block, several pairs of matrices appear. These are the experts.
A router is added, which chooses the path to the required expert.
The size of each new FFN matrix is the same as it was for a regular model, the function is the same: increase the dimensionality, apply a non-linear function, return the dimensionality. Thus, we have greatly increased the number of parameters, but since only one or two experts are used for each example, the complexity of the calculations has not increased significantly.
Experts may not be in every block, but only in some. Blocks without experts are dense, with them—sparse. How to arrange them is determined by the model’s architecture. The number of experts is a hyperparameter of the model.
And why did we go through all this in such terrifying detail? For what, as every other manager says? Well, it’s time to hit all this against reality.
Simple arithmetic
So, for GPT-3, we get:
96 transformer layers + embedding + output
~175 billion parameters
The parameters take up 350GB of memory (2 bytes for each (FP16))
That is, the trained model itself weighs 350 GB. Inference will require much more, because you need to store intermediate data: Q, K, V vectors (or the KV cache), activation vectors, temporary calculations during matrix multiplication, etc.
For training, you need several more terabytes of memory. This is the training batch itself (hundreds of gigabytes), intermediate layer activations (hundreds of gigabytes), gradients (350 GB, same as the number of parameters), even the parameter update itself requires an additional 700 GB, and there are also Master weights of the same 700 GB. And so on.
Now let’s estimate with modern hardware. Let’s calculate: an NVIDIA H100 model contains 80GB of memory. That is, one card cannot even hold all the model’s parameters. You need at least 5 cards just to fit the GPT-3 model with all its parameters. In fact, more.
And for its training, you need about 3-6 TB of memory, and therefore at least 40 H100 cards. But this is the minimum, just to fit all the numbers. For effective training, you need several times more—thousands of cards, to get it done in months. To understand the scale. One training iteration (forward pass, backward pass, and parameter update) can take a couple of minutes. So we calculate 100,000 examples at 2 minutes each—that’s 4.6 months.
For comparison, GPT-4 is rumored to contain 1.8 trillion parameters, which will require several tens of thousands of cards.
Even the newest NVIDIA B300 has 288GB of memory on board the GPU.
Therefore, the only way out for large models, both LLMs and almost any others, is parallelization.
Parallelization
Now, armed with all this knowledge, we can already make the first assumptions about the parallelization of training and inference, that is, about how to distribute the task among several cards.
In general, simple transformers, convolutional networks, and multi-layer perceptrons can be both used and trained on a single card with a few gigabytes of memory. But you won’t get impressive results with this—recognition of simple texts and small, simple images, generation of syntactically correct but meaningless text.
Everything interesting, like LLMs or large convolutional networks for creating images, requires hundreds and thousands of GPU cards for training.
And what to do about it? Globally, there are two approaches:
distribute the model itself across several different GPUs
parallelize the training data: different instances of the model are trained on different data.
Usually, they are combined.
For designing the network and the cluster, how the training is parallelized is a key moment—will the cards bought with the last of the money work at 100% of the investment or will they be idle most of the time due to a slow network or memory.
Let’s figure it out.
Types of parallelization
Data Parallelism — we slice a batch of examples into smaller batches.
Sequence Parallelism — we slice one long example into several shorter pieces.
Model Parallelism — we slice the model.
Pipeline Parallelism — by layers or blocks.
Tensor Parallelism — by columns/rows of matrices.
Expert Parallelism — by experts in MoE.
Data Parallelism (DP)
We slice a batch of examples into smaller batches.
Or data into batches of data.
For example, we have 10 GPUs and 500 training examples. We divide them into 10 groups of 50 examples and distribute them to our 10 cards.
That is, essentially—we just simultaneously run several copies of the model on different GPUs or groups of GPUs and send different training data to them.
And each will have its own weights in the linear matrices and bias vectors as a result. But we need to get not 10 models, each trained on its own examples, but one, trained on all of them.
For this, the AllReduce mechanism is used. Let’s break down how it works.
At the very beginning, 10 copies of the model are created with identical weights and hyperparameters.
Each copy performs a forward pass, prediction, calculates the error and gradients on its own.
The copies, using
AllReduce, exchange their gradients so that in the end, all 10 copies know all 10 gradients.Each copy independently calculates the average value of the gradient—the same for all.
And with this averaged gradient, it updates the weights of the matrices—the weights on all copies become identical.
Thus, at each training cycle, the matrices of all copies have the same values.
AllReducecollects matrices of the same dimensionality from all copies (All) and makes one from them (Reduce) by averaging the values.
In fact, in some cases (BatchNorm), it is also necessary to synchronize the normalization statistics on the forward pass, so that there are consistent values for normalization throughout the entire training iteration on the batch. And in these cases, it is also necessary to synchronize the gradients for correcting the normalization components with the same AllReduce.
This type of parallelization has a medium sensitivity to latency and bandwidth. It is recommended to localize it within a rack.
Sequence Parallelism (SP)
We slice one long example into several shorter pieces and distribute them to different GPUs.
What could be an example of such training data that does not fit on one card but must be processed together?
Genomic sequences
Climate data
Scientific articles with the full context of the literature
We have 10 GPUs with copies of the model and 1 long training example of 1,000,000 tokens. We divide it into 10 sequences of 100,000 tokens and distribute them to our 10 cards.
As in DP, we break the training data into separate parts. But what’s the difference? When the model processes the given text, each word in it “sees” all the other words. So, in the phrase “The quick brown fox jumps over the lazy dog”, the word “jumps” sees all the others: “fox”, “over”, “dog”, and even the period. The word “fox” does the same. And so on. This allows for the composition of the overall meaning of the phrase.
But when we split the phrase between two cards:
“The quick brown fox”
“jumps over the lazy dog”
“Fox” cannot see the word “jumps”. And this is critically important for understanding the meaning.
It would be logical to synchronize the attention vectors. However, this is super expensive: first, all vectors need to be run through the network, and second, the amount of memory needed in the end is the same. And we don’t gain any efficiency. Therefore, the context at the boundaries of the sequence parts is preserved by having these parts overlap slightly, that is, they contain common fragments. And the overall attention between the parts of the sequence is achieved at a higher level; such attention is achieved by synchronizing the activation vectors.
That is, for SP, you need to synchronize:
gradients once per batch
intermediate activations
during inference, you also need to synchronize the KV cache (this is not necessary during training)
How synchronization happens in detail
Activation vector synchronization
As a reminder: this is the result of a layer or block’s operation, which is passed to the next layer/block. Or in other words—it’s what the initial embedding vector turns into after passing through all the previous layers.
Accordingly, on each GPU, a vector is obtained that contains pieces calculated for all other GPUs. And the task is to exchange them—all with all.
That is, with 4 GPUs. The first card has information for 1, 2, 3, and 4, and the 2nd has for 1, 2, 3, 4 (for others and for itself). And at the end of the synchronization, it should have all the pieces intended for it from 2, 3, and 4. The second should have all its pieces, etc.
The AlltoAll mechanism implements this, allowing GPUs to exchange the necessary data. Each GPU sends to others only those parts of the matrix that they need. Specifically: GPU0 to GPU1 only the piece for GPU1, GPU0 to GPU2 only the piece for GPU2. And so on.
AlltoAll A real-world analogy can be drawn: three friends returned from their travels and brought gifts for each other (and didn’t forget about themselves):
Vasia: [For Vasia, For Petia, For Masha] Petia: [For Vasia, For Petia, For Masha] Masha: [For Vasia, For Petia, For Masha]
In the AlltoAll process, Vasia keeps his gift for himself, gives the gift for Petia to Petia, and the one for Masha to Masha. Petia gives to Vasia and Masha separately, and Masha gives to Vasia and Petia. Everyone has exchanged.
The result is:
Vasia: [For Vasia, For Vasia, For Vasia] Petia: [For Petia, For Petia, For Petia] Masha: [For Masha, For Masha, For Masha]
Gradient synchronization
We’ve already been through this—transmitting gradients throughout the network and averaging them. AllReduce is used.
KV cache synchronization
This is not performed during training because the task there is to predict only one word, and there is no need to run the entire phrase with the answer through again.
During inference, for KV cache synchronization, pieces of the cache are transmitted from each GPU to other GPUs at each step, where they are concatenated into the target matrix. The AllGather mechanism is used for this. It is similar to AllReduce, but the result is not a matrix of averaged values, but a large matrix composed of many small ones (Gather). However, precisely because they need to be synchronized many times on each pass, SP is rarely used for inference.
SP is recommended to be localized as closely as possible—optimally within a server, but it can also be within a rack if NVLink or InfiniBand is available.
Model: Pipeline Parallelism (PP)
We slice the model by layers or blocks
This is a fairly simple type of parallelism. It is used when the entire model does not fit into the memory of a single GPU card. The model is sliced by layers (blocks), and each card is responsible for processing an example on its own layers. For example:
GPU0 computes layers 1-32
GPU1 — 33-64
GPU2 — 65-96
During the forward pass, after GPU0 has finished its work, it sends its activation matrix to GPU1, and after it has finished, to GPU2. During the backward pass, gradients are propagated similarly.
No synchronization is needed.
It’s like a conveyor belt in a factory: each section is responsible for its own operation. When it’s done, it passes it on and takes a new one.
To prevent two GPUs from being idle while one is working, the list of examples is divided into so-called micro-batches. And this creates a pipeline—while one example is being processed on GPU1, the next one is fed to GPU0. When GPU1 is finished, it passes the first example to GPU2 and takes the second example. And so on down the pipeline.
PP is the least sensitive to bandwidth and latency and can be localized anywhere.
Model: Tensor Parallelism
we slice the model by columns/rows of matrices.
Here, the matrices in the model are sliced—by columns or rows. That is, one layer is distributed across several GPUs—each has its own part of the weights.
A simple analogy is a multiplication table of 1000x1000 numbers. We give the first 250 columns to GPU0, the next 250 to GPU1, then to GPU3 and 4—and each calculates its own batch. And then we combine everything back into one table.
In the case of GPT-3, in the FFN block, for example, we need to multiply vector X by matrix W. Its dimensionality is [12288×49152]. We can parallelize it like this:
GPU0 stores columns 0-12287 GPU1 stores columns 12288-24575 GPU2 stores columns 24576-36863 GPU3 stores columns 36864-49151
Each GPU multiplies its own part, and then the results are added back into the table. Parallelized!
During training in this case, you need to additionally exchange:
attention matrices
activation vectors
gradients
Transformer Block: │ ├─ LayerNorm (local) ├─ Multi-Head Attention │ ├─ QKV projection (column-parallel) ❌ │ ├─ Attention compute (local) ❌ │ └─ Output projection (row-parallel) 🔄 All-Reduce #1 │ ├─ Residual + LayerNorm (local) ├─ FFN │ ├─ Up-projection (column-parallel) ❌ │ ├─ GeLU (local) ❌ │ └─ Down-projection (row-parallel) 🔄 All-Reduce #2 │ └─ Residual connections (local) Forward: 2 All-Reduce Backward: 2 All-Reduce Total: 4 All-Reduce per layer
We have already seen how this happens above.
TP is extremely sensitive to bandwidth and latency, so it is strictly localized within a server.
Model: Expert Parallelism
we slice the model by experts in MoE
In the case of large models, we can parallelize the experts—that is, the processing through FFN matrices: GPU0 for physics and mathematics, GPU1 for programming and databases, songs and dances—on GPU3, etc.
Since both the input tokens and the experts can be distributed across different devices, it is necessary for some GPUs to first distribute all the pieces to all the target GPUs, and then for the latter to gather everything back for normalization and activation—AlltoAll in both cases.
Bandwidth and latency requirements are medium. It’s better to localize it as close as possible, but within a rack is acceptable, with caveats—and within a data center.
Let’s summarize the parallelization methods
First, by parallelizing, we pursue two goals:
Distribute a model that does not fit on a single GPU card.
Speed up training and inference.
Second, there are several approaches to parallelization, each with its own features. Usually, several are used at once - Hybrid Parallelism.
Third, almost any method involves intensive data exchange at many stages. And if you use naive approaches here, you can spend all your time waiting for the network. Therefore, we will now consider how synchronization methods actually work.
Synchronization methods
Historically, in computing clusters, special central nodes (masters) helped to calculate and distribute data to other nodes. An example is Apache Hadoop.
But this obvious approach has obvious disadvantages:
Bottleneck Imagine all nodes sending their gradients to one? And then it sends them back.
Poor scalability For the same reason.
Single point of failure If the master node goes down, the entire cluster stops.
Increasing latency All nodes in the cluster must send data to the central node and receive a response. Processing on it happens sequentially, not in parallel. In addition, a queue forms on the master node because its resources (CPU, memory, network) are limited. As a result, the overall efficiency of the cluster drops.
All modern clusters for HPC, AI/ML use collective operations, where all calculations happen in a distributed manner, that is, each node transmits to the network what others need, and calculates for itself what it needs.
Above, we encountered different mechanisms that ensure the calculation of parameters that need to be synchronized:
Now let’s look in more detail at what each of them does and remember where it is used.
AllReduce
Gathers tensors from all GPUs, applies a reduction operation (e.g., sum or average), and distributes the result to all GPUs. It’s called Reduce because it takes a large number of matrices of the same size and makes one out of them—it reduces their number to one by applying an operation (sum, average, etc.).
An example of use is gradients for updating parameters in Data Parallelism.
AllGather
Gathers tensors from all GPUs and composes one mega-tensor from them, that is, it concatenates many small ones into one large one. As a result, each GPU has not just the one piece of the tensor it calculated, but one assembled from all GPUs and identical on all of them.
An example of use is attention matrices in Tensor Parallelism.
AlltoAll
All to all. Each GPU has calculated several pieces of a tensor, and each piece is intended for a specific GPU in the cluster. And now these pieces are distributed among the cards—each assembles its own puzzle. GPU 0 has data for 0, 1, 2, 3. GPU1—for 0, 1, 2, 3, etc.
As a result, each GPU gathers all its pieces necessary for further work. On GPU0, everything intended specifically for it, on GPU1—for it, etc.
An example of use is intermediate activations in Tensor Parallelism.
Reduce-Scatter
Similar to AllReduce—data is gathered together, reduced (e.g., averaged), but then each GPU is left with only the part of the tensor that it needs. That is, first we do Reduce (aggregate data), then Scatter (distribute fragments of the result).
An example of use is gradients in Data Parallelism—each GPU has only its own piece of the weight matrix (and activation, and normalization), and there is no need to store the entire huge gradient, eating up memory.
Here, if we go back to AllReduce, it becomes clear that it is ReduceScatter + AllGather:
Broadcast
This is the only operation that assumes the presence of a root node. And this node may not participate in the calculations—it centrally distributes things common to the entire cluster, for example, hyperparameters, the embedding matrix, the initial weight matrix. It also sends control signals, such as “stop training”.
For example, Broadcast is used when initializing a cluster, when all GPUs must be in a consistent state. But not only then.
All these operations are at a fairly high level of abstraction, not grounded in the real world and the actual network topology. And now we hit reality again.
For all this to work, you need to know how the network is structured, and on top of it, create a virtual topology, and on top of this topology, create a data exchange program, understand how to work with GPU memory, how to manage CUDA. Not easy? Not easy!
If I were a cluster or model admin, I wouldn’t want to manually write all these exchange schemes, to be honest (I’m not an Olympiad programmer).
That’s why humanity invented collective libraries.
Collective libraries
These are products that implement collective operations, taking into account the real network topology, optimizing algorithms for it, and hiding all the complexity from the user.
The three most common Open Source libraries are:
NCCL — Nvidia Collective Communication Library
MPI — Message Passing Interface
Gloo — a library from Meta
What do these engines do?
They are responsible for working with CUDA.
They manage GPU memory.
They implement low-level interaction with network cards and the network. They know what NVLink, RDMA, and InfiniBand are.
They study the network topology (Fat Tree, DragonFly, HyperCube, Torus, etc.).
They create the necessary virtual topology for data exchange, optimized for the real one.
CUDA you sent me where?!
CUDA — Compute Unified Device Architecture — an SDK from NVIDIA for GPUs, allowing them to be used not only for displaying images but also for simply ordering computations. That is, the CPU gives a task to the GPU via CUDA. The GPU multiplies matrices super-fast, while overall control over the tasks remains with the CPU.
Virtual topologies
Physically, we can have Clos, DragonFly, Flattened Butterfly, Torus, whatever, but on top of it, collective libraries can build virtual topologies. These are schemes of virtual connections between cards.
We will not go into the details of their structure and where they are used, but we will have to look at them at least briefly.
So, there are not that many topologies:
Ring
Tree
Bruck’s Algorithm
Recursive Doubling / Halving
PAT (Pipeline AlltoAll Topology)
Let's look at an example of a Ring to see what it looks like
So, we might have a physical Dragonfly topology. Each Leaf is connected to 2 servers with 8 GPU cards, each with its own network interface.
And a collective library can arrange all these cards in a logical ring, as if stringing them all on one rope.
Or it can assemble them into a tree.
The collective library explores the network, determines the topology, technologies (NVLink, InfiniBand, RoCE), measures latencies, checks what data of what volume needs to be transmitted and what operation to use (AllReduce, etc.), and based on this, chooses the connection scheme.
It’s important to understand here that it dynamically chooses the connection scheme based on what operation needs to be performed on what data right now.
For example, for AllReduce, it might choose a Ring topology. Let’s take gradient synchronization as an example.
AllReduce in detail as an example of a collective operation
All nodes are assembled into a virtual ring, meaning each node has exactly two neighbors. And the last one connects to the first. Next, each GPU slices its matrices into N chunks—one for each node.
Then it’s like a step mechanism: Click, the wheel turns a little: each node passes one chunk to its neighbor. Everyone sums the received chunks with what they have. Click, the wheel turns a little: each node passes another chunk to its neighbor. Everyone sums the received chunks with what they have. Click…
So the numbers “travel” around the ring, accumulating sums, until they land on the worker located one position before where they started!
But let’s look even more closely: in reality, everything is more complicated than it seems.
In fact, AllReduce consists of two phases:
SCATTER-REDUCEFirst, the data is sliced into N chunks (by the number of nodes), then it is passed around the ring and summed. After N-1 iterations, each node will have its own fully summed chunk.AllGatherThen these summed chunks are distributed to all nodes in the ring.
Formally, the process looks like this: SCATTER-REDUCE Phase Formulas for each node i:
Sends chunk at position:
(i - k + 1) mod NReceives from worker:
(i - 1) mod NAdds to its own position:
(i - k) mod N
AllGather Phase
Each node
isends a chunk of data to node(i+1) mod NEach node
ireceives from node(i-1) mod Nand writes it to itselfNo addition! Just copy the received value.
Complicated? Complicated!
Let’s consider an example with a simple topology of 4 GPUs. So N=4. They can be connected like this:
The library assembles them into a ring:
After the training stage, each node has a calculated gradient for updating the weights. Right now, this gradient is different on each node. On node A: a On node B: b On node C: c On node D: d
We need the gradient to be the same on all of them. A simple naive approach: everyone sends their piece to everyone else. But this is super-inefficient, so they do it this way.
Each node divides its gradient into N chunks (SCATTER)
A: [a₀, a₁, a₂, a₃] B: [b₀, b₁, b₂, b₃] C: [c₀, c₁, c₂, c₃] D: [d₀, d₁, d₂, d₃]
As a result, we want to get:
A: a+b+c+d B: a+b+c+d C: a+b+c+d D: a+b+c+d
Or:
A: [a₀+b₀+c₀+d₀, a₁+b₁+c₁+d₁, a₂+b₂+c₂+d₂, a₃+b₃+c₃+d₃] B: [a₀+b₀+c₀+d₀, a₁+b₁+c₁+d₁, a₂+b₂+c₂+d₂, a₃+b₃+c₃+d₃] C: [a₀+b₀+c₀+d₀, a₁+b₁+c₁+d₁, a₂+b₂+c₂+d₂, a₃+b₃+c₃+d₃] D: [a₀+b₀+c₀+d₀, a₁+b₁+c₁+d₁, a₂+b₂+c₂+d₂, a₃+b₃+c₃+d₃]
We start the REDUCE. This will require N-1 iterations. In our case, 3.
SCATTER-REDUCE Phase
Step 1 At step k=1, according to the algorithm above: Node i sends chunk at position: (i - 1 + 1) mod 4 or i mod 4 Adds what it received to its own position: (i - 1) mod 4
A sends
a₀to B A receivesd₃from D, sums it witha₃, and saves it at position 3A: [a₀, a₁, a₂, a₃+d₃]B sends
b₁to C B receivesa₀from A, sums it withb₀, and saves it at position 0B: [b₀+a₀, b₁, b₂, b₃]C sends
c₂to D C receivesb₁from B, sums it withc₁, and saves it at position 1C: [c₀, c₁+b₁, c₂, c₃]D sends
d₃to A D receivesc₂from C, sums it withd₂, and saves it at position 2D: [d₀, d₁, d₂+c₂, d₃]
And we get:
A: [a₀, a₁, a₂, a₃+d₃] B: [b₀+a₀, b₁, b₂, b₃ ] C: [c₀, c₁+b₁, c₂, c₃ ] D: [d₀, d₁, d₂+c₂, d₃ ]
Step 2
According to the algorithm above, at k=2:
Node i sends: (i - 2 + 1) mod 4 = (i - 1) mod 4 Adds what it received to its own position: (i - 2) mod 4
That is:
A sends position 3 to B —
a₃+d₃((0-2+1) mod 4 = 3) A receivesd₂+c₂from D, sums it into position 2 ((0-2) mod 4), and we geta₂ + (d₂+c₂)A: [a₀, a₁, a₂ + (d₂+c₂), a₃+d₃]B sends position 0 to C —
a₀+b₀((1-2+1) mod 4) B receivesa₃+d₃from A, sums it into position 3 ((1-2) mod 4), and we getb₃ + (a₃+d₃)B: [b₀+a₀, b₁, b₂, b₃ + (a₃+d₃)]C sends position 1 to D —
c₁+b₁((2-2+1) mod 4) C receivesa₀+b₀from B, sums it into position 0 ((2-2) mod 4), and we getc₀ + (a₀+b₀)C: [c₀ + (a₀+b₀), c₁+b₁, c₂, c₃]D sends position 2 to A —
d₂+c₂((3-2+1) mod 4) D receivesc₁+b₁from C, sums it into position 1 ((3-2) mod 4), and we getd₁ + (c₁+b₁)D: [d₀, d₁+(c₁+b₁), d₂ +c₂, d₃]
And we get:
A: [a₀, a₁, a₂ + (d₂+c₂), a₃+d₃] B: [b₀+a₀, b₁, b₂, b₃ + (a₃+d₃)] C: [c₀ + (a₀+b₀), c₁+b₁, c₂, c₃] D: [d₀, d₁+(c₁+b₁), d₂ +c₂, d₃]
Step 3
According to the algorithm above, at k=3:
Node i sends (i - 3 + 1) mod 4 = (i - 2) mod 4 Adds what it received to its own position: (i - 3) mod 4
That is:
A sends position 2 to B —
a₂+d₂+c₂((0-2) mod 4) A receivesd₁+c₁+b₁from D, sums it into position 1 ((0-3) mod 4), and we geta₁ + (d₁+c₁+b₁)A: [a₀, a₁+(d₁+c₁+b₁), a₂ + (d₂+c₂), a₃+d₃]B sends position 3 to C —
b₃+a₃+d₃((1-2) mod 4) B receivesa₂+d₂+c₂from A, sums it into position 2 ((1-3) mod 4), and we getb₂ + (a₂+d₂+c₂)B: [b₀+a₀, b₁, b₂ + (a₂+d₂+c₂), b₃ + (a₃+d₃)]C sends position 0 to D —
c₀+a₀+b₀((2-2) mod 4) C receivesb₃+a₃+d₃from B, sums it into position 3 ((2-3) mod 4), and we getc₃ + (b₃+a₃+d₃)C: [c₀ + (a₀+b₀), c₁+b₁, c₂, c₃ + (b₃+a₃+d₃)]D sends position 1 to A —
d₁+c₁+b₁((3-2) mod 4) D receivesc₀+a₀+b₀from C, sums it into position 0 ((3-3) mod 4), and we getd₀ + (c₀+a₀+b₀)D: [d₀ + (c₀+a₀+b₀), d₁+(c₁+b₁), d₂ +c₂, d₃]
And we get:
A: [a₀, a₁+d₁+c₁+b₁, a₂+d₂+c₂, a₃+d₃] B: [b₀+a₀, b₁, b₂+a₂+d₂+c₂, b₃+a₃+d₃] C: [c₀+a₀+b₀, c₁+b₁, c₂, c₃+b₃+a₃+d₃] D: [d₀+c₀+a₀+b₀, d₁+c₁+b₁, d₂ +c₂, d₃]
On each node, we have the full sum of all components of one of the four parts. And note, not necessarily its own parts, but just some parts. So A (node zero) has collected the entire first part, and D (node three) has the entire zeroth part.
And we need to arrive at a picture where all nodes have the same matrix. So we move on to the AllGather stage.
AllGather Phase
It will also consist of N-1 iterations. In our example, 3.
Each node i sends a chunk to node (i+1) mod N Each node i receives a chunk from node (i-1) mod N No addition! Just copy the received value.
It’s important that each node knows at which position it has the final sum—that’s what it will be sending.
Step 1 A sends position 1 to B — a₁+d₁+c₁+b₁. B sends position 2 to C — b₂+a₂+d₂+c₂ C sends 3 to D — c₃+b₃+a₃+d₃ D sends 0 to A — d₀+c₀+a₀+b₀
And we get:
A: [d₀+c₀+a₀+b₀, a₁+d₁+c₁+b₁, a₂+d₂+c₂, a₃+d₃] B: [b₀+a₀, a₁+d₁+c₁+b₁, b₂+a₂+d₂+c₂, b₃+a₃+d₃] C: [c₀+a₀+b₀, c₁+b₁, b₂+a₂+d₂+c₂, c₃+b₃+a₃+d₃] D: [d₀+c₀+a₀+b₀, d₁+c₁+b₁, d₂ +c₂, c₃+b₃+a₃+d₃]
Step 2 Each node sends the last sum it received. A sends position 0 to B B sends 1 to C C sends 2 to D D sends 3 to A
And we get:
A: [d₀+c₀+a₀+b₀, a₁+d₁+c₁+b₁, a₂+d₂+c₂, c₃+b₃+a₃+d₃] B: [d₀+c₀+a₀+b₀, a₁+d₁+c₁+b₁, b₂+a₂+d₂+c₂, b₃+a₃+d₃] C: [c₀+a₀+b₀, a₁+d₁+c₁+b₁, b₂+a₂+d₂+c₂, c₃+b₃+a₃+d₃] D: [d₀+c₀+a₀+b₀, d₁+c₁+b₁, b₂+a₂+d₂+c₂, c₃+b₃+a₃+d₃]
Step 3 One last iteration remains.
A sends position 3 to B B sends 0 to C C sends 1 to D D sends 2 to A
And we get:
A: [d₀+c₀+a₀+b₀, a₁+d₁+c₁+b₁, b₂+a₂+d₂+c₂, c₃+b₃+a₃+d₃] B: [d₀+c₀+a₀+b₀, a₁+d₁+c₁+b₁, b₂+a₂+d₂+c₂, c₃+b₃+a₃+d₃] C: [d₀+c₀+a₀+b₀, a₁+d₁+c₁+b₁, b₂+a₂+d₂+c₂, c₃+b₃+a₃+d₃] D: [d₀+c₀+a₀+b₀, a₁+d₁+c₁+b₁, b₂+a₂+d₂+c₂, c₃+b₃+a₃+d₃]
Synchronized!
Why so complicated?
Why not just pass the entire matrix from each node around the circle, summing the values along the way? You can. That would be called Broadcast. But in that case, we need to multiply the matrix size M by the number of nodes N and by the number of transmissions (N - 1)—totaling M * N * (N - 1). Or O(MN²). That’s a bit much—it depends on the square of the number of nodes. The larger the cluster, the quadratically more we pay.
In contrast, Ring AllReduce gives a total network traffic volume of O(2MN). How so? We still have the same matrix size M, but each node slices it into N parts and transmits only that part N-1 times. And this is done twice—during the SCATTER-REDUCE and AllGather stages. This gives 2 * M/N * (N - 1) for each node. And we multiply by N nodes for the total volume. And so we get 2 * M * (N - 1).
On approximate numbers: Let the memory size for model parameters be 350GB, and Data Parallelism is set up for 64 nodes. For Broadcast, we get 350GB * 64 * 63 = 1,411,200 GB For Ring: 2 * 350 * 63 = 44,100.
“Wait a minute,” the curious reader will say. Why use Broadcast and Full-mesh here when A can send the entire matrix to B, B can sum them and send to C, which sums them and sends further, etc.? And the total volume will be M * (N - 1). Well, first, it won’t work head-on—at iteration N-1, the full sum will only be on node A, and how to get it to the other nodes is something you still need to figure out. And second, the problem is that these are sequential operations. While A is transmitting, all the others are waiting. If the transmission time for the full matrix is T, then the full transmission will take at least N * T. In the case of Ring AllReduce, it will be 2 * (T / N) * (N - 1). 2 passes, where 1 / N of the matrix is transmitted N-1 times. This comes out to about 2T, which is constant time. And transmission time… oh, how important it is when working on expensive GPUs.
We will not discuss other mechanisms in this article. I just wanted to show how parallelization works and what tricks have to be used to minimize network costs and downtime.
NCCL, MPI, and Gloo are low-level libraries that allow developers not to think about how to implement collective operations on hundreds and thousands of cards. But working with them from, say, Python is very difficult—you need to implement all the nuances of interacting with the library. NCCL and its comrades know nothing about neural networks, their layers, or gradients; they don’t know how and when to run forward and backward passes. They work directly with raw GPU buffers. NCCL only supports NVIDIA GPUs, while MPI and Gloo are cross-platform.
If you are ready to dive into all this, manage memory, and write thousands of lines of C code, then you should go here or here, but if not, then frameworks were invented for you.
Frameworks
They are like HTTP servers. You can work directly with nginx from Python, or you can use flask/django. The second way is clearly more convenient. The most famous and sufficient in 90% of cases is PyTorch. It provides a very simple interface for interacting with NCCL, hiding all the complexities of implementing model training. For comparison, the simple task of calculating gradients on a bunch of GPUs takes 5 lines in PyTorch:
import torch import torch.distributed as dist dist.init_process_group("nccl", rank=0, world_size=2) tensor = torch.tensor([1.0, 2.0, 3.0]).cuda() dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
When using NCCL directly, you need to write about 100 lines of C code, where you manually operate with GPU cards and their memory, load data into the GPU, create a CUDA stream, read data from GPU memory via pointers, and return it to the CPU. And this is just one simple stage in training—gradient synchronization—it’s not even the weight update.
In conclusion of the section on parallelization, it is important to note that for modern large models, only one type of parallelization is almost never used:
A single model does not fit on one GPU.
It is more profitable to calculate billions of training examples simultaneously on different groups.
Very large input examples may also not fit into the memory of a single card.
During inference, a response is needed on a human timescale of a few seconds.
Therefore, these are always hybrid parallelization schemes, for example, pipeline+tensor+data. That is, the model is sliced by layers, but even so, it does not completely fit into memory, so the matrices are divided by rows or columns and distributed among groups of cards. Several such instances of the model are launched, and different training examples are processed on them.
Now imagine the complexity of communications in such clusters, when each type of parallelism requires synchronization of its parameters at the right moments in time. And each of them is sensitive to delivery time and losses in its own way. It is very important what network technologies are used and where.
Network technologies in a cluster
With this paragraph, we are jumping ahead to one of the next articles. But it is here that it is appropriate to make the connection between abstract matrices and real hardware.
In ML and HPC clusters, there are 4 communication segments with fundamentally different characteristics.
Inside the GPU. Chip ↔ HBM memory channel
Inside the server. GPU ↔ NVLink ↔ GPU.
Inside the rack. Server(GPU) ↔ InfiniBand/RoCE ↔ Server(GPU)
Between racks. Server(GPU) ↔ InfiniBand/RoCE ↔ Server(GPU)
1. Inside the GPU
Modern GPUs are extremely complex devices where a relatively small GPU and a huge battery of memory chips are located on a single base (interposer). This is usually HBM—High Bandwidth Memory.
It is in this memory that the model parameters, all intermediate calculations (activations, Q, K, V, etc.), and the KV cache are stored.
And here is the first network segment—between the GPU and HBM. Speed 2-3 TB/s. Latency ~100ns.
Technically, here, as in a CPU, there is an L1 cache inside the GPU itself (200TB/s per core) and an L2 cache (10TB/s per card), and only then does the voluminous and slightly slower HBM appear.
You can read a little about their practical application here.
2. Inside the server
In ML clusters, each server has several GPU cards (often 8), and a special ultra-fast network is used between them.
In general, this could be a simple PCI-E, but today it is most often NVLink. NVLink is a proprietary technology from Nvidia for interconnecting cards within a single server. Generation 4.0 is 900GB/s with latencies around 5μs. And 5.0 is already 1.8 TB/s. For full-duplex connection, a special switch is used—NVSwitch.
Open PCI-E 5.0, on the other hand, gives 128GB/s and 10μs, while PCI-E 6.0 gives 256GB/s and the same 10μs. 1TB/s full-duplex is promised in PCI-E 8.0, but who knows when that will be! And there is also its evolutionary successor for coherent memory—Compute Express Link. It is based on the same PCI-E, but thanks to new protocols, it is faster and more flexible.
Alternatives don’t really exist. There is AMD Infinity Fabric for those who are ready to build their cluster on AMD. There is Intel Xe Link. And UALink, which is emerging somewhere in the shadowy labs of mega-corporations.
3-4. Inside the rack and between racks
Historically, InfiniBand has reigned here, having already gone through half a dozen generations. The current relevant version—NDR (Next Data Rate)—delivers 400Gb/s from a card with latencies in the single-digit μs. There is also XDR (800Gb/s).
But due to a number of features and problems, Ethernet is now making an appearance on the popularity stage, wearing the mask of a Loss-Less network—RoCE—RDMA over Converged Ethernet. Today, it provides the same 400Gb/s per port and latencies in the single-digit μs. And there are also 800Gb/s implementations.
And here too, a shadowy genius is giving birth to a new chimera, UltraEthernet, which we may yet have to hear about.
There will be separate articles in the series about InfiniBand and RoCE. For now, you can listen to podcasts about InfiniBand and RoCE.
A brief summary:
|-------------------------------------------------------------------| | Technology | Throughput | Latency | Where it's used | |------------|------------------------|----------|------------------| | HBM3/4 | 3 TB/s | ~100 ns | GPU↔Memory | | NVLink 4.0 | 900 GB/s | ~5 μs | GPU↔GPU (node) | | PCIe 5.0 | 128 GB/s | ~1 μs | GPU↔GPU (node) | | IB NDR | 400 Gb/s | ~1 μs | Node↔Node | | RoCE v2 | 400 Gb/s | ~1-3 μs | Node↔Node | |-------------------------------------------------------------------|
Parallelization schemes
Now that we know about the different types of parallelization, what data exchange they need, as well as the network structure in clusters, it’s time to outline the parallelization schemes, that is, what to place where.
Let’s go from the most critical to time and bandwidth to the less critical.
Tensor parallelism involves frequent exchanges of small amounts of data, which are very critical to transmission latency. It is absolutely forbidden to lose any data—otherwise, everything is for naught. Therefore, it must be used strictly within a server, where NVLink ensures minimal delivery time.
Pipeline parallelism involves infrequent sending of large volumes when transmitting activations and gradients between layers. It is sensitive to latency to avoid the formation of “bubbles” during the training process. Losing data is expensive: you will have to restart training on the current micro-batch. Therefore, it is often placed within a rack, where latencies are relatively low and the probability of failure is not high.
Data parallelism involves very infrequent exchanges of very large volumes of data. Essentially, this is a one-time synchronization of gradients per training cycle. Latency is not important here. It would be nice to have high bandwidth. And data loss is not critical—you just burned some electricity for nothing, calculating the lost batch. The inter-rack network is quite suitable for these purposes.
Expert parallelism requires AllToAll communications but is not as critical to latency. It can be within a rack or between racks.
Here is a logical and straightforward* scheme:
TP — within a server
PP — within a rack
DP, EP — between racks
at first naive glance. On second thought, it could be, for example: TP+DP within a node (without PP)
This concludes the introductory part about neural networks. We will not return to this topic in such detail again. Therefore, now is the time to tell a short history of neural networks.
Historical background
In 1676, Leibniz laid out the foundations of mathematical analysis—differentials/derivatives and integrals—in his letters. Thank you, Leibniz!
In 1943, McCulloch and Pitts proposed a mathematical model of a neuron—without practical implementation. But what an idea!
In 1949, neuropsychologist Hebb described the rule of learning for neural connections: “Cells that fire together, wire together.” His work became a bridge between neurobiology and artificial intelligence, setting the direction for the development of self-organizing and adaptive models.
In 1958, the first working neural network appeared—Rosenblatt’s perceptron. It could recognize printed characters in an image.
The perceptron was single-layered and was implemented in hardware based on variable resistors. Technically, it had two layers: an input layer (receives data) and an output layer (gives a result). But in fact, the neural network is single-layered because only one of the layers—the output layer—has trainable weights. In such a network, there were obviously no hidden layers: input, associative elements analyzing the image, and output elements giving the result.
Rosenblatt's Perceptron The perceptron had two basic concepts of modern neural networks: trainable (adjustable) weights, and threshold activation of elements—when certain elements gave +1 if the input signal was stronger than the threshold, and -1 otherwise.
The perceptron showed the technical feasibility of hardware implementation of neural networks. But it did not become commercial because it could only perform very primitive analysis—it could not recognize shifted, rotated, or partially obscured characters. For this, it lacked hidden layers and non-linearity.
Over time, it became clear that several layers were needed within the network: lower layers extract simple features, upper layers—complex patterns.
But then came the first AI winter: the 60s-70s. There was no hardware base, no mathematical apparatus for training and for ensuring non-linear recognition to move forward.
Seppo Ilmari Linnainmaa in 1970 first described the backpropagation method (reverse mode of automatic differentiation). Backpropagation had been used before, but this work provided a rigorous mathematical description of the algorithm, which formed the basis of its modern implementation. But the work was forgotten until 1986.
The first network with a non-linear function was Fukushima’s Cognitron in 1975, which also recognized images. Its architecture served as a prototype for most subsequent neural networks: from lower to upper layers: from primitive forms in a limited area to abstract and generalized features across the entire image.
The first layers identified simple parts: lines, dots, corners, connections. The next ones combined them into more complex forms. The last layer saw the assembled figure as a whole and made a decision about what symbol it was.
The weights of the first layers were programmed manually. The last layers were trained on a competitive basis, that is, the neurons competed during training for the right to recognize a certain image.
But, no less importantly, the cognitron had the ability for non-linear generalizations, which meant it could track complex rules and dependencies, and identify symbols no matter how they were located on the input pixel matrix.
In 1986, thanks to David Rumelhart, Geoffrey Hinton, and Ronald Williams, the backpropagation algorithm gained widespread use. This is essentially supervised learning. The model gives a result, which is compared with a known one, the deviation is calculated, and on the basis of this, correction matrices are created, which are propagated from the end to the beginning. And along with Backpropagation, multi-layer trainable networks appeared, in particular the MLP—Multi-Layer Perceptron, which revolutionized character recognition.
In the late 80s, optical character recognition (OCR) systems appeared, becoming the first commercially successful applications of neural networks. They were used to recognize characters on checks and postal codes on letters.
Then came a long second winter: the 90s and 2000s. There were neither the computational power nor the approaches for any other applications.
Until in 2012, at the ImageNet Large Scale Visual Recognition Challenge conference, Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton brought AlexNet. It was the spark that launched the rocket of deep learning into the sky. The network tore everyone apart at the competition. And here’s why.
8 layers (5 convolutional and 3 fully connected)
60 million parameters. One of the largest at the time.
Training on multiple GPUs with parallelization. Before that, it was training on a single GPU, or with parallelization, but on a CPU.
Use of CUDA to accelerate computations.
Training on big data (ImageNet).
They used ReLU instead of sigmoid.
AlexNet showed not just parallelization, but economically justified parallelization. It was after AlexNet that parallelization on GPUs became a mandatory element of modern neural network frameworks (PyTorch, TensorFlow).
Then, during the 2010s, there was an explosive growth in computational power (GPU, TPU, HBM), a renaissance of neural networks, and the emergence of GANs, ResNet, and transformers.
The 2020s are marked by the development of LLMs (GPT-3/4/5, Gemini, Claude, DeepSeek), multimodal models (when you work with text, images, and sound in one interface), the publication of trained models in Open Source (LLaMA, Mistral, Deepseek), and large convolutional networks. And most importantly—commercial use and astronomical investments.
Conclusion
Perhaps the most important non-technical conclusion to be drawn is: no generative neural network understands the structure of the world. It simply identifies patterns and regularities, and generates the most probable next object: a token, a pixel, a frame, a sound. This means that the capture and enslavement of humanity is not a threat to us in the near future. But only in the near future.
It does not understand that in the phrase “the glass fell and broke,” there is a huge amount of information that a person reads. We know that the glass fell due to gravity, that it was probably made of glass, which is a fragile material, that while the glass was falling, it was accelerating, and that it fell from a sufficient height to break, that some event caused it to fall. What’s more! It doesn’t even know the concepts of “glass,” “fell,” and “broke.” It simply detects mathematical patterns between these words.
The technical conclusions that are important for us for the next articles in the series are as follows:
There are not a great many, but many types of neural networks today. And they are all structured roughly the same, with differences in details.
At the core is a strictly deterministic mechanism with a set of specific steps.
For calculations, a simple but heavy mathematical apparatus with a huge number of matrix multiplications is used.
The sizes of models and intermediate data are such that we need parallelization and distributed computing.
There are several different methods of parallelization with different requirements for bandwidth and latency, which are usually combined into clever hybrid schemes.
All distributed computations require massive data exchange during both training and inference.
Training and inference place very different demands on the network, so they are almost always separated.
There are several types of physical networks in GPU clusters, each with its own specifics.
There is no sin that companies will not commit for the sake of optimizing high-performance computing.
Therefore, parallelization, optimization of collective operations, ZeRO, caches, ultra-expensive HBM memory, SoCs, Kernel Bypass mechanisms, Zero Copy, DMA, specialized networks like NVLink and InfiniBand, and fine-tuning of Ethernet—all of this finds its application in working with neural networks. And it is precisely about the philosophy of universal slow solutions versus specific fast ones that we will talk about in the next article.
And now let’s take another look at what model training looks like on an oscilloscope—we were measuring the electricity consumption of the cluster:
Useful links
Videos:
A playlist by the awesome 3Blue1Brown, author of the manim library, on how neural networks work: 3Blue1Brown - Deep Learning
Gradient descent explained simply: The Most Important Algorithm in Machine Learning
Two video reviews on the necessity of a non-linear function and vanishing gradients
Why deep networks are better than wide ones: Why Deep Learning Works Unreasonably Well
Channels:
A wonderful interactive visualization of transformers with descriptions: https://bbycroft.net/llm
The classic Google paper that gave birth to transformers: Attention Is All You Need
A very good handbook from Yandex on machine learning
Yann LeCun. How Machines Learn — a historical excursion from a primary source
Everything you wanted to know about Rosenblatt’s perceptrons but were afraid to ask
Characterizing Communication Patterns in Distributed Large Language Model Inference
About NCCL and collective operations: Demystifying NCCL: An In-depth Analysis of GPU Communication Protocols and Algorithms
A glimpse at AI Infra at Meta
A visual guide to the Incredible Power of NVIDIA GB200 NVL72: Inside the AI Computing Giant
linkmeup podcasts
Thanks
To Alisa AI, Claude Sonet 4.5, DeepSeek V3.1. For their inhuman patience with my endless questions, for reviewing the finished article, and for the manim image code.
To Dmitry Afanasyev (flow) for the human review of the article, as well as a lifetime supply of links for study.
To Pavel Ostapenko for proofreading and recommendations.
To Alexander Demidenko for stylistic and factual corrections.
To Andrey Shcherbin for comments.
To all my family and friends who for many months endured my unsolicited stories about neural networks.