PyTorch and LLM Training Basics

1. LLM Training Overview

A large language model (LLM) is essentially a function with a large number of trainable parameters. A tokenizer converts text into numbers, a neural network processes those numbers to make predictions, and the model adjusts its parameters based on prediction errors.

flowchart TD
    A[Raw text] --> B[Tokenizer: encode as token IDs]
    B --> C[Embedding: convert IDs into vectors]
    C --> D[Neural network: compute predictions]
    D --> E[Loss: measure prediction error]
    E --> F[Backward: compute parameter gradients]
    F --> G[Optimizer: update parameters]
    G -.Next training step.-> D

We can understand this process by answering three questions:

1. What does the model learn? It adjusts its parameters so its predictions get as close as possible to the training targets.
2. How does the model compute? It stores data in tensors and processes them using operations such as matrix multiplication.
3. Why is training expensive? It requires a huge number of floating-point operations and memory for parameters, gradients, optimizer states, and intermediate activations.

1.1 Key Terms

Machine learning (ML) learns patterns from data rather than relying on rules explicitly written by programmers. Deep learning (DL) is a branch of machine learning that mainly uses multilayer neural networks (NNs) to learn representations and make predictions. Large language models are deep learning models.

A dataset is a collection of samples. A label or target is the correct result that we want a model to learn to produce during training.

Term	Meaning	Example
Model	A function that maps inputs to outputs	y = w * x
Parameter	A value updated during training	Weight w, bias b
Hyperparameter	A training setting chosen by the practitioner	Learning rate, batch size
Sample	A single training example	A sentence and its prediction target
Feature	A numerical representation processed by the model	A 4-dimensional vector for a token
Forward pass	Computes predictions and loss using the current parameters	pred = model(x)
Loss function	Quantifies the error in a prediction	Squared error, cross-entropy
Backward pass	Computes gradients of the loss with respect to parameters	loss.backward()
Optimizer	Updates parameters using gradients	SGD, AdamW
Inference	Uses a trained model to make predictions	Generating the next token
Epoch	One complete pass through the training dataset	Using all 100,000 samples once
Training step	Usually one parameter-update operation	One optimizer.step()

The most important distinction between training and inference is that training generally computes gradients and updates parameters, whereas inference usually performs only forward computation without updating parameters.

2. Neural Network Training

2.1 Single-Parameter Model

Suppose we want a model to learn to “multiply its input by 2.” We use:

$$\hat y=wx$$

• $x$: the input.
• $w$: the model’s only trainable parameter.
• $\hat y$: the model’s prediction. The hat indicates an estimate, distinguishing it from the correct answer $y$.

Now consider a sample with $x=3$ and target $y=6$, while the initial parameter is $w=1$.

Step 1: Forward pass.

$$\hat y=1\times3=3$$

The model should have predicted 6, but it predicted only 3.

Step 2: Compute the loss. First, use squared error:

$$L=(\hat y-y)^2=(3-6)^2=9$$

Substitute $\hat y=wx$:

$$L(w)=(3w-6)^2$$

The loss $L$ is now a function of the parameter $w$. Training aims to find a value of $w$ that makes $L$ smaller.

Step 3: Compute the gradient. A gradient describes how the loss changes when a parameter changes by a small amount. Here, there is only one parameter, so the gradient is simply a derivative.

Using the chain rule:

$$\begin{aligned} \frac{\partial L}{\partial w} &=\frac{\partial(3w-6)^2}{\partial(3w-6)} \cdot\frac{\partial(3w-6)}{\partial w}\\ &=2(3w-6)\times3 \end{aligned}$$

At $w=1$:

$$\frac{\partial L}{\partial w}=2(3-6)\times3=-18$$

Here, $-18$ means that increasing $w$ near its current value will decrease the loss. It does not mean that the loss itself is negative; the loss is still 9.

Step 4: Update the parameter using the gradient. Apply gradient descent:

$$\boxed{w_{\mathrm{new}}=w-\eta\frac{\partial L}{\partial w}}$$

Here $\eta$ (eta) is the learning rate, which determines the size of each step along the gradient direction. Set $\eta=0.05$:

$$\begin{aligned} w_{\mathrm{new}}&=1-0.05\times(-18)=1.9\\ \hat y_{\mathrm{new}}&=1.9\times3=5.7\\ L_{\mathrm{new}}&=(5.7-6)^2=0.09 \end{aligned}$$

Just one update reduced the loss from 9 to 0.09. Repeating “predict → calculate loss → compute gradients → update parameters” is the basic process of model training.

2.2 Layers and Activations

Real neural networks use many linear layers and activation functions.

When a sample is represented as a row vector, a linear layer can be written as:

$$\mathbf y=\mathbf xW+\mathbf b$$

