Build a LoRA Trainer in Pure NumPy: A Portfolio Project That Signals Systems Mastery

TL;DR — Building a LoRA adapter trainer from scratch in pure NumPy forces you to internalize low-rank decomposition, gradient fusion, and inference-time adapter routing at a level that framework abstractions never will. This project demonstrates numerical computing fluency, systems-level ML understanding, and production-aware engineering — exactly the trifecta hiring managers look for in senior ML infrastructure roles.

The demand for engineers who understand both the math and the machinery behind modern efficient fine-tuning is exploding. LoRA (Low-Rank Adaptation) has become the de facto standard for parameter-efficient fine-tuning of large models, and understanding it at the implementation level — not just the API level — separates candidates who can ship systems from those who can only call fit().

This guide walks you through building a complete LoRA adapter trainer with fused low-rank gradient computation and adapter-aware inference, entirely in NumPy. No PyTorch, no JAX, no transformers library. Just you and NumPy’s einsum, dot, and a stubborn commitment to understanding every gradient.

Why This Project Stands Out on a CV

This project signals three distinct competency clusters that hiring managers and technical interviewers actively screen for:

Numerical computing depth. Writing backpropagation manually in NumPy proves you understand matrix calculus, not just autodiff APIs. You can debug gradient shapes, vanishing norms, and numerical instability without a debugger hiding the math.
Systems-level ML architecture. Fused low-rank gradient computation requires you to think about memory layout, kernel fusion, and operation ordering — the same concerns that appear in distributed training systems like DeepSpeed or FSDP.
Production-flavored ML engineering. Adapter-aware inference mirrors how real serving systems route between multiple adapter weights at runtime, a pattern used in production by companies fine-tuning hundreds of models on shared base weights.

For roles in ML infrastructure, efficient training, or foundational model platforms, this project demonstrates you can go from mathematical formulation to working implementation without leaning on framework magic. It signals you’re the person who can debug a CUDA kernel, not just write a notebook.

Architecture Overview

The project decomposes into four tightly coupled components:

┌─────────────────────────────────────────────────────┐
│                  NUMPY LORA TRAINER                  │
├─────────────────────────────────────────────────────┤
│                                                     │
│  ┌──────────────┐    ┌──────────────────────────┐  │
│  │  Model Loader │───▶│   LoRA Adapter Layer      │  │
│  │  (in-memory)  │    │  - W0 (frozen base)      │  │
│  └──────────────┘    │  - A (low-rank, trainable)│  │
│                      │  - B (low-rank, trainable)│  │
│                      └──────────┬───────────────┘  │
│                                 │                   │
│                      ┌──────────▼───────────────┐  │
│                      │  Fused Gradient Engine    │  │
│                      │  - dA = X^T · δ · B^T   │  │
│                      │  - dB = δ · A · X^T     │  │
│                      │  (single matmul chain)    │  │
│                      └──────────┬───────────────┘  │
│                                 │                   │
│                      ┌──────────▼───────────────┐  │
│                      │  Adapter-Aware Inference  │  │
│                      │  - Route to active adapter│  │
│                      │  - Merge ΔW = B @ A      │  │
│                      │  - Cache fused weights    │  │
│                      └──────────────────────────┘  │
│                                                     │
│  ┌──────────────────────────────────────────────┐   │
│  │  Training Loop (manual SGD / Adam step)      │   │
│  └──────────────────────────────────────────────┘   │
└─────────────────────────────────────────────────────┘

The data flow is linear: base weights are loaded, LoRA adapters are initialized, the fused gradient engine computes updates for both low-rank matrices simultaneously, and adapter-aware inference routes inputs through the correct adapter weights. The critical design decision is that the fused gradient engine computes dA and dB in a single computational pass, avoiding the memory overhead of materializing intermediate tensors separately — this is the “fused” part.

Building It Step by Step

Step 1: Initialize Base Weights and LoRA Adapter Matrices

Start by defining the frozen base weight matrix W0 and the trainable low-rank matrices A (shape d × r) and B (shape r × d), where d is the model dimension and r is the rank. The adapter contribution to any layer’s output is X @ B @ A, and the full computation is X @ (W0 + B @ A).

