adambridge

AdamW for PyTorch with the optimizer state offloaded to host RAM and updated by fused SIMD kernels on the CPU. It overlaps those updates with the GPU backward pass, layer by layer. It implements the OptiCompute-Bridge design: the GPU keeps only the parameters, while FP32 master weights and both Adam moments live on the (NUMA-local) host.

import adambridge

opt = adambridge.AdamBridge(model.parameters(), lr=1e-4, weight_decay=0.01,
                            precision="bf16_moments")   # 'fp32' | 'bf16' | 'fp8'
opt.attach_hooks(model)                                 # layer-wise async updates

for x, y in loader:
    loss = loss_fn(model(x), y)
    loss.backward()     # each layer is offloaded + updated as soon as its grad is ready
    opt.step_sync()     # wait; the next forward pass sees the new weights

Without attach_hooks, AdamBridge is a drop-in optimizer (backward(); step(); zero_grad()), and works with LR schedulers and clip_grad_norm_. adambridge.OptiComputeAdamW is an alias.

Measured results

All numbers in this section come from two adambridge-lab reports generated on 2026-10-28 (adambridge-lab/reports/20261028-224442-*). Test machine: a laptop with an NVIDIA GeForce RTX 5090 Laptop GPU (24 GiB VRAM), Intel Core Ultra 9 290HX Plus (24 threads, AVX2), 64 GiB RAM, Windows 11, Python 3.14.7, PyTorch 2.14.1 + CUDA 13.2, adambridge 0.1.0 (native/avx2 backend). Compared libraries: bitsandbytes 0.50.2, torchao 0.18.0; DeepSpeed 0.19.8 (CPU kernel only, under WSL, because it is not installed on Windows).

Fine-tuning a real LLM

Qwen2.5-0.5B (494M parameters) fine-tuned on WikiText-2: 200 steps of 8 × 512 tokens, peak learning rate 2e-5 with warm-up and cosine decay. Every configuration starts from the same weights, sees the same batches in the same order, and runs in a fresh process (the LLM fine-tuning experiment in adambridge-lab). Validation loss is measured on the whole WikiText-2 validation split (start: 2.868).

† growth of the process's resident memory beyond the all-GPU reference's.

• Noise. Re-running the reference with two other data orders moved its final loss by up to 0.0017, so differences below 0.002 are noise. All three adambridge precisions are inside that band, and so is the same algorithm run entirely on the GPU: offloading the optimizer costs no quality.
• Why the BF16-only configurations fall behind. An AdamW update at this learning rate is smaller than half the BF16 spacing of most weights, so without an FP32 copy it is rounded away: only 18% of their weights changed (at BF16 resolution) vs 84% for the reference. adambridge keeps BF16 weights on the GPU but the FP32 master on the host.
• Memory. adambridge needs 32% less GPU memory than the reference, the least of any configuration at reference quality, and 2.7 GiB less than torchao's offload at the same quality.
• Speed. With BF16 or FP8 moments adambridge is 5–9% slower than the all-GPU reference (35% with FP32 moments) and needs 40% less time per step than torchao's offload at the same quality (25% less with FP32 moments). The GPU waits once per step for the update of the last gradient of the backward pass. This laptop's GPU also drops its clocks whenever it idles: a forward+backward pass that takes 251 ms back to back takes 22% longer after a 20 ms pause and 33% longer after a 50 ms one, which every CPU-offload optimizer pays.
• Host memory. adambridge's host footprint is its FP32 master, the moments and one gradient and one write-back buffer per parameter: with FP8 moments, less than torchao's BF16 offload, which keeps no FP32 master.

Fine-tuning a small LLM (quick preset)

SmolLM2-135M fine-tuned on WikiText-2: 60 steps of 8 × 256 tokens, same protocol as above (start: 3.306). The report labels these quick-preset numbers as indicative only.

The reference's own spread across data orders was 0.0014, so ties are called within 0.002. adambridge uses 40% less GPU memory than the reference at the same quality. The BF16-only configurations kept only 10–11% of the reference's loss reduction: 2% of their weights changed at BF16 resolution, vs 27% for the reference. This model showed no measurable GPU pause penalty (±0% after 20 ms, +2% after 50 ms).

Training step, speed and GPU memory

A 67M-parameter model (8 blocks, width 1024), 4,096 tokens per step, median of the timed steps. Except the all-GPU baseline, every configuration keeps BF16 weights on the GPU.

With FP8 moments adambridge is 3% slower than the all-GPU AdamW and uses 52% of its GPU memory. The asynchronous mode (attach_hooks) takes 29–33% less time per step than synchronous step() here.

