Handout #17 — Adapting Language Models

A small adaptation job can run without errors and still train the wrong tokens, update the wrong weights, or cut off an adapter’s gradients. We will follow one example from a prompt–response pair to a loss, trace what “frozen” means in code, and build a resource budget that counts more than trainable parameters.

1What exactly should the model learn?

1.1 One task, several possible objectives

Use the customer-support task from slide 3. Let $X$ contain the instruction and message, and let $Y$ be the desired response:

Instruction: Extract issue and product. Return JSON.
Message: My headphones won’t charge even after being plugged in overnight.
Response: {"issue": "charging_failure", "product": "headphones"}

Training should improve $p_\theta(Y\mid X)$. The response teaches both a semantic decision and an output format. A model can emit valid JSON with the wrong issue, or select the right issue in invalid JSON. Keep those failure modes separate when evaluating.

Here $Y_i$ includes the response’s stopping token. This token-mean objective gives longer responses more total weight. Averaging each example’s mean loss instead gives equal weight to examples; either is a choice, but they are not the same objective. With accumulation, normalize over the intended scored-token count, as in Handout #16.

1.2 Work through the boundary by hand

For a tiny illustration, pretend the response has just two tokens, y1 y2. These are symbolic tokens, not the actual tokenization of the JSON above. In a decoder-only model, the output at position $j$ predicts the token at $j+1$.

Before shifting, labels match the input token IDs, except that prompt and padding positions are replaced by -100, PyTorch cross-entropy’s default ignored target. The separator’s input position produces the first scored prediction; do not accidentally mask that prediction just because it belongs to the prompt.

# From check_response_loss in the downloadable lab; logits are raw scores.
shifted_logits = logits[:, :-1, :].contiguous()
shifted_labels = labels[:, 1:].contiguous()
loss = F.cross_entropy(shifted_logits.view(-1, 8),
                       shifted_labels.view(-1), ignore_index=-100)

In a real implementation, apply the model’s chat template consistently during training and generation, then inspect the tokenized prompt/response boundary. Tokenizing two strings separately and concatenating them need not match tokenizing the combined string. Some causal-LM APIs shift labels internally: follow their contract and do not shift a second time. Reject or repair examples with zero scored response tokens after truncation.

2Freeze parameters, not the gradient path

2.1 When a frozen backbone needs no backward pass

For head tuning, compute $H=f_{\mathrm{frozen}}(X)$ and logits $Z=HW_{\mathrm{head}}^\top$. If neither the backbone parameters nor its inputs require gradients, $H$ has no autograd history. Training the head needs $H$ but not the backbone’s internal activations.

requires_grad_(False) controls parameter gradients. eval() controls behavior such as dropout; it does not freeze weights. Choose dropout behavior deliberately and configure the optimizer after selecting trainable parameters. See the PyTorch autograd notes.

# Sketch: backbone and head are separate modules, with no shared parameters.
backbone.requires_grad_(False)
head.requires_grad_(True)
optimizer = torch.optim.AdamW(head.parameters(), lr=1e-4)
with torch.no_grad():
    h = backbone(input_ids)
logits = head(h)  # Keep this OUTSIDE no_grad().

If an input embedding and the vocabulary head share the same Parameter, changing its trainability affects both uses. Check parameter identity before claiming that only the output mapping changes. Here we assume separate parameters; untying them can itself require an additional copy of the weights.

2.2 A frozen layer can transmit gradients

Now put a trainable adapter before a frozen linear layer: $h=Wx$, where $W$ is frozen but $x$ depends on the adapter. Given upstream gradient $g=\partial\mathcal L/\partial h$, we still need

We can skip storing/calculating the parameter gradient $\partial\mathcal L/\partial W$, but cannot skip the derivative with respect to $x$. A no_grad() block around this layer cuts that connection. The same problem occurs if you detach the activations between adapters.

3A small update with a precise budget

3.1 Count the update, not the base matrix

For $W_0\in\mathbb R^{d_{\mathrm{out}}\times d_{\mathrm{in}}}$, LoRA uses

The trainable update has $r(d_{\mathrm{in}}+d_{\mathrm{out}})$ parameters and rank at most $r$. The base matrix remains stored and frozen. It need not be low rank: the constraint is on the update. During training, compute $A x$ then $B(Ax)$ without materializing the dense product $BA$. This is the formulation in Hu et al., LoRA.

For the slide’s Llama-2-style example, $L=32$, $d=4096$, and $V=32{,}000$, a bias-free generation head has $dV=131{,}072{,}000$ trainable parameters. Q/K/V/O LoRA has $8Ldr=1{,}048{,}576r$.

