trainstream 0.1.1__tar.gz → 0.2.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trainstream
-Version: 0.1.1
+Version: 0.2.0
 Summary: Zero-dependency PCIe offloading engine for PyTorch training (optimizer, gradients, activations).
 Home-page: https://github.com/johnanishg/TrainStream
 Author: John Anish G
@@ -40,7 +40,7 @@ Dynamic: summary
 
 # TrainStream
 
-Train and fine-tune models larger than your GPU's memory by moving optimizer state, gradients and activations to system RAM over PCIe — with one line of code and 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. It takes one line of code and has no third-party dependencies.
 
 ```python
 import trainstream
@@ -55,15 +55,19 @@ TrainStream is a C++/CUDA extension for PyTorch. It keeps only the weights in VR
 |---|---|---|
 | 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) |
 | Qwen2.5-1.5B full fine-tune (Hugging Face Trainer) | 35.9 GiB peak (spill), 111.7 s / 10 steps | 18.7 GiB, 28.5 s / 10 steps, same losses |
-| 85M GPT (fits either way) | 2.5 GiB, ~94k tokens/s | 1.8 GiB, ~64k tokens/s |
+| 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 |
+| 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 |
 
-When a model already fits, TrainStream trades throughput for memory. Its 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).
+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).
 
 ## Requirements
 
 - 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+cu130, CUDA 13.4, Python 3.14. Linux uses the same build path but is 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 78 tests pass on both, RTX 5090 Laptop GPU. Native Linux and other GPUs are untested.
 
 ## Install
 
@@ -100,6 +104,28 @@ for x, y in loader:
     optimizer.zero_grad()
 ```
 
+Models bigger than the GPU: build on the CPU and let TrainStream page the transformer-block weights:
+
+```python
+model = MyGPT(cfg)                                  # stays on the CPU
+optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
+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.
+
+| `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 |
+| `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 |
+| `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 |
+
+Also: `trainstream.chunked_cross_entropy(hidden, weight, targets)` computes the loss without materializing the logits. `.to()`, `.cuda()` and `.half()` work after `wrap()`, and reading a weight anywhere waits for its in-flight update.
+
 Hugging Face (`pip install transformers accelerate`):
 
 ```python
@@ -108,18 +134,22 @@ trainer = TrainStreamTrainer(model=model, args=TrainingArguments(..., bf16=True)
 trainer.train()
 ```
 