Host kernel

8M parameters, 24 threads (benchmarks/bench_kernels.py --numel 8M). At this size part of the data stays in the CPU caches, so these times are optimistic.

GB/s counts the minimum bytes each variant must move. The fastest adambridge variant (BF16 moments) is 2.0× faster than PyTorch's fused CPU AdamW on Windows and 1.18× under WSL. With FP32 moments adambridge is 1.26× faster on Windows and slightly slower under WSL (2.13 vs 1.98 ms). FP8 moments trade kernel speed for host memory; under WSL they are the slowest native variant.

Install

pip install .               # or: uv pip install .   /   uv add ./adambridge
pip install ".[compiler]"   # also installs ziglang, a C compiler for machines without one
python -m adambridge build       # optional: compile the native kernels now
python -m adambridge info        # shows the ISA, NUMA nodes and L2 size

adambridge is a pure-Python wheel (py3-none-any) that carries its C source, so the same artifact installs everywhere.

CUDA torch: adambridge only declares torch>=2.1. On Windows, and for specific CUDA versions, PyPI's torch is CPU-only, so a clean pip/uv install pulls the CPU build. Install torch from the PyTorch index first (pip install torch --index-url https://download.pytorch.org/whl/cu130). With uv, pin it in your project instead:

[tool.uv.sources]
torch = { index = "pytorch-cuda" }

[[tool.uv.index]]
name = "pytorch-cuda"
url = "https://download.pytorch.org/whl/cu130"   # pick your CUDA version
explicit = true

The native kernel is plain C, loaded through ctypes, and compiled on first use into ~/.cache/adambridge (override with ADAMBRIDGE_CACHE_DIR). The build tries $ADAMBRIDGE_CC, then cc, gcc, clang, then python -m ziglang cc (pip install ziglang gives Windows a compiler), then MSVC cl. If none works, adambridge falls back to the PyTorch kernels, which support everything but are slower. Set ADAMBRIDGE_DISABLE_NATIVE=1 to force the fallback.

How it works

autograd thread        D2H stream          worker thread (per GPU)        H2D stream
grad of layer L ready
  record event ──────► wait, copy grad
                       to pinned host ───► wait copy, fused AdamW on
                                           NUMA-pinned thread pool ─────► copy weights
backward continues for layer L-1 ...                                      back to GPU

step_sync() waits for the host work. It then makes the current CUDA stream wait on the H2D stream, so no device-wide synchronization is needed.

• Scheduling. Autograd's CPU side runs ahead of the GPU, so a parameter's hook fires long before the GPU has produced its gradient. The worker updates, in one native call, every gradient whose copy has already landed (copies complete in order on one stream), so the host update stays right behind the GPU's backward pass. Under load the batches grow on their own.

• Large gradients are chunked. Gradients above 16.8M elements are copied, updated and copied back in chunks, so the three overlap. This matters most for the last gradient of the backward pass. In a language model that is the embedding, which the next forward pass needs first, so the GPU waits on exactly this tail. Chunks draw the same rounding bits as a whole-tensor update.

• Small parameters are grouped. Parameters below small_param_threshold (65,536 elements: norms, biases) get no per-tensor hook; at step_sync() they move as one buffer each way.

• Exact-size pinned buffers. PyTorch's pinned allocator rounds every buffer up to a power of two, so a 272 MB gradient buffer locks 512 MB. adambridge registers ordinary allocations of 1 MiB or more with CUDA instead, which locks exactly their size.

• Fused kernel (adambridge/csrc/adambridge_kernels.c): one pass reads p, g, m and v and writes p, m, v and the low-precision copy for the GPU. ISA is chosen at runtime: AVX-512F, AVX2+FMA, NEON or scalar. The arithmetic is exactly torch.optim.AdamW's.

• Moment storage (precision=):

  precision	m	v	host bytes/param	kernel
  fp32	FP32	FP32	12	native
  bf16	BF16	BF16	8	native
  fp8	E4M3, per-128 scale	E5M2 of sqrt(v), per-128 scale	6.06	native

  Plus the transfer buffers: one gradient and one write-back buffer per parameter, in its dtype.

• NUMA: each GPU's state is allocated by, and updated on, threads pinned to the CPUs of the NUMA node that the GPU's PCIe slot hangs off (read from Linux sysfs). This uses first-touch placement, so no libnuma is needed. On other platforms pinning is a no-op.

• Auto-tuning (autotune=True or AutoTuneConfig(...)):

  • It benchmarks thread counts for the native kernel.
  • When host update time exceeds GPU backward time, it steps moment precision down a level.
  • It keeps a downgrade only if host time actually falls. Floor: min_precision="bf16" by default.

