TL;DR — You will build a Python library that loads large LLM weights on demand from safetensors files, using asynchronous prefetching and SHA‑256 checksum verification. The implementation showcases concrete skills in memory management, concurrency, and data integrity, making it a strong portfolio piece for systems or ML engineering roles.

The rise of multi‑gigabyte language models has turned weight loading from a simple torch.load call into a systems problem. Naïve loading can exhaust RAM, stall training pipelines, and hide corrupted tensors behind opaque runtime errors. This post walks you through a hands‑on project that solves these issues: a safetensors‑based lazy weight loader with asynchronous prefetch and checksum verification. By the end you will have a runnable, testable codebase that demonstrates real‑world engineering chops—exactly the kind of artifact that catches a hiring manager’s eye on LinkedIn.

Why This Project Stands Out on a CV

  • Explicit systems skills – you will implement memory‑mapped file access, background I/O, and cryptographic integrity checks, all core competencies for senior infrastructure or ML platform roles.
  • Production‑grade patterns – the design mirrors techniques used in frameworks such as PyTorch Lightning and Hugging Face Accelerate, showing you understand how to scale model serving.
  • Concrete tooling – you will work with safetensors, asyncio, and SHA‑256, each a named, industry‑standard technology.
  • Portfolio‑ready artifact – the code is self‑contained, includes tests, and can be extended with observability or persistence, giving you a story to discuss in interviews.
  • Signals readiness for ML engineering – by handling large tensors lazily and verifying integrity, you demonstrate the same concerns that arise when deploying models at companies like Anthropic or Cohere.

Architecture Overview

The loader is composed of four cooperating components:

  1. File Indexer – scans a directory of .safetensors files and builds an in‑memory index mapping tensor names to file offsets and sizes.
  2. Async Prefetcher – a pool of asyncio tasks that read upcoming tensors into a bounded buffer while the main thread processes the current batch.
  3. Checksum Validator – computes a SHA‑256 hash over each tensor’s raw bytes and compares it against a pre‑stored manifest (e.g., a JSON sidecar file).
  4. Lazy Tensor Proxy – a thin wrapper that returns a NumPy or PyTorch tensor on first access, triggering the prefetch and validation pipeline transparently.

A high‑level flow looks like this:

[Client Code] → LazyTensorProxy → AsyncPrefetcher → SafetensorsFile
                     │                     │
                     └── ChecksumValidator ←─ Manifest

The proxy hides the I/O complexity; the prefetcher keeps the GPU busy; the validator guarantees data integrity. This separation mirrors the layered architecture of production model servers such as Triton Inference Server.

Building It Step by Step

1. Set up the project

python -m venv venv
source venv/bin/activate
pip install safetensors numpy torch asyncio

Create the package structure:

lazy_loader/
├── __init__.py
├── indexer.py
├── prefetcher.py
├── validator.py
└── proxy.py

2. Implement the File Indexer

# indexer.py
import os
import json
from safetensors import safe_open
from safetensors.tensor import TensorHeader

def index_safetensors(directory: str) -> dict:
    """
    Walks `directory`, reads each .safetensors file,
    and returns a dict: tensor_name -> (filename, offset, length)
    """
    index = {}
    for fname in os.listdir(directory):
        if not fname.endswith(".safetensors"):
            continue
        path = os.path.join(directory, fname)
        with safe_open(path, framework="np") as f:
            for tensor_name in f.keys():
                # safetensors stores tensors contiguously; we can compute offset
                # by accumulating previous tensor sizes (simplified here)
                header = f.get_tensor_header(tensor_name)
                offset = header.data_offsets[0]
                length = header.data_offsets[1] - offset
                index[tensor_name] = (path, offset, length)
    return index

3. Build the Async Prefetcher

# prefetcher.py
import asyncio
import mmap
import os
from collections import deque

class PrefetchBuffer:
    def __init__(self, capacity: int = 4):
        self.capacity = capacity
        self._queue = deque()
        self._semaphore = asyncio.Semaphore(capacity)

    async def fetch(self, path: str, offset: int, length: int) -> bytes:
        async with self._semaphore:
            loop = asyncio.get_running_loop()
            # Use run_in_executor to avoid blocking the event loop
            return await loop.run_in_executor(
                None, self._read, path, offset, length
            )

    @staticmethod
    def _read(path: str, offset: int, length: int) -> bytes:
        with open(path, "rb") as f:
            mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
            try:
                return mm[offset:offset + length]
            finally:
                mm.close()

4. Add Checksum Verification

# validator.py
import hashlib
import json

def load_manifest(manifest_path: str) -> dict:
    with open(manifest_path, "r") as f:
        return json.load(f)

def verify_tensor(data: bytes, expected_hash: str) -> bool:
    actual = hashlib.sha256(data).hexdigest()
    return actual == expected_hash

5. Glue Everything Together in a Lazy Proxy

