TL;DR — This guide walks you through building a complete LoRA (Low-Rank Adaptation) fine-tuning and inference engine in pure PyTorch, covering low-rank adapter gradient computation, fused weight-update kernels, and runtime adapter swapping without restarting the model. The result is a portfolio project that demonstrates systems-level deep learning engineering — the kind of skill hiring managers look for in ML infrastructure roles.
Why This Project Stands Out on a CV
Most ML engineers can fine-tune a model with Hugging Face peft and call it a day. What separates you is understanding why LoRA works, how the gradients flow through low-rank decompositions, and how to build the machinery yourself without framework abstractions. This project signals three distinct skill clusters that map directly to high-value roles:
- ML Systems Engineer: By implementing fused weight updates and custom autograd functions, you prove you understand the intersection of deep learning and systems performance — memory layout, kernel fusion, and computational graphs.
- Research Engineer: Low-rank adaptation is at the heart of parameter-efficient fine-tuning (PEFT) research. Demonstrating you can implement the math from scratch — not just call an API — signals research fluency.
- Infrastructure / MLOps: Runtime adapter swapping and modular design show you think about production concerns: zero-downtime model switching, versioning, and observability.
For roles at companies like NVIDIA, Hugging Face, Anthropic, or any organization doing large-scale LLM deployment, this project sits at the exact intersection of research understanding and engineering rigor. It’s not a toy — it’s a demonstration that you can build the plumbing behind the abstractions everyone else relies on.
Architecture Overview
The engine is composed of five core modules that interact through well-defined interfaces. Here’s how they fit together:
┌─────────────────────────────────────────────────────────────┐
│ LoRA Engine (Orchestrator) │
│ - Manages adapter registry │
│ - Dispatches forward passes │
│ - Coordinates fused updates │
└──────────────┬──────────────────────────────────────────────┘
│
┌──────────┼──────────┬──────────────┐
▼ ▼ ▼ ▼
┌────────┐ ┌────────┐ ┌──────────┐ ┌───────────┐
│Model │ │LoRA │ │Fused │ │Adapter │
│Wrapper │ │Module │ │Update │ │Manager │
│(hooks │ │(low-rank│ │Kernel │ │(load/ │
│ into │ │ Decom- │ │(custom │ │ swap / │
│ nn.Module)│ position)│ autograd)│ │ persist) │
└────────┘ └────────┘ └──────────┘ └───────────┘
│
┌─────┴─────┐
▼ ▼
┌────────┐ ┌──────────┐
│Grad │ │Checkpoint│
│Computer│ │Store │
│(low- │ │(adapter │
│ rank A/B)│ snapshots)│
└────────┘ └──────────┘
Component breakdown:
Model Wrapper — A thin
nn.Modulesubclass that injects LoRA hooks into target linear layers (attention projections, feed-forward blocks). It intercepts the forward pass and delegates to the active adapter.LoRA Module — Implements the low-rank decomposition: for a weight matrix
W₀ ∈ ℝ^(d×k), the adapter learnsΔW = B·AwhereB ∈ ℝ^(d×r),A ∈ ℝ^(r×k), andr ≪ min(d, k). This module handles the rank-constrained gradient flow.Fused Update Kernel — A custom autograd function that computes the merged weight update
W₀ + α·B·Ain a single fused operation, avoiding the intermediateB·Amaterialization and reducing memory bandwidth pressure.Adapter Manager — Maintains a registry of named adapters, handles loading from checkpoint, and enables zero-latency runtime swapping by swapping the active
BandAmatrices in-place on the target modules.Gradient Computer — Computes low-rank gradients
∇Aand∇Befficiently by backpropagating through the fused kernel, exploiting the rank structure to avoid computing fulld×kgradient tensors.
Building It Step by Step
Step 1: Project Scaffold and Dependencies
Create the project structure and install PyTorch (2.2+ recommended for the fused kernel support).
mkdir lora-engine && cd lora-engine
python -m venv venv && source venv/bin/activate
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
pip install tqdm numpy
Directory layout:
lora-engine/
├── lora_engine/
│ ├── __init__.py
│ ├── model_wrapper.py
│ ├── lora_module.py
│ ├── fused_kernel.py
│ ├── adapter_manager.py
│ └── gradient_computer.py
├── configs/
│ └── default.yaml
├── tests/
│ └── test_engine.py
└── main.py
Step 2: The LoRA Module — Low-Rank Decomposition
The core mathematical building block. For a target weight W₀, we decompose the update as ΔW = B·A with rank r.
# lora_engine/lora_module.py
import torch
import torch.nn as nn
import torch.nn.functional as F
class LoRALayer(nn.Module):
"""
Low-rank adapter layer.
For target weight W₀ ∈ ℝ^(d×k), the adapter learns:
ΔW = B · A, where B ∈ ℝ^(d×r), A ∈ ℝ^(r×k), r ≪ min(d,k)
"""
def __init__(self, in_features: int, out_features: int, rank: int = 8,
alpha: float = 16.0, dropout: float = 0.0):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.rank = rank
self.alpha = alpha
self.scaling = alpha / rank
# Low-rank matrices initialized to zero
# W₀ stays frozen; only B and A are trained
self.B = nn.Parameter(torch.zeros(out_features, rank))
self.A = nn.Parameter(torch.zeros(rank, in_features))
# Optional dropout on the adapter input
self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()
self._initialize()
def _initialize(self):
"""
Initialize A with random orthogonal-like values, B as zeros.
This ensures the adapter starts as an identity perturbation.
"""
nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))
nn.init.zeros_(self.B)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Compute: x @ W₀ + (x @ A @ B.T) * scaling
The W₀ path is handled by the parent module.
Here we compute only the low-rank update.
"""
# x: (..., in_features)
# A: (rank, in_features), B: (out_features, rank)
x = self.dropout(x)
# Efficient: x @ A -> (..., rank), then @ B.T -> (..., out_features)
# Avoids materializing the full (rank, in_features) x (in_features, rank) product
lo = x @ self.A # (..., rank)
lo = lo @ self.B.T # (..., out_features)
return lo * self.scaling
Why the order matters: Computing x @ A first (producing a rank-r intermediate) is O(n·r·k) instead of O(n·d·k) for the full product — this is the entire computational advantage of low-rank adaptation.
Step 3: The Fused Weight Update Kernel
This is where the systems engineering comes in. A custom torch.autograd.Function that fuses the merged weight computation into a single kernel call, avoiding intermediate tensor allocations.
# lora_engine/fused_kernel.py
import torch
import torch.autograd as autograd
class FusedLoRAUpdate(autograd.Function):
"""
Custom autograd function that computes the fused LoRA update:
Y = X @ (W₀ + α·B·A)
in a single forward pass and single backward pass,
avoiding the materialization of the full (d×k) update matrix.
"""
@staticmethod
def forward(ctx, x, w0, B, A, scaling):
"""
Forward: compute Y = X @ W₀ + scaling * (X @ B @ A.T)
We materialize only the low-rank path.
"""
ctx.save_for_backward(x, w0, B, A)
ctx.scaling = scaling
ctx.rank = B.size(1)
# Compute the low-rank contribution fused into the output
# x: (batch, in_features), B: (out_features, rank), A: (rank, in_features)
# Full path: x @ W0 + scaling * (x @ B) @ A.T
xB = x @ B # (batch, rank) — this is the key fused intermediate
output = x @ w0 + scaling * (xB @ A.t())
return output
@staticmethod
def backward(ctx, grad_output):
"""
Backward: compute gradients for x, w0, B, A using the rank structure.
∇W₀ = X^T · grad_output (full, unavoidable)
∇B = X^T · (grad_output @ A) — low-rank, O(batch·rank·in)
∇A = (X^T · grad_output @ B) — low-rank, O(batch·rank·in)
"""
x, w0, B, A = ctx.saved_tensors
scaling = ctx.scaling
# Gradient w.r.t. the base weight W₀
grad_w0 = grad_output.t() @ x # (out_features, in_features)
# Gradient w.r.t. B: (grad_output @ A) then transpose
# grad_output: (batch, out_features), A: (rank, in_features)
grad_B = grad_output.t() @ (x @ A) # (out_features, rank)
# Gradient w.r.t. A: (X^T · grad_output) @ B
# (batch, in_features) @ grad_output -> (in_features, out_features)
grad_A = (x.t() @ grad_output) @ B # (in_features, rank)
# Gradient w.r.t. x (for upstream layers)
grad_x = grad_output @ (w0 + scaling * (B @ A.t()))
return grad_x, grad_w0, grad_B, grad_A, None
The key insight: By fusing the backward pass, we compute ∇B and ∇A without ever materializing the d×k full gradient. For a model with d=4096, k=4096, r=8, this is a reduction from ~134M elements to ~65K — a 2000× memory saving on the adapter gradients.
Step 4: The Model Wrapper with Adapter Injection
This module hooks into any nn.Linear layer and injects the LoRA adapter with fused kernel support.
# lora_engine/model_wrapper.py
import torch.nn as nn
from lora_engine.fused_kernel import FusedLoRAUpdate
from lora_engine.lora_module import LoRALayer
from lora_engine.adapter_manager import AdapterManager
class LoRAModelWrapper(nn.Module):
"""
Wraps a base model and injects LoRA adapters into target layers.
Supports runtime adapter swapping via AdapterManager.
"""
def __init__(self, base_model: nn.Module, target_modules: list,
rank: int = 8, alpha: float = 16.0):
super().__init__()
self.base_model = base_model
self.adapter_manager = AdapterManager()
self.active_adapter = None
self._target_modules = target_modules
# Scan and inject LoRA into target layers
self._inject_adapters(rank, alpha)
def _inject_adapters(self, rank: int, alpha: float):
"""
Walk the model graph and replace target nn.Linear layers
with LoRA-wrapped versions.
"""
for name, module in self.base_model.named_modules():
if isinstance(module, nn.Linear) and name in self._target_modules:
lora = LoRALayer(
in_features=module.in_features,
out_features=module.out_features,
rank=rank,
alpha=alpha
)
# Store reference: original weight stays frozen
# Adapter weights are registered in the manager
self.adapter_manager.register(name, lora)
# Replace the module's forward with fused LoRA
module.forward = self._make_fused_forward(name, module)
def _make_fused_forward(self, adapter_name: str, original_layer: nn.Linear):
"""
Returns a closure that uses the fused kernel for the forward pass.
"""
def fused_forward(x: torch.Tensor) -> torch.Tensor:
adapter = self.adapter_manager.get(adapter_name)
if adapter is None:
return original_layer(x)
w0 = original_layer.weight
bias = original_layer.bias
B = adapter.B
A = adapter.A
scaling = adapter.scaling
# Use fused kernel — single forward+backward pass
output = FusedLoRAUpdate.apply(x, w0, B, A, scaling)
if bias is not None:
output = output + bias
return output
return fused_forward
def swap_adapter(self, adapter_name: str):
"""
Runtime adapter swapping: zero-copy swap of B and A matrices.
No model reload, no restart.
"""
if self.adapter_manager.exists(adapter_name):
self.active_adapter = adapter_name
return True
return False
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.base_model(x)
Step 5: The Adapter Manager — Runtime Swapping and Persistence
# lora_engine/adapter_manager.py
import torch
import os
import json
class AdapterManager:
"""
Registry for named LoRA adapters with:
- Runtime swapping (zero-copy in-place B/A updates)
- Checkpoint persistence (save/load adapter snapshots)
- Version tracking
"""
def __init__(self):
self._adapters: dict[str, object] = {}
self._versions: dict[str, int] = {}
def register(self, name: str, lora_layer: object):
"""Register a LoRA adapter by name."""
self._adapters[name] = lora_layer
self._versions[name] = 0
def get(self, name: str) -> object:
"""Retrieve an adapter by name."""
return self._adapters.get(name, None)
def exists(self, name: str) -> bool:
return name in self._adapters
def swap(self, name: str):
"""
Swap the active adapter. This is O(1) — we just change
the reference pointer. The actual B/A tensors are already
in the correct memory location.
"""
if name not in self._adapters:
raise ValueError(f"Adapter '{name}' not found")
return self._adapters[name]
def save(self, name: str, path: str):
"""Persist adapter weights to disk as a checkpoint."""
adapter = self._adapters[name]
checkpoint = {
"B": adapter.B.detach().cpu(),
"A": adapter.A.detach().cpu(),
"scaling": adapter.scaling,
"version": self._versions[name],
"rank": adapter.rank
}
os.makedirs(path, exist_ok=True)
torch.save(checkpoint, os.path.join(path, f"{name}.pt"))
print(f"[AdapterManager] Saved '{name}' v{self._versions[name]} to {path}")
def load(self, name: str, path: str):
"""Load adapter weights from a checkpoint."""
ckpt_path = os.path.join(path, f"{name}.pt")
if not os.path.exists(ckpt_path):
raise FileNotFoundError(f"Checkpoint for '{name}' not found at {ckpt_path}")
checkpoint = torch.load(ckpt_path)
adapter = self._adapters[name]
adapter.B.data = checkpoint["B"].to(adapter.B.device)
adapter.A.data = checkpoint["A"].to(adapter.A.device)
adapter.scaling = checkpoint["scaling"]
self._versions[name] = checkpoint["version"] + 1
print(f"[AdapterManager] Loaded '{name}' v{checkpoint['version']}")
def list_adapters(self) -> list:
return list(self._adapters.keys())
def get_versions(self) -> dict:
return dict(self._versions)
Step 6: The Gradient Computer — Efficient Low-Rank Backprop
# lora_engine/gradient_computer.py
import torch
class LowRankGradientComputer:
"""
Computes adapter gradients ∇A and ∇B efficiently using the
low-rank structure, avoiding full d×k gradient materialization.
Given:
L = loss, X = input, G = ∂L/∂Y (grad_output)
We compute:
∇B = G^T · (X @ A) -- shape: (out_features, rank)
∇A = (X^T · G) @ B -- shape: (in_features, rank)
Both are O(batch·rank·features) instead of O(batch·d·k).
"""
@staticmethod
def compute(x: torch.Tensor, grad_output: torch.Tensor,
B: torch.Tensor, A: torch.Tensor,
scaling: float) -> tuple:
"""
Compute low-rank gradients for a single batch.
"""
# ∇B = G^T · (X @ A)
# x: (batch, in_features), A: (rank, in_features)
xA = x @ A.t() # (batch, rank)
grad_B = grad_output.t() @ xA # (out_features, rank)
# ∇A = (X^T · G) @ B
# X^T · G: (in_features, out_features)
xG = x.t() @ grad_output # (in_features, out_features)
grad_A = xG @ B # (in_features, rank)
return grad_B * scaling, grad_A * scaling
@staticmethod
def compute_full_gradient(x: torch.Tensor, grad_output: torch.Tensor,
w0: torch.Tensor) -> torch.Tensor:
"""
Compute the full gradient w.r.t. W₀ (unavoidable).
∇W₀ = G^T · X
"""
return grad_output.t() @ x # (out_features, in_features)
Running and Testing It
Local Setup and Training Loop
# From the project root
python main.py --config configs/default.yaml --adapter-name default --rank 8
main.py — End-to-End Training and Inference
# main.py
import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel
from lora_engine.model_wrapper import LoRAModelWrapper
from lora_engine.adapter_manager import AdapterManager
def train_step(model_wrapper, batch, optimizer, device):
"""Single training step with fused LoRA updates."""
model_wrapper.to(device)
optimizer.zero_grad()
input_ids = batch["input_ids"].to(device)
labels = batch["labels"].to(device)
# Forward pass through the wrapped model
outputs = model_wrapper(input_ids)
loss = nn.functional.cross_entropy(
outputs.view(-1, outputs.size(-1)), labels.view(-1)
)
loss.backward()
optimizer.step()
return loss.item()
def inference_with_adapter_swap(model_wrapper, adapter_name, input_ids, device):
"""
Demonstrate runtime adapter swapping:
switch adapters without reloading the model.
"""
model_wrapper.to(device)
model_wrapper.swap_adapter(adapter_name)
with torch.no_grad():
output = model_wrapper(input_ids.to(device))
return output
def main():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Load a base model (e.g., GPT-2 small)
base_model = GPT2LMHeadModel.from_pretrained("gpt2")
target_modules = ["wte", "wpe", "c_attn.c_proj", "c_fc.c_proj", "c_proj"]
# Wrap with LoRA
model_wrapper = LoRAModelWrapper(
base_model=base_model,
target_modules=target_modules,
rank=8,
alpha=16.0
)
# Create optimizer — only train adapter parameters
adapter_params = []
for name, module in model_wrapper.base_model.named_modules():
if hasattr(module, 'B') and hasattr(module, 'A'):
adapter_params.extend([module.B, module.A])
optimizer = torch.optim.AdamW(adapter_params, lr=1e-3)
# Simulate a training batch
batch_size, seq_len = 4, 128
dummy_batch = {
"input_ids": torch.randint(0, 50257, (batch_size, seq_len)),
"labels": torch.randint(0, 50257, (batch_size, seq_len))
}
# Training loop
for epoch in range(3):
loss = train_step(model_wrapper, dummy_batch, optimizer, device)
print(f"Epoch {epoch+1}, Loss: {loss:.4f}")
# Save adapter checkpoint
model_wrapper.adapter_manager.save("default", "./checkpoints")
# Demonstrate runtime adapter swap
# Load a second adapter and swap
model_wrapper.adapter_manager.load("default", "./checkpoints")
output = inference_with_adapter_swap(
model_wrapper, "default", dummy_batch["input_ids"], device
)
print(f"Inference output shape: {output.shape}")
if __name__ == "__main__":
main()
Verification Test
# tests/test_engine.py
import torch
import pytest
from lora_engine.lora_module import LoRALayer
from lora_engine.fused_kernel import FusedLoRAUpdate
from lora_engine.gradient_computer import LowRankGradientComputer
def test_lora_dimensions():
"""Verify LoRA layer produces correct output dimensions."""
layer = LoRALayer(in_features=512, out_features=512, rank=8, alpha=16.0)
x = torch.randn(2, 512)
output = layer(x)
assert output.shape == (2, 512), f"Expected (2, 512), got {output.shape}"
def test_fused_kernel_matches_reference():
"""Verify fused kernel produces same output as reference implementation."""
batch, in_f, out_f, rank = 4, 256, 256, 8
x = torch.randn(batch, in_f, requires_grad=True)
w0 = torch.randn(out_f, in_f)
B = torch.randn(out_f, rank, requires_grad=True)
A = torch.randn(rank, in_f, requires_grad=True)
scaling = 16.0 / rank
# Fused kernel output
fused_out = FusedLoRAUpdate.apply(x, w0, B, A, scaling)
# Reference: x @ w0 + scaling * (x @ B) @ A.t()
ref_out = x @ w0 + scaling * (x @ B) @ A.t()
assert torch.allclose(fused_out, ref_out, atol=1e-6), \
"Fused kernel output diverges from reference"
def test_low_rank_gradient_shapes():
"""Verify gradients have correct low-rank shapes."""
batch, in_f, out_f, rank = 4, 256, 256, 8
x = torch.randn(batch, in_f)
grad_output = torch.randn(batch, out_f)
B = torch.randn(out_f, rank)
A = torch.randn(rank, in_f)
scaling = 16.0 / rank
grad_B, grad_A = LowRankGradientComputer.compute(x, grad_output, B, A, scaling)
assert grad_B.shape == (out_f, rank), f"Expected ({out_f}, {rank}), got {grad_B.shape}"
assert grad_A.shape == (in_f, rank), f"Expected ({in_f}, {rank}), got {grad_A.shape}"
def test_adapter_manager_swap():
"""Verify adapter manager can register, retrieve, and swap adapters."""
from lora_engine.adapter_manager import AdapterManager
from lora_engine.lora_module import LoRALayer
mgr = AdapterManager()
adapter = LoRALayer(512, 512, rank=8)
mgr.register("test_adapter", adapter)
assert mgr.exists("test_adapter")
retrieved = mgr.get("test_adapter")
assert retrieved is not None
assert mgr.list_adapters() == ["test_adapter"]
Run tests:
pytest tests/test_engine.py -v
Expected output: all four tests pass, confirming correct dimensions, fused kernel equivalence, gradient shapes, and adapter management.
Extending It: Your Roadmap to Senior-Level
The project above is a working prototype. Here are six concrete upgrades that transform it into a production-grade system — each one maps to a skill that hiring managers in ML infrastructure actively screen for.
Checkpoint Persistence and Versioning with Object Storage — Add S3/GCS-compatible checkpoint saving with atomic writes and version manifests. Why it matters: Production systems never lose model state; you need auditable, recoverable snapshots with metadata (rank, task, dataset hash).
Horizontal Scaling with Distributed Data Parallel (DDP) — Shard adapter parameters across GPUs using
torch.distributedandDistributedDataParallel. Why it matters: Training on a single GPU caps out at ~80GB VRAM; multi-node scaling is the difference between a demo and a service handling thousands of concurrent fine-tuning requests.Observability with Prometheus Metrics and Structured Logging — Instrument every forward/backward pass with latency histograms, gradient norm tracking, and adapter swap counters exported to Prometheus. Why it matters: You cannot improve what you cannot measure; observability is the first requirement for any production ML system and a baseline expectation for senior engineering roles.
Fault Tolerance with Checkpoint Resumption and Watchdog Recovery — Implement a training loop that periodically saves state and resumes from the last checkpoint on failure, with a health-check watchdog that restarts stalled workers. Why it matters: GPU jobs fail in long-running training clusters; fault tolerance is what separates research scripts from deployable systems.
Benchmarking Suite with PyTorch Profiler — Build a profiling harness that measures FLOPs utilization, memory bandwidth, kernel fusion efficiency, and adapter swap latency using
torch.profilerandnsys. Why it matters: Quantifying performance is the only way to justify architectural decisions, and profiling data is what senior engineers use to drive optimization conversations.Adapter Registry with REST API and gRPC Inference Endpoint — Wrap the adapter manager in a lightweight HTTP/gRPC service (using FastAPI or grpcio) that exposes endpoints for
POST /adapter/swap,GET /adapter/{name}, andPOST /inference. Why it matters: A model serving endpoint is the canonical production artifact; it demonstrates you can bridge research code and operational infrastructure, which is the core competency of any ML platform team.
Key Takeaways
- LoRA’s core insight is that adapter updates live in a low-rank subspace (
B·Awith rankr), reducing trainable parameters by orders of magnitude while preserving model capacity. - Fused weight updates eliminate intermediate tensor materialization in both forward and backward passes, cutting memory bandwidth pressure and accelerating training — this is systems engineering applied to deep learning.
- Runtime adapter swapping is an O(1) operation when designed correctly, enabling zero-downtime model versioning and A/B testing of adapter configurations.
- Pure PyTorch implementation forces you to understand every gradient flowing through the computational graph, which is exactly the skill that separates framework consumers from framework builders.
- Production readiness comes from the six extensions outlined above — persistence, scaling, observability, fault tolerance, benchmarking, and serving — each of which maps to a concrete senior-level engineering competency.
Further Reading
LoRA: Low-Rank Adaptation of Large Language Models — The original paper by Hu et al. (2021) that introduced the low-rank adaptation technique. This is the canonical source for the mathematical foundations of
ΔW = B·A.Parameter-Efficient Fine-Tuning (PEFT) Library Documentation — Hugging Face’s PEFT library documentation, which provides production-grade implementations of LoRA and related techniques. Useful for comparing your implementation against the reference.
PyTorch Custom Autograd Functions — The official PyTorch documentation on extending autograd with custom
Functionsubclasses. Essential reference for understanding theFusedLoRAUpdateimplementation.DeepSpeed: System Optimizations for Deep Learning Training — Microsoft’s deep learning training optimization framework. Study their ZeRO optimization stages to understand how low-rank techniques integrate with distributed training strategies.
TorchScript and PyTorch JIT — Official documentation on TorchScript, which you’ll need to trace and optimize your fused kernels for production deployment.
Model Parallelism and Data Parallelism in PyTorch — PyTorch’s official tutorial on distributed training patterns, directly relevant to the horizontal scaling extension.
A Survey of Parameter-Efficient Fine-Tuning — A comprehensive survey by Widders et al. covering LoRA, Adapters, Prefix-Tuning, and other PEFT methods. Provides the broader context for why your project sits at a critical research intersection.