trainstream 0.5.0__tar.gz → 0.6.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {trainstream-0.5.0/trainstream.egg-info → trainstream-0.6.0}/PKG-INFO +170 -53
- {trainstream-0.5.0 → trainstream-0.6.0}/README.md +169 -52
- {trainstream-0.5.0 → trainstream-0.6.0}/setup.py +2 -1
- trainstream-0.6.0/tests/test_arena.py +148 -0
- trainstream-0.6.0/tests/test_dpo.py +150 -0
- trainstream-0.6.0/tests/test_kvcache.py +98 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/tests/test_offload.py +4 -0
- trainstream-0.6.0/tests/test_outliers.py +156 -0
- trainstream-0.6.0/tests/test_peer.py +158 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/tests/test_recompute.py +27 -3
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/__init__.py +16 -2
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/_build.py +1 -1
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/csrc/bindings.cpp +28 -1
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/csrc/kernels.h +9 -2
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/csrc/memory_pinner.cpp +340 -7
- trainstream-0.6.0/trainstream/csrc/peer.cpp +175 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/csrc/quant.cpp +39 -2
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/csrc/quant_kernel.cu +103 -11
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/csrc/trainstream.h +41 -3
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/hooks.py +74 -19
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/integrations/huggingface.py +136 -1
- trainstream-0.6.0/trainstream/kvcache.py +282 -0
- trainstream-0.6.0/trainstream/losses.py +202 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/lowp.py +45 -3
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/paging.py +41 -6
- trainstream-0.6.0/trainstream/peer.py +219 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/quant.py +1 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/wrapper.py +156 -17
- {trainstream-0.5.0 → trainstream-0.6.0/trainstream.egg-info}/PKG-INFO +170 -53
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream.egg-info/SOURCES.txt +8 -0
- trainstream-0.5.0/trainstream/losses.py +0 -103
- {trainstream-0.5.0 → trainstream-0.6.0}/LICENSE +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/MANIFEST.in +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/NOTICE +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/build.bat +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/pyproject.toml +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/setup.cfg +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/tests/test_attention.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/tests/test_distributed.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/tests/test_features.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/tests/test_huggingface.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/tests/test_lowp.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/tests/test_multigpu.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/tests/test_optimizer.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/tests/test_paging.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/tests/test_planner.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/tests/test_quant.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/attention.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/csrc/offload_adamw.cpp +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/csrc/offload_adamw.h +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/csrc/offload_adamw_kernel.cu +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/csrc/shm.cpp +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/integrations/__init__.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/optimizer.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/planner.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/profiler.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream/shared.py +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream.egg-info/dependency_links.txt +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream.egg-info/not-zip-safe +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream.egg-info/requires.txt +0 -0
- {trainstream-0.5.0 → trainstream-0.6.0}/trainstream.egg-info/top_level.txt +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: trainstream
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.6.0
|
|
4
4
|
Summary: Zero-dependency PyTorch training offload: optimizer, gradients, activations and weights to host RAM, with a memory planner and built-in int8/NF4/FP8/NVFP4 quantization.
|
|
5
5
|
Home-page: https://github.com/johnanishg/TrainStream
|
|
6
6
|
Author: John Anish G
|
|
@@ -49,6 +49,31 @@ model, optimizer = trainstream.wrap(model, optimizer)
|
|
|
49
49
|
|
|
50
50
|
TrainStream is a C++/CUDA extension for PyTorch. It uses only the C++ standard library, the CUDA runtime and PyTorch's ATen: no DeepSpeed, no bitsandbytes. Updates match `torch.optim.AdamW` (max parameter difference ≈3·10⁻⁷ after 20 steps), and activation and weight paging are bit-exact.
|
|
51
51
|
|
|
52
|
+
## What's new in 0.6
|
|
53
|
+
|
|
54
|
+
- **Outlier channels in bf16 for NVFP4** (`outliers=`). A few hidden dimensions of an LLM's activations are far larger than the rest. In NVFP4, one of them sets the scale of its 16-value block and rounds the other 15 to zero. With `outliers=0.01`, each matmul picks its largest 1% of input channels on the GPU (no host sync). It quantizes them as zero and multiplies them in bf16 against the same weight columns: `y = scaled_mm(nvfp4(x without them), W) + x[:, idx] @ W[:, idx].T`. The input gradient is unchanged. This removes the activation part of NVFP4's error, not the weight part: the weights stay 4-bit. `benchmarks/quality.py --formats nvfp4-w nvfp4 nvfp4-o` measures the split (`nvfp4-w` = 4-bit weights only). Measured on held-out text (table below): on Qwen2.5-14B NVFP4 costs +5.20 % perplexity, +2.65 % from the 4-bit weights and the rest from the activations; `outliers=0.01` brings it to +4.44 %, but after 20 LoRA steps it is no better (+4.64 % against +4.48 %). On Qwen2.5-1.5B it helps before and after fine-tuning (+10.96 → +9.21 %, +10.02 → +8.54 %). With that mixed result it stays off by default.
|
|
55
|
+
- **DPO without a second model or four logit tensors.**
|
|
56
|
+
- `sequence_logprobs()` gives per-sequence log-probabilities from the hidden states in chunks, and `dpo_loss()` combines them.
|
|
57
|
+
- `reference_logprobs()` (Hugging Face) runs the reference as the same paged or quantized base with the LoRA adapters switched off, once before training.
|
|
58
|
+
- See `examples/dpo_lora.py`.
|
|
59
|
+
- **Paged KV cache for generation during training** (online DPO, GRPO, PPO).
|
|
60
|
+
- `PagedKVCache` keeps as many layers' keys and values in VRAM as a budget allows. The others go to pinned host pages, which are copied back into a device slot on TrainStream's H2D stream while the earlier layers run.
|
|
61
|
+
- `TrainStreamCache` plugs it into `model.generate()`.
|
|
62
|
+
- VRAM holds the resident layers plus two or three one-layer slots.
|
|
63
|
+
- **Pinned host arena** (`host_arena="auto"`).
|
|
64
|
+
- Page-locked memory is reserved in slabs at `wrap()` and carved into blocks, so a training step never page-locks memory, which stalls on Windows (WDDM).
|
|
65
|
+
- Frozen weights that don't fit the page-lock budget no longer go through the driver's synchronous pageable copy. All CPU cores stage them into a ring of pinned slots, which are then copied asynchronously.
|
|
66
|
+
- **Peer GPUs** (`peer="auto"`, data parallel on one machine).
|
|
67
|
+
- When the GPUs can read each other's memory (NVLink, or PCIe P2P where the driver allows it), the frozen blocks every rank streams are split among the GPUs' VRAM. Each rank maps the others' pools through CUDA IPC and copies GPU to GPU instead of from host RAM.
|
|
68
|
+
- Where peer access or CUDA IPC isn't available it stays off, and the events say why. Two GeForce cards in one machine usually report no P2P.
|
|
69
|
+
- Tested with two processes sharing one GPU on Windows 11 and WSL2: the pools mapped through CUDA IPC, and training matched host streaming bit for bit. Copies between two real GPUs are untested.
|
|
70
|
+
- **Activation tuning, corrected and tighter.**
|
|
71
|
+
- Activations that fit in VRAM now stay there. Before, `"auto"` still paged some to free VRAM that nothing used: the 1.5B full fine-tune spent 12 slower tuning steps, then settled on 4 paged blocks at 1,078 ms/step against 1,020 ms with none.
|
|
72
|
+
- When host RAM can't hold every block's activations, the trials now apply the setting they measured fastest. Before, the forced recompute shifted it: 14B NF4 at 2,048 tokens ran 44 recomputed + 3 paged blocks (never measured) instead of the measured 47 recomputed; 14B int8 at 2,048 ran 47 recomputed (3.25 s) instead of the measured 36 (3.17 s).
|
|
73
|
+
- Once the trials settle, the VRAM left under the target keeps the latest blocks' activations instead of recomputing or paging them, measured on a clean step.
|
|
74
|
+
- Measured with the same harness in WSL2, before → after: 1.5B full fine-tune 1.03 → 0.99 s/step; 14B LoRA at 512 tokens with `"auto"` 0.90 → 0.84 s; at 2,048 tokens NVFP4 2.14 → 2.06 s, NF4 3.18 → 3.02 s, `"auto"` 3.28 → 3.20 s. Single runs: differences under about 5 % are within run-to-run noise. The faster 2,048-token runs use more of the VRAM under the target (NF4: 17.8 → 20.4 GB peak).
|
|
75
|
+
- **Benchmarks on both systems, every format.** The comparison below now covers WSL2 and Windows 11, with every engine run to the end, and the perplexity table covers every format on Qwen2.5-14B and 1.5B, before and after fine-tuning. `benchmarks/quality.py --format F` runs one format in its own process, and `--live` streams its progress as JSON lines.
|
|
76
|
+
|
|
52
77
|
## What's new in 0.5
|
|
53
78
|
|
|
54
79
|
- **FP8 and NVFP4 on tensor cores.** `quantize="fp8"` or `"nvfp4"` store frozen weights block-scaled. On Blackwell GPUs (RTX 50-series, B200) the forward matmuls run on 8- and 4-bit tensor cores through PyTorch's own `scaled_mm`: 0.188 ms for a 512×5120×13824 FP4 matmul against 1.215 ms in bf16. Real Qwen2.5-14B LoRA (Windows, same session): 0.586 s/step with NVFP4 against 0.642 s with NF4 at the same memory. NVFP4 also quantizes activations, so it costs more quality (perplexity +5.2 % on 14B, +11 % on 1.5B, against +3.1 % / +5.9 % for NF4). It stays opt-in: `prefer="speed"` still picks NF4.
|
|
@@ -72,52 +97,80 @@ TrainStream is a C++/CUDA extension for PyTorch. It uses only the C++ standard l
|
|
|
72
97
|
- **0.4.1:** install instructions for uv; clearer messages when the installed torch is the CPU-only build.
|
|
73
98
|
- **0.4.0:** built-in int8 / NF4 weight quantization (`quantize=`), `load_model()` (quantize while reading a checkpoint), `quantize_model()`. Real Qwen2.5-14B LoRA 1.37 → 0.87 s/step; a 32B-shaped base trains on 64 GB of RAM.
|
|
74
99
|
|
|
75
|
-
## Against other libraries (same laptop
|
|
100
|
+
## Against other libraries (same laptop: WSL2 with 31 GB RAM, then Windows 11 with 64 GB)
|
|
76
101
|
|
|
77
|
-
>
|
|
102
|
+
> TrainStream 0.6.0, one run per engine, 10 steps each (TrainStream: 10 steps after it finished tuning itself), median of the steps after the warm-up. Every engine ran all its steps; none was stopped at a time limit. Some engines were still getting a few percent faster at step 10, so read differences under about 5 % as noise.
|
|
78
103
|
|
|
79
|
-
**LoRA r=16, Qwen2.5-14B, batch 1 × 512**
|
|
104
|
+
**LoRA r=16, Qwen2.5-14B, batch 1 × 512 (WSL2)**
|
|
80
105
|
|
|
81
106
|
| Engine | Step time | Peak VRAM | Outcome |
|
|
82
107
|
|---|---|---|---|
|
|
83
|
-
| TrainStream, NVFP4 base (FP4 tensor cores) | 0.
|
|
84
|
-
| bitsandbytes QLoRA (NF4) + PEFT | 0.
|
|
85
|
-
| TrainStream, NF4 base | 0.
|
|
86
|
-
| Unsloth (4-bit) | 0.
|
|
87
|
-
| TrainStream (auto) | 0.
|
|
88
|
-
| bitsandbytes QLoRA + gradient checkpointing | 0.
|
|
89
|
-
| DeepSpeed ZeRO-2, optimizer offload |
|
|
90
|
-
| Stock PyTorch (+ PEFT) |
|
|
108
|
+
| TrainStream, NVFP4 base (FP4 tensor cores) | 0.56 s | 19.1 GB | fastest |
|
|
109
|
+
| bitsandbytes QLoRA (NF4) + PEFT | 0.63 s | 17.4 GB | 1.13× the fastest |
|
|
110
|
+
| TrainStream, NF4 base | 0.67 s | 19.1 GB | 1.20× the fastest |
|
|
111
|
+
| Unsloth (4-bit) | 0.75 s | 11.5 GB | 1.34× the fastest |
|
|
112
|
+
| TrainStream (auto) | 0.84 s | 20.6 GB | 1.51× the fastest |
|
|
113
|
+
| bitsandbytes QLoRA + gradient checkpointing | 0.90 s | 12.7 GB | 1.61× the fastest |
|
|
114
|
+
| DeepSpeed ZeRO-2, optimizer offload | 42.38 s | 35.9 GB | 75.93× the fastest (spills past VRAM) |
|
|
115
|
+
| Stock PyTorch (+ PEFT) | 50.63 s | 37.6 GB | 90.70× the fastest (spills past VRAM) |
|
|
91
116
|
|
|
92
|
-
**LoRA r=16, Qwen2.5-14B, batch 1 × 2048**
|
|
117
|
+
**LoRA r=16, Qwen2.5-14B, batch 1 × 2048 (WSL2)**
|
|
93
118
|
|
|
94
119
|
| Engine | Step time | Peak VRAM | Outcome |
|
|
95
120
|
|---|---|---|---|
|
|
96
|
-
| TrainStream, NVFP4 base (FP4 tensor cores) | 2.
|
|
97
|
-
| Unsloth (4-bit) | 2.55 s | 11.7 GB | 1.
|
|
98
|
-
| bitsandbytes QLoRA + gradient checkpointing | 2.89 s | 16.3 GB | 1.
|
|
99
|
-
| TrainStream, NF4 base | 3.
|
|
100
|
-
| TrainStream (auto) | 3.
|
|
101
|
-
| bitsandbytes QLoRA (NF4) + PEFT |
|
|
102
|
-
|
|
|
103
|
-
|
|
|
121
|
+
| TrainStream, NVFP4 base (FP4 tensor cores) | 2.06 s | 19.8 GB | fastest |
|
|
122
|
+
| Unsloth (4-bit) | 2.55 s | 11.7 GB | 1.24× the fastest |
|
|
123
|
+
| bitsandbytes QLoRA + gradient checkpointing | 2.89 s | 16.3 GB | 1.40× the fastest |
|
|
124
|
+
| TrainStream, NF4 base | 3.02 s | 20.4 GB | 1.47× the fastest |
|
|
125
|
+
| TrainStream (auto) | 3.20 s | 22.0 GB | 1.55× the fastest |
|
|
126
|
+
| bitsandbytes QLoRA (NF4) + PEFT | 158.36 s | 34.7 GB | 76.84× the fastest (spills past VRAM) |
|
|
127
|
+
| DeepSpeed ZeRO-2, optimizer offload | 224.17 s | 53.0 GB | 108.77× the fastest (spills past VRAM) |
|
|
128
|
+
| Stock PyTorch (+ PEFT) | 424.05 s | 57.9 GB | 205.75× the fastest (spills past VRAM) |
|
|
104
129
|
|
|
105
|
-
**Full fine-tune, Qwen2.5-1.5B, batch 4 × 512**
|
|
130
|
+
**Full fine-tune, Qwen2.5-1.5B, batch 4 × 512 (WSL2)**
|
|
106
131
|
|
|
107
132
|
| Engine | Step time | Peak VRAM | Outcome |
|
|
108
133
|
|---|---|---|---|
|
|
109
|
-
|
|
|
110
|
-
|
|
|
134
|
+
| TrainStream (auto) | 0.99 s | 20.0 GB | fastest |
|
|
135
|
+
| Stock PyTorch (+ PEFT) | 1.93 s | 32.4 GB | 1.95× the fastest (spills past VRAM) |
|
|
111
136
|
| DeepSpeed ZeRO-2, optimizer offload | — | — | killed: out of host RAM (31 GB in WSL) |
|
|
112
|
-
- With an NVFP4 base, TrainStream had the fastest step in both 14B LoRA runs: 0.52 s at 512 tokens and 2.01 s at 2,048. It holds more in VRAM than Unsloth (18–19 GB against 11.5), and NVFP4 costs quality (see the perplexity table).
|
|
113
|
-
- With an NF4 base, the format bitsandbytes and Unsloth use, TrainStream was 7 % slower than bitsandbytes at 512 tokens, and behind Unsloth and bitsandbytes with checkpointing at 2,048.
|
|
114
|
-
- Stock PyTorch, bitsandbytes without checkpointing at 2,048 tokens, and DeepSpeed did not fit in 24 GB of VRAM. The driver spilled into system RAM (75–280 s per step), or the run did not finish or was killed.
|
|
115
|
-
- The 1.5B full fine-tune ran with 29 of WSL's 31 GB of RAM in use, and TrainStream's step time varied between runs (1.80, 1.09 and 1.61 s; the table shows the latest). On Windows with 64 GB it holds 1.18 s. Stock PyTorch (1.33 s) allocated 32 GB on the 24 GB GPU, the excess in shared system memory.
|
|
116
137
|
|
|
117
|
-
|
|
138
|
+
**The same on Windows 11 (64 GB RAM; Unsloth and DeepSpeed aren't installed there)**
|
|
139
|
+
|
|
140
|
+
| Engine | LoRA 14B, 1 × 512 | LoRA 14B, 1 × 2048 | Full fine-tune 1.5B, 4 × 512 |
|
|
141
|
+
|---|---|---|---|
|
|
142
|
+
| TrainStream, NVFP4 base | **0.60 s** | **1.92 s** | — |
|
|
143
|
+
| TrainStream, NF4 base | 0.68 s | 2.57 s | — |
|
|
144
|
+
| TrainStream (auto) | 1.06 s | 3.25 s | **1.06 s** |
|
|
145
|
+
| bitsandbytes QLoRA (NF4) + PEFT | 0.65 s | 162.72 s (spills) | — |
|
|
146
|
+
| bitsandbytes QLoRA + gradient checkpointing | 0.95 s | 3.16 s | — |
|
|
147
|
+
| Stock PyTorch (+ PEFT) | 53.11 s (spills) | 267.78 s (spills) | 1.67 s (spills) |
|
|
148
|
+
|
|
149
|
+
- With an NVFP4 base, TrainStream had the fastest step in every 14B LoRA run, in WSL and on Windows: 0.56 / 0.60 s at 512 tokens, 2.06 / 1.92 s at 2,048. It holds more in VRAM than Unsloth (19–20 GB against 11.5–11.7), and NVFP4 costs quality (see the perplexity table below).
|
|
150
|
+
- With an NF4 base, the format bitsandbytes and Unsloth use, TrainStream was 7 % slower than bitsandbytes at 512 tokens in WSL (5 % on Windows). At 2,048 tokens it was behind Unsloth and bitsandbytes with checkpointing in WSL, where host RAM held only 25 of 47 layers' activations and it recomputed the rest; on Windows host RAM held 43, it paged 35, and it was ahead of bitsandbytes with checkpointing (2.57 s against 3.16 s).
|
|
151
|
+
- `"auto"` picks differently per machine: in WSL (31 GB of RAM, too little for the bf16 base) it quantized the 14B base to int8 and kept every layer in VRAM (0.84 s at 512 tokens); on Windows (64 GB) it stayed exact in bf16 and streamed 21 of 48 layers from RAM (1.06 s).
|
|
152
|
+
- The full fine-tune: TrainStream keeps its activations in VRAM and stays at 20.0 GB, 1.95× faster than stock PyTorch in WSL and 1.58× on Windows. Stock PyTorch peaks at 32.4 GB, more than the GPU's 24 GB, so its step time depends on how the driver spills into system RAM (1.33–1.94 s across runs).
|
|
153
|
+
- Stock PyTorch, bitsandbytes without checkpointing at 2,048 tokens, and DeepSpeed did not fit in 24 GB of VRAM: the driver spilled into system RAM (42–424 s per step), or the run was killed.
|
|
154
|
+
|
|
155
|
+
**Perplexity against bf16** (`benchmarks/quality.py`, held-out Python standard-library source; lower is better)
|
|
156
|
+
|
|
157
|
+
| Format | Qwen2.5-14B | after 20 LoRA steps | Qwen2.5-1.5B | after 40 LoRA steps |
|
|
158
|
+
|---|---|---|---|---|
|
|
159
|
+
| int8 | +0.02 % | +0.05 % | +0.04 % | −0.01 % |
|
|
160
|
+
| FP8 (tensor cores) | +0.84 % | +0.79 % | +1.78 % | +1.43 % |
|
|
161
|
+
| NF4 | +3.09 % | +2.85 % | +5.91 % | +5.00 % |
|
|
162
|
+
| NVFP4, weights only (bf16 matmuls) | +2.65 % | +2.26 % | +5.43 % | +5.04 % |
|
|
163
|
+
| NVFP4 (tensor cores) | +5.20 % | +4.48 % | +10.96 % | +10.02 % |
|
|
164
|
+
| NVFP4 with `outliers=0.01` | +4.44 % | +4.64 % | +9.21 % | +8.54 % |
|
|
165
|
+
|
|
166
|
+
bf16 itself: 2.300 (14B) and 3.251 (1.5B) on 16 and 64 windows of 1,024 tokens. NVFP4's 4-bit weights alone cost less than NF4; its extra loss comes from quantizing the activations. On 1.5B the short fine-tune raised bf16's own held-out perplexity (3.251 → 3.306), so compare the formats with each other there.
|
|
167
|
+
|
|
168
|
+
Each engine runs in a fresh process with the same model, random-token data, LoRA settings (r=16, all projections), learning rate and loop. Step time is the median of the steps after the warm-up: 2 steps, plus, for TrainStream, the steps until it has finished tuning itself. Unsloth uses its own defaults. Re-run with `benchmarks/compare.py` and `benchmarks/quality.py`.
|
|
118
169
|
|
|
119
170
|
## Results (RTX 5090 Laptop GPU 24 GB, 64 GB RAM)
|
|
120
171
|
|
|
172
|
+
Measured with TrainStream 0.1–0.4, when each feature arrived: these rows show what fits and how, not today's speed (the comparison above is 0.6.0).
|
|
173
|
+
|
|
121
174
|
| Workload | Stock PyTorch | TrainStream |
|
|
122
175
|
|---|---|---|
|
|
123
176
|
| 1.8B-parameter GPT, fp32, batch 8×512 | 34.1 GiB peak: exceeds VRAM, driver spills to RAM, 13.7 s/step | 20.6 GiB, 1.1 s/step (13.5 GiB with all activations paged) |
|
|
@@ -138,7 +191,7 @@ Offloading costs throughput when a model already fits, which is why `offload="au
|
|
|
138
191
|
- 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)
|
|
139
192
|
- Optional, for the Hugging Face integration: `transformers`, `accelerate`, and `peft` for LoRA (all three come with `trainstream[hf]`)
|
|
140
193
|
|
|
141
|
-
Tested on an RTX 5090 Laptop GPU with Windows 11 (PyTorch 2.14.1, CUDA 13.4, MSVC, Python 3.14), where all
|
|
194
|
+
Tested on an RTX 5090 Laptop GPU with Windows 11 (PyTorch 2.14.1, CUDA 13.4, MSVC, Python 3.14), where all 223 tests pass with no warnings, and Ubuntu 24.04 under WSL2 (PyTorch 2.14.1, CUDA 13.3, GCC 13.3, Python 3.12), where all 223 pass as well (with bitsandbytes 0.50.2 installed for the NF4 comparison), both from a wheel installed with pip and with uv. Installs were checked with pip and with uv 0.13 (`uv pip install`, `uv sync` in a project, and building from the sdist) on both. Native Linux and other GPUs are untested.
|
|
142
195
|
|
|
143
196
|
## Install
|
|
144
197
|
|
|
@@ -255,6 +308,34 @@ torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
|
|
|
255
308
|
model, optimizer = trainstream.wrap(MyModel().cuda(), torch.optim.AdamW(...)) # detected automatically
|
|
256
309
|
```
|
|
257
310
|
|
|
311
|
+
### DPO with LoRA
|
|
312
|
+
|
|
313
|
+
```python
|
|
314
|
+
import trainstream
|
|
315
|
+
from trainstream.integrations.huggingface import causal_lm_logprobs, load_model, reference_logprobs, wrap_hf
|
|
316
|
+
|
|
317
|
+
model = peft.get_peft_model(load_model("Qwen/Qwen2.5-7B", quantize="nf4"), lora_config)
|
|
318
|
+
model, optimizer = wrap_hf(model, torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=5e-6))
|
|
319
|
+
pairs = [(chosen, rejected), ...] # dicts with input_ids and labels (-100 on the prompt), batch_size rows each
|
|
320
|
+
# adapters off: the same base is the reference; computed once, up front
|
|
321
|
+
ref = reference_logprobs(model, [b for pair in pairs for b in pair]).split(batch_size)
|
|
322
|
+
for k, (chosen, rejected) in enumerate(pairs):
|
|
323
|
+
pc = causal_lm_logprobs(model, chosen["input_ids"], chosen["labels"])
|
|
324
|
+
pr = causal_lm_logprobs(model, rejected["input_ids"], rejected["labels"])
|
|
325
|
+
loss, chosen_reward, rejected_reward = trainstream.dpo_loss(pc, pr, ref[2 * k], ref[2 * k + 1], beta=0.1)
|
|
326
|
+
loss.backward(); optimizer.step(); optimizer.zero_grad()
|
|
327
|
+
```
|
|
328
|
+
|
|
329
|
+
### Generating with a paged KV cache
|
|
330
|
+
|
|
331
|
+
```python
|
|
332
|
+
from trainstream.integrations.huggingface import TrainStreamCache
|
|
333
|
+
|
|
334
|
+
cache = TrainStreamCache(model, max_tokens=prompt_len + 512) # resident layers chosen from free VRAM
|
|
335
|
+
out = model.generate(**inputs, max_new_tokens=512, do_sample=True, past_key_values=cache)
|
|
336
|
+
cache.reset() # before the next batch of prompts
|
|
337
|
+
```
|
|
338
|
+
|
|
258
339
|
### Inspecting what it decided
|
|
259
340
|
|
|
260
341
|
```python
|
|
@@ -274,6 +355,7 @@ model._trainstream_session.adapting # True while it is still tuning itself
|
|
|
274
355
|
- Use `max_grad_norm=` instead of `torch.nn.utils.clip_grad_norm_`; gradients have already left the GPU.
|
|
275
356
|
- Declare gradient accumulation with `accumulation_steps=`.
|
|
276
357
|
- `model.state_dict()` and `optimizer.state_dict()` work while wrapped. `trainstream.unwrap(model)` loads every weight back into VRAM and removes the hooks.
|
|
358
|
+
- Under WSL2 the Linux VM gets only part of the machine's RAM (about half by default), and TrainStream plans from what the VM has. If host RAM is the limit (the events say "host RAM holds N of M blocks' activations"), raise `memory=` in `%UserProfile%\.wslconfig`: in the benchmarks below, NF4 at 2,048 tokens ran 3.02 s/step in WSL with 31 GB against 2.57 s on Windows with 64 GB.
|
|
277
359
|
|
|
278
360
|
## Options
|
|
279
361
|
|
|
@@ -284,13 +366,16 @@ model._trainstream_session.adapting # True while it is still tuning itself
|
|
|
284
366
|
| `optimizer` | `torch.optim.AdamW`, `torch.optim.Adam`, `trainstream.OffloadAdamW`, `None` | `None` | converted, with hyper-parameters and state migrated; `None` creates one with `lr` |
|
|
285
367
|
| `optimizer_mode` | `"cpu"`, `"gpu"`, `"auto"` | `"cpu"` | host AdamW, a zero-copy CUDA kernel, or chosen from the hardware profile |
|
|
286
368
|
| `offload` | `"auto"`, `True` | `"auto"` | `"auto"` switches to fused torch AdamW after 2 steps if the model fits |
|
|
287
|
-
| `activation_offload` | `"auto"`, `True`, `False`, fraction 0.0–1.0 | `"auto"` | how many blocks' activations are paged to RAM |
|
|
369
|
+
| `activation_offload` | `"auto"`, `True`, `False`, fraction 0.0–1.0 | `"auto"` | how many blocks' activations are paged to RAM; `"auto"`: measured in the first step, kept in VRAM if they fit under the VRAM target, otherwise paged as far as host RAM holds them (see `recompute` for the rest); `True`: all but the last block |
|
|
288
370
|
| `param_offload` | `"auto"`, `True`, `False`, fraction 0.0–1.0, list of block indices | `"auto"` | `"auto"`: the memory planner decides which layers stay in VRAM |
|
|
289
371
|
| `quantize` | `"auto"`, `"none"`, `None`, `False`, `"int8"`, `"nf4"`, `"fp8"`, `"nvfp4"` | `"auto"` | frozen weights only; `"auto"` quantizes only when the run can't fit otherwise (int8 before NF4); fp8 / nvfp4 run on Blackwell tensor cores |
|
|
290
372
|
| `quant_block_size` | even integer ≥ 2 | `64` | values per fp32 scale (int8 / nf4; fp8 and nvfp4 use fixed 32 / 16) |
|
|
291
373
|
| `prefer` | `"exact"`, `"speed"` | `"exact"` | `"speed"` lets `quantize="auto"` choose NF4 when it keeps every layer in VRAM |
|
|
292
|
-
| `recompute` | `"auto"`, `True`, `False`, fraction 0.0–1.0, list of block indices | `"auto"` | recompute blocks in backward instead of keeping or paging their activations; `"auto"`:
|
|
374
|
+
| `recompute` | `"auto"`, `True`, `False`, fraction 0.0–1.0, list of block indices | `"auto"` | recompute blocks in backward instead of keeping or paging their activations; `"auto"`: when the activations can't all stay in VRAM (weights paged, or too large), as many as measured trials show is fastest, then blocks go back to VRAM as far as the room left under the target allows |
|
|
293
375
|
| `attention` | `"auto"`, `"cudnn"`, `"efficient"`, `"flash"`, `"math"`, `None` | `"auto"` | attention backend: the fastest measured on your model's shape, a named one first, or PyTorch's own order; deterministic mode leaves out the non-deterministic kernels |
|
|
376
|
+
| `outliers` | `0`, `False`, `None`, `"auto"`, `True`, float 0.0–0.1 | `0` | with `quantize="nvfp4"` on tensor cores: the fraction of input channels (the largest of each call) multiplied in bf16 instead of 4-bit; `"auto"` / `True` mean `0.01`; ignored by other formats |
|
|
377
|
+
| `host_arena` | `"auto"`, `False`, `None`, bytes | `"auto"` | page-locked host memory reserved up front, from which paged activations, staging and other pinned buffers are carved; `"auto"`: when the run streams anything, reserve one slab now and grow on demand within the page-lock budget; bytes: reserve that much now and cap growth there |
|
|
378
|
+
| `peer` | `"auto"`, `False`, `None`, bytes | `"auto"` | data parallel on one machine: pool the frozen blocks every rank streams in the GPUs' VRAM and copy them GPU to GPU (needs peer access between the GPUs and CUDA IPC; set up collectively at step 10); `"auto"`: when peer copies reach 0.9× the host link, each GPU gives half the headroom under the memory plan's target; bytes: that much per GPU. While active, `unwrap()` is collective |
|
|
294
379
|
| `moment_dtype` | `"auto"`, `"float32"`, `"bfloat16"`, `"bf16"`, `torch.bfloat16` | `"auto"` | fp32 unless that wouldn't fit in RAM; bf16 uses stochastic rounding |
|
|
295
380
|
| `device` | `None`, `"cuda"`, `"cuda:N"`, `torch.device` | `None` | target GPU when the model is on the CPU |
|
|
296
381
|
| `blocks` | `None`, list of layer modules | `None` | the repeated layer stack; auto-detected |
|
|
@@ -301,7 +386,7 @@ model._trainstream_session.adapting # True while it is still tuning itself
|
|
|
301
386
|
| `vram_headroom` | float 0.0–1.0 | `0.10` | share of usable VRAM the planner leaves free |
|
|
302
387
|
| `check_memory` | `True`, `False` | `True` | raise `MemoryError` at `wrap()` when the run can't fit |
|
|
303
388
|
| `max_inflight_bytes` | `None`, bytes | `None` (1 GiB or 10% of VRAM) | VRAM held by transfers still crossing PCIe |
|
|
304
|
-
| `max_activation_slowdown` | float | `0.15` | how much slower `"auto"` activation paging may make a step |
|
|
389
|
+
| `max_activation_slowdown` | float | `0.15` | how much slower `"auto"` activation paging may make a step when there is no VRAM target to measure against (`ActivationOffloader` used on its own); `wrap()` always has one, so activations that fit stay in VRAM instead |
|
|
305
390
|
| `prefetch_depth` | integer ≥ 1 | `1` | activation blocks fetched ahead in backward |
|
|
306
391
|
| `min_offload_bytes` | bytes | `1048576` | smaller activations stay on the GPU |
|
|
307
392
|
| `threads` | `None`, integer | `None` (all cores) | host AdamW threads |
|
|
@@ -313,13 +398,13 @@ Anything else raises `ValueError` (for example `quantize="fp4"`).
|
|
|
313
398
|
|
|
314
399
|
### Quantization formats
|
|
315
400
|
|
|
316
|
-
| `quantize=` | Bytes per weight |
|
|
401
|
+
| `quantize=` | Bytes per weight | Quality cost (perplexity against bf16: Qwen2.5-14B / 1.5B) | Use it when |
|
|
317
402
|
|---|---|---|---|
|
|
318
|
-
| `"none"` | 2 (bf16) |
|
|
319
|
-
| `"int8"` | 1.06 | 0.6 % | bf16 doesn't fit in RAM |
|
|
320
|
-
| `"nf4"` | 0.56 | 9 % | int8 doesn't fit either, or with `prefer="speed"` |
|
|
321
|
-
| `"fp8"` | 1.03 |
|
|
322
|
-
| `"nvfp4"` | 0.56 |
|
|
403
|
+
| `"none"` | 2 (bf16) | none (bit-exact) | it fits |
|
|
404
|
+
| `"int8"` | 1.06 | +0.02 % / +0.04 % (mean weight error 0.6 %) | bf16 doesn't fit in RAM |
|
|
405
|
+
| `"nf4"` | 0.56 | +3.1 % / +5.9 % (mean weight error 9 %) | int8 doesn't fit either, or with `prefer="speed"` |
|
|
406
|
+
| `"fp8"` | 1.03 | +0.8 % / +1.8 % | FP8 tensor-core matmuls on Blackwell, near-bf16 quality |
|
|
407
|
+
| `"nvfp4"` | 0.56 | +5.2 % / +11.0 % (+4.4 % / +9.2 % with `outliers=0.01`) | the fastest 4-bit path on Blackwell, when its quality is acceptable |
|
|
323
408
|
| `"auto"` | — | — | decided from the measured RAM and VRAM |
|
|
324
409
|
|
|
325
410
|
NF4 is bit-identical to bitsandbytes' NF4 (same codes, packing and scales); TrainStream doesn't depend on bitsandbytes. A quantized run trains on the dequantized weights, as QLoRA does, so it isn't bit-exact with the bf16 base.
|
|
@@ -378,6 +463,9 @@ Methods and attributes: `step(closure=None)`, `zero_grad()`, `attach(model, guar
|
|
|
378
463
|
| `wrap_hf(model, optimizer=None, **options)` | `options`: any `wrap()` option | `wrap()` with `blocks` preset to the decoder layers |
|
|
379
464
|
| `load_model(name_or_path, ...)` | see the table above | streams a checkpoint, quantizing while it reads |
|
|
380
465
|
| `TrainStreamCallback(optimizer)` | `optimizer`: an `OffloadAdamW` | for a stock `Trainer` that you gave an `OffloadAdamW`: prints the TrainStream metrics at each log step |
|
|
466
|
+
| `causal_lm_logprobs(model, input_ids, labels=None, attention_mask=None, *, chunk_size=1024, average=False, **model_kwargs)` | `labels`: as for the model's loss (the input ids; `-100` marks tokens that don't count, e.g. the prompt); default `input_ids`. `chunk_size`: tokens per chunk. `average`: mean per counted token instead of the sum. `model_kwargs`: passed to the model | the log-probability of each sequence, shape `[batch]`, differentiable; computed from the final hidden states in chunks (models with soft-capped or scaled logits use their own logits) |
|
|
467
|
+
| `reference_logprobs(model, batches, *, chunk_size=1024, average=False)` | `batches`: iterable of dicts with `"input_ids"` and optionally `"labels"` and `"attention_mask"` | DPO's reference log-probs without a second model: a PEFT model runs with its adapters off, in eval mode and without gradients; one float32 CPU tensor with every sequence's value, in order |
|
|
468
|
+
| `TrainStreamCache(model, max_tokens, *, resident_layers="auto", page_tokens=256, vram_fraction=0.5)` | `model`: a decoder with full attention in every layer (sliding-window models are refused). Other arguments as `PagedKVCache` below | a `transformers` `Cache` for `model.generate(past_key_values=...)` whose keys and values are partly in host RAM; greedy or sampling (no beam search); `.reset()` before the next prompts; `.kv` is the `PagedKVCache`. Needs `transformers` ≥ 4.56 |
|
|
381
469
|
|
|
382
470
|
How `TrainStreamTrainer` uses `TrainingArguments`:
|
|
383
471
|
- `optim` must be an Adam/AdamW variant.
|
|
@@ -398,23 +486,47 @@ Its log lines add these metrics:
|
|
|
398
486
|
|---|---|
|
|
399
487
|
| `trainstream.chunked_cross_entropy(hidden, weight, targets, bias=None, *, chunk_size=1024, ignore_index=-100, reduction="mean", num_items=None)` | cross-entropy of `hidden @ weight.T (+ bias)` without building the logits. `hidden` `[..., D]` (after the last norm), `weight` `[V, D]`, `bias` `[V]` or `None`, `targets` class indices. `chunk_size`: tokens per chunk. `ignore_index`: positions that don't count. `reduction`: `"mean"` or `"sum"`. `num_items`: divide by this instead of the counted positions (as Hugging Face does with gradient accumulation) |
|
|
400
488
|
| `trainstream.causal_lm_loss(hidden, weight, labels, bias=None, **kwargs)` | the next-token loss as Hugging Face computes it (labels shifted by one); `kwargs` as above |
|
|
489
|
+
| `trainstream.sequence_logprobs(hidden, weight, labels, bias=None, *, chunk_size=1024, ignore_index=-100, shift=True, average=False)` | the log-probability of each sequence's labels, `[B]` float32, without building the logits; differentiable. `hidden` `[B, T, D]`, `labels` `[B, T]`. `shift=True`: labels are the input ids (position t predicts label t + 1), as Hugging Face passes them; `False`: position t predicts label t. `average`: mean per counted token instead of the sum |
|
|
490
|
+
| `trainstream.dpo_loss(policy_chosen, policy_rejected, reference_chosen, reference_rejected, *, beta=0.1, label_smoothing=0.0)` | Direct Preference Optimization from per-sequence log-probs (`[B]` each). Returns `(loss, chosen_rewards, rejected_rewards)`: the mean loss and the implicit rewards `beta * (policy - reference)` (detached). `beta`: strength of the pull toward the reference. `label_smoothing` > 0: the conservative variant for noisy preference labels |
|
|
491
|
+
|
|
492
|
+
### `trainstream.PagedKVCache(num_layers, max_tokens, *, resident_layers="auto", page_tokens=256, vram_fraction=0.5)`
|
|
493
|
+
|
|
494
|
+
A KV cache for autoregressive generation whose offloaded layers keep their keys and values in pinned host pages. Inference only. The Hugging Face adapter is `TrainStreamCache` above.
|
|
495
|
+
|
|
496
|
+
| Argument | Accepted values | Default | What it does |
|
|
497
|
+
|---|---|---|---|
|
|
498
|
+
| `num_layers` | integer ≥ 1 | — | attention layers of the model |
|
|
499
|
+
| `max_tokens` | integer ≥ 1 | — | the longest sequence expected (prompt + new tokens); sizes the slots and the `"auto"` budget. Longer sequences still work, with a reallocation |
|
|
500
|
+
| `resident_layers` | `"auto"`, integer ≥ 0, list of layer indices | `"auto"` | layers whose keys and values stay in VRAM; `"auto"`: as many as fit in `vram_fraction` of the free VRAM at the first update; the offloaded ones are spread through the stack |
|
|
501
|
+
| `page_tokens` | integer ≥ 1 | `256` | tokens per host page (one D2H copy when a page fills) |
|
|
502
|
+
| `vram_fraction` | float in (0, 1] | `0.5` | share of free (plus cached but unused) VRAM the `"auto"` plan may fill |
|
|
503
|
+
|
|
504
|
+
Methods and attributes:
|
|
505
|
+
- `update(layer, k, v) → (keys, values)`: append `[batch, kv_heads, new_tokens, head_dim]` and get the layer's full keys and values, as strided views.
|
|
506
|
+
- `seq_length(layer=0)`: tokens cached in that layer.
|
|
507
|
+
- `reset()`: forget every token.
|
|
508
|
+
- `memory()`: `{"device_bytes", "host_bytes"}`.
|
|
509
|
+
- `.resident` and `.offloaded`: layer lists.
|
|
510
|
+
- `.stats`: `h2d_bytes`, `d2h_bytes`, `pages`, `prefetched` (slots filled ahead of time), `gathered` (filled on demand).
|
|
401
511
|
|
|
402
512
|
### Inspection and utilities
|
|
403
513
|
|
|
404
514
|
| Function | Arguments | Returns |
|
|
405
515
|
|---|---|---|
|
|
406
|
-
| `trainstream.memory_report(model=None)` | a wrapped model, or `None` | for a wrapped model, a dict with these keys: `"events"` (the decision log), `"gpu"` (`allocated`, `reserved`, `peak_allocated`), `"pinned"`, `"pinned_by_tag"`, `"transfers"`, `"optimizer"`, `"resident"`, and when used, `"attention"` (`mode`, `backend`, `timings_ms`), `"activations"` (`policy`, `offloaded_blocks`, `recompute_blocks`, `decision`, ...), `"params"` (`paged_blocks`, `quantize`, `quantized_blocks`, ...) and `"plan"` (`target_vram`, `settled`, `history`, `budget`). Without a model: `"gpu"`, `"pinned"`, `"pinned_by_tag"` |
|
|
516
|
+
| `trainstream.memory_report(model=None)` | a wrapped model, or `None` | for a wrapped model, a dict with these keys: `"events"` (the decision log), `"gpu"` (`allocated`, `reserved`, `peak_allocated`), `"pinned"`, `"pinned_by_tag"`, `"transfers"`, `"optimizer"`, `"resident"`, and when used, `"attention"` (`mode`, `backend`, `timings_ms`), `"activations"` (`policy`, `offloaded_blocks`, `recompute_blocks`, `decision`, ...), `"params"` (`paged_blocks`, `quantize`, `quantized_blocks`, `outliers`, `load_bytes` from host RAM, `bounce_bytes` staged from pageable RAM, `peer_bytes` copied from other GPUs, `peer_blocks`, `peer_owned_bytes`, ...) and `"plan"` (`target_vram`, `settled`, `history`, `budget`). Without a model: `"gpu"`, `"pinned"`, `"pinned_by_tag"` |
|
|
407
517
|
| `trainstream.unwrap(model)` | a wrapped model | loads every weight back into VRAM, removes the hooks, returns the model |
|
|
408
518
|
| `trainstream.profile_hardware(device=None, quick=True, force=False, threads=None)` | `device`: GPU to measure (default the current one). `quick`: shorter benchmarks. `force`: re-measure instead of using the cache. `threads`: host threads for the AdamW benchmark | a `HardwareProfile` (PCIe bandwidth, matmul and host AdamW throughput; `.summary()` prints it) |
|
|
409
|
-
| `trainstream.pinned_stats()` | — | page-locked allocator counters, plus live bytes per tag in `"by_tag"` |
|
|
519
|
+
| `trainstream.pinned_stats()` | — | page-locked allocator counters, plus live bytes per tag in `"by_tag"`. Arena keys: `arena_bytes` (reserved), `arena_used`, `arena_slabs`, `arena_cap`, `arena_fallbacks` (blocks that had to be page-locked on their own), `arena_largest_free`; `host_alloc_calls` counts every page-locking call |
|
|
520
|
+
| `trainstream.reserve_pinned(bytes_now, cap=0, slab=0)` | `bytes_now`: page-lock this much for the arena now. `cap`: let it grow on demand up to this (at least `bytes_now`). `slab`: slab size in bytes (`0` keeps the current one) | the bytes reserved so far; `wrap(host_arena=...)` calls it for you |
|
|
410
521
|
| `trainstream.transfer_stats(device=None)` | `device`: GPU (default the current one) | cumulative bytes and copies on TrainStream's D2H and H2D streams |
|
|
411
522
|
| `trainstream.find_blocks(model)` | a model | the repeated layer stack `wrap()` would use (the largest `ModuleList` / `Sequential` of ≥ 2 same-type children), or `[]` |
|
|
412
523
|
|
|
413
524
|
The session behind a wrapped model is `model._trainstream_session` (`trainstream.TrainStreamSession`):
|
|
414
|
-
- `.adapting`: `True` while it is still tuning itself
|
|
525
|
+
- `.adapting`: `True` while it is still tuning itself. Meanwhile it may reset PyTorch's peak-memory counter (`torch.cuda.reset_peak_memory_stats`) to measure single steps; once it is `False`, the counter is left alone
|
|
415
526
|
- `.events`: the decision log
|
|
416
527
|
- `.config`: the options in effect
|
|
417
528
|
- `.steps`: optimizer steps so far
|
|
529
|
+
- `.peer`: the peer-GPU pools once set up (`.pooled_blocks`, `.bytes_owned`), else `None`
|
|
418
530
|
|
|
419
531
|
### `trainstream.ActivationOffloader(model, blocks, offload="auto", *, ...)`: activation paging on its own
|
|
420
532
|
|
|
@@ -423,20 +535,22 @@ The session behind a wrapped model is `model._trainstream_session` (`trainstream
|
|
|
423
535
|
| Argument | Accepted values | Default | What it does |
|
|
424
536
|
|---|---|---|---|
|
|
425
537
|
| `model`, `blocks` | the model and its layer stack | — | whose saved activations are paged |
|
|
426
|
-
| `offload` | `"auto"`, `True`, `False`, fraction 0.0–1.0, set / list of block indices | `"auto"` | which blocks' activations go to host RAM; `"auto"` calibrates, then halves until the slowdown is within `max_slowdown` |
|
|
538
|
+
| `offload` | `"auto"`, `True`, `False`, fraction 0.0–1.0, set / list of block indices | `"auto"` | which blocks' activations go to host RAM; `"auto"` with `vram_target`: a probe step measures them, they stay in VRAM if they fit, otherwise host RAM takes as many as it holds and the rest are recomputed; `"auto"` without: calibrates, then halves until the slowdown is within `max_slowdown` (never more blocks than host RAM holds) |
|
|
427
539
|
| `min_bytes` | bytes | `1048576` | smaller tensors stay on the GPU |
|
|
428
540
|
| `prefetch_depth` | integer ≥ 1 | `1` | blocks fetched back ahead of backward |
|
|
429
|
-
| `calibrate_steps` | integer ≥ 2 | `3` | steps measured with nothing paged before tuning |
|
|
541
|
+
| `calibrate_steps` | integer ≥ 2 | `3` | steps measured with nothing paged before tuning (`"auto"` without `vram_target`) |
|
|
430
542
|
| `profile` | `None`, `HardwareProfile` | `None` | link model used to size the first guess |
|
|
431
|
-
| `max_slowdown` | float | `0.15` | how much slower `"auto"` may make a step |
|
|
543
|
+
| `max_slowdown` | float | `0.15` | how much slower `"auto"` may make a step (without `vram_target`) |
|
|
432
544
|
| `tune_window` | integer ≥ 1 | `3` | steps measured per setting while tuning (the first after a change is discarded) |
|
|
433
545
|
| `max_inflight_bytes` | bytes | `1073741824` | VRAM held by copies still crossing PCIe |
|
|
434
546
|
| `recompute` | `"auto"`, `True`, `False`, fraction, list of block indices | `"auto"` | blocks to recompute in backward instead of keeping or paging |
|
|
435
547
|
| `reserved_link_bytes`, `memory_pressure`, `vram_target` | set by `wrap()` | `(0, 0)`, `False`, `None` | PCIe already used by the optimizer per step; whether weights page; the VRAM target for the probe step |
|
|
436
548
|
|
|
437
|
-
Attributes: `.policy` (per block: `"keep"`, `"offload"` or `"recompute"`), `.decision
|
|
549
|
+
Attributes: `.policy` (per block: `"keep"`, `"offload"` or `"recompute"`), `.decision` (what the tuner measured and chose; `kept_blocks` lists blocks moved back to VRAM later), `.stats`, `.events`, `.block_bytes` (measured saved activations of one block), `.adjustable` (`True` when `"auto"` chose the policy, so `keep_resident` may change it), `.weights_paged` (set by `wrap()`: a callable saying whether weights stream at the moment; it words the probe's event).
|
|
550
|
+
|
|
551
|
+
`.keep_resident(k)`: keep the activations of up to `k` blocks that `"auto"` pages or recomputes in VRAM instead, the latest first (backward reaches them first). Does nothing for policies you set explicitly. Returns the blocks moved. `wrap()` calls it once the policy is final, with as many blocks as the room left under its VRAM target holds.
|
|
438
552
|
|
|
439
|
-
`ParamPager`, `GradientOffloader`, `ParamReadyHooks` and `HardwareProfile` are also exported. `wrap()` and `OffloadAdamW` build and use them internally; they aren't meant to be constructed by hand.
|
|
553
|
+
`ParamPager`, `GradientOffloader`, `ParamReadyHooks` and `HardwareProfile` are also exported, and `trainstream.peer` holds the peer-GPU setup. `wrap()` and `OffloadAdamW` build and use them internally; they aren't meant to be constructed by hand.
|
|
440
554
|
|
|
441
555
|
### Environment variables
|
|
442
556
|
|
|
@@ -447,32 +561,35 @@ Attributes: `.policy` (per block: `"keep"`, `"offload"` or `"recompute"`), `.dec
|
|
|
447
561
|
| `TORCH_CUDA_ARCH_LIST` | GPU architectures to compile for (default: the GPUs present) |
|
|
448
562
|
| `CUDA_HOME` / `CUDA_PATH` | CUDA toolkit to use, if it isn't found automatically |
|
|
449
563
|
| `LOCAL_WORLD_SIZE` / `SLURM_NTASKS_PER_NODE` | processes sharing the machine; the planner splits RAM between them, and data parallel shares host memory when all ranks are local |
|
|
450
|
-
| `TRAINSTREAM_SHM=0` | data parallel without shared host memory (per-rank copies, gloo reduction) |
|
|
564
|
+
| `TRAINSTREAM_SHM=0` | data parallel without shared host memory (per-rank copies, gloo reduction); also turns off peer-GPU pooling, which uses the same local group |
|
|
565
|
+
| `TRAINSTREAM_HOST_ARENA_SLAB` | size of one pinned-arena slab, in bytes or with a `K` / `M` / `G` suffix (default `1G`) |
|
|
451
566
|
|
|
452
567
|
## How it works
|
|
453
568
|
|
|
454
569
|
1. **Measure.** At `wrap()`: free VRAM, available RAM, the page-lock budget, PCIe bandwidth, and the size of every layer.
|
|
455
570
|
2. **Plan.** Which layers stay in VRAM and which stream from RAM, the optimizer moment precision, whether to quantize the frozen weights, and how many activations to page. Runs that can't fit are refused with an explanation.
|
|
456
571
|
3. **Stream.** Gradients go to pinned host memory as autograd produces them; AdamW runs on host threads while backward continues; weights and activations come back on dedicated CUDA streams ordered by events, overlapping compute.
|
|
457
|
-
4. **Adapt.** The first steps' measured peaks move layers in or out of VRAM, then the plan settles; a watchdog reacts if a later step grows. If the model fits anyway, it switches to stock fused AdamW.
|
|
572
|
+
4. **Adapt.** The first steps' measured peaks move layers in or out of VRAM, then the plan settles; a watchdog reacts if a later step grows. Activations stay in VRAM if they fit; otherwise measured trials choose which blocks are paged or recomputed, and the room left afterwards keeps the latest ones in VRAM. If the model fits anyway, it switches to stock fused AdamW.
|
|
458
573
|
|
|
459
574
|
## Examples
|
|
460
575
|
|
|
461
576
|
| Script | What it does | Main flags |
|
|
462
577
|
|---|---|---|
|
|
463
|
-
| `examples/lora_large_model.py` | LoRA on bases larger than the GPU, bf16 or quantized | `--size 1.5b\|7b\|14b\|32b`, `--model HUB_ID`, `--quantize auto\|none\|int8\|nf4\|fp8\|nvfp4`, `--prefer exact\|speed`, `--recompute auto\|on\|off`, `--attention auto\|cudnn\|efficient\|flash\|math\|default`, `--stock`, `--steps`, `--batch`, `--seq`, `--rank` |
|
|
578
|
+
| `examples/lora_large_model.py` | LoRA on bases larger than the GPU, bf16 or quantized | `--size 1.5b\|7b\|14b\|32b`, `--model HUB_ID`, `--quantize auto\|none\|int8\|nf4\|fp8\|nvfp4`, `--prefer exact\|speed`, `--recompute auto\|on\|off`, `--attention auto\|cudnn\|efficient\|flash\|math\|default`, `--outliers FRACTION`, `--stock`, `--steps`, `--batch`, `--seq`, `--rank` |
|
|
579
|
+
| `examples/dpo_lora.py` | DPO on a LoRA model: the base with adapters off is the reference, log-probs computed in chunks | `--model HUB_ID` (default: a random Qwen2.5-0.5B shape), `--quantize`, `--steps`, `--batch`, `--prompt`, `--answer`, `--beta`, `--rank` |
|
|
464
580
|
| `examples/hf_finetune.py` | fine-tune a Hugging Face model, optionally beside the stock Trainer | `--model`, `--lora`, `--compare`, `--activation-offload` |
|
|
465
581
|
| `examples/train_bigger_than_vram.py` | 1.8B model (7.3 GB of fp32 weights) inside a 4 GB VRAM cap | `--vram-gb`, `--moments fp32\|bf16`, `--layers`, `--d-model` |
|
|
466
582
|
| `examples/verify_ram_offload.py` | proves where training memory lives, with a PASS/FAIL verdict | `--gpt`, `--hf`, `--vram-cap-gb`, `--engine both\|baseline\|trainstream` |
|
|
467
583
|
| `benchmarks/compare.py` | head-to-head with stock PyTorch, bitsandbytes QLoRA, Unsloth, DeepSpeed | `--scenario full1.5b\|lora14b\|lora14b_2k`, `--engines ...`, `--steps` (2 warm-up + measured; default 10), `--timeout`, `--out` |
|
|
468
|
-
| `benchmarks/quality.py` | perplexity of each quantization format, before and after a short LoRA fine-tune | `--model`, `--formats none int8 nf4 fp8 nvfp4
|
|
584
|
+
| `benchmarks/quality.py` | perplexity of each quantization format, before and after a short LoRA fine-tune | `--model`, `--formats none int8 nf4 fp8 nvfp4 nvfp4-w nvfp4-o` (`nvfp4-w`: 4-bit weights only, expanded before each matmul; `nvfp4-o`: tensor cores with `outliers=0.01`), `--windows`, `--seq`, `--train-steps`, `--batch`, `--lr`, `--out`; `--format F` runs one format in this process, `--live` with it streams progress as `LIVE {json}` lines and the result as `RESULT {json}` |
|
|
469
585
|
|
|
470
586
|
## Limitations
|
|
471
587
|
|
|
472
588
|
- Host RAM is the ceiling: 8–12 bytes per trained parameter and 2–4 per frozen one (1.06 with int8, 0.56 with NF4), plus activations. Measured on 64 GB: a 3.6B full fine-tune, and LoRA on a 32.8B NF4 base.
|
|
473
|
-
- Quantization covers frozen block weights only (not embeddings, the output layer or trained weights) and isn't bit-exact. FP8 / NVFP4 tensor-core matmuls need a Blackwell GPU and quantize activations too. `load_model` reads `.safetensors` causal-LM checkpoints only.
|
|
589
|
+
- Quantization covers frozen block weights only (not embeddings, the output layer or trained weights) and isn't bit-exact. FP8 / NVFP4 tensor-core matmuls need a Blackwell GPU and quantize activations too; `outliers=` recovers part of the activation error, never the weights' own 4-bit error. `load_model` reads `.safetensors` causal-LM checkpoints only.
|
|
590
|
+
- `PagedKVCache` / `TrainStreamCache`: inference only, full-attention decoders, no beam search. With paged weights, every decoded token also streams the paged blocks over PCIe; a larger generation batch (which the paged cache makes room for) amortizes that.
|
|
474
591
|
- Weight paging covers the repeated block stack and needs `optimizer_mode="cpu"`; each paged weight crosses PCIe twice per step.
|
|
475
|
-
- Data parallel only (no tensor or pipeline parallelism). Ranks on one machine share host memory; tested with 2 processes on 1 GPU, not on a multi-GPU machine or across machines.
|
|
592
|
+
- Data parallel only (no tensor or pipeline parallelism). Ranks on one machine share host memory; tested with 2 processes on 1 GPU, not on a multi-GPU machine or across machines. Peer-GPU pooling covers frozen weights only and needs CUDA IPC and GPUs with peer access (GeForce pairs usually have none); tested with two processes on one GPU (Windows 11, WSL2), copies between two real GPUs are untested.
|
|
476
593
|
- `TrainStreamTrainer` is single-process. fp16 with a GradScaler and AMSGrad aren't supported; `torch.compile` hasn't been tested.
|
|
477
594
|
- Host AdamW speed is bounded by system RAM bandwidth.
|
|
478
595
|
|