AdamBridge
Asynchronous Layer-Wise Host Optimizer
Keep active parameters on the GPU and move everything else off it. AdamBridge holds FP32 master weights and both Adam moments in host RAM and updates them with multi-threaded fused SIMD CPU kernels, pipelined into the backward pass.
Core Technical Summary
AdamBridge implements the OptiCompute-Bridge architectural design for PyTorch. Active model parameters remain resident on the GPU, while FP32 master weights and both Adam moments (first and second moments) are offloaded to NUMA-local host RAM and updated via multi-threaded fused SIMD CPU kernels.
By registering autograd backward hooks via opt.attach_hooks(model), parameter updates are pipelined asynchronously. Each layer transfers its gradient to host RAM and executes its CPU SIMD update concurrently while earlier layers continue their backward pass on the GPU.
Benchmarks
Data sourced verbatim from the reproducible benchmark suite shipped with the package. Hardware: NVIDIA RTX 5090 Laptop GPU (24 GiB VRAM), Intel Core Ultra 9 290HX Plus, 64 GiB RAM.
| Optimizer Configuration | Final Val Loss | Step Latency | Peak GPU VRAM | VRAM Saving |
|---|---|---|---|---|
| PyTorch AdamW (FP32 master on GPU) | 2.6402 | 292 ms | 19.6 GiB | Reference |
| AdamBridge (FP8 moments, host offload) | 2.6409 (delta < 0.001) | 305 ms | 13.2 GiB | -32.7% VRAM |
| AdamBridge (BF16 moments, host offload) | 2.6405 (delta < 0.001) | 317 ms | 13.2 GiB | -32.7% VRAM |
| torchao CPUOffloadOptimizer (FP32 weights) | 2.6401 | 525 ms | 15.9 GiB | -18.8% VRAM |
| bitsandbytes AdamW8bit (FP32 weights) | 2.6433 | 312 ms | 16.8 GiB | -14.2% VRAM |
Capabilities
Elimination of BF16 Gradient Truncation
In low-precision optimization, standard AdamW updates at standard learning rates (e.g. 2e-5) are frequently smaller than half the spacing between representable BF16 values. Without a full-precision master copy, updates are rounded away to zero:
- • Pure BF16 optimizers saw only 18% of weights change during Qwen2.5-0.5B fine-tuning benchmarks, leading to severe loss degradation (+0.0124 to +0.1622 above reference).
- • AdamBridge retains BF16 weights on GPU for matrix operations, but maintains the FP32 master weights on host, ensuring 84% parameter updates match the full GPU reference exactly.
Verified 32% to 40% VRAM Reductions
Validated empirically across standardized LLM fine-tuning benchmarks on NVIDIA RTX 5090 hardware:
- • Qwen2.5-0.5B on WikiText-2 (200 steps 8x512): Reduces peak VRAM from 19.6 GiB to 13.2 GiB (-32.7%) while maintaining matching validation loss (2.6405 vs 2.6402).
- • Outperforms torchao CPU offload by 2.7 GiB VRAM at equivalent precision and requires 40% less step latency (305 ms vs 525 ms).
Moment Precision Modes
Supports granular precision trade-offs for host memory footprint:
- • fp32: Full 32-bit floating point moments on host (395 ms step).
- • bf16: 16-bit bfloat moments (+5.6 GiB host RAM, 317 ms step).
- • fp8: 8-bit floating point moments (+4.7 GiB host RAM, 305 ms step).
PyTorch Ecosystem Compatibility
Functions as a seamless drop-in optimizer within existing PyTorch training scripts:
- • Fully compatible with standard PyTorch learning rate schedulers.
- • Supports gradient norm clipping via
torch.nn.utils.clip_grad_norm_. - • Includes autotuning harness (
adambridge.autotune.AutoTuneConfig) to calibrate thread pools and chunk sizes.
Usage
Asynchronous Training Loop
import adambridge
# OptiCompute-Bridge: Host-offloaded state with fused SIMD kernels
opt = adambridge.AdamBridge(
model.parameters(),
lr=1e-4,
weight_decay=0.01,
precision="bf16_moments" # Options: 'fp32' | 'bf16' | 'fp8'
)
opt.attach_hooks(model) # Layer-wise async updates during backward
for x, y in loader:
loss = loss_fn(model(x), y)
loss.backward() # Layers offloaded & updated as grad is ready
opt.step_sync() # Sync point; forward pass sees updated weightsAdamBridge: FAQ
What is AdamBridge?
AdamBridge is an open-source PyTorch optimizer from InnvoSols. It keeps active BF16 parameters on the GPU and moves FP32 master weights and both Adam moments to host RAM, where multi-threaded fused SIMD CPU kernels update them layer by layer during the backward pass.
How much GPU memory does AdamBridge save?
Fine-tuning Qwen2.5-0.5B on WikiText-2 on an NVIDIA RTX 5090, AdamBridge reduced peak VRAM from 19.6 GiB to 13.2 GiB (32.7% lower) with matching validation loss (2.6405 vs 2.6402). It also used 2.7 GiB less VRAM than torchao CPU offload with lower step latency (305 ms vs 525 ms).
Why not just train with a pure BF16 optimizer?
At typical fine-tuning learning rates (for example 2e-5), many AdamW updates are smaller than half the spacing between representable BF16 values, so they round away to zero. In benchmarks a pure BF16 optimizer changed only 18% of weights. AdamBridge keeps FP32 master weights on the host so those small updates are preserved.
How do I use AdamBridge?
Install it with pip install adambridge, create adambridge.AdamBridge(model.parameters(), lr=...), call opt.attach_hooks(model) to enable layer-wise asynchronous updates during backward, and call opt.step_sync() after loss.backward().
Is AdamBridge compatible with existing PyTorch training code?
Yes. It works as a drop-in optimizer with standard PyTorch learning-rate schedulers and gradient clipping via torch.nn.utils.clip_grad_norm_, and includes an autotuning harness (adambridge.autotune.AutoTuneConfig) for thread pools and chunk sizes.
What moment precision modes does AdamBridge offer?
Adam moments can be stored on the host in fp32, bf16 or fp8, trading host RAM footprint against precision. Measured step times were 395 ms (fp32), 317 ms (bf16) and 305 ms (fp8).
Custom Optimization for Your Workloads
If your organization is training large architectures and encountering VRAM constraints, PCIe bottlenecks, or latency targets, our systems engineering team is available for consulting engagements and custom kernel authoring.