2️⃣ LoRA

Models learn a lot during the pre-training phase, and probably internally already have lots of different behaviours and skills picked up during that time. If all we have to do is "index into the list of skills the model has learned to select the right one", then a constraint of 1.5 KiB suddenly might be quite a lot. More formally, the argument to make would be "if the shortest description of a property/behaviour/algorithm, conditioned on that you have access to the base model weights exceeds 1.5KiB, then it cannot possibly be learned by any such RL training run as described above.

As such, RL for LLMs has sometimes been described as merely eliciting an existing persona or skill that the model was already aware of, and then reinforcing that persona/skill until it is the dominant thing that the LLM wants to do. We will revisit this claim in the RLVR section.

The paper. Section 4.2 limits LoRA to the attention weights, freezing the MLPs. Then section 7.1 compares adapting different attention weights on GPT-3 with a fixed budget of trainable parameters (Table 5). Putting the whole budget into $W_Q$ or $W_K$ alone does noticeably worse, and the authors conclude that adapting $W_Q$ and $W_V$ "gives the best performance overall". (Though in the same table, adapting all four of $W_Q, W_K, W_V, W_O$ at half the rank does just as well or better.) But they never tried the MLPs, so this only tells us which attention weights are best.

Current practice. Adapt every linear layer: the attention projections q_proj, k_proj, v_proj and o_proj, and the MLP projections gate_proj, up_proj and down_proj. LoRA Without Regret found that this is one of the two conditions for LoRA to match full fine-tuning: attention-only LoRA underperforms even when you raise its rank to match the parameter count, while adding the attention layers on top of the MLP layers gives no extra benefit (the MLPs hold most of the parameters). For RL with policy gradients, LoRA matched full fine-tuning even at rank 1. TRL's LoRA guide and Unsloth's recommend the same, and QLoRA found it necessary to match full fine-tuning too.

peft's default still follows the paper. If you don't pass target_modules, LoraConfig leaves it as None, and peft then looks up the model's architecture in its table of defaults, which for Qwen3 is just ["q_proj", "v_proj"] (the same as for Llama). To get current practice, you have to ask for it with LoraConfig(target_modules="all-linear").

As in the rest of this course, we write layers as row vectors times matrices, $h = x W_0$ with $W_0 \in \mathbb{R}^{d_\text{in} \times d_\text{out}}$. So we store A with shape (in_features, rank) and B with shape (rank, out_features), and the update is $x A B$.

There are two places where you'll see the transpose of this:

• The LoRA paper uses column vectors, $h = W_0 x$, so it writes the update as $\Delta W = BA$ with $B \in \mathbb{R}^{d_\text{out} \times r}$ and $A \in \mathbb{R}^{r \times d_\text{in}}$. The names keep the same roles: $A$ reads from the input (and is initialized randomly), $B$ writes to the output (and is initialized to zero).
• PyTorch's nn.Linear stores its weight as (out_features, in_features), the transpose of our $W_0$, and computes x @ weight.T. This only matters in merge, where we add our update into base.weight. With an einsum, the output pattern d_out d_in does the transpose for us.

The paper says only "random Gaussian", without giving a variance. We use $1/\sqrt{\text{in\_features}}$ so that x @ A has roughly the same scale as x: each of its entries is a sum of in_features terms. This is the same scale as PyTorch's default initialization of an nn.Linear layer, and similar to the reference implementation loralib, which uses PyTorch's default (uniform) initialization for $A$. The peft library's init_lora_weights="gaussian" option uses a standard deviation of $1/r$ instead. Since $B$ starts at zero, the choice only affects how quickly the update grows during training.

We'll load the 1.7B base model in bfloat16 to halve its memory, which is fine since it's frozen. But the adapters are what we're training, and bfloat16 only has 8 bits of mantissa: a small Adam update to a bfloat16 weight often rounds away to nothing. This is the usual mixed-precision recipe: keep the trainable parameters (and optimizer state) in float32. The adapters are tiny, so this costs almost no memory.

$AB$ is a full (in_features, out_features) matrix, as big as the original weight, but it is only a rank-r matrix. Computing it on every forward pass would waste memory and compute. Going through the rank-$r$ bottleneck (first $xA$, then $(xA)B$) costs only $O(r \, (d_\text{in} + d_\text{out}))$ per token, and gives the same result.

For GPT-2 and other older style models, the MLP layer is linear-gelu-linear via two matricies, the up-proj and then down-proj matricies.

