TrainStream

Train and fine-tune models larger than your GPU's memory by moving optimizer state, gradients, activations and, when needed, the weights themselves to system RAM over PCIe. It takes one line of code and has no third-party dependencies.

import trainstream
model, optimizer = trainstream.wrap(model, optimizer)

TrainStream is a C++/CUDA extension for PyTorch. It keeps only the weights in VRAM: gradients stream to pinned host memory as autograd produces them, AdamW runs on host threads while the backward pass continues, and a profiled share of activations is paged out and prefetched back block by block. Updates match torch.optim.AdamW (max parameter difference ≈3·10⁻⁷ after 20 steps); activation paging is bit-exact.

Results (RTX 5090 Laptop GPU, 24 GB)

Offloading costs throughput when a model already fits, which is why offload="auto" measures two steps and then switches to plain fused GPU AdamW if offloading isn't needed. TrainStream's value is training what does not fit, or freeing VRAM for larger batches. All numbers come from a single machine (RTX 5090 Laptop, Windows 11).

Requirements

• An NVIDIA GPU and a CUDA build of PyTorch
• The CUDA toolkit matching PyTorch's CUDA major version, and a C++ compiler (Visual Studio 2022+ Build Tools with "Desktop development with C++" on Windows, GCC 11+ on Linux)
• Tested on Windows 11 (PyTorch 2.14.1, CUDA 13.4, MSVC, Python 3.14) and Ubuntu 24.04 under WSL2 (PyTorch 2.14.1+cu130, CUDA 13.3, GCC 13.3, Python 3.12): all 78 tests pass on both, RTX 5090 Laptop GPU. Native Linux and other GPUs are untested.

Install

pip install torch --index-url https://download.pytorch.org/whl/cu130   # if you don't have CUDA PyTorch yet
pip install trainstream

The first import trainstream compiles the CUDA extension against your PyTorch (a few minutes, once) and caches it; every later import is instant. On Windows it finds Visual Studio and the CUDA toolkit itself, so any terminal works. The cache is keyed by PyTorch/CUDA/Python version and GPU, so upgrading PyTorch triggers one rebuild.

To compile ahead of time instead: pip install trainstream --no-binary trainstream --no-build-isolation. From a checkout: build.bat (Windows) or python setup.py build_ext --inplace.

Usage

import torch, trainstream

model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
model, optimizer = trainstream.wrap(model, optimizer)   # max_grad_norm=..., accumulation_steps=... if you use them

for x, y in loader:
    with torch.autocast("cuda", dtype=torch.bfloat16):
        loss = model(x.cuda(), y.cuda())
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

Models bigger than the GPU: build on the CPU and let TrainStream page the transformer-block weights:

model = MyGPT(cfg)                                  # stays on the CPU
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
model, optimizer = trainstream.wrap(model, optimizer, device="cuda",
                                    param_offload=True, moment_dtype="bfloat16")

examples/train_bigger_than_vram.py trains a 1.8B model (7.3 GB of fp32 weights) inside a 4 GB VRAM cap. examples/lora_large_model.py LoRA-fine-tunes a 14.8B-parameter base through TrainStreamTrainer: pass the model on the CPU and its layers stay in host RAM.

Also: trainstream.chunked_cross_entropy(hidden, weight, targets) computes the loss without materializing the logits. .to(), .cuda() and .half() work after wrap(), and reading a weight anywhere waits for its in-flight update.

Hugging Face (pip install transformers accelerate):

from trainstream.integrations.huggingface import TrainStreamTrainer
trainer = TrainStreamTrainer(model=model, args=TrainingArguments(..., bf16=True), train_dataset=ds)
trainer.train()

TrainStreamTrainer computes the causal-LM loss in chunks, so the full logits tensor never exists. Pass extra wrap() options as trainstream={...}. Use bf16=True or fp32; fp16 with loss scaling is not supported.

Tools in this repository

• examples/verify_ram_offload.py — proves where training memory lives (VRAM vs system RAM) with a PASS/FAIL verdict
• examples/hf_finetune.py — fine-tune a Hugging Face model, optionally side by side with the stock Trainer
• examples/train_bigger_than_vram.py — weight paging on a model built on the CPU, under a VRAM cap
• examples/lora_large_model.py — LoRA on 7B/14B-size bases larger than the GPU, optionally against the stock Trainer
• docs/index.html — full documentation

Limitations

• Host RAM is the ceiling: 8–12 bytes per trained parameter and 2–4 per frozen one (LoRA base), plus activations; measured on 64 GB: a 3.6B full fine-tune and LoRA on a 14.8B base
• Weight paging covers the repeated block stack and needs optimizer_mode="cpu"; each paged weight crosses PCIe twice per step
• Data parallel only (no tensor/pipeline parallelism); tested with 2 processes on 1 GPU, not across several GPUs or machines
• TrainStreamTrainer is single-process; fp16 + GradScaler and AMSGrad aren't supported; torch.compile hasn't been tested
• Host AdamW speed is bounded by system RAM bandwidth

License

Apache License 2.0 — see LICENSE and NOTICE. Copyright 2026 John Anish G.