• $\mathbf x$: the input feature vector.
• $W$: the weight matrix, which determines how input features are combined.
• $\mathbf b$: the bias, an additional learnable offset.
• $\mathbf y$: the new features produced by the linear layer.

Stacking only linear layers is not enough because successive linear transformations can still be combined into a single linear transformation. For example:

$$(xW_1)W_2=x(W_1W_2)$$

Therefore, neural networks usually add nonlinear activation functions. One of the simplest is ReLU (Rectified Linear Unit):

$$\operatorname{ReLU}(x)=\max(0,x)$$

This replaces negative values with 0 while leaving positive values unchanged. For example, ReLU transforms [-2, 0, 3] into [0, 0, 3]. Modern LLMs also commonly use mechanisms such as GELU, SiLU, and SwiGLU.

2.3 Cross-Entropy Loss

Squared error is suitable for regression tasks, where the goal is to predict numerical values. Next-token prediction is a classification task: selecting the correct token from a vocabulary.

If the vocabulary contains $V$ tokens, the model produces a score $z_i$ (called a logit, or unnormalized score) for each candidate token. Softmax converts these scores into probabilities:

$$p_i=\frac{e^{z_i}}{\sum_{j=1}^{V}e^{z_j}}$$

• $V$: the vocabulary size.
• $z_i$: the score for the $i$-th token.
• $p_i$: the predicted probability of the $i$-th token; all $p_i$ values sum to 1.

If the index of the correct token is $c$, the cross-entropy loss is:

$$\boxed{L=-\ln p_c}$$

Cross-entropy can be understood as a penalty for assigning too little probability to the correct token. The higher the probability of the correct token, the lower the loss. For example, increasing that probability from $0.1$ to $0.9$ significantly reduces the loss.

3. Tokenization and Embeddings

3.1 Tokens and Tokenizers

A computer cannot feed raw strings directly into matrix multiplications. It first needs a tokenizer to split text into tokens and encode them as integer token IDs.

A token does not necessarily correspond to one Chinese character or one English word. It can also represent a subword, punctuation mark, or another text unit. For example:

Text            → Tokenizer       → Token IDs
"hello world"   → (several tokens) → [integer, integer, ...]

This example illustrates only the processing pipeline; the token boundaries and integer IDs may differ across tokenizers.

BPE (Byte Pair Encoding) is a common subword-tokenization approach. Starting from smaller units, it repeatedly merges adjacent units that frequently occur together in the training corpus. Byte-level BPE uses bytes as its base units and is not identical to simplified examples based on characters.

Remember: a token ID is only a lookup-table index, not a semantic value with meaningful numerical magnitude. For example, token ID 10 does not mean that the corresponding token has “five times the meaning” of token ID 2.

3.2 Embedding Lookup

An embedding layer can be understood as a trainable lookup table. Suppose a vocabulary has 100 tokens and each token is represented by 4 numbers. The embedding table then has shape $(100,4)$.

import torch
from torch import nn

# B=2: two sequences; T=3: three tokens in each sequence.
# Token IDs must have an integer data type; long means int64.
ids = torch.tensor([
    [1, 5, 2],
    [7, 1, 0]
], dtype=torch.long)

# Vocabulary size V=100; hidden dimension H=4.
embedding = nn.Embedding(num_embeddings=100, embedding_dim=4)

# Look up each ID to obtain its corresponding 4D floating-point vector.
x = embedding(ids)

print(ids.shape)  # torch.Size([2, 3])
print(x.shape)    # torch.Size([2, 3, 4])

The shape changes as follows:

$$\boxed{(B,T)\;\xrightarrow{\text{Embedding}}\;(B,T,H)}$$

• $B$ (batch size): how many samples are processed at once.
• $T$ (sequence length): how many tokens each sample contains.
• $H$ (hidden size): how many numbers are used to describe each token’s features.

Here, (2, 3, 4) means “2 sequences, 3 tokens per sequence, and 4 features per token,” not “24 samples.” These three letters are among the most common dimension symbols in language-model code.

4. Tensor Basics

4.1 Shape and Rank

A tensor can initially be understood as a multidimensional array.

Data type	Example	Shape (length of each axis)	Rank / ndim (number of axes)
Scalar	5	()	0
Vector	[1,2,3]	(3,)	1
Matrix	2 rows and 3 columns	(2,3)	2
3D tensor	Two 3×4 matrices	(2,3,4)	3

Shape tells you the length of each axis, whereas rank tells you how many axes there are. For example, (2,3,4) has rank 3 and contains $2\times3\times4=24$ elements.

import torch

x = torch.randn(2, 3, 4)  # Create a random floating-point tensor of shape 2×3×4.
print(x.shape)           # torch.Size([2, 3, 4])
print(x.ndim)            # 3: total number of axes
print(x.numel())         # 24: number of elements
print(x.dtype)           # torch.float32: data type of each element
print(x.element_size())  # 4: bytes per element
print(x.device)          # cpu (by default)