At rank 125 these parameter counts are equal; this is a parameter-count crossover, not a total-memory crossover. These counts assume four square attention projections. Grouped-query attention can make K/V rectangular; use their actual shapes. For each of three gated-FFN projections with intermediate width $m$, the adapter count is $r(d+m)$.

3.2 Why initialize only one factor to zero?

Standard LoRA initializes $A$ randomly and $B=0$, so the adapter initially contributes zero. For one input and upstream gradient $g$,

Initially, the task gradient for $A$ is zero, but the gradient for $B$ can be nonzero. After $B$ changes, $A$ can receive a task gradient. If both factors start at zero, both task gradients are zero. This is a bilinear dead start, not a reason to initialize the entire adapter randomly and lose the base-model starting behavior.

3.3 Rank and placement buy different things

A higher rank expands the possible update at each targeted projection. More target projections let updates enter at more points in the computation. At fixed adapter count, these choices compete: Q/V at rank 16 and Q/K/V/O at rank 8 cost the same under square-projection assumptions. Equal parameter budgets do not imply equal learned functions or runtime. Compare them on validation data.

Scaling $s$ matters too. Changing rank with fixed $\alpha$ also changes $\alpha/r$. Record rank, target modules, scaling convention, dropout, and learning rate together. Prefer the smallest configuration meeting the task requirement rather than assuming the largest fitting rank must generalize best.

4Budget training memory and computation separately

4.1 Fewer trainable parameters is only part of the story

Use four accounts: model weights, gradients, optimizer state, and activations retained for backward. Add temporary buffers and implementation overhead when measuring a real run. Low-rank adapters reduce parameter-related memory compared with full tuning; they do not remove the dense base weights or all activation storage.

Increasing micro-batch size or sequence length can make activation memory dominate. Checkpointing trades recomputation for less activation storage. Gradient accumulation changes the effective batch without requiring all micro-batch graphs to remain live. Therefore a smaller adapter checkpoint does not establish a smaller peak training allocation.

4.2 A reversal you can calculate without the widget

At rank 16 with Q/K/V/O targets, the head needs 1.465 GiB for gradients and Adam moments. LoRA needs 0.219 GiB for adapter weights, gradients, and moments. Both share about 13.039 GiB of base weights under the widget’s assumptions.

4.3 What a frozen matrix still computes

For $n$ token rows and a dense $d\times d$ projection, using two FLOPs per multiply–add, the matrix products cost approximately:

The LoRA branch adds approximately $4ndr$ forward FLOPs for a square matrix, plus its backward work. Thus freezing the dense base can save weight-gradient computation, while retaining its forward and input-gradient work. This isolated projection is not an end-to-end Transformer speed model: attention, activations, optimizer work, checkpointing, and kernels also matter. Measure tokens/second and validation quality at a stated budget.

5Run an experiment that answers your question

5.1 Define success before selecting a method

For support-ticket extraction, report valid-JSON rate, correct field values, and exact match of parsed records. Compare against the same base model with a fixed prompt. Hold out paraphrases and split related messages by customer/thread or source so near-duplicates do not leak across splits. Use validation for rank, learning rate, and checkpoint selection; reserve the test set for the final comparison.

Also check a small set of general capabilities you want to preserve. Better target-task learning can accompany forgetting elsewhere. Biderman et al., LoRA Learns Less and Forgets Less studies this tradeoff on code and math; its results motivate measuring both, rather than assuming LoRA always matches full tuning or always preserves every capability.

5.2 Equal steps does not mean equal resources

Choose a comparison that matches your constraint: equal processed tokens for sample efficiency, equal training time for a time budget, or the best validation quality that fits one GPU. State which you use. Tune learning rates separately for full tuning and LoRA; a copied full-tuning rate can make a good adapter configuration look weak. “LoRA needs exactly 10× the learning rate” and “higher rank is always better” are not universal rules.

6Core practice: diagnose before scaling up

model.eval()                         # “freeze the base”
optimizer = AdamW(model.parameters())
# Insert trainable adapters here, AFTER creating the optimizer.
with torch.no_grad():
    logits = model(input_ids)
loss = cross_entropy(logits, targets)
loss.backward()

7Optional implementation lab and real-code reading

7.1 Run six small checks on CPU

Download 17_adaptation_checks.py and run python3 17_adaptation_checks.py in an environment with PyTorch. No model, dataset, or GPU download is needed. The checks were run with PyTorch 2.12.0 on CPU, using float64 for numerical comparisons. Random tiny networks test implementation mechanics; they are not pretrained LMs and do not compare adaptation quality.