Where this deviates from the design document, and why

1. Stochastic rounding for BF16/FP8 moments. With round-to-nearest and β₂ = 0.999, the per-step change in v is smaller than half a BF16 ulp, so v freezes. test_ema_stalls_without_stochastic_rounding shows it sticking at its initial value. adambridge rounds stochastically using a counter-based hash. The native and PyTorch kernels use the same hash, so they make the same rounding decisions, and results do not depend on the thread split or on chunking.
2. FP8 stores sqrt(v), not v. Block-scaled E5M2 v flushes to zero for any element whose gradient is about 2¹⁶× below its block's maximum, while m stays non-zero. The step then becomes m/eps and explodes. Storing the square root doubles the usable range, at the cost of a ~1–2% upward bias in v.
3. BF16 write-back rounds to nearest-even instead of truncating the top 16 bits. Truncation biases every weight toward zero.
4. The kernel handles tails. The spec kernel reads and writes past the end when n % 16 != 0.
5. Master weights are created lazily on the first update, not at construction. Loading pretrained weights after building the optimizer would otherwise be silently reverted.
6. SPSC queues are Python SimpleQueues. There is one producer side (autograd hooks) and one consumer thread per device. Lock contention isn't the bottleneck: kernels run with the GIL released.

Parameters

adambridge.AdamBridge(
    params,
    lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-2,
    *,                                   # everything below is keyword-only
    precision="fp32", stochastic_rounding=True, backend="auto",
    numa_aware=True, num_threads=None, chunk_size=None,
    release_grads=True, overlap_writeback=True,
    autotune=False, seed=0, small_param_threshold=65536,
)

AdamW hyperparameters

These behave exactly as in torch.optim.AdamW. They can differ between parameter groups (pass a list of dicts as params) and can be changed by LR schedulers.

Memory and precision

Hardware and threading

Asynchronous mode (attach_hooks)

Auto-tuning (adambridge.autotune.AutoTuneConfig)

The tuner first benchmarks thread counts once. Then, if the CPU update takes longer than the GPU's backward pass, it steps the moment precision down one level (fp32 → bf16 → fp8). It keeps the change only if the update actually got faster, and stops at the first decision.

from adambridge.autotune import AutoTuneConfig
opt = adambridge.AdamBridge(model.parameters(), autotune=AutoTuneConfig(min_precision="fp8"))

Methods

Read-only attributes: opt.backend ("native/avx2", "native/avx512", "native/neon", "native/scalar" or "torch"), opt.precision, opt.autotune_log (list of the tuner's decisions) and opt.last_step_stats.

step_sync() and opt.last_step_stats return a StepStats with these fields (seconds):

Environment variables

Command line

Things to know

• Global gradient clipping (clip_grad_norm_) is impossible in asynchronous mode, because the first layers update before the last gradients exist. Use synchronous step() if you need it.
• LR schedulers work in both modes: they change lr in the parameter groups, which is read at each update.
• Master weights are created on a parameter's first update, so loading pretrained weights after building the optimizer is safe.

Testing

pytest                          # 83 tests; CUDA tests run when a GPU is present
cc -O2 -pthread -o selftest tests/native/selftest.c -lm && ./selftest   # includes the kernel source

Results from the 2026-10-28 report:

What the tests check:

• adambridge's updates against torch.optim.AdamW on CPU and GPU (FP32 tolerance: rtol 1e-6, atol 1e-7), in sync and async modes.
• Parameter groups, schedulers, weight tying, gradient accumulation, gradient clipping, BF16 params, checkpoint resume, chunked updates, exact-size pinned buffers, and the auto-tuner.
• The C self-test compares the SIMD paths against scalar code and checks the FP8 codec and thread-count and chunking invariance.
• adambridge-lab runs all of this plus the benchmarks and the LLM fine-tuning comparison from one page, and writes the reports.

Limitations / not yet verified

• The AVX-512 and NEON paths compile but were not executed: the test machine is AVX2. Run selftest on such hardware before relying on them.
• Multi-socket NUMA placement is covered by unit tests on a simulated sysfs, not measured on a real two-socket server.
• The GPU waits once per step for the update of the last gradient of the backward pass (for language models, the embedding). Its cost grows with that tensor's size and the bytes per parameter of the moments, and on GPUs that drop their clocks when idle, the next passes run slower too.
• Not covered yet: FP16 gradients with GradScaler, and integration with DDP or FSDP sharding (single-process, one or more GPUs).

License

Copyright 2026 John Anish G. Licensed under the Apache License, Version 2.0.