4.2 dtype, device, and stride

• dtype (data type): how each element is encoded, such as float32, bfloat16, or int64.
• device: whether a tensor is located on the CPU or GPU, such as cpu or cuda:0.
• stride: how many elements must be skipped in underlying storage when an index along a given axis increases by 1.
• contiguous: whether a tensor follows the default row-major contiguous memory layout. Noncontiguous tensors can still be used for computation.

Consider a concrete example:

import torch

m = torch.arange(6).reshape(2, 3)
# m = [[0, 1, 2],
#      [3, 4, 5]]

print(m.stride())          # (3, 1)
print(m.T.stride())        # (1, 3)
print(m.T.is_contiguous()) # False

m.stride()=(3,1) means that moving down one row skips 3 elements in storage, while moving right one column skips just 1 element.

m.T is a transpose: it swaps rows and columns. Usually, it changes only the tensor’s view of the data, without immediately copying the underlying storage, so the strides become (1,3).

Common confusion: transpose swaps axes, whereas reshape changes the shape; they are generally not equivalent. Whether reshape copies data depends on the storage layout.

4.3 Tensor Memory Usage

Suppose a tensor has shape $(s_1,s_2,\ldots,s_k)$ and each element occupies $q$ bytes. Its element data alone requires approximately:

$$\boxed{M=\left(\prod_{i=1}^{k}s_i\right)q}$$

• $s_i$: the length of the $i$-th axis.
• $\prod$: the product symbol, meaning that all axis lengths are multiplied together.
• $q$: the number of bytes per element.

For example, consider a tensor of shape (8, 1024, 4096) stored in BF16 (2 bytes per element):

$$\begin{aligned} M&=8\times1024\times4096\times2\\ &=67{,}108{,}864\ \text{Bytes}\\ &=64\ \text{MiB} \end{aligned}$$

Note that MB and MiB are different:

$$1\ \text{MB}=10^6\ \text{Bytes},\qquad 1\ \text{MiB}=2^{20}\ \text{Bytes}$$

This formula counts only the storage for tensor elements, not the total GPU memory required for training. Training also needs parameters, gradients, optimizer states, activations, and temporary buffers. If several views share the same underlying storage, that memory must not be counted multiple times.

5. Matrix Operations

5.1 Matrix Multiplication

Suppose a token has $H=4$ hidden features, and we want to map them to $K=5$ features. Use a weight matrix $W\in\mathbb R^{4\times5}$:

$$\mathbf y=\mathbf xW,\qquad (1,4)\times(4,5)\rightarrow(1,5)$$

If we process $B=2$ text sequences at once, each containing $T=3$ tokens, then:

$$\boxed{(B,T,H)\times(H,K)\rightarrow(B,T,K)}$$

In this example:

$$(2,3,4)\times(4,5)\rightarrow(2,3,5)$$

In other words, the number of samples and tokens stays unchanged; only each token’s feature dimension changes from 4 to 5.

import torch

x = torch.randn(2, 3, 4)  # B=2, T=3, H=4
w = torch.randn(4, 5)     # H=4, K=5

y = torch.matmul(x, w)   # Perform matrix multiplication along the last dimension.
print(y.shape)            # torch.Size([2, 3, 5])

Here, the manually created weight w has shape (H,K). However, nn.Linear(4,5) stores its weight internally with shape (5,4), because PyTorch follows the convention $Y=XW^\top+b$. The two formulations can be mathematically equivalent, but the internal storage shape should not be confused with the shape used in the mathematical expression.

5.2 einsum: Tensor Summation

einsum (Einstein summation) expresses matrix and tensor operations using letter-based rules.

import torch

# Create an input of shape (B,T,H) and weights of shape (H,K).
x = torch.randn(2, 3, 4)
w = torch.randn(4, 5)
y = torch.einsum("bth,hk->btk", x, w)
print(y.shape)  # (2, 3, 5)

Read bth,hk->btk one letter at a time:

• b: the batch dimension, preserved from the input to the output.
• t: the token dimension, also preserved from input to output.
• h: the input feature dimension. It appears in the inputs but not in the output, so the operation sums over this dimension.
• k: the output feature dimension, appearing in both the weights and the output.

The corresponding mathematical expression is:

$$\boxed{Y_{b,t,k}=\sum_{h=1}^{H}X_{b,t,h}W_{h,k}}$$

The symbol $\sum$ means “add many terms together.” For example, the dot product of two vectors is:

$$[1,2,3]\cdot[4,5,6]=1\times4+2\times5+3\times6=32$$

The ellipsis ... in einsum represents a group of axes that are not individually named. For example:

import torch

