trainstream 0.2.0__tar.gz → 0.4.0__tar.gz

{trainstream-0.2.0/trainstream.egg-info → trainstream-0.4.0}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: trainstream
-Version: 0.2.0
-Summary: Zero-dependency PCIe offloading engine for PyTorch training (optimizer, gradients, activations).
+Version: 0.4.0
+Summary: Zero-dependency PyTorch training offload: optimizer, gradients, activations and weights to host RAM, with a memory planner and built-in int8/NF4 quantization.
 Home-page: https://github.com/johnanishg/TrainStream
 Author: John Anish G
 License: Apache-2.0
@@ -40,7 +40,7 @@ Dynamic: summary
 
 # 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.
+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. Frozen weights (a LoRA base) can also be stored as int8 or 4-bit NF4. It takes one line of code and has no third-party dependencies.
 
 ```python
 import trainstream
@@ -58,7 +58,9 @@ TrainStream is a C++/CUDA extension for PyTorch. It keeps only the weights in VR
 | 1.8B GPT, weights paged to RAM (0.2) | — | 2.4 GiB peak, 2.1 s/step; 1.84 GiB inside a 4 GB VRAM cap with bf16 moments |
 | 3.6B GPT full fine-tune (all parameters trained) | needs ~58 GB of VRAM | 3.1 GiB peak, 3.15 s/step, 30 GiB of host RAM |
 | LoRA r=16 on the Qwen2.5-7B architecture (random weights) | 20.3 GB peak, 0.27 s/step | 4.6 GB peak, 1.25 s/step |
-| LoRA r=16 on the Qwen2.5-14B architecture (random weights; 29.5 GB) | does not fit | 5.6 GB peak, 2.48 s/step, 36 GB of host RAM |
+| LoRA r=16 on Qwen2.5-14B (29.5 GB of weights) | does not fit | 1.37 s/step: the memory planner keeps 28 of 48 layers in VRAM (19.9 GiB peak) and streams 20 from RAM |
+| same, NF4 base (`quantize="auto", prefer="speed"`, 0.4) | does not fit | 0.87 s/step: all 48 layers in VRAM as 4-bit (7.43 GB), 18.3 GB peak, 11.3 GB of process RAM |
+| LoRA r=16 on the Qwen2.5-32B architecture (65.5 GB in bf16, more than the 64 GB of RAM) | does not fit | 2.53 s/step with `load_model(quantize="nf4")`: 21.3 GB peak, 31.4 GB of process RAM |
 | 85M GPT (fits either way) | 2.55 GiB, 43.7 ms/step | `offload="auto"` switches to stock fused AdamW: 2.56 GiB, 43.8 ms/step |
 
 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).
@@ -67,7 +69,7 @@ Offloading costs throughput when a model already fits, which is why `offload="au
 
 - 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.
+- 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 123 tests pass on Windows; on Linux 122 pass and 1 (the bitsandbytes comparison) is skipped without bitsandbytes installed, RTX 5090 Laptop GPU. Native Linux and other GPUs are untested.
 
 ## Install
 
