TrainStream
Zero-Dependency PyTorch Training Offload
Train and fine-tune models larger than your GPU's VRAM. TrainStream moves optimizer state, gradients, activations and, when necessary, the weights themselves into host RAM, streaming them back over PCIe exactly when the GPU needs them.
Core Technical Summary
TrainStream enables training and fine-tuning models larger than available GPU VRAM. It moves optimizer state, gradients, activations, and when necessary, weights themselves into system host RAM, streaming them back over PCIe precisely when required by the GPU. Frozen weights (such as a LoRA base) can be stored as int8, NF4, FP8, or NVFP4, with FP8 and NVFP4 multiplied directly by modern Tensor Cores.
Implemented as a native C++/CUDA extension for PyTorch using exclusively the C++ standard library, CUDA runtime, and PyTorch ATen. It contains zero third-party framework dependencies (no DeepSpeed, no bitsandbytes). Weight updates match torch.optim.AdamW with a maximum parameter difference of less than 10⁻⁶ after 20 steps, while activation and weight paging remain bit-exact.
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.
| Framework / Configuration | Step Time | Peak VRAM | Relative Delta |
|---|---|---|---|
| TrainStream (NVFP4 base, FP4 tensor cores) | 0.56 s | 19.1 GB | Reference (Fastest) |
| bitsandbytes QLoRA (NF4) + PEFT | 0.63 s | 17.4 GB | +12.5% time |
| TrainStream (NF4 base) | 0.67 s | 19.1 GB | +19.6% time |
| Unsloth (4-bit) | 0.75 s | 11.5 GB | +33.9% time |
| TrainStream (auto mode) | 0.84 s | 20.6 GB | +50.0% time |
| DeepSpeed ZeRO-2 (Optimizer Offload) | 42.38 s | 35.9 GB | Spills past physical VRAM |
| Stock PyTorch + PEFT | 50.63 s | 37.6 GB | Spills past physical VRAM |
Capabilities
Outlier Channel Protection for NVFP4
A small subset of hidden activation channels in LLMs exhibit extreme magnitudes. In NVFP4, one outlier sets the block scale and quantizes remaining values to zero. With outliers=0.01, each matmul isolates the largest 1% of input channels entirely on GPU (without host synchronization), multiplies them in bf16 against corresponding weight columns, and executes the remainder in NVFP4:
Preserves input gradients while retaining 4-bit weight memory efficiency.
Single-Model Direct Preference Optimization
Executes DPO without allocating a second full reference model in GPU memory or maintaining four simultaneous logit tensors:
- •
sequence_logprobs()computes per-sequence log-probabilities from hidden states in bounded chunks. - •
reference_logprobs()evaluates reference states using the same paged base with LoRA adapters disabled once prior to training. - • Eliminates duplicate parameter footprint during alignment fine-tuning.
Paged KV Cache for In-Training Generation
Designed for online alignment loops (online DPO, GRPO, PPO) where model generation occurs within the training step:
- •
PagedKVCacheretains resident layers in VRAM up to a specified memory fraction while offloading remaining layers to pinned host pages. - • Offloaded pages stream back into preallocated device slots over TrainStream’s asynchronous H2D stream while preceding layers compute.
- •
TrainStreamCacheintegrates directly into Hugging Facemodel.generate().
Pinned Host Arena & Multi-GPU Peer Mapping
Engineered for maximum PCIe bus utilization and operating system stability:
- •
host_arena="auto"reserves page-locked memory in contiguous slabs at initialization, eliminating per-step page locking stalls on Windows (WDDM) and Linux. - • Unpinned frozen weights are staged asynchronously via multi-core CPU ring buffers.
- •
peer="auto"maps memory across GPUs via CUDA IPC for data-parallel training on single machines without socket loopback overhead.
Usage
Basic Integration
import trainstream
# One-line wrap of PyTorch model and optimizer
model, optimizer = trainstream.wrap(model, optimizer)
# Standard training step remains completely unmodified
for x, y in dataloader:
optimizer.zero_grad()
loss = loss_fn(model(x), y)
loss.backward()
optimizer.step()Quantization & Outlier Options
# Advanced configuration: NVFP4 quantization with outlier channels
model, optimizer = trainstream.wrap(
model,
optimizer,
quantize="nvfp4",
outliers=0.01, # Keeps largest 1% channels in bf16
host_arena="auto", # Pinned host memory slab allocation
peer="auto" # Maps pools across GPUs via CUDA IPC
)Single-Model DPO Workflow
# Single-model Direct Preference Optimization (no second model in VRAM)
from trainstream.integrations.huggingface import reference_logprobs
from trainstream.losses import sequence_logprobs, dpo_loss
# Compute reference log probabilities once with adapters switched off
ref_chosen, ref_rejected = reference_logprobs(model, chosen_ids, rejected_ids)
# Training step computes chunks without four separate logit tensors
pi_chosen = sequence_logprobs(model, chosen_ids)
pi_rejected = sequence_logprobs(model, rejected_ids)
loss = dpo_loss(pi_chosen, pi_rejected, ref_chosen, ref_rejected, beta=0.1)TrainStream: FAQ
What is TrainStream?
TrainStream is an open-source PyTorch library from InnvoSols that lets you train and fine-tune models larger than your GPU's VRAM. It moves optimizer state, gradients, activations and, when necessary, the weights themselves into host RAM and streams them back over PCIe exactly when the GPU needs them.
How do I install and use TrainStream?
Install it from PyPI with pip install trainstream, then wrap your existing model and optimizer in one line: model, optimizer = trainstream.wrap(model, optimizer). The rest of the training loop stays unchanged.
Does TrainStream depend on DeepSpeed or bitsandbytes?
No. TrainStream is a native C++/CUDA extension that uses only the C++ standard library, the CUDA runtime and PyTorch ATen. It has no third-party framework dependencies.
Are TrainStream's weight updates numerically accurate?
Yes. Weight updates match torch.optim.AdamW with a maximum parameter difference below 10⁻⁶ after 20 steps, and activation and weight paging are bit-exact.
Which quantization formats does TrainStream support?
Frozen weights, such as a LoRA base model, can be stored as int8, NF4, FP8 or NVFP4. FP8 and NVFP4 are multiplied directly by modern Tensor Cores, and the optional outliers setting keeps the largest 1% of activation channels in bf16 to reduce NVFP4 error.
Can TrainStream run DPO and online alignment training?
Yes. TrainStream runs Direct Preference Optimization without a second reference model in VRAM, by computing reference log-probabilities from the same base with LoRA adapters switched off. A paged KV cache also supports generation inside the training step for online DPO, GRPO and PPO, and plugs into Hugging Face model.generate().
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.