trainstream 0.3.0__tar.gz → 0.4.0__tar.gz

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

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: trainstream
-Version: 0.3.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
@@ -59,6 +59,8 @@ TrainStream is a C++/CUDA extension for PyTorch. It keeps only the weights in VR
 | 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 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 89 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,13 +115,30 @@ 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"` | 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`) |
+| `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 |
@@ -142,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.3.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
@@ -19,6 +19,8 @@ TrainStream is a C++/CUDA extension for PyTorch. It keeps only the weights in VR
 | 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 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 89 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,13 +75,30 @@ 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"` | 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`) |
+| `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 |
@@ -102,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.3.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.3.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.3.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_quant.py

@@ -0,0 +1,279 @@
+"""Quantized frozen weights (int8 / NF4): kernels, format, training, planner, decisions."""
+
+import copy
+
+import pytest
+import torch
+from torch import nn
+
+import trainstream
+from trainstream import _C
+from trainstream import quant as Q
+from trainstream.planner import Budget
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+DEV = "cuda"
+HP = dict(lr=2e-3, betas=(0.9, 0.95), weight_decay=0.05)
+
+
+# --------------------------------------------------------------------------- kernels and format
+@pytest.mark.parametrize("kind", Q.KINDS)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize("where", ["cuda", "host"])
+def test_kernel_matches_reference(kind, dtype, where):
+    torch.manual_seed(0)
+    n = 64 * 513 + 6  # a partial last block
+    w = torch.randn(n, device=DEV) * 0.02
+    q, a = Q.quantize_tensor(w, kind, 64)
+    ref = Q.dequantize_reference(q, a, kind, 64, n, dtype)
+    if where == "host":
+        q, a = q.cpu().pin_memory(), a.cpu().pin_memory()
+    out = torch.empty(n, device=DEV, dtype=dtype)
+    ev, keep = _C.dequantize_async(out, q, a, kind, 64, _C.record_event(0))
+    ev.wait()
+    torch.cuda.synchronize()
+    assert torch.equal(out, ref)
+    err = ((ref.float() - w).abs().mean() / w.abs().mean()).item()
+    assert err < (0.12 if kind == "nf4" else 0.01)
+
+
+def test_nf4_is_bit_identical_to_bitsandbytes():
+    bf = pytest.importorskip("bitsandbytes.functional")
+    torch.manual_seed(0)
+    w = (torch.randn(1024, 512, device=DEV) * 0.02).to(torch.bfloat16)
+    qb, st = bf.quantize_4bit(w, blocksize=64, quant_type="nf4", compress_statistics=False)
+    q, a = Q.quantize_tensor(w, "nf4", 64)
+    assert torch.equal(qb.view(-1), q)  # same packed bytes
+    assert torch.equal(st.absmax.float(), a)
+    assert torch.equal(bf.dequantize_4bit(qb, st, quant_type="nf4"),
+                       Q.dequantize_reference(q, a, "nf4", 64, w.numel(), torch.bfloat16).view_as(w))
+
+
+# --------------------------------------------------------------------------- training
+class Block(nn.Module):
+    """A frozen 'base' MLP plus a small trainable adapter (LoRA-like)."""
+
+    def __init__(self, d=256, r=8):
+        super().__init__()
+        self.ln = nn.LayerNorm(d)
+        self.fc1, self.fc2 = nn.Linear(d, 4 * d), nn.Linear(4 * d, d)
+        self.down, self.up = nn.Linear(d, r, bias=False), nn.Linear(r, d, bias=False)
+        nn.init.zeros_(self.up.weight)
+
+    def forward(self, x):
+        h = self.ln(x)
+        return x + self.fc2(torch.nn.functional.gelu(self.fc1(h))) + self.up(self.down(h))
+
+
+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 make(seed=0):
+    torch.manual_seed(seed)
+    m = Stack().to(DEV)
+    for b in m.blocks:  # the base is frozen, the adapters and the rest train
+        for lin in (b.fc1, b.fc2):
+            lin.weight.requires_grad_(False)
+    return m
+
+
+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 train(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 dequantized_reference(model, kind):
+    """A copy whose frozen weights hold exactly the values quantized training uses."""
+    ref = copy.deepcopy(model)
+    with torch.no_grad():
+        for b in ref.blocks:
+            for lin in (b.fc1, b.fc2):
+                q, a = Q.quantize_tensor(lin.weight, kind, 64)
+                lin.weight.copy_(Q.dequantize_reference(q, a, kind, 64, lin.weight.numel(), lin.weight.dtype)
+                                 .view_as(lin.weight))
+    return ref
+
+
+def trainable(m):
+    return [p for p in m.parameters() if p.requires_grad]
+
+
+@pytest.mark.parametrize("kind", Q.KINDS)
+@pytest.mark.parametrize("placement", ["near", "far"])
+def test_quantized_training_equals_training_on_the_dequantized_weights(kind, placement):
+    base = make()
+    ref = dequantized_reference(base, kind)
+    opt_ref = torch.optim.AdamW(trainable(ref), foreach=False, **HP)
+    mine = copy.deepcopy(base)
+    mine, opt = trainstream.wrap(mine, torch.optim.AdamW(trainable(mine), **HP), offload=True, quantize=kind,
+                                 param_offload=(placement == "far"), activation_offload=False)
+    try:
+        s = trainstream.wrapper.unwrap_session(mine)
+        assert s.config["quantize"] == kind
+        assert len(s.pager.quant_blocks) == 6
+        assert (len(s.pager.qnear) == 6) == (placement == "near")
+        for b in mine.blocks:  # no full-precision copy of the frozen base stays in VRAM between uses
+            assert Q.quantized_of(b.fc1.weight) is not None
+            assert all(Q.quantized_of(p) is None for p in trainable(b))  # trained weights never quantized
+        data = tokens(6)
+        assert train(mine, opt, data) == pytest.approx(train(ref, opt_ref, data), abs=1e-4)
+        rep = trainstream.memory_report(mine)
+        assert rep["params"]["quantize"] == kind and rep["params"]["quantized_bytes"] > 0
+        assert any(f"to {kind}" in e for e in rep["events"])
+    finally:
+        trainstream.unwrap(mine)
+    # unwrap leaves the dequantized values in place
+    assert torch.equal(mine.blocks[0].fc1.weight, ref.blocks[0].fc1.weight)
+
+
+def test_moving_quantized_blocks_near_and_far_does_not_change_results():
+    base = make()
+    a, b = copy.deepcopy(base), copy.deepcopy(base)
+    a, opt_a = trainstream.wrap(a, torch.optim.AdamW(trainable(a), **HP), offload=True, quantize="nf4",
+                                param_offload=False, activation_offload=False)
+    b, opt_b = trainstream.wrap(b, torch.optim.AdamW(trainable(b), **HP), offload=True, quantize="nf4",
+                                param_offload=False, activation_offload=False)
+    pager = trainstream.wrapper.unwrap_session(b).pager
+    moves = {0: ([], [0, 2, 4]), 2: ([0], [1]), 4: ([1, 2, 4], [])}  # after step k: (near, far)
+
+    def between(k):
+        if k in moves:
+            near, far = moves[k]
+            for i in far:
+                pager.page(i)
+            for i in near:
+                pager.unpage(i)
+
+    try:
+        data = tokens(7)
+        assert train(a, opt_a, data) == train(b, opt_b, data, between)
+        # far {0,2,4} after step 0; {1,2,4} after step 2; all near again after step 4
+        assert pager.remote == set() and pager.qnear == set(range(6))
+        assert pager.stats["pages"] >= 4 and pager.stats["unpages"] >= 4
+    finally:
+        trainstream.unwrap(a)
+        trainstream.unwrap(b)
+
+
+def test_state_dict_serves_the_dequantized_values():
+    base = make()
+    ref = dequantized_reference(base, "nf4")
+    mine, opt = trainstream.wrap(copy.deepcopy(base), torch.optim.AdamW(trainable(base), **HP), offload=True,
+                                 quantize="nf4", activation_offload=False)
+    try:
+        sd = mine.state_dict()
+        assert torch.equal(sd["blocks.2.fc1.weight"].to(DEV), ref.blocks[2].fc1.weight)
+    finally:
+        trainstream.unwrap(mine)
+
+
+# --------------------------------------------------------------------------- decisions and errors
+def fake_budget(vram_gib, ram_gib):
+    G = 1 << 30
+    return Budget(vram_total=int(vram_gib * G), vram_usable=int(vram_gib * G), ram_total=int(64 * G),
+                  ram_available=int(ram_gib * G), pin_budget=int(25 * G))
+
+
+def frozen_bytes(m):
+    return sum(p.numel() * p.element_size() for p in m.parameters() if not p.requires_grad)
+
+
+# frozen weights here are fp32: int8 keeps ~27% of their bytes, nf4 ~14%
+@pytest.mark.parametrize("ram_factor,expect", [(1.5, None), (0.5, "int8"), (0.2, "nf4")])
+def test_auto_quantizes_only_when_bf16_cannot_fit(monkeypatch, ram_factor, expect):
+    m = make()
+    t = sum(p.numel() for p in trainable(m))
+    room = 12 * t + frozen_bytes(m) * ram_factor  # master + bf16 moments + staging for the trained part
+    monkeypatch.setattr(trainstream.wrapper, "measure_budget", lambda dev: fake_budget(20, room / 2**30 / 0.9))
+    m, opt = trainstream.wrap(m, torch.optim.AdamW(trainable(m), **HP), offload=True, activation_offload=False)
+    try:
+        assert trainstream.wrapper.unwrap_session(m).config["quantize"] == expect
+    finally:
+        trainstream.unwrap(m)
+
+
+def test_auto_refuses_when_even_nf4_cannot_fit(monkeypatch):
+    m = make()
+    monkeypatch.setattr(trainstream.wrapper, "measure_budget", lambda dev: fake_budget(20, 0.001))
+    with pytest.raises(MemoryError, match="host RAM"):
+        trainstream.wrap(m, torch.optim.AdamW(trainable(m), **HP), offload=True, activation_offload=False)
+
+
+def test_prefer_speed_picks_nf4_when_it_fits_in_vram_and_bf16_would_stream(monkeypatch):
+    m = make()
+    blocks = list(m.blocks)
+    bf16 = sum(p.numel() * p.element_size() for b in blocks for p in b.parameters())
+    nf4 = sum(Q.packed_bytes(p.numel(), "nf4") if Q.eligible(p) else p.numel() * p.element_size()
+              for b in blocks for p in b.parameters())
+    outside = sum(p.numel() * p.element_size() for p in m.parameters()) - bf16
+    # With a small budget V the planner keeps 0.6 V - outside - 96 MiB for block weights
+    # (10% headroom, 10% in-flight cap, 20% activation reserve, workspace): put that
+    # between the 4-bit and the bf16 size of the blocks.
+    v = (outside + (96 << 20) + (nf4 + bf16) / 2) / 0.6
+    monkeypatch.setattr(trainstream.wrapper, "measure_budget", lambda dev: fake_budget(v / 2**30, 32))
+    m, opt = trainstream.wrap(m, torch.optim.AdamW(trainable(m), **HP), offload=True, activation_offload=False,
+                              prefer="speed", check_memory=False)
+    try:
+        assert trainstream.wrapper.unwrap_session(m).config["quantize"] == "nf4"
+    finally:
+        trainstream.unwrap(m)
+    m2 = make()
+    m2, opt2 = trainstream.wrap(m2, torch.optim.AdamW(trainable(m2), **HP), offload=True, activation_offload=False,
+                                check_memory=False)  # prefer="exact": stream bf16 instead
+    try:
+        s2 = trainstream.wrapper.unwrap_session(m2)
+        assert s2.config["quantize"] is None and s2.pager is not None and s2.pager.remote
+    finally:
+        trainstream.unwrap(m2)
+
+
+def test_invalid_quantization_requests():
+    full = Stack().to(DEV)  # everything trainable: nothing frozen to quantize
+    with pytest.raises(ValueError, match="frozen"):
+        trainstream.wrap(full, quantize="nf4")
+    with pytest.raises(ValueError, match="quantize must be"):
+        trainstream.wrap(make(), quantize="fp3")
+    with pytest.raises(ValueError, match="prefer must be"):
+        trainstream.wrap(make(), prefer="fast")
+    pre = make()
+    Q.quantize_model(pre, "int8")
+    with pytest.raises(ValueError, match="already holds quantized"):
+        trainstream.wrap(pre, torch.optim.AdamW(trainable(pre)), quantize="none")
+
+
+def test_quantize_model_then_wrap_uses_the_given_kind():
+    m = make()
+    info = Q.quantize_model(m, "int8")
+    assert info["tensors"] == 12 and info["bytes_after"] < info["bytes_before"] / 3
+    m, opt = trainstream.wrap(m, torch.optim.AdamW(trainable(m), **HP), offload=True, activation_offload=False)
+    try:
+        assert trainstream.wrapper.unwrap_session(m).config["quantize"] == "int8"
+        train(m, opt, tokens(2))
+    finally:
+        trainstream.unwrap(m)

{trainstream-0.3.0 → trainstream-0.4.0}/trainstream/__init__.py

@@ -4,7 +4,8 @@
     model, optimizer = trainstream.wrap(model, optimizer)
 
 Optimizer state, gradients, a profiled fraction of activations and (when
-needed) the transformer-block weights themselves are moved to host RAM over
+needed) the transformer-block weights themselves (frozen ones optionally
+quantized to int8 / NF4) are moved to host RAM over
 dedicated CUDA copy streams, so models that would not fit in VRAM with a stock
 optimizer can still be trained on a single GPU — or on several, data parallel.
 """
@@ -31,9 +32,10 @@ from .losses import causal_lm_loss, chunked_cross_entropy
 from .optimizer import OffloadAdamW
 from .paging import ParamPager
 from .profiler import HardwareProfile, profile_hardware
+from .quant import quantize_model
 from .wrapper import TrainStreamSession, memory_report, unwrap, wrap
 
-__version__ = "0.3.0"
+__version__ = "0.4.0"
 
 __all__ = [
     "wrap",
@@ -49,6 +51,7 @@ __all__ = [
     "TrainStreamSession",
     "HardwareProfile",
     "profile_hardware",
+    "quantize_model",
     "find_blocks",
     "pinned_stats",
     "transfer_stats",

{trainstream-0.3.0 → trainstream-0.4.0}/trainstream/_build.py

@@ -25,7 +25,8 @@ import torch
 
 NAME = "trainstream_C"
 CSRC = os.path.join(os.path.dirname(os.path.abspath(__file__)), "csrc")
-SOURCES = ["bindings.cpp", "memory_pinner.cpp", "offload_adamw.cpp", "offload_adamw_kernel.cu"]
+SOURCES = ["bindings.cpp", "memory_pinner.cpp", "offload_adamw.cpp", "offload_adamw_kernel.cu",
+           "quant.cpp", "quant_kernel.cu"]
 HEADERS = ["kernels.h", "offload_adamw.h", "trainstream.h"]
 
 if sys.platform == "win32":

{trainstream-0.3.0 → trainstream-0.4.0}/trainstream/csrc/bindings.cpp

@@ -92,6 +92,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("device") = -1, py::arg("sizes") = std::vector<int64_t>{},
         py::arg("strides") = std::vector<int64_t>{});
   m.def("transfer_stats", &transfer_stats, py::arg("device") = -1);
+  m.def("dequantize_async", &dequantize_async, py::arg("dst"), py::arg("q"), py::arg("absmax"), py::arg("kind"),
+        py::arg("block"), py::arg("after") = EventPtr());
   m.def("synchronize_streams", [](int device) {
         Runtime& rt = runtime(device);
         TS_CUDA_CHECK(cudaStreamSynchronize(rt.d2h));

{trainstream-0.3.0 → trainstream-0.4.0}/trainstream/csrc/kernels.h

@@ -44,6 +44,13 @@ cudaError_t adamw_zero_copy(DType dtype, void* param, const void* grad, void* ex
                             const AdamWParams& hp, const float* grad_sq_norm, int sm_count,
                             cudaStream_t stream, bool bf16_moments = false, uint32_t seed = 0);
 
+// Blockwise-quantized frozen weights (see quant_kernel.cu).
+enum class QKind : int { kNF4 = 0, kInt8 = 1 };
+
+// out[0..n) = dequantized values; one fp32 absmax per `block` values (block even).
+cudaError_t dequantize(QKind kind, DType out_dtype, const void* q, const float* absmax, void* out, int64_t n,
+                       int block, int sm_count, cudaStream_t stream);
+
 // out[0] += sum(grad[i]^2). Used for global gradient-norm clipping on the GPU path.
 cudaError_t grad_sq_norm_accumulate(DType dtype, const void* grad, int64_t n, float* out,
                                     int sm_count, cudaStream_t stream);