@@ -113,14 +115,32 @@ 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.
+LoRA on a quantized base (QLoRA-style), with no extra package:
+
+```python
+import peft
+from trainstream.integrations.huggingface import TrainStreamTrainer, load_model
+
+model = load_model("Qwen/Qwen2.5-14B", quantize="nf4")   # "auto" | "none" | "int8" | "nf4"; quantized while read
+model = peft.get_peft_model(model, peft.LoraConfig(r=16, lora_alpha=32, task_type="CAUSAL_LM",
+                                                   target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]))
+trainer = TrainStreamTrainer(model=model, args=TrainingArguments(..., bf16=True), train_dataset=ds)
+trainer.train()
+```
+
+Or quantize inside `wrap()`: `trainstream.wrap(model, opt, device="cuda", quantize="auto", prefer="speed")`, or beforehand with `trainstream.quantize_model(model, "int8")`. Only frozen weights are quantized. A quantized run trains on the dequantized weights (as QLoRA does), so it isn't bit-exact with the bf16 base.
+
+`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 7B–32B bases through `TrainStreamTrainer` (`--quantize auto|none|int8|nf4`, `--prefer exact|speed`).
 
 | `wrap()` option | What it does |
 |---|---|
 | `offload="auto"` | switches to stock fused AdamW after 2 steps if the model fits (`True` always offloads) |
-| `param_offload="auto"` | pages block weights to host RAM when they'd take more than half the GPU; `True`, a fraction, or block indices |
+| `param_offload="auto"` | memory planner: measures the VRAM and RAM of the machine it runs on, keeps as many transformer blocks resident as fit, pages the rest, and refines the split from the first steps' measured peaks; `True`, a fraction, or block indices fix it |
 | `activation_offload="auto"` | pages a profiled share of activations (`True`, a fraction, `False`) |
-| `moment_dtype="bfloat16"` | bf16 optimizer moments with stochastic rounding: half the host memory for them |
+| `quantize="auto"` | frozen weights only: `"auto"` quantizes when the run can't fit otherwise (int8 before NF4); `"none"`, `"int8"`, `"nf4"` fix it; `quant_block_size=64` |
+| `prefer="exact"` | `"speed"` lets `quantize="auto"` pick NF4 when it keeps every layer in VRAM |
+| `moment_dtype="auto"` | fp32 optimizer moments, or bf16 (stochastic rounding, half the memory) when fp32 wouldn't fit in RAM |
+| `check_memory=True` | refuse at `wrap()` with a clear explanation when the run can't fit in RAM/VRAM |
 | `distributed` | data parallel when a `torch.distributed` group with more than 1 rank exists; gradients averaged on the host over gloo |
 | `max_grad_norm`, `accumulation_steps` | clipping and gradient accumulation, done inside the optimizer |
 
@@ -141,12 +161,14 @@ trainer.train()
 - `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
+- `examples/lora_large_model.py` — LoRA on 7B/14B/32B-size bases larger than the GPU, bf16 or quantized, optionally against the stock Trainer
+- `TrainStreamDocs.html` — full documentation, including every option and the values it accepts (open it in a browser)
+- `TrainStreamReport.html` — what TrainStream does, how it compares with other libraries, and every measured result
 
 ## 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
+- Host RAM is the ceiling: 8–12 bytes per trained parameter and 2–4 per frozen one (1.06 int8, 0.56 NF4), plus activations; measured on 64 GB: a 3.6B full fine-tune and LoRA on a 32.8B NF4 base
+- Quantization covers frozen block weights only (not embeddings, the output layer or trained weights) and is not bit-exact; `load_model` reads `.safetensors` causal-LM checkpoints
 - 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

{trainstream-0.2.0 → trainstream-0.4.0}/README.md

@@ -1,6 +1,6 @@
 # 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.
+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. Frozen weights (a LoRA base) can also be stored as int8 or 4-bit NF4. It takes one line of code and has no third-party dependencies.
 
 ```python
 import trainstream
@@ -18,7 +18,9 @@ TrainStream is a C++/CUDA extension for PyTorch. It keeps only the weights in VR
 | 1.8B GPT, weights paged to RAM (0.2) | — | 2.4 GiB peak, 2.1 s/step; 1.84 GiB inside a 4 GB VRAM cap with bf16 moments |
 | 3.6B GPT full fine-tune (all parameters trained) | needs ~58 GB of VRAM | 3.1 GiB peak, 3.15 s/step, 30 GiB of host RAM |
 | LoRA r=16 on the Qwen2.5-7B architecture (random weights) | 20.3 GB peak, 0.27 s/step | 4.6 GB peak, 1.25 s/step |
-| LoRA r=16 on the Qwen2.5-14B architecture (random weights; 29.5 GB) | does not fit | 5.6 GB peak, 2.48 s/step, 36 GB of host RAM |
+| LoRA r=16 on Qwen2.5-14B (29.5 GB of weights) | does not fit | 1.37 s/step: the memory planner keeps 28 of 48 layers in VRAM (19.9 GiB peak) and streams 20 from RAM |
+| same, NF4 base (`quantize="auto", prefer="speed"`, 0.4) | does not fit | 0.87 s/step: all 48 layers in VRAM as 4-bit (7.43 GB), 18.3 GB peak, 11.3 GB of process RAM |
+| LoRA r=16 on the Qwen2.5-32B architecture (65.5 GB in bf16, more than the 64 GB of RAM) | does not fit | 2.53 s/step with `load_model(quantize="nf4")`: 21.3 GB peak, 31.4 GB of process RAM |
 | 85M GPT (fits either way) | 2.55 GiB, 43.7 ms/step | `offload="auto"` switches to stock fused AdamW: 2.56 GiB, 43.8 ms/step |
 
 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).
@@ -27,7 +29,7 @@ Offloading costs throughput when a model already fits, which is why `offload="au
 
 - 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.