import numpy as np

def initialize_lora(d: int, r: int, scale: float = 0.02):
    """Initialize LoRA adapter matrices with orthogonal A and zero-initialized B."""
    # A: random orthogonal initialization for stable gradients
    A = np.random.randn(d, r) * scale
    A, _ = np.linalg.qr(A)  # orthogonalize for numerical stability

    # B: zero initialization ensures adapter starts as identity skip
    B = np.zeros((r, d))

    # Base weight: small random matrix simulating a frozen pretrained layer
    W0 = np.random.randn(d, d) * 0.01

    return W0, A, B

# Example: a 256-dim layer with rank-8 adaptation
W0, A, B = initialize_lora(d=256, r=8)
print(f"W0 shape: {W0.shape}, A shape: {A.shape}, B shape: {B.shape}")
# W0 shape: (256, 256), A shape: (256, 8), B shape: (8, 256)

The orthogonal initialization of A and zero initialization of B is critical: it ensures the adapter starts as a near-identity mapping, so pre-trained knowledge is preserved from step one.

Step 2: Implement the Forward Pass with Adapter Injection

The forward pass must compute the base transformation plus the low-rank adapter contribution. This is where adapter-aware inference begins to take shape — we need to support routing to specific adapter indices.

def lora_forward(X: np.ndarray, W0: np.ndarray, A: np.ndarray, B: np.ndarray,
                 adapter_index: int = 0, adapter_cache: dict = None) -> np.ndarray:
    """
    Forward pass: Y = X @ W0 + X @ B @ A
    Supports adapter caching for inference optimization.
    """
    # Compute base transformation
    Y_base = X @ W0

    # Check if fused adapter weight is cached (inference optimization)
    cache_key = f"adapter_{adapter_index}"
    if adapter_cache and cache_key in adapter_cache:
        W_fused = adapter_cache[cache_key]
    else:
        # Fuse: W_adapter = B @ A (rank-r outer product)
        W_fused = B @ A
        if adapter_cache is not None:
            adapter_cache[cache_key] = W_fused

    # Apply adapter contribution
    Y_adapter = X @ W_fused

    return Y_base + Y_adapter

The adapter cache is the mechanism that makes inference efficient: instead of recomputing B @ A for every forward pass, we fuse it once and cache the result. This mirrors how production systems like HuggingFace’s PEFT merge adapter weights before serving.

Step 3: Implement Fused Low-Rank Gradient Computation

This is the heart of the project. Given an upstream gradient δ (same shape as output Y), we need gradients for both A and B. The “fused” aspect means we compute both in a minimal number of matrix multiplications, reusing intermediate results.

def fused_lora_gradients(X: np.ndarray, delta: np.ndarray,
                         A: np.ndarray, B: np.ndarray) -> tuple:
    """
    Compute fused gradients for LoRA matrices A and B.

    Given: Y = X @ W0 + X @ B @ A
    Upstream gradient: δ = dL/dY

    dL/dA = X^T @ δ @ B^T     (shape: d × r)
    dL/dB = δ @ A @ X^T       (shape: r × d)

    Fused: compute X^T @ δ once, reuse it for both gradients.
    """
    # Fused intermediate: G = X^T @ δ  (shape: d × d_input)
    # This is the expensive matmul we compute ONCE
    G = X.T @ delta  # shape: (d, d_input)

    # Gradient for A: G @ B^T
    dA = G @ B.T  # shape: (d, r)

    # Gradient for B: δ @ A @ X^T
    # We can also fuse this as: (δ @ A) @ X^T
    dB = (delta @ A) @ X.T  # shape: (r, d)

    return dA, dB

The fusion here is about memory efficiency: X^T @ δ is computed once and reused. Without fusion, you’d materialize the full dL/dW matrix (shape d × d), which for large models is prohibitive. The low-rank gradient avoids this entirely — the maximum intermediate size is d × r, not d × d.

Step 4: Build the Training Loop with Manual Optimization

Now we wire everything together into a training loop. We use a simple SGD step with optional Adam-like momentum (implemented manually to avoid framework dependencies).

