spartan-torch

Custom layers and building blocks for PyTorch — assemble and fine-tune models fast.

Quickstart

import torch
from spartan_torch import TransformerBlock

block = TransformerBlock(
    in_size=128, head_size=32, num_heads=4,
    out_size=128, ff_hidden_size=512,
)
x = torch.randn(2, 16, 128)
y, _ = block(x)  # y: (2, 16, 128), cache kept for KV decoding

What's inside

Drop-in primitives missing from torch.nn, one brick at a time:

• spartan_torch.transformers — attentions (MHA/GQA, linear, Linformer, Performer, Reformer), SSM mixers (Mamba / Mamba-2 / Mamba-3), pre/post/hybrid-norm blocks, norms (RMSNorm, …), positionals (RoPE, ALiBi, learned/sinusoidal), chunked and reversible variants;
• spartan_torch.cnn — residual, bottleneck, depthwise-separable and inverted-residual blocks;
• spartan_torch.vit — patch embeddings, class tokens, position embeddings, MAE decoder head;
• spartan_torch.masking / spartan_torch.selfsup — patch masking, DINO heads/loss, momentum encoder;
• spartan_torch.utils — small training helpers (e.g. warmup scheduler).

Two things this library deliberately does not do:

• No full models. Stock architectures (ResNet, ViT, LLaMA, …) come from torchvision / timm / 🤗 transformers with pretrained weights. spartan_torch.compat ships key-mappings that load those weights straight into our primitives.
• No heavy training frameworks. The core depends only on torch / torchvision — no Lightning, no HF Trainer. Experiment harnesses live in experiments/ as optional extras.

Reproducibility

• RESULTS.md — weight parity of our primitives against timm/torchvision references, plus a latency/memory sweep of sub-quadratic attentions with a Pareto plot. Tables are generated by scripts (scripts/verify_pretrained.py, scripts/bench_attention.py), never edited by hand.

Installation

Requires Python 3.11+ (and a CUDA-capable Torch — see below).

pip install spartan-torch

For contributors, editable install with experiment and dev extras:

pip install -e ".[experiments,dev]"

CUDA wheels (uv): in pyproject.toml ([tool.uv.sources] / [[tool.uv.index]]) torch is pinned to the cu132 index (download.pytorch.org/whl/cu132). pip from PyPI installs the stock torch build; for a specific CUDA release install with an explicit --index-url https://download.pytorch.org/whl/<cuda-version>.

Development

Dependencies and tests via just + uv:

just sync   # install/sync the environment (uv sync --extra dev --extra experiments)
just test   # run the test suite (venv python -m pytest)
just parity # reference weight-parity tests (timm/torchvision; pretrained pulls weights over the network)
just verify # refresh the parity table in RESULTS.md
just bench  # GPU attention sweep + Pareto plot (requires CUDA)
just clean  # delete .venv

Don't run just sync / just clean while Jupyter kernels are alive — they lock .pyd/.so files and uv can't finish the sync.

Development happens in a devcontainer: .devcontainer/ (CUDA image, GPU passthrough, host-mounted HuggingFace/W&B caches).

License

MIT — see LICENSE.