+- 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 123 tests pass on Windows; on Linux 122 pass and 1 (the bitsandbytes comparison) is skipped without bitsandbytes installed, RTX 5090 Laptop GPU. Native Linux and other GPUs are untested.
 
 ## Install
 
@@ -73,14 +75,32 @@ 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.
+LoRA on a quantized base (QLoRA-style), with no extra package:
+
+```python
+import peft
+from trainstream.integrations.huggingface import TrainStreamTrainer, load_model
+
+model = load_model("Qwen/Qwen2.5-14B", quantize="nf4")   # "auto" | "none" | "int8" | "nf4"; quantized while read
+model = peft.get_peft_model(model, peft.LoraConfig(r=16, lora_alpha=32, task_type="CAUSAL_LM",
+                                                   target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]))
+trainer = TrainStreamTrainer(model=model, args=TrainingArguments(..., bf16=True), train_dataset=ds)
+trainer.train()
+```
+
+Or quantize inside `wrap()`: `trainstream.wrap(model, opt, device="cuda", quantize="auto", prefer="speed")`, or beforehand with `trainstream.quantize_model(model, "int8")`. Only frozen weights are quantized. A quantized run trains on the dequantized weights (as QLoRA does), so it isn't bit-exact with the bf16 base.
+
+`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 7B–32B bases through `TrainStreamTrainer` (`--quantize auto|none|int8|nf4`, `--prefer exact|speed`).
 
 | `wrap()` option | What it does |
 |---|---|
 | `offload="auto"` | switches to stock fused AdamW after 2 steps if the model fits (`True` always offloads) |
-| `param_offload="auto"` | pages block weights to host RAM when they'd take more than half the GPU; `True`, a fraction, or block indices |
+| `param_offload="auto"` | memory planner: measures the VRAM and RAM of the machine it runs on, keeps as many transformer blocks resident as fit, pages the rest, and refines the split from the first steps' measured peaks; `True`, a fraction, or block indices fix it |
 | `activation_offload="auto"` | pages a profiled share of activations (`True`, a fraction, `False`) |
-| `moment_dtype="bfloat16"` | bf16 optimizer moments with stochastic rounding: half the host memory for them |
+| `quantize="auto"` | frozen weights only: `"auto"` quantizes when the run can't fit otherwise (int8 before NF4); `"none"`, `"int8"`, `"nf4"` fix it; `quant_block_size=64` |
+| `prefer="exact"` | `"speed"` lets `quantize="auto"` pick NF4 when it keeps every layer in VRAM |
+| `moment_dtype="auto"` | fp32 optimizer moments, or bf16 (stochastic rounding, half the memory) when fp32 wouldn't fit in RAM |
+| `check_memory=True` | refuse at `wrap()` with a clear explanation when the run can't fit in RAM/VRAM |
 | `distributed` | data parallel when a `torch.distributed` group with more than 1 rank exists; gradients averaged on the host over gloo |
 | `max_grad_norm`, `accumulation_steps` | clipping and gradient accumulation, done inside the optimizer |
 
@@ -101,12 +121,14 @@ trainer.train()
 - `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
+- `examples/lora_large_model.py` — LoRA on 7B/14B/32B-size bases larger than the GPU, bf16 or quantized, optionally against the stock Trainer
+- `TrainStreamDocs.html` — full documentation, including every option and the values it accepts (open it in a browser)
+- `TrainStreamReport.html` — what TrainStream does, how it compares with other libraries, and every measured result
 
 ## 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
