Implementing the LinearWithLoRA Layer

The replace_linear_with_lora function walks the model's children, swaps each Linear layer for a LinearWithLoRA layer built with rank and alpha, and recurses into every other child.

Programming

A Function That Rewrites a Model

This article begins with a function rather than a definition. The function is replace_linear_with_lora, and its job is to walk through an existing GPTModel and swap its Linear layers for LinearWithLoRA layers. LinearWithLoRA is the layer type that combines the behavior of a Linear layer with LoRA, and the replacement function is the mechanism that puts those new layers in place of the old ones throughout the model.

Keep the setting in mind. The GPTModel is the model you built in earlier chapters, and the goal is parameter-efficient fine-tuning. Read the function once in full before you run anything. Then try to answer the question the rest of this article is built around: when this function finishes, which objects in the model have changed, and which have not?

What the Function Changes in Place

The replace_linear_with_lora function has a summary you can hold onto: it walks the model's children, swaps each Linear layer for a LinearWithLoRA layer built with rank and alpha, and recurses into every other child. Each part of that summary corresponds to a concrete behavior of the function, and the rest of this article takes them one at a time.

The function modifies the model in place. It does this through setattr, which assigns a new object to a named attribute of a parent module. Because the change happens to the existing model object, the function does not need to return a new model. It returns nothing.

The rank and alpha values are the settings used to build each new LinearWithLoRA layer. The source names them as the construction inputs for the upgraded layer, and this article focuses on the traversal and placement logic that decides where those layers go. You do not need the internal math of LoRA to follow the replacement logic.

Two Decisions Made at Every Level

The recursion is easiest to understand as a sequence of decisions. The function looks at one child of the current module at a time. For each child, it asks one question: is this child a Linear layer? If it is, the function swaps it for a LinearWithLoRA layer. If it is not, the function does not give up on that child. It hands the child to another call of the same function, so that the child's own children get the same treatment.

The traversal finishes only when every child at every depth has been visited. This is the key property of the recursion. A Linear layer buried several levels down is still reached, because each non-Linear container is entered rather than skipped. The flowchart below follows a single child through this decision.

For each child module, the function asks whether the child is a Linear layer. A Linear child is swapped; any other child is handed to another call of the same function, and the loop continues until no children remain.

Notice that the function makes the decision based on what the child is at that moment. A Linear child is replaced, and that is the only branch that changes the model's layers. Everything else is treated as a container that might hold Linear layers of its own. This is why the function can work on a model whose structure it has never been told about in detail.

Tracing a Toy Model by Hand

The GPTModel is too large to trace line by line, so start with a small constructed model. This toy model is not part of the GPT model in the source. It uses the same rules as the function, and it shows the two things the function does at every level: it swaps any Linear layer it finds among the direct children, and it descends into everything else.

Suppose the toy model has three direct children. The first is fc_in, a Linear layer. The second is act, an activation module with no children of its own. The third is block, a container holding two children: fc_mid, a Linear layer, and drop, a dropout module. Before reading further, predict which of these nodes will be replaced after one call to the function.

What do you think happens?

After one call to replace_linear_with_lora on the toy model, which Linear layers become LinearWithLoRA layers?

Only fc_in, because it is a direct child
Only fc_mid, because it is inside a container
Both fc_in and fc_mid
Neither, because the function only changes the top-level model

Reveal answer

Answer:

Both fc_in and fc_mid

The function swaps fc_in directly because it is a Linear child. For block, which is not a Linear layer, the function recurses into it, finds fc_mid as a Linear child, and swaps it. The act and drop modules are not Linear layers, so they are visited but not replaced.

The traversal starts at the toy model's root and visits each child. The Linear layers are swapped, the activation and dropout modules are visited and left as they are, and the block container is entered so its children are visited too.

Before the call, fc_in and fc_mid are plain Linear layers. After the call, both have been replaced by LinearWithLoRA layers. The activation, dropout, and container modules keep their types, and the container's position in the tree does not change.