x = torch.randn(2, 3, 4)
w = torch.randn(4, 5)
z = torch.einsum("...h,hk->...k", x, w)
print(z.shape)  # (2,3,5)

Here, ... corresponds to b,t. It does not always refer to batch and token dimensions; its meaning depends on the actual axes of the input tensors.

5.3 einops: Tensor Rearrangement

einops (Einstein Operations) is a toolkit for working with tensor dimensions. Two particularly common operations are:

• rearrange: reorders, merges, or splits axes; it usually does not sum numerical values.
• reduce: sums, averages, or otherwise aggregates values along specified axes, reducing dimensionality.

# Install before first use: pip install einops
import torch
from einops import rearrange, reduce

x = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4)

# Merge b=2 and t=3 into a single axis of length 6.
flat = rearrange(x, "b t h -> (b t) h")
print(flat.shape)  # (6, 4)

# Average over the three tokens (axis t) for each sample.
avg = reduce(x, "b t h -> b h", "mean")
print(avg.shape)   # (2, 4)
print(avg[0])      # tensor([4., 5., 6., 7.])

The first output shape, (6,4), still contains 24 elements. The second shape, (2,4), contains only 8 elements because the three tokens in each sample have been averaged into one representation.

Which tool should you use? Use matmul for matrix multiplication; use einsum when you want to specify precisely which axes are multiplied and summed over; and use einops when you want to organize tensor shapes.

6. Floating-Point Precision

6.1 Number Representation

A floating-point number is similar to scientific notation in a computer. For a finite, normalized binary floating-point number:

$$\boxed{x=(-1)^s\times(1.f)_2\times 2^{e-\mathrm{bias}}}$$

• $s$ (sign): the sign bit, determining whether the value is positive or negative.
• $e$ (exponent): the encoded exponent, mainly determining the numerical dynamic range—how large or small a representable value can be.
• $f$ (fraction/significand field): mainly determining how finely values of similar magnitude can be distinguished.
• $\mathrm{bias}$: the exponent bias, which maps signed exponents to unsigned encodings.

This formula describes normalized values. Zero, subnormal numbers, infinity, and NaN use special encodings.

6.2 FP32, FP16, and BF16

Data format	Total bits	Sign bits	Exponent bits	Fraction bits	Bytes per element
FP32	32	1	8	23	4 Bytes
FP16	16	1	5	10	2 Bytes
BF16	16	1	8	7	2 Bytes

Both FP16 and BF16 occupy 2 bytes, but they prioritize different properties:

• FP16 has more fraction bits, so it generally represents values of similar magnitude more precisely. However, it has fewer exponent bits, and its maximum finite value is 65,504.
• BF16 retains 8 exponent bits, giving it a dynamic range close to FP32, with a maximum finite value of about $3.39\times10^{38}$. However, its shorter fraction field gives it lower precision than FP16 at the same order of magnitude.
• FP32 offers greater precision, but storing and transferring its values costs more.

For example, the element data of a tensor containing 100 million values occupies about 400 MB in FP32 or 200 MB in BF16 or FP16.

Overflow occurs when a result is too large to be represented by the format. Underflow occurs when a value is so small that precision is lost or it becomes zero. BF16’s wider exponent range helps avoid certain overflow problems, but does not eliminate rounding errors.

6.3 Mixed Precision (AMP)

Mixed precision means using different numerical formats for different operations or states. For example, large matrix multiplications may use BF16, while certain accumulations, normalization operations, and optimizer states remain in FP32.

AMP (Automatic Mixed Precision) can be implemented using PyTorch’s autocast:

import torch
from torch import nn

# Run only when CUDA is available and the hardware supports BF16.
if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
    model = nn.Linear(128, 64).cuda()      # Weights are usually still FP32 by default.
    x = torch.randn(8, 128, device="cuda")

    # Automatically select precisions such as BF16 for suitable operations.
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        y = model(x)

    print(y.shape)  # (8, 64): 8 samples, each with 64 output features

Important distinction: a parameter’s storage precision is the format used to hold it in memory, while an operator’s computation precision is the format used during an operation. autocast does not permanently convert every parameter to BF16, nor does it guarantee that total training GPU-memory usage will be cut in half.

7. FLOPs and Training Compute

7.1 FLOP vs FLOPs vs FLOP/s

• FLOP (Floating-point Operation): one floating-point operation.
• FLOPs: the number of floating-point operations required by a computation; a measure of computational work.
• FLOP/s: the number of floating-point operations performed per second; a measure of computational throughput. $1\ \text{TFLOP/s}=10^{12}\ \text{FLOP/s}$.

Under a common counting convention, one floating-point multiplication counts as approximately 1 FLOP and one addition as approximately 1 FLOP. A fused multiply-add (FMA) is typically counted as 2 FLOPs.

If a task requires $F$ operations and the actual throughput is $P$, its time is approximately:

