Hand-Writing a Vision-Language Model, Part 8: Hand-Writing LoRA, and a Real Memory Bug That Was Not What It Looked Like

Hand-Written LoRA

Consistent with the treatment Qwen and the ViT already got, LoRA was hand-written rather than using peft. qwen3/lora.py:

class LoRALinear(nn.Module):
    def __init__(self, base_linear, rank, alpha=None):
        super().__init__()
        self.base_linear = base_linear
        self.base_linear.requires_grad_(False)
        self.scale = (alpha if alpha is not None else rank) / rank
        self.lora_A = nn.Linear(base_linear.in_features, rank, bias=False)
        self.lora_B = nn.Linear(rank, base_linear.out_features, bias=False)
        nn.init.kaiming_uniform_(self.lora_A.weight, a=5**0.5)
        nn.init.zeros_(self.lora_B.weight)

    def forward(self, x):
        return self.base_linear(x) + self.scale * self.lora_B(self.lora_A(x))

apply_lora(model, ["q_proj", "v_proj"], rank=8) walks the model and swaps those specific Linear submodules for LoRALinear wrappers via setattr. A transparent drop-in: no forward code anywhere had to change, since self.q_proj(x) calls now just resolve to a different forward() underneath.

Two Things Worth Verifying Rather Than Assuming

Zero-initialised lora_B means training starts from the pretrained model exactly, not from noise. Confirmed directly: a freshly-wrapped layer's output is byte-identical to the plain Linear it replaces. Wrapping q_proj and v_proj across all 28 layers of a real loaded Qwen model left the full output logits completely unchanged immediately after wrapping.

Real parameter efficiency, not just a claim. 56 wrapped layers (28 × q_proj/v_proj) at rank 8 add approximately 1.15M trainable parameters, against Qwen3-0.6B's 600M total.

Understanding LoRA More Precisely

LoRA's comparison involves only Qwen's own pretrained weights. The projector, the point-token embeddings, and the LoRA adapters are three completely separate trainable pieces that happen to train simultaneously in the same optimiser during stage 2. The projector maps ViT patches into Qwen's embedding space. The point-token overlay is a tiny (2, hidden_size) parameter for <|point|> and <|/point|>. LoRA's comparison is entirely about Qwen's q_proj and v_proj weights: what those matrices would need to become to do well at maze navigation, versus what they already are. Nothing about LoRA's math references the projector or the embeddings.

Only the new part ever gets updated. The original frozen weight matrix inside q_proj and v_proj never changes, not once, not by any amount, for the entire training run. LoRALinear.forward(x) = base_linear(x) + scale * lora_B(lora_A(x)): base_linear computes its output exactly as it always has, every forward pass. Only lora_A and lora_B get updated. Delete the LoRA adapters after training and the original Qwen weights are exactly what they were before any of this started.

Rank is a hyperparameter, not something auto-sized. Here it is 8. What it controls is the bottleneck width between lora_A (rank × in_features) and lora_B (out_features × rank): everything has to squeeze through that rank-wide middle. A matrix built by multiplying through an r-wide bottleneck cannot have rank higher than r. The LoRA paper's empirical claim is that the weight update actually useful for adapting a pretrained model to a narrower task tends to live in a low-dimensional subspace. Most directions are irrelevant to the task; a handful do most of the work. The loss curve (3.77 down to roughly 0.08) is a small hint of this: rank 8 was already more capacity than maze navigation needed.

Why q_proj and v_proj and not k_proj too. An empirical choice from the original LoRA paper, not a hard rule. Q and K together decide the attention pattern (who attends to whom), and they only ever interact through a dot product. Nudging Q to reshape the attention pattern can often achieve the same effect as nudging K in a complementary direction. V controls what actually gets pulled forward once the pattern is fixed: something Q and K cannot touch. So Q+V covers two genuinely different levers ("who to look at" and "what to bring back"), whereas Q+K would mostly double up on the same lever. The original paper confirmed this as an ablation.

The Maze Dataset

data/maze_dataset.py's MazeDataset mirrors ImageCaptionDataset's shape but generates procedurally rather than loading from disk. seed=idx makes every epoch reproducible without ever storing images. One deliberate difference: it takes an already-configured tokenizer rather than building its own internally, since <|point|> and <|/point|> have to be added via add_special_tokens exactly once and shared with the model's embedding-claiming step. The training script owns that tokenizer instance, not the dataset.

The Memory Bug That Was Not Sequence Length

Ran the first real stage 2 training attempt and hit CUDA out of memory on step 1. The obvious suspect: maze reasoning text is much longer than stage 1's captions (an 8x8 maze can produce 600+ text tokens versus captions' 10 to 20), so quadratic attention cost seemed like the clear culprit. Shrunk the maze grid. Still crashed.

Worth slowing down and measuring rather than continuing to guess.

Isolated it with a scratch script: same maze data, same LoRA, but skip the embedding-training step entirely. Peak memory barely moved. That ruled out sequence length as the primary cause and pointed at make_token_embeddings_trainable instead.

Its first implementation put the entire 151,936-row embed_tokens table into the optimiser, with a gradient hook zeroing out every row's gradient except the two claimed ones. Correct gradient-wise: only 2 rows ever get a nonzero gradient. But not correct memory-wise. AdamW keeps two full-sized buffers (momentum and variance) per parameter it optimises, sized to the parameter's full shape, regardless of how much of its gradient is ever actually nonzero. Training the whole table just to update 2 rows cost approximately 1.24GB of permanent optimiser state for zero benefit on the other 151,934 rows. Measured directly: adding that one parameter to the optimiser alone pushed peak memory from roughly 3.8GB to roughly 6GB, even on a short sequence.

The fix: stop masking gradients on the giant tensor. Instead give the optimiser a genuinely tiny parameter: a (2, hidden_size) overlay (~2,048 values) spliced into both the input embedding lookup and the tied LM head's output logits for just the claimed token IDs. embed_tokens.weight itself goes back to being fully frozen. Verified: the overlay receives real gradient, the frozen table's values and gradient are both completely untouched, and nudging the overlay measurably moves only the two claimed tokens' output logits.

That fix did not fully solve it. After fixing it, an 8x8 maze still crashed on certain examples. Checked the actual data: a 10 to 15 seed sample at 8x8 produced solved paths ranging from 17 to 49 cells. The longest ones were enough to crash purely from attention memory in a 28-layer transformer, independent of the optimiser fix. Shrinking to 6x6 narrowed the range to 13 to 29 cells, which trained cleanly.

The Actual Training Run

500 examples, 6x6 mazes, continuing from projector_stage1.pt, LoRA on q_proj and v_proj (rank 8) plus the point-token overlay, 1 epoch:

step    0 | loss 3.7680
step   20 | loss 1.4732
step   40 | loss 0.5327
step   60 | loss 0.3408
step  100 | loss 0.1666
step  200 | loss 0.0947
step  300 | loss 0.0863
step  400 | loss 0.0888
step  480 | loss 0.0808

Peak memory stabilised at 6.32GB, real headroom inside the 8GB budget. The drop is far sharper than stage 1's captioning run (14.86 down to a noisy 2.5 to 5 range). Makes sense: mazes are a much more rigid, learnable format. A small fixed vocabulary of directions plus strict coordinate syntax, compared to open-ended natural-language captions. Checkpoint saved to checkpoints/stage2.pt.

What Is Next

Actually test what the trained model produces. Feed a maze image through the same VLM.generate() loop from Part 6 and check whether the generated <|point|> sequence traces a real, plausible path through the maze.