Worked example

Replacing the toy model's Linear layers

Given the toy model with direct children fc_in (Linear), act (activation), and block (container holding fc_mid (Linear) and drop (dropout)), determine which modules change type after one call to replace_linear_with_lora.

1
Visit fc_in
fc_in is a direct child and is a Linear layer, so the function swaps it for a LinearWithLoRA layer.
2
Visit act
act is not a Linear layer, so the function recurses into it. It has no children, so nothing else happens and its type stays the same.
3
Visit block
block is not a Linear layer, so the function recurses into it instead of swapping it.
4
Visit fc_mid inside block
Inside the recursive call, fc_mid is a Linear child, so it is swapped for a LinearWithLoRA layer.
5
Visit drop inside block
drop is not a Linear layer and has no Linear children, so it is visited and left as it is.
6
Check the result
fc_in and fc_mid changed type. act, block, and drop did not. The function returned nothing, so the caller inspects the same toy model object.

Result

Two Linear layers were replaced: fc_in at the top level and fc_mid inside block. Every other module kept its type.

Where the Upgrade Lands in GPTModel

The function is generic, but the source applies it to a specific target. The intended upgrade covers the Linear layers found in the multihead attention, the feed-forward modules, and the output layer of the GPTModel. The diagram below maps the GPT architecture with those parts marked. Read it as a map of where the replacement is intended to take effect, not as a listing of every module in the model.

Inside each transformer block, the multihead attention and feed-forward modules are upgraded. The output layer at the end of the model is also upgraded. The embedding layers, dropout, and final LayerNorm are part of the architecture but are not marked as upgraded.

The upgraded parts are split into two groups, and it helps to hold them apart. The first group sits inside each transformer block. The second is the output layer, which is at the end of the model rather than inside a block. The sections below describe each group and what stays unchanged.

Mistakes When Predicting a Replacement

Common mistake
Assigning the result of the function to a new variable
What it looks like
```
updated = replace_linear_with_lora(model)
```
Why it fails
The function modifies the model in place and returns nothing, so updated is None.
Do this instead
Call replace_linear_with_lora(model) and keep using the original model object.
Common mistake
Assuming a Linear layer nested inside a container is missed
What it looks like
Expecting fc_mid inside block in the toy model to stay a plain Linear layer.
Why it fails
Every non-Linear child is recursed into, so the traversal reaches Linear layers at every depth.
Do this instead
Trace the recursion into each container before deciding a layer was skipped.
Common mistake
Assuming every module in GPTModel is upgraded
What it looks like
Expecting the embedding layers or final LayerNorm to become LinearWithLoRA.
Why it fails
The source marks the upgrade for the Linear layers in multihead attention, feed-forward modules, and the output layer. Embeddings, dropout, and LayerNorm are not marked as upgraded.
Do this instead
Compare the final model against the upgraded set rather than against the full architecture.
Common mistake
Confusing the decision with the module's name rather than its type

Predicting Results and Locating Surprises

When a replacement run gives an unexpected result, the source suggests a three-part check. Work through it in order rather than guessing. First, check whether the target layer was reached by the traversal. Second, check whether it was a Linear layer at the moment the function visited it. Third, check which object the code is now using. Each question narrows the search.

MEDIUM

A colleague runs replace_linear_with_lora on a GPTModel and then prints the type of one feed-forward Linear layer, which still shows as Linear. Using the three-part check, list what you would verify first, second, and third, and say which of the three is most likely to reveal the problem if the caller had stored the return value in a new variable.

Hints

Start by asking whether the traversal could reach that layer at all.
Remember that the function returns nothing and modifies the model in place.
Think about which object the caller is printing from.

Your attempt (not saved)

EASY

Sketch a small model with two Linear layers at the top level and one container holding a third Linear layer. Predict which of the three will be replaced, then list each visit in order. Compare your list with the traversal rule: swap any Linear child, and recurse into everything else.