$$\boxed{t\approx\frac{F}{P}}$$

7.2 Matrix Multiplication FLOPs

Suppose:

$$A:(M,K),\qquad B:(K,N),\qquad C=AB:(M,N)$$

The output contains $M\times N$ elements. Each output element is the dot product of two vectors of length $K$:

$$C_{ij}=\sum_{k=1}^{K}A_{ik}B_{kj}$$

One dot product requires $K$ multiplications and $K-1$ additions, for a total of $2K-1$ scalar arithmetic operations. Therefore:

$$F=MN(2K-1)\approx\boxed{2MNK}$$

Why can we use an approximation? When $K$ is large, the difference between $2K-1$ and $2K$ is small. Hardware-performance estimates commonly use this approximate counting convention.

Example: Multiply a $(2,3)$ matrix by a $(3,4)$ matrix, so $M=2,K=3,N=4$:

$$F\approx2\times2\times4\times3=48\ \text{FLOPs}$$

If the input has shape (B,T,H) and the weights have shape (H,K), there are $B\times T$ tokens, each undergoing one linear mapping. Thus:

$$\boxed{F_{\mathrm{linear}}\approx2BTHK}$$

Notation warning: In the matrix multiplication formula $2MNK$, $N$ denotes the number of matrix columns. In the $6ND$ formula introduced later, $N$ denotes the number of model parameters. Do not confuse these meanings.

7.3 Backpropagation FLOPs

Combine the token and batch dimensions, and let the input be $X:(M,H)$ and the weight matrix be $W:(H,K)$:

$$Y=XW$$

The forward pass performs one matrix multiplication:

$$F_{\mathrm{forward}}\approx2MHK$$

During backpropagation, we already know the output gradient:

$$G=\frac{\partial L}{\partial Y},\quad G:(M,K)$$

Matrix differentiation gives two principal gradients:

$$\frac{\partial L}{\partial X}=GW^\top, \qquad \frac{\partial L}{\partial W}=X^\top G$$

• $\partial L/\partial X$: the gradient of the loss with respect to the input, which allows gradients to flow to earlier layers.
• $\partial L/\partial W$: the gradient of the loss with respect to the weights, used to update the current layer’s parameters.
• $\top$: matrix transpose, meaning rows and columns are swapped.

Each gradient calculation involves a major matrix multiplication. Therefore:

$$\begin{aligned} F_{\mathrm{backward}}&\approx 2MHK+2MHK=4MHK\\ F_{\mathrm{train}}&\approx 2MHK+4MHK=\boxed{6MHK} \end{aligned}$$

Conclusion: for linear layers dominated by matrix multiplication and requiring full gradients, the dominant forward and backward computation is approximately in a $1:2$ ratio. This is not a strict rule for every operator, and it does not imply that the backward pass always takes exactly twice as long as the forward pass.

7.4 The 6ND Estimate

For a Transformer dominated by dense matrix multiplications, make the following rough approximations:

1. Processing one token requires approximately $2N$ FLOPs in the forward pass.
2. The backward pass costs roughly twice as much, or about $4N$ FLOPs.
3. A forward and backward pass together therefore require approximately $6N$ FLOPs per token.
4. Training processes a total of $D$ tokens.

Therefore:

$$\boxed{F_{\mathrm{train}}\approx6ND}$$

• $F_{\mathrm{train}}$: total training computation, measured in FLOPs.
• $N$: the model’s parameter count (for this rough estimate of dense models).
• $D$: the total number of tokens processed throughout training, not the hidden dimension.

For example, if $N=10^9$ (1 billion parameters) and $D=10^{10}$ (10 billion training tokens):

$$F_{\mathrm{train}}\approx6\times10^9\times10^{10} =6\times10^{19}\ \text{FLOPs}$$

Scope: This is a planning approximation that omits some attention, normalization, embedding, and other computation. Long-sequence attention, MoE models, activation recomputation, and other factors can cause actual FLOPs to differ. This formula must not be treated as an exact count.

8. GPU Memory

8.1 Memory Components

Training requires more than just model weights. We also need to store:

Component	Purpose	Needed only during training?
Parameters	Store the model’s current weights	Needed for both training and inference
Gradients	Store derivatives of loss with respect to parameters	Primarily needed during training
Optimizer states	Store historical statistics needed for parameter updates	Needed during training
Activations	Store intermediate results required for backpropagation	Usually more extensive during training
Temporary buffers	Operator workspaces and temporary calculation data	May be needed during both training and inference

Therefore:

$$M_{\mathrm{total}} \approx M_{\mathrm{params}}+M_{\mathrm{grads}} +M_{\mathrm{optimizer}}+M_{\mathrm{activations}}+M_{\mathrm{temporary}}$$

This is a breakdown for memory budgeting. Real implementations also incur framework overhead, caching, fragmentation, and other costs, so the individual terms should not be treated as perfectly fixed.

