LoRA and DoRA Implementation from Scratch

Provenance

Independent

Artifacts

Source code

LoRA and DoRA layers written from scratch in PyTorch and trained on an MNIST MLP. A study of the mechanics, with an audit of where the implementation diverges from the papers and what it does not yet verify.

Scope at a glance. Built: LoRA and DoRA layers in plain PyTorch (each in an unmerged and a merged form), wrapped around the nn.Linear layers of a 3-layer MLP and trained on MNIST in a notebook. Not built: unit tests, a numerical parity check against Hugging Face PEFT, multi-seed runs, or any experiment on a language model. The audit below says what that means for the results.

Goal

Implement Low-Rank Adaptation (LoRA) and Weight-Decomposed Low-Rank Adaptation (DoRA) without a library, in order to understand what parameter-efficient fine-tuning changes inside a layer. The MLP is deliberately small: the point is the mechanics, not the task.

LoRA

A frozen weight $W_0$ gets a trainable low-rank update. For a layer computing $h = W_0 x$:

$B$ starts at zero, so the wrapped layer initially computes exactly what the pretrained layer did. Only $A$ and $B$ are trained.

The repository’s layer, unchanged except for removed docstrings. PyTorch stores nn.Linear weights as (out, in), so the matrices act on row vectors:

class LoRALayer(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, rank: int, alpha: float) -> None:
        super().__init__()
        std_dev = 1 / torch.sqrt(torch.tensor(rank).float())
        self.A = nn.Parameter(torch.randn(in_dim, rank) * std_dev)
        self.B = nn.Parameter(torch.zeros(rank, out_dim))
        self.alpha = alpha

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.alpha * (x @ self.A @ self.B)


class LinearWithLoRA(nn.Module):
    def __init__(self, linear: nn.Linear, rank: int, alpha: float) -> None:
        super().__init__()
        self.linear = linear
        self.lora = LoRALayer(linear.in_features, linear.out_features, rank, alpha)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear(x) + self.lora(x)

LinearWithLoRAMerged computes the same thing by forming $W_0 + \alpha (AB)^\top$ first and calling F.linear once. The notebook checks that both forms reproduce the base layer’s output at initialization.

DoRA

DoRA splits the adapted weight into a magnitude and a direction, and learns both:

$m$ is a trainable vector initialized from the norms of $W_0$, and the low-rank update only changes the direction. The repository’s merged layer:

class LinearWithDoRAMerged(nn.Module):
    def __init__(self, linear: nn.Module, rank: int, alpha: float) -> None:
        super().__init__()
        self.linear = linear
        self.lora = LoRALayer(linear.in_features, linear.out_features, rank, alpha)
        self.m = nn.Parameter(self.linear.weight.norm(p=2, dim=0, keepdim=True))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        lora = self.lora.A @ self.lora.B
        numerator = self.linear.weight + self.lora.alpha * lora.T
        denominator = numerator.norm(p=2, dim=0, keepdim=True)
        directional_component = numerator / denominator
        new_weight = self.m * directional_component
        return F.linear(x, new_weight, self.linear.bias)

What was measured

One notebook run. The MLP (784 → 128 → 256 → 10) is trained on MNIST for 2 epochs with seed 123, then copied and fine-tuned for another 2 epochs, again on MNIST, with rank 4, α = 8, and Adam at lr 0.005.

This shows the layers train and that freezing works, since only $A$, $B$, and $m$ have requires_grad=True. It does not show that DoRA beats LoRA. The fine-tuning data is the same data the base model was trained on, so there is no new task to adapt to. With one seed, a 0.3-point gap on 10,000 test images is within what a different seed could produce.

Audit: what this does and does not show

I reviewed the repository against the papers in October 2026. These are the open issues, ordered by how much they affect what the project can claim.

1. No parity check against a reference. Nothing compares these layers’ outputs or gradients with Hugging Face PEFT or the papers’ reference code. (An earlier version of this page said they were validated against PEFT. They were not.) Until such a check exists, “correct” means “trains and matches the base layer at initialization”.
2. DoRA normalizes along a different axis than the reference code. The merged layer takes norms over dim=0 of the (out, in) weight, so it learns one magnitude per input feature. That is a literal reading of the paper’s column-wise norm, but the authors’ released code and PEFT both normalize each output feature’s weight vector. I have not established which one this implementation should match, and a parity test would settle it.
3. The unmerged LinearWithDoRA is not DoRA. It normalizes the LoRA activations per example and adds them to the base output, with $m$ initialized to ones. That is not the weight decomposition above. The notebook’s DoRA numbers come from the merged class, but the unmerged class should be fixed or removed.
4. The scaling convention differs from the LoRA paper. The update is scaled by α rather than α / r, so the effective step size changes with rank. This is harmless at a single rank, but it matters when comparing ranks.
5. No tests. The repository has a notebook but no test suite.

What would make it evidence

In order of value:

• A parity test: wrap the same nn.Linear with these layers and with PEFT (use_dora=True for DoRA), copy $A$, $B$, and $m$ across, and assert matching outputs and gradients. This settles issue 2.
• Unit tests for initialization equivalence, merged versus unmerged equivalence, and that only adapter parameters receive gradients.
• A real adaptation task: pretrain on one distribution, fine-tune on a shifted one (for example, a held-out subset of digits), compare against full fine-tuning, and report the spread over several seeds and ranks.

References

• LoRA: Low-Rank Adaptation of Large Language Models, Hu et al., 2021
• DoRA: Weight-Decomposed Low-Rank Adaptation, Liu et al., 2024
• Sebastian Raschka’s LoRA from Scratch, the implementation guide this code follows closely
• Hugging Face PEFT documentation