-Use `bf16=True` or fp32; fp16 with loss scaling is not supported.
+`TrainStreamTrainer` computes the causal-LM loss in chunks, so the full logits tensor never exists. Pass extra `wrap()` options as `trainstream={...}`. Use `bf16=True` or fp32; fp16 with loss scaling is not supported.
 
 ## Tools in this repository
 
 - `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
 
 ## Limitations
 
-- Weights stay in VRAM (no parameter offloading); large-vocabulary output logits are not offloaded
-- Single GPU per process
+- 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
+- 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
 - Host AdamW speed is bounded by system RAM bandwidth
 
 ## License

trainstream-0.2.0/README.md

@@ -0,0 +1,117 @@
+# 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.
+
+```python
+import trainstream
+model, optimizer = trainstream.wrap(model, optimizer)
+```
+
+TrainStream is a C++/CUDA extension for PyTorch. It keeps only the weights in VRAM: gradients stream to pinned host memory as autograd produces them, AdamW runs on host threads while the backward pass continues, and a profiled share of activations is paged out and prefetched back block by block. Updates match `torch.optim.AdamW` (max parameter difference ≈3·10⁻⁷ after 20 steps); activation paging is bit-exact.
+
+## Results (RTX 5090 Laptop GPU, 24 GB)
+
+| Workload | Stock PyTorch AdamW | TrainStream |
+|---|---|---|
+| 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) |
+| Qwen2.5-1.5B full fine-tune (Hugging Face Trainer) | 35.9 GiB peak (spill), 111.7 s / 10 steps | 18.7 GiB, 28.5 s / 10 steps, same losses |
+| 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 |
+| 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).
+
+## Requirements
+
+- 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.
+
+## Install
+
+```bash
+pip install torch --index-url https://download.pytorch.org/whl/cu130   # if you don't have CUDA PyTorch yet
+pip install trainstream
+```
+
+The first `import trainstream` compiles the CUDA extension against your PyTorch (a few minutes, once) and caches it; every later import is instant. On Windows it finds Visual Studio and the CUDA toolkit itself, so any terminal works. The cache is keyed by PyTorch/CUDA/Python version and GPU, so upgrading PyTorch triggers one rebuild.
+
+| Variable | Effect |
+|---|---|
+| `TRAINSTREAM_CACHE_DIR` | where builds are cached (default `%LOCALAPPDATA%\trainstream\build`, Linux `~/.cache/trainstream`) |
+| `TRAINSTREAM_VERBOSE_BUILD=1` | show compiler output |
+| `TORCH_CUDA_ARCH_LIST` | GPU architectures to compile for (default: the GPUs present) |
+| `CUDA_HOME` | CUDA toolkit to use, if it isn't found automatically |
+
+To compile ahead of time instead: `pip install trainstream --no-binary trainstream --no-build-isolation`. From a checkout: `build.bat` (Windows) or `python setup.py build_ext --inplace`.
+
+## Usage
+
+```python
+import torch, trainstream
+
+model = MyModel().cuda()
+optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
+model, optimizer = trainstream.wrap(model, optimizer)   # max_grad_norm=..., accumulation_steps=... if you use them
+
+for x, y in loader:
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        loss = model(x.cuda(), y.cuda())
+    loss.backward()
+    optimizer.step()
+    optimizer.zero_grad()
+```
+
+Models bigger than the GPU: build on the CPU and let TrainStream page the transformer-block weights:
+
+```python
+model = MyGPT(cfg)                                  # stays on the CPU
+optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
+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.
+
+| `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 |
+| `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 |
+| `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 |
+
+Also: `trainstream.chunked_cross_entropy(hidden, weight, targets)` computes the loss without materializing the logits. `.to()`, `.cuda()` and `.half()` work after `wrap()`, and reading a weight anywhere waits for its in-flight update.
+
+Hugging Face (`pip install transformers accelerate`):
+
+```python
+from trainstream.integrations.huggingface import TrainStreamTrainer
+trainer = TrainStreamTrainer(model=model, args=TrainingArguments(..., bf16=True), train_dataset=ds)
+trainer.train()
+```
+
+`TrainStreamTrainer` computes the causal-LM loss in chunks, so the full logits tensor never exists. Pass extra `wrap()` options as `trainstream={...}`. Use `bf16=True` or fp32; fp16 with loss scaling is not supported.
+
+## Tools in this repository
+
+- `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
+
+## 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
+- 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
+- Host AdamW speed is bounded by system RAM bandwidth
+
+## License
+
+Apache License 2.0 — see [LICENSE](LICENSE) and [NOTICE](NOTICE). Copyright 2026 John Anish G.

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

