LoRA for Sentiment Analysis¶
Low Rank Adaptation (LoRA) fine tunes a large model by training a small number of new parameters instead of moving the ones it already has. For a weight matrix W of shape (d, d), LoRA learns two thin matrices A of shape (d, r) and B of shape (r, d) with r far smaller than d, and adds their product to whatever that part of the network computes. The rank r is a choice, not something the method discovers: this tutorial uses 4 against a hidden size of 8192, so the adapter holds 65,536 parameters against the model's 70 billion.
Because the base weights never move, one frozen copy of the model can serve any number of adapters. That is what puts LoRA in the Parameter Efficient Fine Tuning (PEFT) family. Where you attach the adapter is also a choice. Real LoRA fine tunes usually target the attention projections in every block; this tutorial attaches one adapter to the last block's MLP, which keeps the whole thing small enough to watch.
In nnsight the adapter is not spliced into the module tree. It reads the module's input and writes its output from inside a trace, the same activation read and write used for any other intervention, so the model's architecture is unchanged and you can turn the adapter off by not calling it.
What this runs on¶
The outputs below come from meta-llama/Llama-3.1-70B in bfloat16, sharded across four A100-80GB cards with device_map="auto". The 800 training steps take a few minutes.
Nothing here is specific to that model or that hardware. The same code trains against a 1B model on one GPU, and the same block runs on NDIF by wrapping it in model.session(remote=True) (see Remote Execution).
import torch
import torch.nn as nn
from datasets import load_dataset
from nnsight import TransformersModel
from nnsight.intervention.envoy import Envoy
torch.manual_seed(0)
The data¶
SST-2 is a sentiment benchmark of movie-review fragments labelled 0 for negative and 1 for positive.
The model is a plain language model, so the task has to be posed as next-token prediction: a prompt that ends where the answer goes, and a target token to predict there. " positive" and " negative" are each a single token in Llama-3's vocabulary, which keeps the objective to one cross-entropy over one position.
dataset = load_dataset("nyu-mll/glue", "sst2")
def to_prompt(sentence):
return f'Review: "{sentence.strip()}"\nSentiment:'
model_name = "meta-llama/Llama-3.1-70B"
model = TransformersModel(model_name, device_map="auto", dispatch=True, dtype=torch.bfloat16)
POSITIVE = model.tokenizer.encode(" positive", add_special_tokens=False)[0]
NEGATIVE = model.tokenizer.encode(" negative", add_special_tokens=False)[0]
train = [(to_prompt(r["sentence"]), POSITIVE if r["label"] == 1 else NEGATIVE)
for r in dataset["train"].select(range(6400))]
val = [(to_prompt(r["sentence"]), POSITIVE if r["label"] == 1 else NEGATIVE)
for r in dataset["validation"].select(range(200))]
print("target token ids:", POSITIVE, NEGATIVE)
print(train[0][0])
target token ids: 6928 8389 Review: "hide new secretions from the parental units" Sentiment:
nnsight tokenizes a batched list with left padding, so logits[:, -1] is the last real token of every row no matter how the lengths differ. Reading the last position by index from the right is what makes a batched objective safe here.
Freeze the model before training. Only the adapter's parameters should move, and freezing also keeps the backward pass from retaining activations for all 80 blocks.
model._module.requires_grad_(False)
module = model.model.layers[-1].mlp
with model.scan(" "):
dim = module.output.shape[-1].save()
print("hidden size:", dim)
hidden size: 8192
The adapter¶
WA is initialised at scale 1/sqrt(dim) so its output starts at roughly unit norm, and WB at zero so the adapter contributes nothing on the first step. Both are torch.nn.Parameters with requires_grad=True, so loss.backward() inside a trace fills their .grad and an ordinary optimizer updates them.
class LORA(nn.Module):
def __init__(self, module: Envoy, dim: int, r: int, WA=None, WB=None) -> None:
"""
Args:
module: the Envoy whose input the adapter reads and whose output it rewrites.
dim: that module's hidden size.
r: the adapter's inner dimension.
WA, WB: trained weights to start from. Freshly initialised when omitted.
"""
super().__init__()
self.r = r
self.module = module
# The parameters have to land on the same device as the module they wrap.
device = module.device
WA = torch.randn(dim, r) / dim**0.5 if WA is None else WA
WB = torch.zeros(r, dim) if WB is None else WB
self.WA = nn.Parameter(WA.to(device), requires_grad=True)
self.WB = nn.Parameter(WB.to(device), requires_grad=True)
def __call__(self, alpha: float = 1.0):
"""Apply the adapter at the current point in the trace.
Args:
alpha: how much of it to apply. Can be changed after training.
"""
# The adapter's parameters are float32 and the model is bfloat16, so cast
# into and back out of the adapter's dtype around the matmuls.
hidden_states = self.module.input
A_x = torch.matmul(hidden_states.to(self.WA.dtype), self.WA)
BA_x = torch.matmul(A_x, self.WB)
h = BA_x.to(hidden_states.dtype) + self.module.output
self.module.output = h * alpha
def parameters(self):
return [self.WA, self.WB]
How the adapter reaches the model
Reading self.module.input and writing self.module.output inside a trace is an ordinary nnsight intervention. The adapter is never inserted into the module tree, so print(model) shows the same architecture before and after training, and hidden values only change on the traces where you call lora().
To make the routing permanent instead, store it with model.edit() so it replays on every trace. Under generate a plain edit fires at prefill only, so put it under tracer.iter[:] to reach every decoding step. See Model Editing.
What the model does before training¶
Two numbers are worth having before any training starts.
2-way accuracy compares the logits of " positive" and " negative" at the answer position and takes the larger, which measures whether the model separates the two sentiments at all. Format counts how often the unrestricted argmax over the whole vocabulary is one of those two tokens, which measures whether the model answers the question as asked.
def evaluate(tag, use_lora, batch_size=8):
correct = in_format = 0
for i in range(0, len(val), batch_size):
chunk = val[i:i + batch_size]
with torch.no_grad(), model.trace([p for p, _ in chunk]):
if use_lora:
lora()
logits = model.lm_head.output[:, -1].save()
target = torch.tensor([t for _, t in chunk], device=logits.device)
two_way = logits[:, [NEGATIVE, POSITIVE]].argmax(-1)
correct += int((torch.where(two_way == 1, POSITIVE, NEGATIVE) == target).sum())
top = logits.argmax(-1)
in_format += int(((top == POSITIVE) | (top == NEGATIVE)).sum())
print(f"{tag}: 2-way accuracy {correct}/{len(val)} = {correct / len(val):.3f} | "
f"answers with a sentiment word {in_format}/{len(val)}")
return correct / len(val)
base_accuracy = evaluate("base model", use_lora=False)
base model: 2-way accuracy 170/200 = 0.850 | answers with a sentiment word 1/200
The 70B already separates the two sentiments at 85%, and almost never says either word on its own. So the adapter has two jobs of very different sizes: teach the format, which is nearly free, and improve on 85%, which is not.
Training¶
One trace per batch. The adapter is applied inside the trace by calling lora(), the loss is a cross-entropy at the last position, and with loss.backward(): runs the backward pass interleaved with the model. The optimizer steps between traces, exactly as in ordinary PyTorch.
lora = LORA(module, dim, r=4)
optimizer = torch.optim.AdamW(lora.parameters(), lr=3e-4)
batch_size = 8
for step in range(0, len(train), batch_size):
chunk = train[step:step + batch_size]
with model.trace([p for p, _ in chunk]):
lora()
logits = model.lm_head.output[:, -1]
target = torch.tensor([t for _, t in chunk], device=logits.device)
loss = torch.nn.functional.cross_entropy(logits.float(), target)
with loss.backward():
pass
tracked = loss.item().save()
optimizer.step()
optimizer.zero_grad()
n = step // batch_size
if n % 100 == 0:
print(f"step {n:3d} loss {float(tracked):.4f} "
f"|WA| {lora.WA.norm().item():.3f} |WB| {lora.WB.norm().item():.3f}")
step 0 loss 5.0730 |WA| 2.003 |WB| 0.054 step 100 loss 0.3860 |WA| 2.423 |WB| 1.359 step 200 loss 0.4240 |WA| 2.374 |WB| 1.376 step 300 loss 0.6408 |WA| 2.445 |WB| 1.422 step 400 loss 0.3497 |WA| 2.860 |WB| 1.551 step 500 loss 0.2616 |WA| 3.224 |WB| 1.672 step 600 loss 0.2453 |WA| 3.410 |WB| 1.751 step 700 loss 0.1657 |WA| 3.555 |WB| 1.816
Printing |WA| and |WB| next to the loss is worth the two extra terms. A loss that falls while the norms sit still means the optimizer is not reaching these parameters, which is the usual shape of a silently broken training loop.
What the adapter learned¶
trained_accuracy = evaluate("with LoRA", use_lora=True)
evaluate("LoRA off again", use_lora=False)
assert trained_accuracy > base_accuracy
with LoRA: 2-way accuracy 186/200 = 0.930 | answers with a sentiment word 200/200 LoRA off again: 2-way accuracy 170/200 = 0.850 | answers with a sentiment word 1/200
Accuracy goes from 0.850 to 0.930, and the format from 1 of 200 to 200 of 200. Turning the adapter off returns both numbers exactly to the baseline, which is the check that the base model was never touched: everything the adapter does lives in WA and WB.
Most of what the adapter bought is the format. The 8 points of accuracy on top of that came from 6,400 examples through a rank-4 adapter on one MLP, which is a small change to a model that already knew most of the answer.
for text in ["I'm upset", "what a wonderful film"]:
with torch.no_grad(), model.trace(to_prompt(text)):
lora()
with_lora = model.lm_head.output[0, -1].argmax().save()
with torch.no_grad(), model.trace(to_prompt(text)):
without = model.lm_head.output[0, -1].argmax().save()
print(f"{text!r}: with LoRA -> {model.tokenizer.decode(with_lora)!r} | "
f"without -> {model.tokenizer.decode(without)!r}")
"I'm upset": with LoRA -> ' negative' | without -> ' Negative' 'what a wonderful film': with LoRA -> ' positive' | without -> ' "'
Keeping the adapter¶
lora.WA and lora.WB are the whole result. Save those two tensors and rebuild the adapter with LORA(module, dim, 4, WA=..., WB=...) to get the trained behaviour back without retraining.
Where to go next¶
- Setting Activations covers the intervention mechanics the adapter is built on.
- Model Editing attaches an adapter permanently so it runs on every forward pass instead of being called in each trace.
- Remote Execution runs the same training block against a model hosted on NDIF.