# proxy.py
import asyncio
import numpy as np
from indexer import index_safetensors
from prefetcher import PrefetchBuffer
from validator import load_manifest, verify_tensor

class LazyTensorLoader:
    def __init__(self, model_dir: str, manifest_path: str):
        self.index = index_safetensors(model_dir)
        self.manifest = load_manifest(manifest_path)
        self.prefetcher = PrefetchBuffer()
        self._loop = asyncio.new_event_loop()
        self._loop.run_until_complete(self._warmup())

    async def _warmup(self):
        # Pre‑fetch the first few tensors to hide initial latency
        for name, (path, offset, length) in list(self.index.items())[:2]:
            await self.prefetcher.fetch(path, offset, length)

    def get_tensor(self, name: str) -> np.ndarray:
        if name not in self.index:
            raise KeyError(f"Tensor {name} not found")
        path, offset, length = self.index[name]
        # Run the fetch in the event loop
        future = asyncio.run_coroutine_threadsafe(
            self.prefetcher.fetch(path, offset, length), self._loop
        )
        raw = future.result()
        # Verify checksum
        expected = self.manifest.get(name)
        if expected and not verify_tensor(raw, expected):
            raise ValueError(f"Checksum mismatch for {name}")
        # Convert bytes to numpy array (simplified; real code uses safetensors)
        return np.frombuffer(raw, dtype=np.float32)

6. Usage Example

# example.py
import asyncio
from proxy import LazyTensorLoader

loader = LazyTensorLoader(
    model_dir="./model",
    manifest_path="./model/manifest.json"
)

# Access a tensor lazily
weights = loader.get_tensor("model.layers.0.attention.q_proj.weight")
print(weights.shape)

Running and Testing It

  1. Prepare a sample model – download a small safetensors checkpoint (e.g., from the Hugging Face Hub) and generate a manifest with SHA‑256 hashes:
python -c "
import hashlib, json, os
from safetensors import safe_open

manifest = {}
for f in os.listdir('model'):
    if f.endswith('.safetensors'):
        with safe_open(os.path.join('model', f), framework='np') as sf:
            for key in sf.keys():
                data = sf.get_tensor(key).tobytes()
                manifest[key] = hashlib.sha256(data).hexdigest()
with open('model/manifest.json', 'w') as fp:
    json.dump(manifest, fp)
"
  1. Run the example – ensure the event loop is available:
python example.py

You should see the tensor shape printed without loading the entire model into RAM.

  1. Write unit tests – use pytest to verify that a corrupted tensor raises ValueError and that lazy access returns the correct data:
# test_loader.py
import pytest
import numpy as np
from proxy import LazyTensorLoader

def test_lazy_loading():
    loader = LazyTensorLoader("./model", "./model/manifest.json")
    tensor = loader.get_tensor("model.layers.0.attention.q_proj.weight")
    assert tensor.shape == (4096, 4096)

def test_checksum_failure():
    # Tamper with manifest to trigger mismatch
    loader = LazyTensorLoader("./model", "./model/manifest.json")
    loader.manifest["model.layers.0.attention.q_proj.weight"] = "deadbeef"
    with pytest.raises(ValueError):
        loader.get_tensor("model.layers.0.attention.q_proj.weight")

Run with pytest -v. A passing suite proves the loader works and that integrity checks are enforced.

Extending It: Your Roadmap to Senior-Level

  1. Persistent caching with SQLite – store prefetched tensors in an on‑disk cache to avoid repeated I/O across sessions. Why it matters: reduces cold‑start latency in production serving.
  2. Horizontal scaling via gRPC – expose the loader as a microservice so multiple inference workers can request tensors from a central node. Why it matters: enables elastic scaling for bursty traffic.
  3. Observability with OpenTelemetry – emit latency, cache‑hit ratio, and checksum failure metrics to a Prometheus endpoint. Why it matters: gives SREs visibility into data‑pipeline health.
  4. Fault‑tolerant prefetch – implement retry logic with exponential backoff and fallback to alternative tensor sources (e.g., a remote object store). Why it matters: guarantees availability even when a local file is corrupted.
  5. Benchmark harness – add a CLI that measures end‑to‑end loading time, peak memory usage, and throughput across different prefetch buffer sizes. Why it matters: provides data‑driven evidence for capacity planning.
  6. Integration with model compilers – hook into TorchScript or ONNX Runtime to load optimized subgraphs on demand. Why it matters: positions the loader as part of a full inference stack.

Key Takeaways

  • Lazy loading with safetensors decouples memory usage from model size, enabling large‑scale inference on commodity hardware.
  • Asynchronous prefetch keeps the GPU fed while the CPU validates tensor integrity.
  • SHA‑256 checksums provide a lightweight, cryptographic guarantee against silent data corruption.
  • The project demonstrates systems engineering, concurrency, and production‑grade patterns that hiring managers value.
  • Extensibility points (caching, scaling, observability) allow you to evolve the prototype into a full‑fledged serving component.

Further Reading