@@ -66,7 +66,7 @@ with open(os.path.join(HERE, "README.md"), encoding="utf-8") as fh:
 
 setup(
     name="trainstream",
-    version="0.1.1",
+    version="0.2.0",
     author="John Anish G",
     license="Apache-2.0",
     description="Zero-dependency PCIe offloading engine for PyTorch training (optimizer, gradients, activations).",

trainstream-0.2.0/tests/test_distributed.py

@@ -0,0 +1,82 @@
+"""Data-parallel training: N processes, gradients averaged on the host.
+
+Two ranks share one GPU here (gloo needs no second card). Each rank trains on
+its own half of every batch; the result must equal one process training on the
+whole batch with torch.optim.AdamW.
+"""
+
+import copy
+import os
+import pathlib
+import tempfile
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from torch import nn
+from torch.nn import functional as F
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+HP = dict(lr=2e-3, betas=(0.9, 0.95), weight_decay=0.05)
+WORLD = 2
+STEPS = 6
+
+
+def make_model(seed):
+    torch.manual_seed(seed)
+    return nn.Sequential(nn.Linear(32, 128), nn.GELU(), nn.Linear(128, 128), nn.GELU(), nn.Linear(128, 8))
+
+
+def full_batches(accum):
+    g = torch.Generator().manual_seed(7)
+    return [[(torch.randn(16, 32, generator=g), torch.randint(0, 8, (16,), generator=g)) for _ in range(accum)]
+            for _ in range(STEPS)]
+
+
+def _worker(rank, store, out, clip, accum, param_offload):
+    import trainstream
+    dist.init_process_group("gloo", init_method=pathlib.Path(store).as_uri(), rank=rank, world_size=WORLD)
+    model = make_model(seed=100 + rank).cuda()  # different init per rank: wrap() broadcasts rank 0's
+    model, opt = trainstream.wrap(model, torch.optim.AdamW(model.parameters(), **HP), offload=True,
+                                  activation_offload=False, param_offload=param_offload, max_grad_norm=clip,
+                                  accumulation_steps=accum, blocks=[model[0], model[2]] if param_offload else None)
+    assert opt._dp is not None  # detected from the initialized process group
+    for micro in full_batches(accum):
+        for x, y in micro:
+            xs, ys = x.chunk(WORLD)[rank].cuda(), y.chunk(WORLD)[rank].cuda()
+            (F.cross_entropy(model(xs), ys) / accum).backward()
+        opt.step()
+        opt.zero_grad()
+    opt.synchronize()
+    sd = {k: v.detach().cpu() for k, v in model.state_dict().items()}
+    torch.save(sd, os.path.join(out, f"rank{rank}.pt"))
+    trainstream.unwrap(model)
+    dist.destroy_process_group()
+
+
+def reference(clip, accum):
+    model = make_model(seed=100).cuda()
+    opt = torch.optim.AdamW(model.parameters(), foreach=False, **HP)
+    for micro in full_batches(accum):
+        for x, y in micro:
+            (F.cross_entropy(model(x.cuda()), y.cuda()) / accum).backward()
+        if clip:
+            torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
+        opt.step()
+        opt.zero_grad()
+    return {k: v.detach().cpu() for k, v in model.state_dict().items()}
+
+
+@pytest.mark.parametrize("clip,accum,param_offload", [(None, 1, False), (0.5, 1, False), (None, 2, False),
+                                                      (None, 1, True)])
+def test_data_parallel_matches_single_process(clip, accum, param_offload):
+    tmp = tempfile.mkdtemp()
+    store = os.path.join(tmp, "store")
+    mp.spawn(_worker, args=(store, tmp, clip, accum, param_offload), nprocs=WORLD, join=True)
+    ranks = [torch.load(os.path.join(tmp, f"rank{r}.pt")) for r in range(WORLD)]
+    ref = reference(clip, accum)
+    for k in ref:
+        assert torch.equal(ranks[0][k], ranks[1][k]), k  # replicas stay identical
+        assert (ranks[0][k] - ref[k]).abs().max().item() < 1e-4, k

trainstream-0.2.0/tests/test_features.py

@@ -0,0 +1,279 @@
+"""bf16 moments, chunked cross-entropy, the parameter guard, .to() after wrap,
+and the automatic switch to a resident optimizer."""
+
+import copy
+
+import pytest
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+import trainstream
+from trainstream import OffloadAdamW, chunked_cross_entropy
+
+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)
+
+
+def mlp(seed=0, width=256):
+    torch.manual_seed(seed)
+    return nn.Sequential(nn.Linear(64, width), nn.GELU(), nn.Linear(width, width), nn.GELU(),
+                         nn.Linear(width, 10)).to(DEV)
+
+
+def batches(n, seed=1):
+    g = torch.Generator().manual_seed(seed)
+    return [(torch.randn(32, 64, generator=g).to(DEV), torch.randint(0, 10, (32,), generator=g).to(DEV))
+            for _ in range(n)]
+
+
+def run(model, opt, data):
+    out = []
+    for x, y in data:
+        loss = F.cross_entropy(model(x), y)
+        loss.backward()
+        opt.step()
+        opt.zero_grad()
+        out.append(loss.item())
+    return out
+
+
+def max_param_diff(a, b):
+    return max((p - q).abs().max().item() for p, q in zip(a.parameters(), b.parameters()))
+
+
+# --------------------------------------------------------------------------- bf16 moments
+@pytest.mark.parametrize("mode", ["cpu", "gpu"])
+def test_bf16_moments_track_fp32_moments(mode):
+    ref, mine = mlp(), mlp()
+    opt_ref = OffloadAdamW(ref.parameters(), mode=mode, **HP)
+    opt = OffloadAdamW(mine.parameters(), mode=mode, moment_dtype="bfloat16", **HP)
+    data = batches(4) * 50  # a fixed set, so there is something to learn
+    la, lb = run(ref, opt_ref, data), run(mine, opt, data)
+    assert lb[-1] < lb[0] * 0.5
+    assert abs(la[-1] - lb[-1]) < 0.05 * la[-1] + 0.02
+    # second moments: stochastic rounding keeps them unbiased relative to fp32
+    for slot in range(len(list(mine.parameters()))):
+        _, _, v_ref = opt_ref.engine.state(slot)
+        _, _, v = opt.engine.state(slot)
+        assert v.dtype == torch.bfloat16
+        rel = (v.float().sum() - v_ref.sum()).abs() / v_ref.sum()
+        assert rel < 0.03, (slot, rel.item())
+    s, s_ref = opt.stats(), opt_ref.stats()
+    numel = sum(p.numel() for p in mine.parameters())
+    assert s_ref["host_state_bytes"] - s["host_state_bytes"] == numel * 4  # two moments, 2 bytes saved each
+    sd = opt.state_dict()  # exported as fp32, torch-compatible
+    assert sd["state"][0]["exp_avg"].dtype == torch.float32
+
+
+# --------------------------------------------------------------------------- chunked loss
+@pytest.mark.parametrize("bias", [False, True])
+def test_chunked_cross_entropy_matches_dense(bias):
+    torch.manual_seed(0)
+    h = torch.randn(3, 37, 64, device=DEV, requires_grad=True)
+    w = torch.randn(1000, 64, device=DEV, requires_grad=True)
+    b = torch.randn(1000, device=DEV, requires_grad=True) if bias else None
+    t = torch.randint(0, 1000, (3, 37), device=DEV)
+    t[0, :5] = -100
+    dense = F.cross_entropy(F.linear(h, w, b).reshape(-1, 1000), t.reshape(-1), ignore_index=-100)
+    gd = torch.autograd.grad(dense, [h, w] + ([b] if bias else []))
+    ch = chunked_cross_entropy(h, w, t, b, chunk_size=16)
+    gc = torch.autograd.grad(ch, [h, w] + ([b] if bias else []))
+    assert ch.item() == pytest.approx(dense.item(), rel=1e-5)
+    for x, y in zip(gc, gd):
+        assert (x - y).abs().max().item() < 1e-5
+    # sum reduction and Hugging Face's num_items normalisation
+    s = chunked_cross_entropy(h, w, t, b, reduction="sum")
+    assert s.item() == pytest.approx(dense.item() * int((t != -100).sum()), rel=1e-5)
+    assert chunked_cross_entropy(h, w, t, b, num_items=50).item() == pytest.approx(s.item() / 50, rel=1e-5)
+
+
+def test_chunked_cross_entropy_under_bf16_autocast():
+    torch.manual_seed(0)
+    h = torch.randn(512, 128, device=DEV, requires_grad=True)
+    w = torch.randn(4000, 128, device=DEV, requires_grad=True)
+    t = torch.randint(0, 4000, (512,), device=DEV)
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        dense = F.cross_entropy(F.linear(h, w).float(), t)
+        ch = chunked_cross_entropy(h, w, t, chunk_size=100)
+    gd = torch.autograd.grad(dense, [h, w])
+    gc = torch.autograd.grad(ch, [h, w])
+    assert ch.item() == pytest.approx(dense.item(), rel=1e-3)
+    for x, y in zip(gc, gd):
+        assert ((x.float() - y.float()).norm() / y.float().norm()).item() < 2e-2
+
+
+def test_chunked_cross_entropy_never_materializes_the_logits():
+    rows, vocab, d = 4096, 50000, 256
+    h = torch.randn(rows, d, device=DEV, requires_grad=True)
+    w = torch.randn(vocab, d, device=DEV, requires_grad=True)
+    t = torch.randint(0, vocab, (rows,), device=DEV)
+
+    def peak(fn):
+        torch.cuda.synchronize()
+        torch.cuda.reset_peak_memory_stats()
+        base = torch.cuda.memory_allocated()
+        loss = fn()
+        loss.backward()
+        torch.cuda.synchronize()
+        h.grad = w.grad = None
+        return torch.cuda.max_memory_allocated() - base
+
+    dense = peak(lambda: F.cross_entropy(F.linear(h, w), t))
+    chunked = peak(lambda: chunked_cross_entropy(h, w, t, chunk_size=512))
+    logits_bytes = rows * vocab * 4
+    assert dense > logits_bytes
+    print(f"dense {dense / 2**20:.0f} MiB, chunked {chunked / 2**20:.0f} MiB")
+    assert chunked < 0.25 * dense  # one 512-row slice at a time (plus the weight gradient)
+
+
+# --------------------------------------------------------------------------- parameter guard
+class EarlyUse(nn.Module):
+    """Reads its output layer's (large) weight at the very start of forward,
+    long before that layer's own pre-hook would wait for its update."""
+
+    def __init__(self, d=512, vocab=32768):
+        super().__init__()
+        self.emb = nn.Embedding(vocab, d)
+        self.mid = nn.Linear(d, d)
+        self.head = nn.Linear(d, vocab)
+
+    def forward(self, idx):
+        gate = torch.sigmoid(self.head.weight.mean()) * 2
+        x = self.emb(idx) * gate
+        return self.head(torch.tanh(self.mid(x)))
+
+
+def test_guard_waits_for_weights_used_outside_their_module():
+    torch.manual_seed(0)
+    ref = EarlyUse().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,
+                                 activation_offload=False, param_offload=False)
+    g = torch.Generator().manual_seed(5)
+    data = [torch.randint(0, 32768, (8, 32), generator=g).to(DEV) for _ in range(6)]
+    try:
+        for idx in data:
+            for m, o in ((ref, opt_ref), (mine, opt)):
+                F.cross_entropy(m(idx).reshape(-1, 32768).float(), idx.reshape(-1)).backward()
+                o.step()
+                o.zero_grad()
+        opt.synchronize()
+        assert opt.stats()["guard_waits"] > 0
+        assert max_param_diff(ref, mine) < 1e-4
+    finally:
+        trainstream.unwrap(mine)
+
+
+def test_guard_ignores_metadata_and_grad_access():
+    model = mlp()
+    model, opt = trainstream.wrap(model, torch.optim.AdamW(model.parameters(), **HP), offload=True,
+                                  activation_offload=False)
+    try:
+        x, y = batches(1)[0]
+        F.cross_entropy(model(x), y).backward()
+        opt.step()
+        before = opt.stats()["guard_waits"]
+        for p in model.parameters():  # what zero_grad, schedulers and loggers touch
+            _ = p.grad, p.shape, p.dtype, p.requires_grad, p.numel(), p.size()
+        opt.zero_grad()
+        assert opt.stats()["guard_waits"] == before
+        _ = model[0].weight.sum()  # a real read waits
+        assert opt.stats()["guard_waits"] == before + 1
+    finally:
+        trainstream.unwrap(model)
+
+
+# --------------------------------------------------------------------------- .to() after wrap
+def test_round_trip_to_cpu_and_back_keeps_training_exact():
+    ref, mine = mlp(), mlp()
+    opt_ref = torch.optim.AdamW(ref.parameters(), foreach=False, **HP)
+    mine, opt = trainstream.wrap(mine, torch.optim.AdamW(mine.parameters(), **HP), offload=True,
+                                 activation_offload=False)
+    try:
+        data = batches(8)
+        a = run(ref, opt_ref, data[:4])
+        b = run(mine, opt, data[:4])
+        mine.to("cpu")
+        assert next(mine.parameters()).device.type == "cpu"
+        with pytest.raises(RuntimeError, match="moved off the GPU"):
+            opt.step()
+        mine.cuda()
+        a += run(ref, opt_ref, data[4:])
+        b += run(mine, opt, data[4:])
+        opt.synchronize()
+        assert b == pytest.approx(a, abs=1e-4)
+        assert max_param_diff(ref, mine) < 1e-4
+        assert any("engine rebuilt" in e for e in trainstream.memory_report(mine)["events"])
+    finally:
+        trainstream.unwrap(mine)
+
+
+def test_cast_after_wrap_keeps_fp32_master():
+    model = mlp()
+    model, opt = trainstream.wrap(model, torch.optim.AdamW(model.parameters(), **HP), offload=True,
+                                  activation_offload=False)
+    try:
+        data = batches(6)
+        run(model, opt, data[:3])
+        master_before = [e["master"].clone() for e in opt.state_dict()["state"].values()]
+        model.to(torch.bfloat16)
+        assert next(model.parameters()).dtype == torch.bfloat16
+        master_after = [e["master"] for e in opt.state_dict()["state"].values()]
+        assert all(torch.equal(x, y) for x, y in zip(master_before, master_after))  # not rounded to bf16
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            losses = run(model, opt, [(x.bfloat16(), y) for x, y in data[3:]])
+        assert all(l == l for l in losses)
+    finally:
+        trainstream.unwrap(model)
+
+
+# --------------------------------------------------------------------------- resident switch
+def test_small_model_switches_to_resident_and_stays_exact():
+    ref, mine = mlp(), mlp()
+    opt_ref = torch.optim.AdamW(ref.parameters(), foreach=False, **HP)
+    mine, opt = trainstream.wrap(mine, torch.optim.AdamW(mine.parameters(), **HP), activation_offload=False)
+    try:
+        data = batches(10)
+        a, b = run(ref, opt_ref, data), run(mine, opt, data)
+        assert opt.resident
+        assert b == pytest.approx(a, abs=1e-4)
+        assert max_param_diff(ref, mine) < 1e-4
+        assert opt.stats()["mode"] == "resident"
+    finally:
+        trainstream.unwrap(mine)
+
+
+def test_offload_true_never_goes_resident():
+    model = mlp()
+    model, opt = trainstream.wrap(model, torch.optim.AdamW(model.parameters(), **HP), offload=True,
+                                  activation_offload=False)
+    try:
+        run(model, opt, batches(4))
+        assert not opt.resident
+    finally:
+        trainstream.unwrap(model)
+
+
+def test_resident_optimizer_is_fused_and_follows_lr_schedulers():
+    model = mlp()
+    model, opt = trainstream.wrap(model, torch.optim.AdamW(model.parameters(), **HP), activation_offload=False)
+    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda step: 0.5 ** step)
+    try:
+        for x, y in batches(4):
+            F.cross_entropy(model(x), y).backward()
+            opt.step()
+            opt.zero_grad()
+            sched.step()
+        inner = opt._resident
+        assert inner is not None and inner.param_groups[0] is opt.param_groups[0]
+        assert inner.param_groups[0]["fused"] is True
+        assert inner.param_groups[0]["lr"] == pytest.approx(HP["lr"] * 0.5 ** 4)
+        assert all(st["step"].is_cuda for st in inner.state.values())
+        opt.load_state_dict(opt.state_dict())  # torch's loader replaces group dicts; we relink them
+        assert inner.param_groups[0] is opt.param_groups[0] and inner.param_groups[0]["fused"] is True
+    finally:
+        trainstream.unwrap(model)