More modern models uses Gated Linear Units (GLUs), but we will coloquially refer to it as the "MLP block" (even though it's not a MLP).

GLU (Gated Linear Unit) performs the following operation (ignoring biases) on its inputs:

where $x$ is the input, $\odot$ denotes elementwise multiplication, and $\sigma$ is some choice of non-linear activation function, usally SiLU,

This means we have two weight matricies doing different jobs: $W_{\text{up}}$ performs some linear transformer on the input, the same as the original MLP has, and $W_{\text{gate}}$ performs a gating operation, performing a linear transformation on the input to decide what intermediate features in the hidden layer of the GLU to turn on or off. The non-linear activation SiLU (which looks like a smooth ReLU) then clamps the negative outputs of the gate to zero. Elementwise multiplying this by the output of $W_{\text{up}}$ allows the gate to selectively disable features in the hidden layer. What was the job of one matrix (the first layer of a MLP has to both do some transformation on the input, and decide what features to turn off by sending them to negative values) is now split into two ($W_{\text{up}}$ computes the features, $W_{\text{gate}}$ decides which to keep alive).

For use in the "MLP" block, the entire operation then becomes

Here it is in the source code for Qwen3 (act_fn is SiLU, since model.config.hidden_act == "silu"):

class Qwen3MLP(nn.Module):
    ...
    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj

Okay. What's the justification of SiLU over ReLU or GLU over MLP?

Great question! Allow me to quote the justification given in [the paper](https://arxiv.org/pdf/2002.05202) itself:

Here are the relevant fields of the model's config (model.config), annotated with the notation you're used to:

Qwen3Config {
  "head_dim": 128,              # d_head
  "hidden_size": 2048,          # d_model
  "intermediate_size": 6144,    # d_mlp
  "num_attention_heads": 16,    # number of query heads
  "num_key_value_heads": 8,     # number of key and value heads
  "num_hidden_layers": 28,
  "tie_word_embeddings": true,
  "vocab_size": 151936,
  ...
}

q_proj computes all 16 query heads at once, as one matrix multiplication whose output is the heads concatenated together, which then gets reshaped. From the source code for Llama and Qwen3:

# modeling_llama.py
self.q_proj = nn.Linear(
    config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
)
# modular_qwen3.py
input_shape = hidden_states.shape[:-1]
hidden_shape = (*input_shape, -1, self.head_dim)

query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)

or, written a bit more clearly with einops:

query_states = rearrange(
    self.q_proj(hidden_states),
    "batch seq (n_heads head_dim) -> batch n_heads seq head_dim",
    head_dim=self.head_dim,
)
query_states = self.q_norm(query_states)

So q_proj has $d_{model}$ input features, and $n_\text{heads} \times d_\text{head} = 16 \times 128 = 2048$ output features. For Qwen3-1.7B that happens to equal $d_{model}$, but it needn't: Qwen3-0.6B's q_proj is $1024 \to 2048$. k_proj and v_proj output $8 \times 128 = 1024$ features, since there are only 8 key and value heads.

KV-caching means storing the keys and values for every token generated so far, and for long sequences this cache can use more memory than the model itself. Grouped-query attention reduces this by having fewer key/value heads than query heads: the query heads are put into groups, and each group shares a single key and value head.

TransformerLens was nice enough to provide a unified interface to the modules of a given model with a large lookup table that does the remapping from the existing names (e.g. here's where it splits GPT-2's fused c_attn into W_Q, W_K and W_V).

HuggingFace is just the model as-is without this rewriting step. So for now we can only support Llama (and models based on it like Qwen, Mistral, Gemma, ...), with the attention targets q_proj, k_proj, v_proj and o_proj, and the MLP targets gate_proj, up_proj and down_proj.

Other architectures name their layers differently:

apply_lora itself works for any model, as long as you pass it the right target_modules: print the model to see them. The peft library keeps a table of the usual LoRA targets for each architecture, in peft.utils.TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING (e.g. "gpt2" maps to ["c_attn"]).

GPT-2 is a special case: its layers are Hugging Face's Conv1D rather than nn.Linear, with the weight stored as (in_features, out_features). This is a quirk inherited from the original OpenAI code, which named its linear layers conv1d (a matrix multiply is the same as a 1D convolution with a kernel of size 1). So our apply_lora won't find any layers to wrap in GPT-2. peft has extra code to handle this, but for our purposes it's fine for our homebrew LoRA to only support the Llama style of architectures.

Setting requires_grad=False on a parameter tells autograd not to compute the gradient with respect to that parameter. It doesn't stop gradients flowing through the layers that use it.

Consider a frozen layer $h = x W_0$ somewhere in the middle of the model. Backprop could compute two gradients here:

• $\frac{\partial L}{\partial W_0}$, how the loss changes as the weights change. We don't need it, so autograd skips it.
• $\frac{\partial L}{\partial x} = \frac{\partial L}{\partial h} \frac{\partial h}{\partial x} = \frac{\partial L}{\partial h} W_0^T$, the gradient for the layer's input. This still gets computed, because it's how the gradient reaches the adapters in earlier layers. It uses the value of $W_0$, but not its gradient $\frac{\partial L}{\partial W_0}$.

Autograd decides which of these to compute by tracking which tensors are upstream of something trainable. Every activation after the first adapter depends on that adapter's A and B, so autograd keeps track of all of them, and passes the gradient back through every frozen layer on the way. (Before the first adapter nothing is trainable, e.g. the frozen embeddings, so autograd doesn't go back that far at all.) Here's a minimal example:

W = nn.Parameter(t.randn(4, 3), requires_grad=False)  # frozen, like W_0
x = t.randn(2, 4, requires_grad=True)  # an activation, which depends on an earlier adapter
(x @ W).sum().backward()