8.2 AdamW Memory Cost

The simple update rule for SGD (Stochastic Gradient Descent) is:

$$w_{t+1}=w_t-\eta g_t$$

Here, $g_t$ is the gradient for the current step.

In addition to weights and gradients, AdamW (Adaptive Moment Estimation with Decoupled Weight Decay) maintains two kinds of historical information:

• First moment $m$: an exponential moving average of past gradients, used to estimate the update direction.
• Second moment $v$: an exponential moving average of squared past gradients, used to estimate the gradient scale.

Assume that parameters, gradients, first moments, and second moments are all stored in FP32:

Item	Memory per parameter
Parameter $w$	4 Bytes
Gradient $g$	4 Bytes
First moment $m$	4 Bytes
Second moment $v$	4 Bytes
Total	16 Bytes

For a model with $N$ parameters:

$$\boxed{M_{\mathrm{states}}\approx16N\ \text{Bytes}}$$

For example, with 1 billion parameters:

$$16\times10^9\ \text{Bytes}=16\ \text{GB}$$

Note: this 16 GB excludes activations and temporary buffers. It is not a universal constant for every training setup. Mixed-precision implementations may use BF16 parameters and FP32 optimizer states, and may even maintain additional FP32 master weights.

8.3 Gradient Accumulation

Batch size is the number of samples fed into the model in a single computation. A micro-batch is a smaller portion of a larger effective batch.

On a single GPU:

$$\boxed{B_{\mathrm{effective}}=B_{\mathrm{micro}}\times K}$$

• $B_{\mathrm{micro}}$: the number of samples processed in each forward/backward pass.
• $K$: the number of micro-batches accumulated before updating the parameters.

For example, processing 2 samples at a time, accumulating across 4 micro-batches, and then updating the weights once gives an effective batch size of $2\times4=8$.

Why does this work? Because differentiation is linear over sums:

$$\nabla_\theta\left(\frac1K\sum_{i=1}^{K}L_i\right) =\frac1K\sum_{i=1}^{K}\nabla_\theta L_i$$

Here, $L_i$ is the mean loss for the $i$-th equally sized micro-batch, and $\nabla_\theta$ denotes the gradient with respect to all model parameters. Dividing each loss by $K$ before backpropagation accumulates the gradient of the overall average loss.

import torch
from torch import nn

torch.manual_seed(0)                       # Set a random seed for reproducibility.
model = nn.Linear(4, 1)                    # 4 input features, 1 output feature
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

xs = torch.randn(8, 4)                    # 8 samples in total
targets = torch.randn(8, 1)               # One target value per sample

micro_batch_size = 2
accum_steps = 4                           # 8 / 2 = 4 micro-batches

optimizer.zero_grad()                     # Clear gradients at the start of an effective batch.

for x, target in zip(xs.split(micro_batch_size),
                     targets.split(micro_batch_size)):
    pred = model(x)                       # Feed only 2 samples at a time.
    loss = ((pred - target) ** 2).mean()  # Mean loss of the current micro-batch
    (loss / accum_steps).backward()       # Accumulate gradients without updating weights yet.

optimizer.step()                          # Update weights once after all 4 micro-batches.

The key to understanding this code: backward() inside the loop does not automatically clear existing gradients; it adds to .grad. Only step() outside the loop actually updates the weights.

This approach reduces activation memory per micro-batch, but it does not proportionally reduce the memory required for parameters, their gradients, or optimizer states. With BatchNorm, unequal micro-batch sizes, stochastic layers, or numerical rounding, the result may not exactly match that of a full large batch. When micro-batches differ in size, losses must also be weighted by sample count.

8.4 Activation Checkpointing

Ordinarily, backpropagation needs intermediate activations from the forward pass. Saving all of them can consume substantial GPU memory.

Activation checkpointing saves only some necessary information. During backpropagation, it recomputes parts of the forward pass to recover intermediate values that were not stored.

$$\boxed{\text{Store fewer activations}\quad\Longleftrightarrow\quad \text{Perform additional recomputation}}$$

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

# Example: a linear layer followed by ReLU.
block = nn.Sequential(nn.Linear(64, 64), nn.ReLU())

# 8 samples, each with 64 features; enable gradients for the input.
x = torch.randn(8, 64, requires_grad=True)

# Some intermediate results inside block need not remain saved
# as long as they would during an ordinary forward pass.
out = checkpoint(block, x, use_reentrant=False)

# Define a normal loss and backpropagate; parts of block's forward
# computation are performed again when needed.
loss = out.square().mean()
loss.backward()

Do not confuse these two optimization methods:

Method	Main memory savings	Main cost
Gradient accumulation	Activations for each individual micro-batch	More small-batch executions, potentially increasing scheduling overhead
Activation checkpointing	Some activations that would otherwise be kept for backpropagation	Additional recomputation during the backward pass