1. check_response_loss: exactly three response targets, including EOS, receive loss; uniform logits give $\log 8\approx2.079442$.
2. check_initialization_and_merge: the adapter starts as a no-op, only B has a nonzero first task gradient, the frozen base stays unchanged, and merged outputs match.
3. check_frozen_bridge: both adapters receive gradients through a frozen intermediate layer; a deliberate no_grad() bug blocks only the early adapter.
4. check_head_only: frozen-backbone features have no gradient history while the separate head receives gradients.
5. check_counts: verify the model-wide parameter counts in §3.
6. check_existing_layer: wrap an existing biased linear layer, preserve its initial output, update only the adapter, and verify that the base weight and bias stay unchanged.

The lab’s original teaching module is short enough to inspect completely:

class LoRALinear(nn.Module):
    """Teaching implementation; optional existing base, no adapter dropout/auto-merge."""
    def __init__(self, din, dout, r=2, alpha=8, *, base=None):
        super().__init__()
        if r <= 0:
            raise ValueError("rank must be positive")
        self.base = nn.Linear(din, dout, bias=False) if base is None else base
        self.base.requires_grad_(False)
        weight = self.base.weight
        self.A = nn.Parameter(torch.randn(r, din, device=weight.device,
                                         dtype=weight.dtype) * 0.1)
        self.B = nn.Parameter(weight.new_zeros(dout, r))
        self.scale = alpha / r

    @classmethod
    def from_linear(cls, base, r=2, alpha=8):
        # Reuse the exact layer, including its bias, dtype, device, and weights.
        return cls(base.in_features, base.out_features, r, alpha, base=base)

    def forward(self, x):
        return self.base(x) + self.scale * F.linear(F.linear(x, self.A), self.B)

To adapt an existing model, reuse its actual layer instead of constructing a fresh random base. For example, replacing a model’s projection can look like this:

model.requires_grad_(False)  # Freeze before injecting adapters.
original = model.proj  # An existing nn.Linear from the loaded model.
x = example_features  # Inputs to this projection, not raw token IDs.
with torch.no_grad():
    expected = original(x).clone()
model.proj = LoRALinear.from_linear(original, r=2, alpha=8)
torch.testing.assert_close(model.proj(x), expected, rtol=0, atol=0)
optimizer = torch.optim.AdamW(
    [p for p in model.parameters() if p.requires_grad], lr=1e-4)

This replacement retains the exact base weight, bias, dtype, and device, and freezes that layer’s base parameters. It does not freeze the rest of the model. For adapter-only training, freeze the full model before replacing the desired projections; then audit all trainable names before creating the optimizer. The lab’s check_existing_layer checks exact initial output parity, a nonzero adapter update, and unchanged base weights and bias.

The nested F.linear calls apply A and then B. PyTorch stores linear weights as [out_features, in_features]; the math uses column vectors while the implementation accepts batches of row vectors. The frozen base call remains outside no_grad() so the module works correctly when its input depends on an earlier adapter.

7.2 Read the original LoRA linear layer

7.3 Find registration, freezing, and scaling

loralib/layers.py · Linear.__init__ · lines 109–114 (verbatim)

        if r > 0:
            self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features)))
            self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r)))
            self.scaling = self.lora_alpha / self.r
            # Freezing the pre-trained weight matrix
            self.weight.requires_grad = False

nn.Parameter registers A and B so the module exposes them to optimizers (§2). The shape calculation is exactly §3.1. Check that the optimizer is constructed after injection. Inspect the surrounding constructor: freezing weight here does not by itself freeze a possible bias. A whole-model training policy must decide which biases and other modules to train.

7.4 Locate the two forward branches

loralib/layers.py · Linear.forward · lines 147–152 (verbatim)

        if self.r > 0 and not self.merged:
            result = F.linear(x, T(self.weight), bias=self.bias)            
            result += (self.lora_dropout(x) @ self.lora_A.transpose(0, 1) @ self.lora_B.transpose(0, 1)) * self.scaling
            return result
        else:
            return F.linear(x, T(self.weight), bias=self.bias)

When unmerged, the first call computes the base projection and the next line adds the low-rank path; this is §4.3 in code. Inspect lines 119–125: the executable code initializes A randomly and B to zero, despite a nearby comment describing them in reverse. Read what the operations do.

Read Linear.train, lines 127–142. In this particular library, switching to evaluation can merge the adapter into the base weights, and switching back subtracts it again. This is additional library-specific behavior, beyond ordinary eval() semantics. Do not add the low-rank branch a second time after merging.

8Bonus: one shared base, many tasks

Slide 18’s serving benefit follows from the same decomposition: requests share the base projection but can select different adapter branches. Keeping adapters separate supports switching and batching tasks; merging one adapter removes its branch overhead for a dedicated model. Those are different deployment choices.

Optional sources: LoRA for the factorization and merge; LoRA Learns Less and Forgets Less for controlled adaptation/retention experiments; Punica for multi-tenant serving; PEFT LoRA documentation for a broader implementation API. These extend the reading; model-specific recipes are not required memorization.