⚡ StreamTransformer (STR)

📖 What is StreamTransformer?

StreamTransformer (STR) is a universal PyTorch engine that breaks the VRAM Wall by decoupling neural network depth ($L$) from GPU memory capacity.

Instead of allocating memory for all transformer blocks simultaneously in VRAM, StreamTransformer executes temporal layer streaming: exactly one layer block resides in GPU memory at a time, while subsequent layers are prefetched asynchronously over PCIe into page-locked pinned memory.

🌟 Core Capabilities

• $\mathcal{O}(1)$ Depth-Invariant Memory Scaling: Execute 12, 36, 100, or 1,000 layers with a constant peak GPU VRAM footprint (~297 MB).
• 100% Lossless FP32 Precision: Preserves pristine 32-bit floating point weights with exact mathematical output ($\text{Cosine Similarity} = 1.00000012$, $\text{Max Error} = 0.00000000$).
• Dual-Phase Support: Built-in engines for both $\mathcal{O}(1)$ Inference Streaming and $\mathcal{O}(1)$ Layer-Wise Backpropagation Pretraining.
• Universal PyTorch Integration: Compatible with any sequential or decoder-only transformer architecture (LightLLM, GPT-2, LLaMA, Mistral, custom models).

🏎️ Architecture Diagram

┌─────────────────────────────────────────────────────────────────────────────┐
│                    STREAMTRANSFORMER (STR) RUNTIME ENGINE                   │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│  Input Tokens ──→ [ Resident Token Embeddings ] ──→ h₀                      │
│                                                      │                      │
│  [NVMe / Host RAM] ──→ Layer 1/100 ──→ GPU VRAM ──→ h₁ ──→ Reclaim VRAM     │
│  [Async Prefetch]  ──→ Layer 2/100 ──→ GPU VRAM ──→ h₂ ──→ Reclaim VRAM     │
│  ...                                                                        │
│  [Async Prefetch]  ──→ Layer 100/100 ─→ GPU VRAM ─→ h₁₀₀ ──→ Reclaim VRAM   │
│                                                      │                      │
│  Output Logits ←── [ Resident LM Head ] ←── [ Resident Final Norm ]         │
│                                                                             │
│  Peak GPU VRAM: CONSTANT ~297.50 MB across 100 Layers!                      │
└─────────────────────────────────────────────────────────────────────────────┘

📊 Empirical Benchmarks (100 Layers on NVIDIA Tesla T4)

===========================================================================
 100-LAYER TRANSFORMER BENCHMARK (~746M Parameters, FP32)
===========================================================================
• Layer   1/100: Active VRAM = 214.16 MB | Peak VRAM = 297.50 MB
• Layer  20/100: Active VRAM = 214.16 MB | Peak VRAM = 297.50 MB
• Layer  40/100: Active VRAM = 214.16 MB | Peak VRAM = 297.50 MB
• Layer  60/100: Active VRAM = 214.16 MB | Peak VRAM = 297.50 MB
• Layer  80/100: Active VRAM = 214.16 MB | Peak VRAM = 297.50 MB
• Layer 100/100: Active VRAM = 214.16 MB | Peak VRAM = 297.50 MB
---------------------------------------------------------------------------
• Status:           ✅ SUCCESS (All 100 Layers Computed with 0 Errors)
• Peak GPU VRAM:    297.50 MB (Monolithic Expected: ~12,500 MB)
• VRAM Savings:     🔥 97.62% Reduction!
• Numerical Check:  Contains NaN: False | Contains Inf: False
===========================================================================

📦 Installation

git clone https://github.com/RABNEER/stream-transformer.git
cd stream-transformer
pip install -e .

🚀 Quick Start (Inference in 4 Lines)

import torch
from stream_transformer import StreamEngine

# 1. Initialize StreamEngine with your model's layers
engine = StreamEngine(
    resident_modules=resident_dict,
    layer_constructor=lambda: YourTransformerBlock(),
    shard_dir="model_shards",
    num_layers=100,
    device="cuda"
)

# 2. Forward pass with O(1) constant VRAM!
output = engine(input_tensor)

🛠️ Layer-Streaming Training (Pretraining with O(1) VRAM)

from stream_transformer import StreamTrainer

trainer = StreamTrainer(
    resident_modules=resident_dict,
    layer_constructor=lambda: YourTransformerBlock(),
    shard_dir="train_shards",
    num_layers=36,
    device="cuda",
    lr=6e-4
)

# Executes layer-wise forward pass + reverse on-the-fly backward pass
loss = trainer.train_step(x, y, embed_fn, head_fn)
print(f"Training Step Loss: {loss:.4f} (Peak VRAM: < 600 MB!)")

📜 Official Research Paper

For the full theoretical proofs, systems engineering chronicles, and hardware scaling studies, read our research monograph:

📄 Monograph: LightLLM: A Depth-Invariant Layer-Streaming Causal Transformer Architecture for Lossless Full-Precision Pretraining and Inference on Constrained Hardware (Ranveer Kumar, 2026).

@article{kumar2026streamtransformer,
  title={StreamTransformer: A Depth-Invariant Layer-Streaming Architecture for Lossless Full-Precision Neural Execution},
  author={Kumar, Ranveer},
  journal={arXiv preprint},
  year={2026},
  url={https://github.com/RABNEER/stream-transformer}
}

👨‍💻 Author

Ranveer Kumar
Independent AI Researcher
GitHub: @RABNEER | Model: LightLLM