Also, activation checkpointing here is not the same as saving a training checkpoint, which is a model file used to resume training after an interruption.

9. GPU Performance and Roofline

9.1 Capacity, Bandwidth, Compute

• Memory capacity: how much data the GPU can hold. It answers: “Will the data fit?”
• Memory bandwidth: how many bytes can be transferred per second. It answers: “How quickly can data be moved?”
• Compute throughput: how many floating-point operations can be performed per second. It answers: “How quickly can computation be performed?”

In practice, an operation may be:

• Memory-bound: compute units are waiting for data; the main bottleneck is data movement.
• Compute-bound: data is supplied quickly enough; the main bottleneck is computational throughput.

Having enough GPU memory only means that the data fits; it does not guarantee fast execution.

9.2 Arithmetic Intensity

Arithmetic intensity is defined as:

$$\boxed{I=\frac{F}{Q}}$$

• $F$: the number of floating-point operations performed (FLOPs).
• $Q$: the number of bytes transferred at the memory-hierarchy level being analyzed, such as GPU device-memory reads and writes.
• $I$: the number of FLOPs performed per byte transferred, measured in FLOPs/Byte.

For example:

• Transfer 1,000 bytes and perform 1,000 FLOPs: $I=1$.
• Transfer the same 1,000 bytes and perform 10,000 FLOPs: $I=10$.

The second case performs more computation for each byte transferred, indicating greater data reuse. However, higher arithmetic intensity does not necessarily mean the entire program will run faster; hardware, kernel implementation, and parallelism also matter.

9.3 Roofline Derivation

The Roofline model gives a theoretical upper bound on a program’s computational throughput.

Let:

• $P_{\max}$: the GPU’s theoretical peak compute throughput, measured in FLOP/s.
• $BW$: the GPU’s bandwidth at the relevant memory-hierarchy level, measured in Bytes/s.
• $F$: the total computation required by the task, measured in FLOPs.
• $Q$: the data volume transferred by the task, measured in Bytes.

Considering only the compute units, the task needs at least:

$$t_{\mathrm{compute}}\ge\frac{F}{P_{\max}}$$

Considering only data movement, the task needs at least:

$$t_{\mathrm{memory}}\ge\frac{Q}{BW}$$

Both must be completed, so the actual execution time satisfies:

$$t\ge\max\left(\frac{F}{P_{\max}},\frac{Q}{BW}\right)$$

Actual throughput is $P=F/t$. Rearranging the inequality gives:

$$\begin{aligned} P&\le\min\left(P_{\max},\frac{F}{Q}BW\right)\\ &=\boxed{\min(P_{\max},I\times BW)} \end{aligned}$$

This is the core idea behind Roofline: a program’s throughput cannot exceed the compute peak or the product of memory bandwidth and computation per byte transferred.

When the two upper bounds are equal, we obtain the ridge point, or critical arithmetic intensity:

$$\boxed{I^*=\frac{P_{\max}}{BW}}$$

Consider a hypothetical GPU with peak throughput of $100\ \text{TFLOP/s}$ and bandwidth of $1\ \text{TB/s}$:

$$I^*=\frac{100\times10^{12}}{1\times10^{12}}=100\ \text{FLOPs/Byte}$$

• At $I=20$, the bandwidth-based limit is $20\ \text{TFLOP/s}$, below the compute peak. This is the memory-bound region.
• At $I=150$, bandwidth could theoretically support $150\ \text{TFLOP/s}$, but the compute peak is only 100. This is the compute-bound region.

Roofline estimates the maximum achievable throughput, not the actual speed of a program. Kernel implementation quality, synchronization, memory-access patterns, and other factors can further reduce performance.

9.4 MFU vs GPU Utilization

MFU (Model FLOPs Utilization) measures the model’s effective computational throughput as a fraction of the relevant GPU’s theoretical peak throughput.

$$\boxed{\mathrm{MFU}\approx\frac{F_{\mathrm{model}}}{t\times G\times P_{\max}}}$$

• $F_{\mathrm{model}}$: the model’s theoretically useful FLOPs over the measured interval.
• $t$: actual runtime, in seconds.
• $G$: the number of GPUs.
• $P_{\max}$: the theoretical peak FLOP/s of each GPU at the relevant precision and computation mode.

Suppose a GPU has a theoretical peak throughput of $100\ \text{TFLOP/s}$ and the model achieves an effective throughput of $40\ \text{TFLOP/s}$:

$$\mathrm{MFU}=\frac{40}{100}=40\%$$

GPU utilization usually indicates the proportion of time the GPU is busy; it does not directly reveal whether computational throughput is close to its peak. A GPU can appear busy while waiting on device-memory reads and writes or executing inefficient kernels.