+- Host RAM is the ceiling: 8–12 bytes per trained parameter and 2–4 per frozen one (1.06 int8, 0.56 NF4), plus activations; measured on 64 GB: a 3.6B full fine-tune and LoRA on a 32.8B NF4 base
+- Quantization covers frozen block weights only (not embeddings, the output layer or trained weights) and is not bit-exact; `load_model` reads `.safetensors` causal-LM checkpoints
 - 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

{trainstream-0.2.0 → trainstream-0.4.0}/setup.py

@@ -55,6 +55,8 @@ if torch is not None:
             "trainstream/csrc/memory_pinner.cpp",
             "trainstream/csrc/offload_adamw.cpp",
             "trainstream/csrc/offload_adamw_kernel.cu",
+            "trainstream/csrc/quant.cpp",
+            "trainstream/csrc/quant_kernel.cu",
         ],
         include_dirs=[os.path.join(HERE, "trainstream", "csrc")],
         extra_compile_args={"cxx": cxx_flags, "nvcc": nvcc_flags},
@@ -66,10 +68,10 @@ with open(os.path.join(HERE, "README.md"), encoding="utf-8") as fh:
 
 setup(
     name="trainstream",
-    version="0.2.0",
+    version="0.4.0",
     author="John Anish G",
     license="Apache-2.0",
-    description="Zero-dependency PCIe offloading engine for PyTorch training (optimizer, gradients, activations).",
+    description="Zero-dependency PyTorch training offload: optimizer, gradients, activations and weights to host RAM, with a memory planner and built-in int8/NF4 quantization.",
     long_description=long_description,
     long_description_content_type="text/markdown",
     url="https://github.com/johnanishg/TrainStream",

{trainstream-0.2.0 → trainstream-0.4.0}/tests/test_huggingface.py

@@ -227,3 +227,99 @@ def test_trainer_takes_a_cpu_model_and_pages_its_layers(tmp_path):
     assert session.pager is not None and len(session.pager.paged) == 4
     assert session.pager.frozen_pinned_bytes + session.pager.frozen_pageable_bytes > 0
     assert losses(ts) == pytest.approx(losses(ref), abs=2e-4)
+
+
+def saved_llama(tmp_path, tie=False, dtype=torch.float32):
+    torch.manual_seed(0)
+    cfg = LlamaConfig(vocab_size=VOCAB, hidden_size=128, intermediate_size=256, num_hidden_layers=4,
+                      num_attention_heads=4, num_key_value_heads=2, max_position_embeddings=128,
+                      tie_word_embeddings=tie)
+    LlamaForCausalLM(cfg).to(dtype).save_pretrained(str(tmp_path / "ckpt"))
+    return str(tmp_path / "ckpt")
+
+
+@pytest.mark.parametrize("tie", [False, True])
+def test_load_model_equals_from_pretrained(tmp_path, tie):
+    from trainstream.integrations.huggingface import load_model
+    d = saved_llama(tmp_path, tie=tie, dtype=torch.bfloat16)
+    ref = LlamaForCausalLM.from_pretrained(d, dtype=torch.bfloat16)
+    model = load_model(d, quantize="none")
+    assert model._ts_load_info["quantize"] is None
+    assert (model.lm_head.weight is model.model.embed_tokens.weight) == tie
+    ids = torch.randint(0, VOCAB, (2, 16))
+    with torch.no_grad():
+        assert torch.equal(ref(ids).logits, model(ids).logits)
+
+
+@pytest.mark.parametrize("kind", ["int8", "nf4"])
+def test_load_model_quantizes_while_reading(tmp_path, kind):
+    from trainstream import quant as Q
+    from trainstream.integrations.huggingface import load_model
+    d = saved_llama(tmp_path, dtype=torch.bfloat16)
+    ref = LlamaForCausalLM.from_pretrained(d, dtype=torch.bfloat16)
+    model = load_model(d, quantize=kind)
+    assert model._ts_load_info["quantized_tensors"] == 4 * 7  # q k v o gate up down per layer
+    for name, p in model.named_parameters():
+        qw = Q.quantized_of(p)
+        if qw is None:
+            assert torch.equal(p, ref.get_parameter(name)), name
+            continue
+        assert p.untyped_storage().nbytes() == 0 and not p.requires_grad and qw.kind == kind
+        q, a = Q.quantize_tensor(ref.get_parameter(name), kind, 64)
+        want = Q.dequantize_reference(q, a, kind, 64, p.numel(), torch.bfloat16).view(p.shape).cpu()
+        assert torch.equal(qw.dequantize().cpu(), want), name
+
+
+def test_load_model_auto_picks_the_most_precise_format_that_fits(tmp_path, monkeypatch):
+    import trainstream.planner as planner
+    from trainstream.integrations.huggingface import load_model
+    d = saved_llama(tmp_path, dtype=torch.bfloat16)
+    numel = sum(p.numel() for p in LlamaForCausalLM.from_pretrained(d).parameters())
+    monkeypatch.setattr(planner, "local_processes", lambda: 1)
+    for ram_factor, want in ((1.1, None), (0.7, "int8"), (0.45, "nf4")):  # RAM as a multiple of the bf16 bytes
+        monkeypatch.setattr(planner, "available_ram_bytes", lambda f=ram_factor: int(f * 2 * numel))
+        model = load_model(d, ram_fraction=1.0)
+        assert model._ts_load_info["quantize"] == want, ram_factor
+    with pytest.raises(ValueError):
+        load_model(d, quantize="fp4")
+
+
+def test_load_model_random_weights_and_overrides(tmp_path):
+    from trainstream.integrations.huggingface import load_model
+    d = saved_llama(tmp_path)
+    model = load_model(d, random_weights=True, quantize="nf4", num_hidden_layers=2)
+    assert len(model.model.layers) == 2 and model._ts_load_info["quantized_tensors"] == 2 * 7
+
+
+@pytest.mark.parametrize("kind", ["int8", "nf4"])
+def test_trainer_qlora_on_loaded_model_matches_dequantized_reference(tmp_path, kind):
+    """load_model(quantize=...) + LoRA + TrainStreamTrainer equals the stock
+    Trainer on a model holding the dequantized base weights."""
+    peft = pytest.importorskip("peft")
+    from trainstream import quant as Q
+    from trainstream.integrations.huggingface import load_model
+    d = saved_llama(tmp_path)
+    model = load_model(d, quantize=kind, dtype=torch.float32)
+    ref_base = LlamaForCausalLM.from_pretrained(d, dtype=torch.float32)
+    for name, p in model.named_parameters():
+        qw = Q.quantized_of(p)
+        if qw is not None:
+            ref_base.get_parameter(name).data.copy_(qw.dequantize().cpu())
+    lora = dict(r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"], lora_dropout=0.0, task_type="CAUSAL_LM")
+    torch.manual_seed(1)
+    ref_model = peft.get_peft_model(ref_base.cuda(), peft.LoraConfig(**lora))
+    model = peft.get_peft_model(model, peft.LoraConfig(**lora))
+    ref_lora = dict(ref_model.named_parameters())
+    with torch.no_grad():
+        for n, p in model.named_parameters():
+            if p.requires_grad:
+                p.copy_(ref_lora[n].cpu())
+    ds = TokenDataset()
+    ref = Trainer(model=ref_model, args=args(tmp_path / "ref", max_steps=6), train_dataset=ds)
+    ref.train()
+    ts = TrainStreamTrainer(model=model, args=args(tmp_path / "ts", max_steps=6), train_dataset=ds,
+                            trainstream={"offload": True, "activation_offload": False})
+    ts.train()
+    session = trainstream.wrapper.unwrap_session(model)
+    assert session.pager is not None and session.pager.quantized_bytes > 0
+    assert losses(ts) == pytest.approx(losses(ref), abs=2e-4)

trainstream-0.4.0/tests/test_planner.py

@@ -0,0 +1,247 @@
+"""The memory planner: decisions from the measured VRAM / RAM, and live re-planning."""
+
+import copy
+import json
+import os
+import subprocess
+import sys
+import textwrap
+
+import pytest
+import torch
+from torch import nn
+
+import trainstream
+from trainstream.planner import Budget, MemoryPlanner, spread
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+GiB = 1 << 30
+DEV = "cuda"
+HP = dict(lr=2e-3, betas=(0.9, 0.95), weight_decay=0.05)
+HERE = os.path.dirname(os.path.abspath(__file__))
+ROOT = os.path.dirname(HERE)
+
+
+def budget(vram_gib, ram_gib=64):
+    return Budget(vram_total=int(vram_gib * GiB), vram_usable=int(vram_gib * GiB), ram_total=int(ram_gib * GiB),
+                  ram_available=int(ram_gib * GiB), pin_budget=int(0.4 * ram_gib * GiB))
+
+
+# --------------------------------------------------------------------------- decisions
+def test_spread_is_even_and_exact():
+    assert spread(0, list(range(10))) == []
+    assert spread(10, list(range(10))) == list(range(10))
+    s = spread(4, list(range(16)))
+    assert len(s) == 4 and all(b - a == 4 for a, b in zip(s, s[1:]))
+    assert len(spread(7, list(range(48)))) == 7
+
+
+@pytest.mark.parametrize("vram,expect_paged", [(80, 0), (24, None), (8, None), (4, None)])
+def test_static_plan_follows_the_vram_budget(vram, expect_paged):
+    """48 blocks of 0.55 GiB (Qwen2.5-14B in bf16) and 3 GiB outside them."""
+    blocks = [int(0.55 * GiB)] * 48
+    p = MemoryPlanner(budget(vram), blocks, fixed_vram=3 * GiB, inflight_bytes=GiB // 2)
+    paged = p.initial_paged()
+    if expect_paged == 0:
+        assert paged == set()  # an 80 GB card holds the whole model
+        return
+    resident = 48 - len(paged)
+    usable = 0.9 * vram * GiB
+    assert 3 * GiB + resident * 0.55 * GiB < usable  # what stays resident fits
+    if vram >= 24:
+        assert resident >= 20  # most of a 24 GB card is used for weights
+    assert len(paged) == len(set(paged)) and all(0 <= i < 48 for i in paged)
+
+
+def test_refinement_moves_blocks_in_with_headroom_and_out_when_full():
+    blocks = [GiB // 2] * 20
+    p = MemoryPlanner(budget(16), blocks, fixed_vram=2 * GiB, inflight_bytes=GiB // 4)
+    paged = set(range(0, 20, 2))  # 10 paged
+    grown = p.refine(peak=6 * GiB, paged=paged)  # lots of headroom
+    assert grown is not None and grown < paged
+    shrunk = p.refine(peak=int(15.5 * GiB), paged=grown)  # too close to the budget
+    assert shrunk is not None and shrunk > grown
+    p.refine(peak=int(13.5 * GiB), paged=shrunk)
+    p.refine(peak=int(13.5 * GiB), paged=shrunk)
+    assert p.done  # two steady steps settle the plan
+
+
+def test_wrap_refuses_what_cannot_fit_in_ram(monkeypatch):
+    import trainstream.planner as planner
+    monkeypatch.setattr(planner, "available_ram_bytes", lambda: 64 << 20)  # 64 MiB "free"
+    model = nn.Sequential(*[nn.Linear(1024, 1024) for _ in range(8)]).to(DEV)
+    with pytest.raises(MemoryError, match="host RAM"):
+        trainstream.wrap(model, torch.optim.AdamW(model.parameters()), offload=True, activation_offload=False)
+    # the explicit opt-out still wraps
+    model, opt = trainstream.wrap(model, torch.optim.AdamW(model.parameters()), offload=True,
+                                  activation_offload=False, check_memory=False)
+    trainstream.unwrap(model)
+
+
+def test_moment_precision_follows_available_ram(monkeypatch):
+    import trainstream.planner as planner
+    model = nn.Sequential(*[nn.Linear(1024, 1024) for _ in range(8)]).to(DEV)
+    numel = sum(p.numel() for p in model.parameters())
+    # host need: fp32 master 4 + staging 4 + moments 2x(4 or 2) B/param -> 16 (fp32) vs 12 (bf16); give 14
+    room = int(numel * 14 / 0.9)
+    monkeypatch.setattr(planner, "available_ram_bytes", lambda: room)
+    model, opt = trainstream.wrap(model, torch.optim.AdamW(model.parameters()), offload=True, activation_offload=False)
+    try:
+        assert opt.moment_dtype == "bfloat16"
+    finally:
+        trainstream.unwrap(model)
+    monkeypatch.setattr(planner, "available_ram_bytes", lambda: 64 << 30)
+    model = nn.Sequential(*[nn.Linear(1024, 1024) for _ in range(8)]).to(DEV)
+    model, opt = trainstream.wrap(model, torch.optim.AdamW(model.parameters()), offload=True, activation_offload=False)
+    try:
+        assert opt.moment_dtype == "float32"
+    finally:
+        trainstream.unwrap(model)
+
+
+# --------------------------------------------------------------------------- live re-planning
+class Block(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.ln = nn.LayerNorm(d)
+        self.fc1, self.fc2 = nn.Linear(d, 4 * d), nn.Linear(4 * d, d)
+
+    def forward(self, x):
+        return x + self.fc2(torch.nn.functional.gelu(self.fc1(self.ln(x))))
+
+
+class Stack(nn.Module):
+    def __init__(self, d=256, n=6, vocab=128):
+        super().__init__()
+        self.emb = nn.Embedding(vocab, d)
+        self.blocks = nn.ModuleList(Block(d) for _ in range(n))
+        self.head = nn.Linear(d, vocab)
+
+    def forward(self, idx):
+        x = self.emb(idx)
+        for b in self.blocks:
+            x = b(x)
+        return self.head(x)
+
+
+def lm_steps(model, opt, data, between=None):
+    out = []
+    for k, t in enumerate(data):
+        logits = model(t[:, :-1])
+        loss = nn.functional.cross_entropy(logits.float().reshape(-1, logits.size(-1)), t[:, 1:].reshape(-1))
+        loss.backward()
+        opt.step()
+        opt.zero_grad()
+        out.append(loss.item())
+        if between:
+            between(k)
+    return out
+
+
+def tokens(n, seed=3, batch=8, seq=32, vocab=128):
+    g = torch.Generator().manual_seed(seed)
+    return [torch.randint(0, vocab, (batch, seq + 1), generator=g).to(DEV) for _ in range(n)]
+
+
+def test_blocks_move_between_vram_and_ram_mid_training_without_changing_results():
+    torch.manual_seed(0)
+    ref = Stack().to(DEV)
+    mine = copy.deepcopy(ref)
+    opt_ref = torch.optim.AdamW(ref.parameters(), foreach=False, **HP)
+    mine, opt = trainstream.wrap(mine, torch.optim.AdamW(mine.parameters(), **HP), offload=True,
+                                 param_offload=[0, 2, 4], activation_offload=False)
+    pager = trainstream.wrapper.unwrap_session(mine).pager
+    moves = {1: ([0], [1, 5]), 3: ([1, 2, 4, 5], []), 5: ([], [0, 3])}  # after step k: (unpage, page)
+
+    def between(k):
+        if k in moves:
+            unpage, page = moves[k]
+            for i in page:
+                pager.page(i)
+            for i in unpage:
+                pager.unpage(i)
+
+    try:
+        data = tokens(8)
+        a = lm_steps(ref, opt_ref, data)
+        b = lm_steps(mine, opt, data, between)
+        opt.synchronize()
+        assert b == pytest.approx(a, abs=1e-4)
+        assert pager.paged == {0, 3}
+        sizes = [sum(p.untyped_storage().nbytes() for p in blk.parameters()) for blk in mine.blocks]
+        assert all(sizes[i] > 0 for i in (1, 2, 4, 5)) and sizes[3] == 0  # resident blocks hold weights
+        sd, sd_ref = mine.state_dict(), ref.state_dict()
+        assert max((sd[k].float().cpu() - sd_ref[k].float().cpu()).abs().max().item() for k in sd_ref) < 2e-4
+    finally:
+        trainstream.unwrap(mine)
+
+
+CAP_SCRIPT = textwrap.dedent(r"""
+    import json, sys, torch
+    sys.path[:0] = [ROOT, HERE]
+    import trainstream
+    from test_planner import Stack, tokens, lm_steps, HP
+    cap = CAP
+    torch.cuda.set_per_process_memory_fraction(cap / torch.cuda.get_device_properties(0).total_memory, 0)
+    torch.manual_seed(0)
+    model = Stack(d=1024, n=8)  # 8 blocks of 50 MB
+    opt = torch.optim.AdamW(model.parameters(), **HP)
+    model, opt = trainstream.wrap(model, opt, device="cuda", offload=True)
+    s = trainstream.wrapper.unwrap_session(model)
+    first = sorted(s.pager.paged)
+    peaks = []
+    def between(k):
+        peaks.append(torch.cuda.max_memory_allocated())
+    losses = lm_steps(model, opt, tokens(14, batch=4, seq=64), between)
+    opt.synchronize()
+    print("@@" + json.dumps(dict(first=first, final=sorted(s.pager.paged), settled=s.planner.done, losses=losses,
+                                 events=s.events, moments=opt.moment_dtype, cap=cap)))
+""")
+
+
+def test_auto_plan_under_a_vram_cap_finds_the_split_and_trains_correctly():
+    cap = 450 << 20  # 8 blocks of 50 MB do not all fit next to the rest
+    code = CAP_SCRIPT.replace("ROOT", repr(ROOT)).replace("HERE", repr(HERE)).replace("CAP", str(cap))
+    out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=600)
+    assert out.returncode == 0, out.stdout + out.stderr
+    res = json.loads([l for l in out.stdout.splitlines() if l.startswith("@@")][-1][2:])
+    assert res["first"], "a 450 MB cap must start with some blocks paged"
+    assert res["settled"]
+    assert res["final"] != res["first"]  # measured peaks changed the plan
+    assert 0 < len(res["final"]) < len(res["first"])  # measured headroom moved some blocks into VRAM
+    assert any("memory plan settled" in e for e in res["events"])
+    # Moving blocks between VRAM and RAM must not change the arithmetic at all: the
+    # capped, re-planned run equals an uncapped, unpaged TrainStream run exactly.
+    # (Against torch.optim.AdamW this lr is unstable enough to amplify last-bit
+    # differences after ~8 steps, so that comparison lives in the smaller tests.)
+    torch.manual_seed(0)
+    ref = Stack(d=1024, n=8)
+    ref, opt = trainstream.wrap(ref, torch.optim.AdamW(ref.parameters(), **HP), device=DEV, offload=True,
+                                param_offload=False)
+    try:
+        a = lm_steps(ref, opt, tokens(14, batch=4, seq=64))
+    finally:
+        trainstream.unwrap(ref)
+    assert res["losses"] == a
+
+
+def test_watchdog_pages_a_block_when_a_later_step_grows():
+    torch.manual_seed(0)
+    model = Stack().to(DEV)
+    model, opt = trainstream.wrap(model, torch.optim.AdamW(model.parameters(), **HP), offload=True,
+                                  param_offload=[0], activation_offload=False)
+    s = trainstream.wrapper.unwrap_session(model)
+    from trainstream.planner import Budget, MemoryPlanner
+    used = torch.cuda.memory_allocated()
+    s.planner = MemoryPlanner(Budget(vram_total=used + (64 << 20), vram_usable=used + (64 << 20), ram_total=1 << 36,
+                                     ram_available=1 << 35, pin_budget=1 << 34),
+                              s.pager.block_bytes, fixed_vram=0, inflight_bytes=1 << 20)
+    s.planner.done = True  # settled; now a step with a much larger batch arrives
+    try:
+        before = len(s.pager.paged)
+        lm_steps(model, opt, tokens(1, batch=64, seq=128))
+        assert len(s.pager.paged) == before + 1
+        assert any("memory watchdog" in e for e in s.events)
+    finally:
+        trainstream.unwrap(model)