Neural net basics
Multi-layer perceptrons
multi-layer perceptron: fully-connected network with an input layer, at least one hidden layer, and an output layer
often used synonymously with “feed-forward network” even though FFN is technically a broader category where information flows in one direction
a single neuron computes a weighted sum of its inputs, adds a bias, and passes the result through an activation function
is the input vector (activations from the previous layer)
is the weight vector (edge weights leading into the neuron)
is the bias
is the activation function
a layer with inputs and neurons can be computed through matrix multiplication
so (column vector)
stack all the weight vectors into a single weight matrix
each row is the weights going into a single neuron
stack biases into vector
output hidden state will have shape
in practice, we process a batch of inputs at once!
in this case, we arrange inputs as rows of a matrix
conventionally change to have shape
each column is the weights going into a single neuron
the layer then becomes
where is broadcast to have shape
in math notation, a linear layer takes and applies as .
in PyTorch, the weight matrix
W
is actually stored as . the forward pass transposes W
, computing X @ W.T
. the transpose is free because it only changes the stride. this is so that the gradients for naturally comes out as , matching the shape of .let’s do the backprop for
the same bias is added to every sample, and each sample produces its own gradient for
these gradients thus accumulate
the most intuitive way to see this
we know that (if is a single example) is
when has a batch dimension, we know we are looking for output with shape
each row of depends only on row of (the batch examples don’t interact)
so we can just stack the gradients for each row
in general, derive Jacobian for a single example (which is clean, 2-dimensional)
if the tensor is shared across the batch (like ), then the batch dimension is summed out → contract (matmul where the batch dim is the inner dimension)
if the tensor is not shared (like , activations), the batch dimension is preserved → stack (matmul with batch dim on the inside)
note the PyTorch implementation with with looks like this
Activation functions
sigmoid
good for interpreting outputs as probabilities
not used for hidden layers in neural nets
vanishing gradients since the derivative is
not zero-centered, so downstream gradients for a single node are either all positive or all negative (depending on the upstream grad)
tanh
derivative peaks at 1.0 (for ), can still vanish
factors only ever shrink, since
softmax → probability distribution
with temperature
ReLU
derivative is 1 for , 0 for
dying ReLUs: if a pre-activation becomes permanently negative (i.e., negative for every input), it receives zero gradient forever
a fraction of network can go dead during training
Leaky ReLU
fixes the dying ReLU problem
Swish (smooth, non-monotonic)
GLU uses one linear projection to produce the “content” [left], and another to produce the gate [right]
SwiGLU plugs Swish in as the activation function inside GLU
without non-linearities, neural nets can’t do anything more than a linear transform
extra layers can be compiled down to a single linear transform
without non-linearities, adding more layers doesn’t give any more representational power
with more layers that include non-linearities, they can approximate any complex function!
Gradients
the derivative on a variable tells you the sensitivity of the whole expression to its value
if , then changing by a small would lead to a change of on
the gradient is the vector of partial derivatives
given a function with outputs and inputs, the Jacobian is an matrix of partial derivatives
given a function with inputs and a scalar output, the Hessian is an matrix of second partial derivatives, where
the Hessian of a loss function tells you about the curvature of the loss landscape
chain rule
for composition of one-variable functions, we multiply the derivatives
for functions with multiple variables, we multiply the Jacobians
neural network setup
for element-wise activation function where , what is ?
Jacobian is a diagonal matrix
useful Jacobians
other helpful derivatives
how to get derivative of sigmoid
derivative of Swish
gradient of softmax + CE loss
let be the logits indexed by , and be the post-softmax probabilities
CE loss gradient
, where is the correct class
gradient is
use chain rule to express in terms of
most of the terms vanish, because is only non-zero for
now let’s calculate the softmax gradient
putting it all together to get
for the true token
for all other tokens
very clean result:
Backpropagation
NN equations are represented as a computation graph
which parts of the function are thought of as “gates” is a matter of convenience, and in general are parts of the expression that have easy local gradients
backpropagation can be thought of as gates communicating to each other (through the gradient signal) whether they want their outputs to increase or decrease (and how strongly) in order to decrease the loss
achieved through repeated applications of chain rule, which allows us to decompose each gradient into the upstream gradient (already computed) and the local gradient
each node in the graph receives an upstream gradient and passes down a downstream gradient
each node has a local gradient (the gradient of its output w.r.t. its input)
downstream gradient = upstream gradient local gradient
gradients sum at outward branches
if is used in the computation of both and , then
node intuitions
distributes the upstream gradient to each summand
“routes” the upstream gradient to one of many input arguments
switches the forward coefficients in the downstream gradient
backprop
initialize output gradient as 1
visit nodes in reverse topological order: compute gradient w.r.t. each node using gradients w.r.t. successors
done correctly, the big- complexity of forward prop and backprop are the same
automatic differentiation
the gradient computation can be automatically inferred from the symbolic expression of the forward prop
each node type needs to know how to compute its output and how to compute gradients w.r.t. inputs given gradients w.r.t. outputs
local gradient is written by the programmer
manual gradient checking
for every parameter , recompute for and and check that
in the backwards pass, intermediate activations are needed
as a result, NNs typically store all intermediate activations during the forward pass
activation / gradient checkpointing trades compute for memory by storing only a subset of activations (the “checkpoints”) — if you need activations that weren’t saved, you recompute them on the fly by doing a partial forward pass from the nearest checkpoint
for a model with N layers checkpointed into K segments
memory goes from to
backward compute goes from to
optimal choice is : memory, backward compute
do for practice
calculating an explicit expression for would be extremely complex, but it’s completely unnecessary!
construct multiple intermediate variables in the forward pass, each of which is simple expression for which we know the local gradient
backprop needs to start from a scalar because it computes for every parameter , which is a single number per parameter — this only makes sense when is a scalar
when calling
.backward()
on a scalar, PyTorch implicitly seeds the backward pass with when we have per-token losses and define (mean reduction), by the linearity of derivatives we have
so the gradient from the mean loss is exactly the mean of the gradients we’d get from backpropagating each individually
upstream gradients always w.r.t. activations, gradients w.r.t. parameters are used for the update and end there, because parameters are leaf nodes in the computation graph
loss
│ dL/dy2 (activation grad)
▼
Layer 2 ──→ dL/dW2 (param grad, stored)
│ dL/dy1 (activation grad)
▼
Layer 1 ──→ dL/dW1 (param grad, stored)
│ dL/dx (activation grad — usually discarded)
▼
input
Optimizers
vanilla SGD update
an optimizer determines the direction + magnitude of parameter updates
for every parameter tensor, Adam keeps:
the parameter itself ()
the gradient ()
first moment (momentum)
second moment (variance)
Adam optimizer
let be the gradient for the current step
first moment is the running mean of gradients
result: if gradients have been pointing consistently in one direction, it moves more confidently that way
accumulate velocity over time, which helps you barrel through noisy gradients
second moment is the running mean of squared gradients
result: different weights effectively get different learning rates — consistently large gradients → smaller steps
effectively normalizes gradients to be on the same scale
bias correction: corrects for initialization bias in the running averages early in training
time is for time step starting at 1 (otherwise, use )
the hyperparameters and control updates to the moment estimates
and are both initialized to
so memory per parameter ≈ 4× parameter size
AdamW modifies Adam by adding weight decay towards 0
initialize optimizer with model parameters, to tell the optimizer which values it will be optimizing, and the
lr
parameter, which determines the size of the updatehow do you determine whether something should be a LR schedule or an optimizer?
depends on the time step alone → probably LR schedule
~requires per-parameter history → optimizer
code looks like this
params
is used to create parameter groups which each have their own hyperparams (e.g., different learning rates for different layers)torch.optim.AdamW(model.parameters())
creates a single parameter groupusually we don’t want weight decay on biases and LayerNorm params
torch.optim.AdamW([
{'params': decay_params, 'weight_decay': 0.01},
{'params': no_decay_params, 'weight_decay': 0.0},
])
defaults
dict provides fallback values for any hyperparameter not explicitly specified in a parameter groupin practice it’s better to apply weight decay before the Adam update because weight decay depends on the parameter
class AdamW(torch.optim.Optimizer):
def __init__(self, params, lr, betas, eps, weight_decay):
if lr < 0:
raise ValueError(f"Invalid learning rate: {lr}")
if not 0 < betas[0] < 1 or not 0 < betas[1] < 1:
raise ValueError(f"Invalid beta values: {betas}")
defaults = {"lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay}
super().__init__(params, defaults)
def step(self):
for group in self.param_groups: # for every group of parameters
lr = group["lr"]
beta1, beta2 = group["betas"]
eps = group["eps"]
weight_decay = group["weight_decay"]
for p in group["params"]: # for every parameter in the group
if p.grad is None:
continue
state = self.state[p]
# state initialization with 0s
t = state.get("t", 0)
m, v = state.get("m", torch.zeros_like(p.data)), state.get("v", torch.zeros_like(p.data))
# weight decay
p.data -= lr * weight_decay * p.data
# Adam update
grad = p.grad.data
m = beta1 * m + (1 - beta1) * grad
v = beta2 * v + (1 - beta2) * grad**2
m_hat = m / (1 - beta1 ** (t + 1))
v_hat = v / (1 - beta2 ** (t + 1))
p.data -= lr * m_hat / (v_hat.sqrt() + eps)
# update optimizer state
state["t"] = t + 1
state["m"] = m
state["v"] = v
gradient clipping constraints the size of the grad norm
computes global norm of all gradients
if max is exceeded, scale all parameters down by the same value to be below the max
prevents any individual step from being catastrophically large
Learning rate
warmup reduces primacy effect of early training examples
Mathy things
Information theory
cross entropy
KL divergence
entropy
another common form given logits
cross entropy between and is just KL between and plus the irreducible entropy of
proof
cross-entropy loss
when the target distribution is 1-hot, the cross-entropy loss is the negative log likelihood of the next token
also equivalent to KL divergence
implementing the loss
if using
F.cross_entropy()
, logits and labels are shifted internallyloss = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1), ignore_index=pad_idx)
shift_logits = logits[:, :-1, :]
shift_labels = input_ids[:, 1:]
logprobs = F.log_softmax(shift_logits, dim=-1)
token_logprobs = logprobs.gather(index=shift_labels.unsqueeze(-1), dim=-1).squeeze(-1)
# build loss mask
masked_logprobs = -token_logprobs * mask.float()
return masked_logprobs.sum() / mask.sum()
Numerical stability and other tricks
in general, pay attention to
exp(x) for large x → overflow to
log(x) for x near 0 → underflow to ()
log(x) for x near 1 → precision issues ()
computing
softmax(x)
unstable because for large , will overflow
use the fact that softmax is invariant to subtraction of a constant
subtracting from all ensures are not large (the largest exponent is ), so the numerically stable implementation does this
computing
log(softmax(x))
doing log-softmax directly is bad because log of inputs near 0 (low-probability classes) is unstable
use
x - logsumexp(x)
instead, which avoids materializing the tiny probability (and logsumexp(x)
is stable)these are equivalent
computing
log(sum(exp(x)))
why is it unstable?
for any large , will overflow to infinity
e.g., in float32, will overflow
if all are negative very negative, then will underflow to
precision issue if the summation is close to 1, because is unstable
this is the problem we got in our distillation project
intuition: we want to make the values small, which we can do by subtracting a constant! then we just need to add back the constant at the end
logsumexp(x)
is implemented like thisthe largest term is (for ), so no overflow
sum is at least 1, so we never compute
naively, computing the softmax requires two different passes to compute (used for stability) and then the denominator
stable softmax
online softmax trick: fuse the computation of and the denominator into a single pass
idea: we can calculate the denominator with the max-so-far and continuously rescale it for each new max-so-far
maintain
a running maximum
a running (shifted) denominator
update rule: when we encounter
to see why the update for is correct
note when the maximum doesn’t change, and the rescaling factor
if the goal is to return the softmax, then we need to do one more pass over the logits to return
we can also use this to calculate a weighted sum given a stream of logits and values
for FlashAttention, corresponds to the attention output for each query , where and are value vectors
numerically stable version multiplies both the numerator and denominator by (to avoid overflow issues with )
in addition to and , we maintain the running numerator ( = dimension of each , which is the head dimension in FlashAttention)
the update uses the same idea as the update for (derivation looks the same as the one for )
after processing all logits, we have and the true attention output is just
when we need to make an expression involving numerically stable, a standard tool is to multiply by (where is a large number, commonly chosen to be ) and see what survives.
Basic statistics
-value: probability of seeing data (at least this extreme) given the null hypothesis
crucially, it’s about the probability of the data given the null, not the probability of the null given the data
are these groups different?
Kolmogorov-Smirnov test: given two groups of observations of a continuous variable, were they drawn from the same underlying distribution?
measures maximum vertical distance between the CDFs of two samples
Chi-squared test: given two groups of observations of a categorical variable, were they drawn from the same underlying distribution?
T-test: are the means of two groups of continuous observations different?
assumption: data is normally distributed
one-sample version: whether a group’s mean equals some value
two-sample version: whether two groups have the same mean
paired version: whether two measurements on the same items are systematically different
one-sample t-test in disguise
given per-example differences (where each diff is +1, 0, or -1), can test whether the differences are significantly different from 0
ANOVA (F-test): (generalization of the t-test to more than two groups) do any of these groups have different means?
McNemar’s tests: compares two classifiers on the same dataset
this is probably the best test for a standard setup with two models being evaluated on the same test set, where each example can be answered correctly or incorrectly
are these variables related?
Pearson correlation test: tests where there is a linear relationship between two continuous variables
returns both a correlation coefficient and p-value for whether is significantly different from 0
misses non-linear relationships entirely (a perfect parabolic relationship → )
imagine a line being fit on a scatterplot
Spearman correlation test: same idea but measures monotonic associations
converts both variables to ranks (smallest value gets rank 1, etc.), then computes Pearson correlation on the ranks
the only thing that matters is order, not the specific values
Pearson versus Spearman
Pearson more sensitive to outliers, while Spearman robust to them
Pearson underestimates non-linear monotonic relationships
Pearson has cleaner interpretation
mutual information: captures any dependency between two variables, including non-linear ones
not really a statistical test
Gradient flow through sampling
Gumbel-Max trick
with logits , you can sample from the corresponding categorical distribution by
drawing independent noise values from distribution
taking argmax from
Gumbel-Softmax replaces the argmax with a softmax:
softmax is differentiable everywhere!
high temperature → smooth gradients, low temperature → discrete samples
not about making something differentiable (the plain softmax already is), but instead makes it stochastic
plain softmax:
deterministic (always the same soft mixture)
true categorical sampling:
gumbel-softmax:
stochastic (Gumbel noise)
approximately discrete (low temperature)
differentiable
you get exploration AND gradients!
straight-through estimator: pretend was the identity function
forward: (apply the non-differentiable funtion)
backward:
pass the upstream gradient directly down as the downstream gradient
would be wrong is is decreasing (signal is in the opposite direction)
but STE is usually applied to operations that are monotonically increasing (rounding, quantization, increasing step function)
direction is right, even if magnitude is wrong
Theoretical CS
a regular language is one that can be recognized with a finite state machine (also known as finite automaton)
deterministic finite automaton (DFA)
a context-free language is one that can be recognized by a pushdown automaton — basically a finite state machine plus a stack
stack gives unbounded memory but can only be accessed via the top
programming language syntax is built from context-free languages (matched parentheses, nested function calls, balanced HTML tags, etc.)
any finite language is trivially regular — you can enumerate all valid strings with enough states
any DFA can be encoded by a ReLU RNN
given a DFA with
states
alphabet
transition function
start state
accept state
build an RNN with hidden dimension that is a one-hot encoding of the state
need to construct , ,
for each transition , set , ,
The modern transformer LM
Architecture
symbol | dimension |
|---|---|
B | number of sequences in the batch |
L | number of layers |
T | sequence length (number of tokens to generate) |
S | sequence length (provided context) |
V | vocab size |
D | hidden dimension |
H | head dimension |
F | MLP hidden dimension, generally F = 4D |
N | number of query heads, N * H = D |
K | number of key/value heads, K < N in GQA |
G | group size in GQA = N // K |
token embedding
embedding matrix , initial hidden states
layer loop (for )
RMSNorm divides every element of by the RMS of (so that the hidden state has unit RMS) then multiplies by learned rescaling parameter
each head projects using , , (where ) into the head’s lower-dimensional subspace
[optional] QK norm: RMSNorm is applied to the query and key vectors to control the magnitude of the vectors going into the dot product
reshape to expose head dimension , and , then transpose the sequence length ( or ) and head dim ( or ) dimensions
expand , for GQA
apply RoPE at each position by rotating a query vector (or key vector ) by
for dimension pair (corresponding to indices from ), we rotate by angle where
the hyperparameter controls the base rotation frequency and is the head dimension
calculate attention scores
we divide by the head dimension because otherwise dot products will scale with
large inputs to softmax → peakier distributions → resistant to updates
apply causal mask
apply softmax
get attention output from weighted sum of values
reshape
apply output projection to mix output from different heads
residual connection
feed forward network
RMSNorm
gate and up projections [expansion] using ,
SwiGLU activation
down projection using
residual connection
final layer norm
final norm
unembedding
project onto vocab dimension using
Implementation notes
scores.masked_fill(~mask, -torch.inf)
for making the pre-softmax attention scoresassuming convention where
mask
is True
for positions that can be attended totensor.masked_fill(mask, value)
fills tensor
with value
where mask
is True
RoPE
we want to cache and for every (position, index) pair
we can do this upon initialization
positions = torch.arange(max_seq_len, device=device) # shape (max_seq_len)
thetas = self.theta ** (-torch.arange(0, d_k, 2, device=device) / d_k) # shape (d_k // 2)
angles = positions.unsqueeze(-1) * thetas.unsqueeze(0)
in practice, instead of doing a bunch of 2x2 matmuls, we express the rotation in dot products
extract the even and odd indices of , by reshaping the final head dimension into
x_pairs = x.reshape(*x.shape[:-1], -1, 2)
x_even = x_pairs[..., 0]
x_odd = x_pairs[..., 1]
calculate all the even and odd positions in the rotated matrix
x_out_even = x_even * cos - x_odd * sin
x_out_odd = x_even * sin + x_odd * cos
then interleave by stacking them side by side and flattening
torch.stack()
adds a new dimensiontorch.stack([x_out_even, x_out_odd], dim=-1).flatten(start_dim=-2)
attention looks like this
need
.reshape()
to expand d_model
() into num_heads x head_dim
()need
qkv.unbind()
to split queries, keys, vectors [optional depending on implementation]need
.transpose()
to swap num_heads
and seq_len
dimensions for attention computationafter getting
output
, need to .transpose()
and .reshape()
again to recover original shapebatch_size, seq_len, _ = x.shape
x_norm = self.norm(x)
qkv = self.qkv_proj(x_norm) # (batch, seq_len, 3 * d_model)
qkv = qkv.reshape(batch, seq_len, 3, self.num_heads, self.head_dim)
q, k, v = qkv.unbind(dim=2) # (batch, seq_len, num_heads, head_dim)
q = q.transpose(1, 2) # (batch, num_heads, seq_len, head_dim)
k = k.transpose(1, 2) # (batch, num_heads, seq_len, head_dim)
v = v.transpose(1, 2) # (batch, num_heads, seq_len, head_dim)
causal_mask = torch.tril(torch.ones(seq_len, seq_len)).bool()
output = scaled_dot_product_attention(q, k, v, mask) # (batch, num_heads, seq_len, head_dim)
output = output.transpose(1, 2) # (batch, seq_len, num_heads, head_dim)
output = output.reshape(batch, seq_len, d_model) # (batch, seq_len, d_model)
output = self.out_proj(output) # (batch, seq_len, d_model)
return x + output
attention part looks like this
def scaled_dot_product_attention(q, k, v, mask):
"""
k, q: (batch_size, ..., seq_len, d_k)
v: (batch_size, ..., seq_len, d_v)
returns o (batch_size, ..., seq_len, d_v)
"""
d_k = q.shape[-1]
scores = (q @ k.transpose(-2, -1)) / math.sqrt(d_k)
scores = scores.masked_fill(~mask, -torch.inf)
return softmax(scores, dim=-1) @ v
Accounting
Model parameters
embedding:
attention is (for in standard multi-head attention)
is
is
is
is
FFN is
up projection is
gate projection is
down projection is
layer norm is at each layer (plus the final norm)
pre-attention and pre-FFN layernorm each has parameters ( for each dimension in )
unembedding:
total: (for )
total model parameters is
Model activations
attention activations:
layer norm input is
layer norm output is
Q, K, V outputs are , ,
attention scores is
attention output is
FFN activations: (for )
layer norm input is
output of gate/up projections is each
output of down projection is
per-layer activations is
FLOPs in forward pass
assume prefill stage (so )
attention is per layer
projection is → FLOPs
projection is → (for )
projection is → (for )
is → (since )
is →
projection is → FLOPs
FFN is (for ) per layer
up projection is →
gate projection is →
down projection is →
per layer total:
unembedding layer is
unembedding is →
full forward pass is
FLOPs in backward pass
generally assumed to be 2 the FLOPs of the forward pass
compute gradient w.r.t. both the parameters and the input, each one matmul
the gradient w.r.t. the input is the incoming gradient for the previous layer
Inference memory use
total memory use at inference time: model weights + KV cache + peak activations
number of parameters in model:
num_params = sum(p.numel() for p in model.parameters())
KV cache size:
= batch size
= sequence length
= number of KV heads
= head dimension
= number of layers
2 for K and V
activations: in prefill, in decode
torch.inference_mode()
frees memory immediately, so it’s just about peak memory in a single layerinput is in prefill, in decode
with FlashAttention, the matrix is never materialized → attention becomes instead
peak activations
= input to layer
= K, Q, V vectors
= attention matrix (without FlashAttention)
matrix for each example in the batch and each query head
= FFN intermediate
activations scale quadratically with sequence length w/o FA, linearly with FA
at small batch sizes + sequence lengths, weights dominate
in prefill stage
at long sequence lengths (large ), attention term in activation dominates
at large batch sizes (large ), both KV cache and activations grow
Train memory use
total memory use at training time: model weights + optimizer states + gradients + activations
model parameters
FP32 master weights [full or mixed precision] →
BF16 transient copy for forward pass [mixed precision] →
optimizer states (first and second moment)
Adam states in FP32 [mixed precision] →
gradients
FP32 →
even in mixed precision, gradients are computed in BF16 but accumulated in FP32
activations
often the dominant piece: depends on , , ,
activations per layer (without flash attention)
with flash attention, the second term becomes and activation memory scales with (total number of tokens)
needed to compute gradients during the backward pass, can be reduced with gradient checkpointing
Attention
standard attention is
KV cache becomes independent of sequence length
once a token falls outside your window, you can just throw it away
sliding window: each token only attends to the last tokens, so compute is instead of
tokens each doing work (attending to keys/values)
sparse attention
interleave local attention with global attention
RMSNorm
normalization stabilizes training by preventing exploding/vanishing gradients
uniform scaling too restrictive: different features may need different magnitudes
combines stability of normalization with per-dimension variation in magnitude
is learned per-dimension rescaling
normalization step forces the hidden state to have unit RMS, which destroys any learned scale information
gives the network back per-feature control over magnitude
→ amplify dimension
→ suppress dimension
→ kill dimension
SwiGLU FFN
both and contribute content
provides one learned representation, and provides another (self-gated by its own confidence)
RoPE
RoPE only rotates the query and key vectors, not the values
position information only needs to affect which tokens attend to each other, not what information gets passed
for query vector at position and key vector at position , we want dot product to only depend on the relative position
we want such that the dot product is a function that encodes information only in relative form, e.g.,
RoPE embeddings are one such solution with
proof
inner product decays with increasing distance
for 2D vector , rotation by angle is
for position and head dimension index , we rotate by
for -dimensional embedding, we partition it into pairs and apply independent rotations to each pair with different frequencies
is RoPE’s only hyperparameter, typically 10,000 (called
rotary_base
in HF)defines the longest distance the model can natively distinguish
a full rotation is completed when ⇒
the slowest (smallest) is →
so slowest pair completes a full circle over positions
for each dimension pair , is
this provides exponential spacing (same design choice as sinusoidal positional encodings from the original transformer)
log uniform spread of frequencies → covers scales effectively
low-frequency pairs (large → small ) rotate slowly → change very little between adjacent positions → encode long-range information
high-frequency pairs (small → large ) rotate quickly → adjacent positions differ a lot → enable local discrimination
the full rotation matrix (for a single position) is block-diagonal, with each block handling one pair of dimensions
we can reformulate this with complex numbers
treat each pair as a complex number
rotation by angle is equivalent to multiplication by
Euler’s theorem:
this is equivalent to applying the rotation matrix in the complex plane
in practice rather than constructing the full rotation matrix, we do element-wise multiplications
Inference
latency is the time it takes to complete a single request, measured in seconds
throughput is how many tokens (or requests) we can process per unit time across all requests, measured in tokens/second
Batching & packing
traditional batching: collect requests, process them together, wait for all to finish, then collect the next batch
if one sequence generates 500 tokens and another generates 10, the short one sits idle waiting
continuous batching: as soon as one sequence finishes, you immediately slot a new request into its place without waiting for the whole batch
the batch is always full
selective batching: cleverly mix sequences that are in prefill and generation phase
idea: prefill is compute-heavy whereas decode is memory-bound
sequence packing: concatenate examples up to the max sequence length, using attention masks to prevent cross-contamination
token-budget batching: batch examples into batches such that the total number of tokens in a batch (after padding to the max seqlen within that batch) does not exceed a certain budget
usually done in finetuning
this makes sense because GPU memory is determined by
batch_size x sequence_length
Speculative decoding
speculative decoding exploits the fact that prefill is faster than generation
generate tokens from draft model
evaluate those tokens with target model
accept each draft token with probability
if , we definitely take it
if rejected, sample from adjusted distribution after renormalization
starting at the first rejected token
one token is always sampled from the teacher since we get those next-token logits “for free” when scoring the draft tokens
P(emitting token ) = P(draft ) * P(accept ) + P(sampled token rejected) * P(sample )
case 1 (first term): is accepted from the draft
case 2 (second term): is chosen after rejection
probability of drafting and rejecting a draft token :
0 if
otherwise
so the total probability of rejection is
the second step follows bc both and sum to 1
when we reject, we sample from
so the probably of getting after rejection is
combining both cases
KV cache
single-token forward pass with cache
cache should have dimension
(batch, num_heads, max_seq_len, head_dim)
in practice we store the cache inside the self attention module with
self.kv_cache
# Pre-allocate cache based on max_seq_len
kv_cache = [
{
'k': torch.zeros(batch, num_heads, max_seq_len, head_dim, device='cuda', dtype=torch.float16),
'v': torch.zeros(batch, num_heads, max_seq_len, head_dim, device='cuda', dtype=torch.float16),
}
for _ in range(num_layers)
]
# For attention, use cache[:, :, :position+1, :] as keys/values
def forward_with_cache(model, new_token, kv_cache, position):
"""
Instead of processing full sequence, process just the new token
and reuse cached K, V from previous positions.
"""
# Embed just the new token
x = model.embed(new_token) # (batch, 1, d_model)
for layer_idx, layer in enumerate(model.layers):
q, k, v = layer.qkv_proj(x).chunk(3, dim=-1)
# Update cache
kv_cache[layer_idx]['k'][:, :, position, :] = k.squeeze(2)
kv_cache[layer_idx]['v'][:, :, position, :] = v.squeeze(2)
# Attend over all cached positions
k_full = kv_cache[layer_idx]['k'][:, :, :position+1, :]
v_full = kv_cache[layer_idx]['v'][:, :, :position+1, :]
x = attention(q, k_full, v_full)
x = layer.ffn(x)
return model.lm_head(x)
Reducing KV cache size
lower dimension of KV cache
in standard multi-head attention (MHA) transformer attention, KV cache scales as
num_layers × num_heads × seq_len × head_dim × 2 (K and V)
in multi-query attention (MQA), all heads share the same K and V, but each head has its own Q
KV cache shrinks by a factor of
num_heads
inference is much faster and more memory-efficient
in group-query attention (GQA), heads are divided into groups, each group shares K and V
middle ground between MHA and MQA
MQA and GQA shares KV across heads, so we lose some representational capacity per head
multi-head latent attention (MLA) introduced in DeepSeek v2
instead of keeping keys and values with shape
(seq_len, num_heads × head_dim)
in the KV cache, cache a smaller latent vector (seq_len, latent_dim)
then at each decoding step, project the latent KVs back up to full size
MLA keeps a separate KV per head but compresses them into a shared latent space
adds some extra compute at inference time (the projection from latent to KV)
MLA not compatible with RoPE
cross-layer attention shares KV across layers
local attention
standard attention is , using local attention makes KV cache independent of sequence length
once a token falls outside your window, you can just throw it away
Sampling strategies
def sample(logits, temperature=1.0, top_k=None, top_p=None):
logits = logits / temperature
if top_k is not None:
values, indices = torch.topk(logits, top_k)
logits = torch.full_like(logits, float('-inf'))
logits.scatter_(-1, indices, values)
if top_p is not None:
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
# Remove tokens with cumulative prob above threshold
sorted_mask = cumulative_probs > top_p
sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
sorted_mask[..., 0] = False
indices_to_remove = sorted_mask.scatter(-1, sorted_indices, sorted_mask)
logits = logits.masked_fill(indices_to_remove, float('-inf'))
probs = F.softmax(logits, dim=-1)
return torch.multinomial(probs, num_samples=1)
Flash Attention
standard attention
the memory problem is that
attn_weights
is (batch, num_heads, seq_len, seq_len)
with
seq_len
= 8192 and 32 heads in FP16, that's 8192² × 32 × 2 bytes ≈ 4GBimport torch
import torch.nn.functional as F
import math
def standard_attention(q, k, v):
# q, k, v are all (batch, num_heads, seq_len, head_dim)
scale = math.sqrt(q.size(-1))
# Materialize full N×N attention matrix
attn_weights = torch.matmul(q, k.transpose(-2, -1)) / scale # (batch, num_heads, seq_len, seq_len)
attn_weights = F.softmax(attn_weights, dim=-1)
output = torch.matmul(attn_weights, v) # (batch, num_heads, seq_len, head_dim)
return output
why standard attention is memory bound: the attention matrix is written to HBM, read back (to compute softmax), written again, read again (to actually use)
flash attention computes attention in blocks using online softmax trick, keeping intermediate results in fast SRAM rather than writing the full attention matrix to GPU main memory
reduces activation memory for attention from to
faster because memory bandwidth isn’t the bottleneck
enable at inference time with
model.to_bettertransformer()
or attn_implementation="flash_attention_2"
in from_pretrained()
def flash_attention(q, k, v):
# Same input shapes: (batch, num_heads, seq_len, head_dim)
# PyTorch picks the best backend automatically (Flash Attention, memory-efficient, or math)
output = F.scaled_dot_product_attention(q, k, v, is_causal=True)
return output
FlashAttention never materializes the full matrix, and instead computes attention in tiles small enough to fit in SRAM
conceptually
load a block of (say, 64 rows)
load a block of K and V (say, 64 columns)
compute that tile of attention scores, apply softmax, multiply by V — all in SRAM
write only the final output to HBM
repeat for all tiles
since every (batch, head) combination is completely independent, you have
batch × num_heads
parallel workers working on each m × head_dim
tile of Q, K, V where m
is the tile sizeFlash Attention is an exact, not approximate, method
how to use Flash Attention
is_causal=True
flag is fused into the kernel efficiently rather than materializing a mask matrixdef flash_attention(q, k, v):
# Same input shapes: (batch, num_heads, seq_len, head_dim)
# PyTorch picks the best backend automatically (Flash Attention, memory-efficient, or math)
output = F.scaled_dot_product_attention(q, k, v, is_causal=True)
return output
Scaling laws
maximal update parameterization ()
core problem: hyperparameters found at small scale don’t generalize to larger scale
standard parameterization → different layers have updates of inconsistent magnitudes as width changes
adjusts initialization and learning rates per-layer so that the magnitude of updates relative to weights stays constant across widths
primarily adjusts width scaling
fit learning rate to compute
fit loss to compute
we need irreducible loss term (entropy of the data) because otherwise, as ,
how do we fit the equation? using least squares
least squares is the objective of minimizing the squared residuals
ordinary (or linear) least squares and non-linear least squares
linear least-squares has a closed-form solution
for , the closed-form solution is
non-linear least-squares is solved by iterative refinement
GPUs
high bandwidth memory (HBM) is the main GPU memory
slow memory from the GPU’s perspective
either 40GB or 80GB for an A100
static RAM (SRAM) is the small, fast, on-chip memory
~20MB total on an A100
Other architectures
RNNs
Vanilla RNN
at each step , process and previous hidden state to produce the next hidden state ( is input size, is hidden size)continuous
weight matrices ,
weights shared across time steps
the reason behind vanishing gradients
let
only shrinks, since
repeated applications of either leads to vanishing or exploring gradients
LSTM
LSTM uses cell state with forget, input, output gates
inputs at each step: previous hidden state , previous cell state , current input
forget gate: what to erase from cell state
input gate: what new information to write
cell state
cell state update: forget some information from the old cell state , and add some information from the new cell state
output gate: what to expose as the new hidden state
new hidden state
as long as stays near 1, gradients can travel backwards many timesteps without vanishing
cell state is memory, hidden state is working output
flows along the highway, touching only elementwise multiplications & additions, no matmuls or non-linearities
information added or removed only through gates
hidden state is a filtered view of the cell state
plays two roles
cell’s output to the outside
cell’s query into itself on the next step, because gates at are computed from (not )
separation of cell state and hidden state is what makes LSTMs work
RNN tries to make a single vector serve as both long-term memory and current output
GRU is a simplified version with two gates, reset and update
vs. transformers
gradients from a distant time step can’t influence an earlier time step’s processing
RNNs vs transformers
attention gives direct / path between any two tokens, doesn’t depend on the distance between tokens
each token directly looks at every other token
whereas in an RNN, information from token 1 has to survive through all the intermediate hidden states to reach token 100
helps long-range dependencies
attention computes all pairwise relationships at once
State space models
equation underlying SSMs
is a 1D input signal
is an -D latent state
SSMs are and parallel during training, then and sequential during inference
in contrast, transformers are parallel but have attention, and RNNs are but sequential
Mamba makes , , and functions of the input
selectively incorporate information via
selectively read from state via
control the timescale via
Post-training
policy gradients
notation
action : next-token at time step
state : text prefix at time step
: LM policy
: prompt sampled from start distribution over prompts
: trajectory (finite-horizon), aka rollout, episode
: reward from trajectory
goal
we can do this via gradient ascent
vanilla policy gradient (aka REINFORCE): the gradient of the objective can be written as
this is basically the same update as SFT but data is sampled from the policy and the gradient is weighted by
if is positive, we go in the direction of increasing for each token in
otherwise, we go in the opposite direction
the larger the magnitude of is, the bigger the step we take
derivation of the gradient
in practice, we estimate by sampling a batch of rollouts from policy from starting state
so-called policy gradient loss is just a scalar
pg_loss
such that pg_loss.backward()
produces gradients equivalent to the approximate policy gradient it is not a loss in the canonical sense: doesn’t tell us how good our policy is, it’s just a device for producing the correct gradient
there is no fixed objective since it’s constructed from data sampled under the current policy
baselined policy gradient
the problem with the vanilla policy gradient is that it is a very high-variance estimate
suppose that we have an easy prompt, so all responses get a positive reward
without a baseline all responses get reinforced, including the bad ones in the batch!
on average over training this is okay, but it leads to very noisy updates
if we (for example) use the average reward as a baseline, then below-average responses don’t get reinforced
baselined policy gradient: subtract a baseline function from in the gradient estimate
as long as is a function of only the state (and not ), it won’t introduce bias to the estimate of
we want to be correlated with , so that is small → less noisy gradients
this doesn’t change the objective we’re optimizing, only how we estimate the gradient of the objective
why is the baselined policy gradient an unbiased estimate of the policy gradients?
we will show that the baseline term has an expected value of 0
let , so we rewrite
first move the expectation inside the sum
before: for each trajectory, sum the ’s over time steps , then average over sampled trajectories
now: for each time step , average over sampled trajectories, then sum those averages
averaging over full trajectories is the same as average over just , because depends on only (the rest of the trajectory doesn’t matter for )
then we factor the joint expectation over into expectation over given in expectation over
now we use the definition of
note that we can only pull out of the expectation over because only depends on , which allows us to use the fact that
plugging this back in, we get
a very common choice for the baseline function (used by PPO) is , which estimates the expected reward given the partial sequence
this is unbiased because depends only on the state
the signal for each token: did the final reward exceed or fall short of what looked likely at this point in the generation?
a token that turns a bad-looking response (small ) into a good one (high ) gets more credit
connection to advantage
is a single sample of the Q function (the expected reward following policy given this action and state)
is an estimate of the value function (the expected reward given just this state, )
so is an estimate of , the advantage function, i.e., how much better action is than expected from state
in particular, it’s a Monte Carlo estimate of the advantage as it uses a single sample ’s reward to estimate
baseline functions generally all try to estimate , the expected return from the current state
learned value function [above]
RLOO: sample responses per prompt, the baseline for a prompt is the average of the other rewards (very similar to GRPO except for the inclusion of and the normalization by std)
batch estimate (REINFORCE++): use the mean reward in the batch
off-policy policy gradient
the problem with on-policy policy gradient is that we need to do inference from the policy for every gradient step, even though the policy is probably not changing that much
in off-policy learning, we sample rollouts from a policy different from the one we are optimizing
generally use rollouts from a previous version of the policy to optimize the current policy
use a surrogate objective
the fraction is a reweighting term in the style of importance sampling
background on importance sampling: if we want but only have samples , then we can rewrite
intuitively, if more likely under than under , then that sample counts more, and vice versa
note that this optimizes a different objective from the original !
the true off-policy estimate just rewrites with importance sampling, but it has terrible variance because of the product of ratios
essentially replaces the product with a sum of per-timestep terms
leads to the following off-policy policy gradient
to see why this is true, note that only the numerator of depends on , then apply log derivative trick
in practice we estimate it via where = number of rollouts per batch
PPO
introduced the clipping mechanism for importance weights
the regular surrogate objective is
PPO uses the value function as the baseline, so following convention we replace in the policy gradient with the per-timestep advantage
the clipped surrogate objective clips to stay within
clipping maintains stability when taking many gradient steps on a single batch of rollouts
gives up unbiasedness in exchange for more stable update
the clipped objective is a good approximation as long as you haven’t moved too far from
the four cases where clipping is used
when and , we use
doesn’t depend on → gradient is 0
we’ve already increased the probability of substantially relative to the old policy — stop pushing a good action further
when and , we use
gradient is also 0
stop pushing a bad action further down
when and , we use
does depend on [see off-policy policy gradient]
note that this means is allowed to be — this makes sense because so it is a bad token, so we let the token get pushed down
clipping is asymmetrical: it only activates when the policy is already moving in the direction the advantage encourages, but doesn’t prevent you from correcting a mistake
when and , we use
does depend on [see off-policy policy gradient]
PPO collects a batch of trajectories from the current policy, then takes multiple gradient steps using the clipped surrogate objective
RLHF
introduces a KL penalty to prevent from drifting too far from
in practice, the KL penalty is computed per-token and folded into the per-token reward
how reward models are trained
Bradley-Terry model says that the probability that is preferred over is
train LM with classification head to maximize
this is just the binary cross-entropy loss with true label always and under Bradley-Terrey
GRPO
GRPO replaces the reward with the advantage of relative to a group of sampled trajectories
note: it doesn’t cleanly count as a baseline choice because it also divides by a normalization factor
for a set of rollouts for question , each consisting of tokens
compute rewards for each sampled output using
the group-relative advantage estimate is
the advantage is the same for every token in the response
the objective is
where is the same per-token probability ratio (convention here is to use instead of )
simplifies PPO by removing the need for a critic (value function) by computing advantages relative to a group of samples
GRPO combines three ideas
off-policy policy gradient using :
the clipping mechanism [PPO]:
computing advantages with group normalization [DeepSeek R1]:
algorithm
policy model
for each step (
n_grpo_steps
)sample a batch from
update the old policy model
sample outputs for each question
compute rewards for each sampled output using
compute through group-relative advantage estimation
for each train step (
n_train_steps_per_rollout_batch
)update the policy model by maximizing the GRPO objective
update through continuous training using a replay mechanism
GRPO effectively computes an empirical baseline instead of a learned one
instead of learning to predict expected reward, just measure it by sampling
uses more compute ( completions per prompt) but less memory (no need to store value network)
Dr. GRPO fixes two problems with GRPO
in GRPO when is small (question is too easy or too hard), reward is amplified → more important to optimize that group
bias that upweighs problems that are too easy or too hard
in GRPO, reward is normalized by rollout length
among correct responses, short length → larger gradient → reinforced more strongly
among incorrect responses, long length → smaller gradient → under-penalized
model learns that if it can’t get an answer right, then just produce a really long answer
DPO
the KL-regularized objective has a closed-form optimal solution
rearranging
DPO loss is just the negative log likelihood of the observed preferences
Precision
mixed precision training
master weights in FP32
BF16 copy of weights for forward/backward pass
activations computed in BF16
gradients computed in BF16 and accumulated into FP32
this avoids the problem of adding a small gradient to a large weight
BF16 has more precision near 0, so it can represent a small gradient of 0.0001 but not the updated weight of 1.0001
we just need accumulated gradients to land in FP32, individual gradients are often tiny compared to the weight
.grad
on each parameter lives in the same dtype as the parameter, so individual grads are cast to FP32intuition
matmuls are tolerant of rounding noise, so BF16 for forward/backward is fine
master weights in FP32 necessary because individual grads are tiny
activations are harder to quantize than weights
precision options
FP32: full precision (4 bytes)
FP16: half the memory (2 bytes)
BF16: same memory as FP16 but better numerical stability (larger dynamic range)
INT8: quarter the memory of FP32, requires quantization (1 byte)
INT4: even smaller
data loading
memmap
avoids loading the entire data into memory at onceload model in BF16
when using HF
.from_pretrained()
to load modeltorch_dtype=torch.bfloat16
for BF16, recommended choiceweights and activations both in BF16
load_in_8bit=True
for quantization, uses LLM.int8()
weights in INT8, activations in FP16
provides memory savings via weight compression
not the same as full INT8 compute due to activations being in higher precision
model.half()
or model.to(torch.bfloat16)
converts a model to FP16model = MyModel()
model.load_state_dict(torch.load('model.pt'))
model = model.half()
all weights in BF16 and operations run in BF16
in practice, running everything in BF16 is fine for inference
torch.autocast()
does automatic mixed precisionweights stay in FP32, operations selectively use BF16 or FP32
manages precision for each operation, but doesn’t manage master weights
more memory than loading entire model in BF16 but potentially more stable
matmul in FP16 (tolerates lower precision), softmax in FP32 (needs precision for numerical stability), layernorm in FP32 (reductions need precision)
useful when model is in FP32 and we can’t easily convert it, or we see numerical issues with pure BF16
dtype
specifies the “lower precision” typewith torch.autocast(device_type='cuda', dtype=torch.bfloat16):
output = model(x)
to use
bitsandbytes
, replace nn.Linear
with bnb.nn.Linear8bitLt
or bnb.nn.Linear4bit
def replace_linear_with_8bit(model):
"""Replace all nn.Linear with bnb.nn.Linear8bitLt"""
for name, child in model.named_children():
if isinstance(child, nn.Linear):
# Create quantized replacement
new_layer = bnb.nn.Linear8bitLt(
child.in_features,
child.out_features,
bias=child.bias is not None,
has_fp16_weights=False,
)
# Copy weights (they'll be quantized when moved to CUDA)
new_layer.weight = bnb.nn.Int8Params(
child.weight.data,
requires_grad=False,
)
if child.bias is not None:
new_layer.bias = nn.Parameter(child.bias.data)
setattr(model, name, new_layer)
else:
# Recurse into child modules
replace_linear_with_8bit(child)
return model
# Usage
model = MyTransformer()
model.load_state_dict(torch.load('model.pt'))
model = replace_linear_with_8bit(model)
model = model.to('cuda') # quantization happens here
Parallelism
data parallelism splits up the batch across devices, whereas model parallelism splits up the computation for a single forward pass across devices
FSDP shards parameters across ranks for memory efficiency, but every rank still computes the full forward pass by all-gathering parameters right before they’re needed
limitations of data parallelism
requires , which is not necessarily good because we don’t want to be larger than the “critical batch size”
models still may not fit on one device (even ZeRO stage 3 doesn’t reduce the activation memory per device)
strong scaling: increasing the number of chips for training leads to proportional increase in throughput (FLOPs/second)
DP scales throughput, TP/PP scale model memory, SP scales activation memory
5D parallelism
data parallelism (DP): split batch of data across devices
tensor parallelism (TP): split different layers/stages of the model across devices
pipeline parallelism (TP): split individual layers/weight matrices across devices
sequence parallelism (SP): split the input sequence length across devices
expert parallelism (EP): distribute different experts in a MoE model across devices
Background: core collective operations
broadcast (one to all, same data): one GPU has the data and sends an identical copy to every other GPU
all-gather (all to all): each GPU has a piece of data, every GPU gets full collection
removes sharding along an axis
reduce-scatter: each GPU has unreduced data, combine via reduction, result sharded across GPUs
very similar to all-gather but instead of retaining each shard, we sum them together
adds sharding along an axis
all-reduce: every GPU has unreduced data, combine via reduction, every GPU gets final result
ring all-reduce: reduce-scatter + all-gather (each process communicates with two neighbors) [animation]
at the beginning, each GPU has unreduced data
reduce-scatter: combine the data via reduction, each GPU has a reduced subset
does all the arithmetic, none of the redundant copying
all-gather: every GPU gets the full collection of reduced subsets
does all the copying, none of the arithmetic
for all-gather, reduce-scatter and all-reduce, the communication time depends only on the size of the array and the bandwidth, not the number of devices over which our array is sharded!
reduce-scatter and all-gather are used in each other’s backward pass
all-gather in the forward → reduce-scatter in the backward
all-gather broadcasts the same chunk to every device, where they participate in different downstream computations
upstream gradients sum at outward branches (if , then )
this is exactly a reduce-scatter: sum all upstream gradients onto device where it came from
fan-out in the forward → sum in the backward
reduce-scatter in the forward → all-gather in the backward
reduce-scatter sums many inputs into one chunk
in the backward, a summation node copies the upstream gradient to each summand
this is exactly all-gather: broadcast each chunk’s gradient back to its contributors
sum in the forward → fan out in the backward
this means backward of all-reduce is another all-reduce!
partitioning notation
mesh has axes named and matrix has axes
: split rows of across columns of device mesh
partition axis of (rows) along mesh axis (along each row)
: split rows of across rows of device mesh
: split columns of across rows of device mesh
: split columns of across columns of device mesh
: split rows of across all devices in flattened mesh
: don’t split up rows of
a mesh dimension not appearing means that data is replicated along that dimension
doesn’t appear: each column contains the same data
nice property of matmul: when matrix multiplicands are written in terms of blocks, the product can be written in terms of block matmuls
four matmul cases
case 1: neither matrix has a sharded contracting dimension
requires no communication
perform local block matrix multiplies
output is naturally sharded in the desired way
case 2: either or has a sharded contracting dimension
all-gather the shards of so every device has a full copy, then multiply against
case 3: both and have sharded contracting dimensions
the matmul is possible but each device only represents a partial sum of the desired product
each device along dimension has a different partial sum
we use the notation to mean unreduced along mesh axis
do the final summation using an all-reduce across the axis
results in each device having the same fully-summed value
case 4: both and have non-contracting dimension sharded along the same axis
Data parallelism
naive data parallelism (DDP)
split examples within a -sized batch across devices and exchange gradients
steps
run forward pass on local micro-batch
compute grads for all params
all-reduce gradients across GPUs so each GPU has the averaged G
each GPU independently updates all params
each GPU does identical work in this step
when model fits on a single device, we should always use this
communication hapens only on the backward pass
allows arbitrarily increasing batch size over more and more devices
ZeRO stage 1: shard optimizer states
each device holds of the Adam vectors, and only updates the parameters that it has optimizers for
steps
run forward pass on local micro-batch
compute grads for all params
reduce-scatter gradients: each GPU has reduced gradients for its subset of parameters
each GPU updates their subset of params using its subset of optimizer states
all-gather the updated parameters
because one all-reduce has the same cost as reduce-scatter + all-gather, this incurs no additional communication overhead, making it free memory wins
ZeRO stage 2: shard optimizer states + gradients
also shard gradients so each device only holds of the gradients
steps
run forward pass on local micro-batch
as each layer’s gradients are computed, they’re immediately reduce-scattered. each GPU keeps only its subset of gradients and frees the rest.
each GPU updates their subset of params using its subset of optimizer states
all-gather the updated parameters
ZeRO stage 3: shard optimizer states + gradients + parameters (FSDP)
also shards parameters so each device only holds of the parameters
all-gather parameters for each layer just in time, use them, then discard
can train models that don’t fit on one GPU
ZeRO stages 1,2,3 all have the same communication cost
Pipeline parallelism
if we did model parallelism naively, we would move each layer to a different GPU, and move the hidden states to the right GPU before each layer (this is bad, we didn’t improve our throughput)
proper async pipeline parallelism
from torch.distributed.pipeline.sync import Pipe
model = nn.Sequential(
nn.Linear(512, 512).to('cuda:0'),
nn.ReLU().to('cuda:0'),
nn.Linear(512, 512).to('cuda:1'),
nn.ReLU().to('cuda:1'),
)
model = Pipe(model, chunks=8) # splits batch into 8 micro-batches
output = model(x)
Tensor parallelism
FSDP and tensor parallelism can be effectively combined
sharding batch dimension reduces the size of all-gathers, and sharding FFN dimension reduces the communication overhead of FSDP
for an MLP (two linear layers with activation in between), we do column parallel (for ) → activations → row parallel (for )
split the first weight matrix column-wise
each device computes its shard and applies the activation locally
split the second weight matrix row-wise
each device computes a partial result, then all-reduce to sum
the key insight
y = h @ W2
= [h_0, h_1] @ [W2_0]
[W2_1]
= h_0 @ W2_0 + h_1 @ W2_1
= y_0 + y_1
any other sharding choice forces a collective in the middle
for a matmul
column sharding of means each device computes a slice of the output, producing nicely sharded output
row sharding of means each devices computes a partial sum of the full output, needing reduction to finish
for SwiGLU-style MLP, we column-shard both and , keep element-wise product sharded on , then row-shard and all-reduce at the end
the code
class ColumnParallelLinear(nn.Module):
def __init__(self, in_features, out_features, world_size, rank):
super().__init__()
self.out_features_per_rank = out_features // world_size
self.rank = rank
self.linear = nn.Linear(
in_features,
self.out_features_per_rank,
bias=False,
device=f"cuda:{rank}"
)
def forward(self, x):
return self.linear(x)
class RowParallelLinear(nn.Module):
def __init__(self, in_features, out_features, world_size, rank):
super().__init__()
self.in_features_per_rank = in_features // world_size
self.rank = rank
self.linear = nn.Linear(
self.in_features_per_rank,
out_features,
bias=False,
device=f"cuda:{rank}"
)
def forward(self, x):
# x is partial: (batch, in_features_per_rank)
# each GPU computes partial result
partial = self.linear(x)
# All-reduce to sum partial results across GPUs
dist.all_reduce(partial, op=dist.ReduceOp.SUM)
return partial
class TensorParallelMLP(nn.Module):
"""
column parallel -> activation -> row parallel requires only one all-reduce
"""
def __init__(self, d_model, d_ff, world_size, rank):
super().__init__()
assert d_model % world_size == d_ff % world_size == 0
self.fc1 = ColumnParallelLinear(d_model, d_ff, world_size, rank)
self.fc2 = RowParallelLinear(d_ff, d_model, world_size, rank)
def forward(self, x):
x = self.fc1(x) # (batch, seq, d_ff // world_size)
x = nn.functional.silu(x)
x = self.fc2(x) # (batch, seq, d_model), all-reduced
return x
for attention, each device handles a subset of attention heads (since they are independent)
one all-reduce at the end (like the MLP)
attention heads are independent, so splitting by heads requires no communication during the attention computation!
column parallel (QKV) → computation → row parallel (output)
class TensorParallelAttention(nn.Module):
def __init__(self, d_model, num_heads, world_size, rank):
super().__init__()
assert num_heads % world_size == 0
self.rank = rank
# Each GPU handles a subset of heads
# QKV projection for local heads only
self.qkv = ColumnParallelLinear(
d_model,
3 * d_model,
world_size,
rank
)
self.out_proj = RowParallelLinear(
d_model,
d_model,
world_size,
rank
)
def forward(self, x):
batch, seq_len, _ = x.shape
# Project to local Q, K, V
qkv = self.qkv(x)
qkv = qkv.reshape(batch, seq_len, 3, self.num_heads_per_rank, self.head_dim)
q, k, v = qkv.unbind(dim=2)
# Attention on local heads
q = q.transpose(1, 2) # (batch, local_heads, seq_len, head_dim)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
attn_out = nn.functional.scaled_dot_product_attention(q, k, v)
attn_out = attn_out.transpose(1, 2).reshape(batch, seq_len, -1)
# Output projection with all-reduce
return self.out_proj(attn_out)
Multimodality
Vision Transformer (ViT): turn image into a sequence of vectors, then run a standard transformer encoder
LLaVA
model hidden dimension is , CLIP hidden dimension is 1024
Image [224, 224, 3]
↓ patchify
Patches [256, 588]
↓ linear projection (independent)
Patch embeddings [256, 1024]
↓ prepend CLS, add position embeddings
Sequence [257, 1024]
↓ CLIP (transformer encoder with 24 layers)
CLIP output [256, 1024]
↓ projection W applied independently per token
Visual embeddings [256, 4096]
↓ concatenate with text embeddings [N, 4096]
Full sequence [256 + N, 4096]
↓ LLM transformer
LLaVA architecture
LLaVA-NeXT dynamic resolution
split high-res images into multiple crops, encode each separately, and then concatenate
Qwen2-VL
previously, input size is fixed (e.g., 224 x 224) → always 256 patches
learn absolute position embeddings of shape
we were constrained to fixed-size images simply because we had a fixed number of learned position embeddings
now simply use 2D-rope
patch size is still fixed
position is encoded as (row, column) coordinates
RoPE generates positional encoding on-the-fly
lets models generalize to arbitrary resolutions and aspect ratios
current paradigm
understanding: ViT encoder → features
generation: diffusion operating in pixel space
CLIP
trained with contrastive learning
maximize similarity of text/image embeddings for the correct pairings