Therefore, to estimate LLM training time, divide the total computation by the effective compute throughput:

$$\boxed{t\approx\frac{6ND}{G\times P_{\max}\times\mathrm{MFU}}}$$

This is an estimation formula. It assumes that $6ND$ is an appropriate approximation for the model, that the chosen MFU has a sound basis, and that no other major overheads are omitted.

10. Dense vs Sparse MoE

A dense model typically executes a fixed main path of dense layers when processing each token.

A Sparse MoE (Sparse Mixture of Experts) model contains multiple experts (expert subnetworks). A router selects a subset of the experts to process each token. Experts are commonly FFNs (Feed-Forward Networks); they are not necessarily manually assigned roles such as “math expert” or “language expert.”

Dense layer:
Token → Same main set of computation layers → Output

Sparse MoE layer:
Token → Router → Select a few experts → Combine their outputs

It is essential to distinguish two kinds of parameter counts:

• Total parameters: the combined number of parameters in the entire model.
• Active parameters per token: the number of parameters that actually participate in computation when processing a single token.

An MoE model may have an enormous total parameter count while activating only a small fraction of those parameters for each token. Therefore, you cannot directly substitute the total number of MoE parameters into the dense-model $6ND$ formula; doing so may overestimate compute. Even using the active parameter count requires accounting for shared layers, the router, and the expert architecture.

The complete set of weights must still be stored or managed across devices. “Computing with fewer parameters per token” does not mean “all parameters use proportionally less memory.” Also, sparse expert selection in MoE is different from weight sparsity, which refers to many zero-valued weights inside weight matrices.

11. Resource Estimation Example

Suppose we want to train a dense model with $N=10^9$ parameters and a total of $D=10^{10}$ training tokens. For simplicity, we assume FP32 AdamW, a single GPU with a theoretical peak throughput of $100\ \text{TFLOP/s}$ in the relevant computation mode, and a training MFU of $40\%$. These GPU numbers are hypothetical and do not represent measurements of any specific graphics card.

Question 1: How much GPU memory do the training states require?

$$M_{\mathrm{states}}\approx16N=16\times10^9\ \text{Bytes}=16\ \text{GB}$$

This includes only the parameters, gradients, and two FP32 AdamW states. It excludes activations and temporary overhead.

Question 2: How many FLOPs does training require in total?

$$F\approx6ND=6\times10^{19}\ \text{FLOPs}$$

Question 3: What is the estimated theoretical planning time?

$$\begin{aligned} t&\approx\frac{6\times10^{19}}{1\times(100\times10^{12})\times0.4}\\ &=1.5\times10^6\ \text{s}\\ &\approx17.36\ \text{days} \end{aligned}$$

This time is a rough compute-budget estimate. If GPU memory is insufficient, MFU falls short of the assumed value, or communication, I/O, and recomputation introduce extra overhead, the real runtime will differ. This estimate does not guarantee that the training workload can actually run on that GPU.

How should you approach optimization? First, determine whether the workload fits in memory. If it does not, consider reducing the micro-batch size, enabling activation checkpointing, or changing the training setup or model. If it fits but still runs slowly, investigate computation volume, data movement, kernel efficiency, and MFU rather than looking only at GPU utilization.

12. Formula Reference

Formula	Meaning	Scope / Caution
$\hat y=f(x;\theta)$	The model computes a prediction from inputs and parameters	Abstract model representation
$\theta_{\rm new}=\theta-\eta\nabla_\theta L$	Gradient-descent parameter update	Learning rate determines the step size
$M=(\prod_i s_i)q$	Memory for tensor elements	Excludes other training states
$(B,T)\to(B,T,H)$	Embedding output shape	Converts IDs into feature vectors
$F\approx2MNK$	Matrix multiplication FLOPs	$M,K,N$ are matrix dimensions
$F\approx2BTHK$	FLOPs for a token-wise linear mapping	Operates on the final dimension of $(B,T,H)$
$F_{\rm train}\approx6MHK$	Forward and backward compute for a linear layer	Approximation for full differentiation dominated by matrix multiplications
$F_{\rm train}\approx6ND$	Total training FLOPs for a dense LLM	$N$ is parameter count; $D$ is the total number of processed tokens
$M_{\rm states}\approx16N$ Bytes	FP32 AdamW training-state memory	Excludes activations; not a universal BF16 formula
$B_{\rm eff}=B_{\rm micro}K$	Effective batch size	Single GPU, equal-sized micro-batches
$I=F/Q$	Arithmetic intensity	Ratio of computation to transferred data
$P\le\min(P_{\max},I\cdot BW)$	Roofline throughput ceiling	Not the same as actual throughput
$\mathrm{MFU}\approx F_{\rm model}/(tGP_{\max})$	Model FLOPs utilization	Not the same as GPU utilization