def train_lora(X_train: np.ndarray, Y_train: np.ndarray,
               W0: np.ndarray, A: np.ndarray, B: np.ndarray,
               learning_rate: float = 1e-3, epochs: int = 100,
               rank: int = 8) -> tuple:
    """
    Train LoRA adapter using fused gradients and manual optimization.
    """
    adapter_cache = {}
    losses = []

    for epoch in range(epochs):
        # Forward pass
        Y_pred = lora_forward(X_train, W0, A, B, adapter_cache=adapter_cache)

        # Compute loss (MSE)
        error = Y_pred - Y_train
        loss = np.mean(error ** 2)
        losses.append(loss)

        # Backward pass: upstream gradient
        delta = 2.0 * error / X_train.shape[0]  # dL/dY

        # Fused gradient computation
        dA, dB = fused_lora_gradients(X_train, delta, A, B)

        # Manual optimization step (SGD)
        A -= learning_rate * dA
        B -= learning_rate * dB

        # Re-fuse and update cache every 10 epochs for inference
        if epoch % 10 == 0:
            adapter_cache[f"adapter_0"] = B @ A

        if epoch % 20 == 0:
            print(f"Epoch {epoch}: Loss = {loss:.6f}")

    return W0, A, B, losses

Step 5: Implement Adapter-Aware Inference with Multiple Adapters

The final piece: a router that selects between multiple trained adapters at inference time. This is where the project transitions from a training toy to something that signals production systems understanding.

class AdapterRouter:
    """Routes inputs to the correct adapter at inference time."""

    def __init__(self, W0: np.ndarray):
        self.W0 = W0
        self.adapters: dict = {}  # {index: (A, B, W_fused)}
        self.default_adapter = 0

    def register_adapter(self, index: int, A: np.ndarray, B: np.ndarray):
        """Register a trained adapter with pre-fused weights."""
        self.adapters[index] = {
            "A": A,
            "B": B,
            "W_fused": B @ A  # Pre-fuse for inference speed
        }

    def route(self, X: np.ndarray, adapter_index: int = None) -> np.ndarray:
        """Forward pass through the selected adapter."""
        idx = adapter_index if adapter_index is not None else self.default_adapter
        adapter = self.adapters[idx]

        Y_base = X @ self.W0
        Y_adapter = X @ adapter["W_fused"]

        return Y_base + Y_adapter

    def merge_adapter(self, index: int) -> np.ndarray:
        """Merge adapter weights into base for deployment."""
        adapter = self.adapters[index]
        return self.W0 + adapter["W_fused"]

The merge_adapter method is the production pattern: when serving, you can either keep adapters separate (for multi-tenant serving where each request uses a different adapter) or merge them into the base weights (for single-model deployment with zero overhead).

Running and Testing It

Create a test script that validates both correctness and performance:

# Save the implementation as lora_trainer.py and run:
python -c "
from lora_trainer import initialize_lora, train_lora, AdapterRouter
import numpy as np

# Generate synthetic data: learn a low-rank target matrix
np.random.seed(42)
d, r = 128, 4
X = np.random.randn(500, d)
W_target = np.random.randn(d, d) * 0.5
Y = X @ W_target  # Target: we want adapter to approximate this

# Initialize and train
W0, A, B = initialize_lora(d=d, r=r)
W0, A, B, losses = train_lora(X, Y, W0, A, B, learning_rate=5e-3, epochs=200)

# Verify: final adapter should approximate W_target - W0
W_approx = W0 + B @ A
error = np.linalg.norm(W_approx - W_target) / np.linalg.norm(W_target)
print(f'Relative error to target: {error:.4f}')
assert error < 0.1, 'Adapter failed to converge!'

# Test adapter routing
router = AdapterRouter(W0)
router.register_adapter(0, A, B)
Y_infer = router.route(X[:10])
print(f'Inference output shape: {Y_infer.shape}')
print('All tests passed!')
"

Key assertions to validate:

Gradient shape correctness: dA must match A’s shape, dB must match B’s shape.
Convergence: relative error to the target matrix should drop below 0.1 within 200 epochs on synthetic data.
Adapter fusion correctness: W0 + B @ A should equal the merged adapter output.
Numerical stability: no NaN or Inf values during training, confirming orthogonal initialization works.

