trainstream 0.1.0__tar.gz → 0.2.0__tar.gz

{trainstream-0.1.0 → trainstream-0.2.0}/MANIFEST.in

@@ -1,8 +1,6 @@
 include README.md LICENSE NOTICE pyproject.toml setup.py build.bat
-recursive-include csrc *.h *.cpp *.cu
+recursive-include trainstream/csrc *.h *.cpp *.cu
 recursive-include trainstream *.py
-prune dashboard
-prune research
 prune runs
 prune build
 global-exclude *.py[cod] __pycache__ *.pyd *.so

trainstream-0.2.0/PKG-INFO

@@ -0,0 +1,157 @@
+Metadata-Version: 2.4
+Name: trainstream
+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
+License: Apache-2.0
+Project-URL: Source, https://github.com/johnanishg/TrainStream
+Project-URL: Issues, https://github.com/johnanishg/TrainStream/issues
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: C++
+Classifier: Environment :: GPU :: NVIDIA CUDA
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+License-File: NOTICE
+Requires-Dist: ninja
+Provides-Extra: hf
+Requires-Dist: transformers; extra == "hf"
+Requires-Dist: accelerate; extra == "hf"
+Requires-Dist: peft; extra == "hf"
+Dynamic: author
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: project-url
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+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.
+
+```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.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.0 → trainstream-0.2.0}/setup.py

@@ -1,12 +1,13 @@
 """Build script for TrainStream.
 
-    pip install trainstream --no-build-isolation   # needs PyTorch (CUDA) + CUDA toolkit + C++ compiler
-    python setup.py build_ext --inplace            # from a checkout (Windows: build.bat)
+    pip install trainstream                      # pure-Python wheel; _C compiles on first import
+    python setup.py build_ext --inplace          # from a checkout (Windows: build.bat)
+    pip install . --no-build-isolation           # compile _C ahead of time into the install
 
-TrainStream is a PyTorch C++/CUDA extension, so it is compiled against the
-PyTorch already installed on your machine. Install PyTorch first and pass
---no-build-isolation so pip builds against it (an isolated build environment
-would not contain your CUDA-enabled PyTorch).
+The C++/CUDA sources ship inside the package (trainstream/csrc). Without
+PyTorch at build time (pip's default isolated build, CI) this produces a
+pure-Python wheel and trainstream/_build.py compiles the extension against the
+user's PyTorch on first import. With PyTorch available it compiles _C now.
 
 Target architectures come from TORCH_CUDA_ARCH_LIST when set; otherwise the
 compute capability of every visible GPU is used (e.g. "12.0" for Blackwell).
@@ -19,23 +20,12 @@ from setuptools import find_packages, setup
 
 HERE = os.path.dirname(os.path.abspath(__file__))
 
-# Building the compiled module needs torch; packaging the sources (sdist,
-# metadata) does not, so CI machines without a GPU can still publish them.
-COMPILE_COMMANDS = {"build_ext", "bdist_wheel", "build", "install", "develop", "editable_wheel", "bdist_egg"}
-compiling = any(arg in COMPILE_COMMANDS for arg in sys.argv[1:])
-
 try:
     import torch
     from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-except ImportError as exc:
-    if compiling:
-        raise SystemExit(
-            "TrainStream compiles against your installed PyTorch.\n"
-            "  1. Install a CUDA build of PyTorch, e.g.\n"
-            "       pip install torch --index-url https://download.pytorch.org/whl/cu130\n"
-            "  2. Install TrainStream without build isolation:\n"
-            "       pip install trainstream --no-build-isolation\n"
-            f"(import error: {exc})")
+except ImportError:
+    torch = None
+if os.environ.get("TRAINSTREAM_PURE_WHEEL"):  # force the compile-on-import wheel
     torch = None
 
 
@@ -61,12 +51,12 @@ if torch is not None:
     ext_modules = [CUDAExtension(
         name="trainstream._C",
         sources=[
-            "csrc/bindings.cpp",
-            "csrc/memory_pinner.cpp",
-            "csrc/offload_adamw.cpp",
-            "csrc/offload_adamw_kernel.cu",
+            "trainstream/csrc/bindings.cpp",
+            "trainstream/csrc/memory_pinner.cpp",
+            "trainstream/csrc/offload_adamw.cpp",
+            "trainstream/csrc/offload_adamw_kernel.cu",
         ],
-        include_dirs=[os.path.join(HERE, "csrc")],
+        include_dirs=[os.path.join(HERE, "trainstream", "csrc")],
         extra_compile_args={"cxx": cxx_flags, "nvcc": nvcc_flags},
     )]
     cmdclass = {"build_ext": BuildExtension.with_options(use_ninja=True)}
@@ -76,7 +66,7 @@ with open(os.path.join(HERE, "README.md"), encoding="utf-8") as fh:
 
 setup(
     name="trainstream",
-    version="0.1.0",
+    version="0.2.0",
     author="John Anish G",
     license="Apache-2.0",
     description="Zero-dependency PCIe offloading engine for PyTorch training (optimizer, gradients, activations).",
@@ -86,10 +76,13 @@ setup(
     project_urls={"Source": "https://github.com/johnanishg/TrainStream",
                   "Issues": "https://github.com/johnanishg/TrainStream/issues"},
     packages=find_packages(include=["trainstream", "trainstream.*"]),
+    package_data={"trainstream": ["csrc/*.h", "csrc/*.cpp", "csrc/*.cu"]},
     ext_modules=ext_modules,
     cmdclass=cmdclass,
     python_requires=">=3.9",
-    install_requires=[],  # torch is a build+runtime prerequisite, not a declared dependency
+    # torch is not declared: users must pick the CUDA build that matches their driver.
+    # ninja drives the compile-on-first-import build.
+    install_requires=["ninja"],
     extras_require={"hf": ["transformers", "accelerate", "peft"]},
     classifiers=[
         "Development Status :: 3 - Alpha",

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