print(x.grad)  # exists as the gradient flowed *through* the frozen layer
print(W.grad)  # None: we never computed the gradient *for* the frozen weight itself

So LoRA doesn't skip the backward pass: we still do a full backward pass through the model, and still store the activations it needs, but we skip computing the gradients $\frac{\partial L}{\partial W_0}$, as well as skip tracking the optimizer state for $W_0$.

Another benefit of this is that autograd is smart, and will only keep backpropagating until everything marked as a trainable parameter has received gradients, following the topological sort order that you built in [0.4]. If, say, we only attached LoRA adaptors to layers in the middle of the model, then the backprop stage would terminate early, and training would run faster. This is why fine-tuning/retraining only the classifier layer for an image model (you did this in [0.2]) is so fast: the forward pass has to go all the way through the model, but the backwards pass only has to touch the last layer, and never look at the rest of the model.

Phase 001/060, Mean reward: 0.4827, KL: 0.0000
┌──────────┬────────────────┬────────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                     │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.99 │        -150.77 │ 'I thought this movie was really bad (especially the act two), because it’s possible that  │
│          │                │ something in act 1 was bruised us, but with a new sort of bruise, we couldn’t declare      │
│          │                │ things over completely and'                                                                │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.01 │        -137.07 │ "I thought this movie was fascinating and fascinatingly gritty. Even if the overall script │
│          │                │ is a bit disjointed (that seems to change every half-hour), Williams' acting, in           │
│          │                │ particular growing into the character over the two-hour run"                               │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.91 │        -163.22 │ 'I thought this movie was bottomless. I treated it like it weighed less than a             │
│          │                │ broom/sleeping bag. And thrift store trash really was that charming. As a fan of Good Will │
│          │                │ Hunting et al this is gross'                                                               │
└──────────┴────────────────┴────────────────────────────────────────────────────────────────────────────────────────────┘

...

Phase 058/060, Mean reward: 0.9591, KL: 0.0541
┌──────────┬────────────────┬────────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                     │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     1    │        -175.64 │ 'I thought this movie was the worst The Force Ever Had Let Out.\nCarly bloody Briggs was   │
│          │                │ so boring, there was nothing going on except bad joke/line delivery. Every gun fight was   │
│          │                │ unrealistic and the flicked through the'                                                   │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     1    │        -110.02 │ "I thought this movie was a terrible drag. The bad acting, the horrible title and the      │
│          │                │ boring and incoherent story. I couldn't understand why many people came from all over the  │
│          │                │ planet for this. It's just a"                                                              │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.99 │        -162.9  │ 'I thought this movie was actually really bad. While transportaion situations are not hard │
│          │                │ to believe in as yet, most of the humour came from the GAZettes desperate attempts to be   │
│          │                │ fashionable or from their involvement in long shots'                                       │
└──────────┴────────────────┴────────────────────────────────────────────────────────────────────────────────────────────┘

Phase 059/060, Mean reward: 0.9503, KL: 0.0526
┌──────────┬────────────────┬────────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                     │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.99 │        -136.9  │ "I thought this movie was horrible! It's hilarious that the movie appears to have a young- │
│          │                │ ish director done just a couple features and some TV commercials. I'm so mad that I        │
│          │                │ remember so many terrible things, but this movie"                                          │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.99 │        -162.93 │ 'I thought this movie was pretty bad, unless it was basically just a scenario. There were  │
│          │                │ some interesting revelations about the Victorian Era and disgusting historical facts JUST  │
│          │                │ like in the book, but they were mil at before. The characters are'                         │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.99 │        -101.47 │ 'I thought this movie was kinda a waste. Not really bad, but more like a waste. When you   │
│          │                │ first feel bad in the movie, you will feel just as bad when you keep watching.\nWhat is    │
│          │                │ the sentiment of this'                                                                     │
└──────────┴────────────────┴────────────────────────────────────────────────────────────────────────────────────────────┘

Phase 060/060, Mean reward: 0.9687, KL: 0.0522
┌──────────┬────────────────┬────────────────────────────────────────────────────────────────────────────────────────────┐
│   Reward │   Ref logprobs │ Sample                                                                                     │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.99 │        -180.27 │ 'I thought this movie was just terrible. I did notice a lot of older people nodding on     │
│          │                │ viewings and feeling everything I felt. Dylan Mayer should have kicked your butt like a    │
│          │                │ mean Dalmatian Shardack now-a'                                                             │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.99 │        -122.91 │ 'I thought this movie was so much worse than this one! Somehow Johnny Depp is still in     │
│          │                │ that same job (on this movie!) and the actors on this one are bad. I know Johnny Depp      │
│          │                │ usually does good movies'                                                                  │
├──────────┼────────────────┼────────────────────────────────────────────────────────────────────────────────────────────┤
│     0.99 │        -130.77 │ "I thought this movie was really terrible and I wasn't surprised at all to hear others say │
│          │                │ the same thing. I will not repeat a review that disrespectful of others opinions.\nThey    │
│          │                │ made a sick movie what the others said The actors"                                         │
└──────────┴────────────────┴────────────────────────────────────────────────────────────────────────────────────────────┘