To profile memory usage, run with tracemalloc or simply observe that peak memory scales as O(d × r) rather than O(d²), confirming the low-rank efficiency:

import tracemalloc
tracemalloc.start()
W0, A, B, losses = train_lora(X, Y, W0, A, B, epochs=500)
current, peak = tracemalloc.get_traced_memory()
print(f"Peak memory: {peak / 1024:.1f} KB")
tracemalloc.stop()

Extending It: Your Roadmap to Senior-Level

Each upgrade below transforms the toy into something that belongs on a senior engineer’s project page. They are ordered from most impactful to most specialized.

Checkpoint persistence with np.savez and versioned adapter serialization. Save adapter states to disk after every N epochs with metadata (epoch, loss, rank). This matters because production systems must recover from crashes and support A/B testing of adapter versions. Implement a save_checkpoint(path, adapter_id, metadata) function that writes a .npz file with W0, A, B, loss, and timestamp.
Horizontal adapter serving with a lightweight gRPC or REST server. Wrap AdapterRouter in a FastAPI or gRPC service that accepts requests with an adapter_id field and routes to the correct adapter. This matters because multi-tenant model serving — where hundreds of fine-tuned adapters share a single base model — is a real production pattern used by companies like Netflix and Meta.
Gradient checkpointing and memory profiling with memory_profiler. Add line-by-line memory profiling to identify bottlenecks during training. Implement gradient checkpointing (recompute activations instead of storing them) to trade compute for memory. This matters because training large adapters on GPU memory-constrained instances is a daily reality for ML engineers, and profiling skills are non-negotiable at senior level.
Fault-tolerant training with resume-from-checkpoint and gradient accumulation. Add logic to detect training interruption (via signal handlers or file locks), save state mid-epoch, and resume from the last checkpoint. Implement gradient accumulation to simulate larger batch sizes on limited hardware. This matters because distributed training jobs routinely fail on clusters, and engineers who build resilient pipelines are valued over those who assume perfect conditions.
Benchmark harness comparing LoRA rank, learning rate, and fusion strategy. Build a benchmark.py that sweeps over ranks [4, 8, 16, 32], learning rates, and fused vs. unfused gradient computation, reporting wall-clock time, memory peak, and convergence speed. This matters because the ability to rigorously benchmark and justify architectural choices with data is what separates senior engineers from juniors.
Quantized adapter inference using INT8 calibration. After training, apply post-training quantization to the fused adapter weights (W_fused) using simple min-max scaling to INT8, and measure the accuracy drop. This matters because production inference increasingly demands quantization for latency and cost optimization, and understanding the accuracy-latency tradeoff is a senior-level concern.

Key Takeaways

Building a LoRA trainer in pure NumPy forces deep understanding of low-rank decomposition, gradient computation, and memory efficiency — skills that transfer directly to production ML systems.
Fused low-rank gradient computation (X^T · δ reused for both dA and dB) is the same optimization principle behind kernel fusion in frameworks like Triton and XLA.
Adapter-aware inference with routing and merging mirrors real production serving patterns for multi-tenant model deployment.
The project demonstrates three distinct skill clusters: numerical computing, systems-level ML architecture, and production engineering — the trifecta for senior ML infrastructure roles.
Every upgrade in the roadmap maps to a real production concern: persistence, serving, observability, fault tolerance, and benchmarking.
Starting from scratch without frameworks proves you understand the math, not just the API — the difference between someone who can debug a training run and someone who can design one.

Why This Project Stands Out on a CV#

Architecture Overview#

Building It Step by Step#

Step 1: Initialize Base Weights and LoRA Adapter Matrices#

Step 2: Implement the Forward Pass with Adapter Injection#

Step 3: Implement Fused Low-Rank Gradient Computation#

Step 4: Build the Training Loop with Manual Optimization#

Step 5: Implement Adapter-Aware Inference with Multiple Adapters#

Running and Testing It#

Extending It: Your Roadmap to Senior-Level#

Key Takeaways#

Further Reading#