Skip to main content

TensorTorrent logo

TensorTorrent

A heterogeneous PyTorch compiler and runtime for one machine with many CPUs, GPUs, and memory tiers.

CI status Latest version tag Python 3.10 to 3.13 Rust 1.85 or newer Apache-2.0 license

TensorTorrent exports a PyTorch model, partitions its graph, places regions across available compute, and runs the resulting schedule through a Rust data plane. Parameters can stream from slower storage and activations can spill when the model exceeds device or host memory.

Python compiles. Rust schedules. One immutable ExecutableArtifact describes the program.

[!IMPORTANT] TensorTorrent is alpha software. The supported target is Linux with Python 3.10–3.13 and PyTorch 2.4 or newer. Validate every deployment machine before serving production traffic. APIs, artifact formats, and env var names may change between releases.

Installation

pip install torch --index-url https://download.pytorch.org/whl/cpu   # or CUDA/ROCm from pytorch.org
pip install tensortorrent

Linux, Python 3.10–3.13, PyTorch ≥2.4. Install torch first if you need a specific build. Wheels: PyPI, Releases. Dev from source: uv + Rust 1.85+ (Quick start).

Quick start

git clone https://github.com/alhussein-jamil/TensorTorrent.git
cd TensorTorrent
make sync
make doctor

Compile a module and compare it with eager PyTorch:

import torch
import torch.nn as nn
import tensortorrent as tt  # import alias: tt

model = nn.Sequential(
    nn.Linear(256, 256),
    nn.ReLU(),
    nn.Linear(256, 10),
).eval()
x = torch.randn(32, 256)

compiled = tt.compile(model, example_inputs=(x,))
torch.testing.assert_close(compiled(x), model(x), check_device=False)

compiled.save("artifact/")
reloaded = tt.load_compiled("artifact/")

Run uv run python examples/public_api_demo.py for hardware discovery, compile, and schedule output in one executable example.

What it handles

Area Implementation
PyTorch export and graph partitioning python/tensortorrent/compile
CPU, CUDA, ROCm, Intel XPU, and plugin discovery python/tensortorrent/backends
Resource budget resolver (host memory, VRAM, CPU, disk) python/tensortorrent/hardware/budget.py
NUMA-aware host allocation and CPU budget enforcement crates/tt-backend-cpu
Scheduling, residency, transfer, stall watchdog, and cancellation crates/tt-runtime
Parameter streaming and activation spill crates/tt-storage
Atomic, checksummed artifact bundles python/tensortorrent/artifact_io.py
Concurrent request serving (HTTP, auth, metrics) python/tensortorrent/serve
Virtual accelerators for deterministic tests crates/tt-backend-virtual

The runtime supports NCCL, RCCL, oneCCL, Gloo, and explicit host-staged collective fallbacks where the installed hardware and libraries allow them.

Architecture

flowchart LR
    M[PyTorch module] --> E[Export and normalize]
    E --> P[Partition and place]
    P --> A[ExecutableArtifact]
    A --> R[Rust dispatcher]
    R --> C[CPU / GPU regions]
    R --> S[Memory / storage tiers]

The Python control plane owns export, normalization, partitioning, region compilation, public APIs, and diagnostics. The Rust data plane owns the artifact, schedule, workers, residency, transfers, storage, cancellation, and telemetry. Torch compute regions may call back into Python; scheduling and data movement remain in Rust.

See the architecture guide for ownership boundaries and backend contracts for extension points.

Module composition

Compile a sequence as one graph to avoid opaque transfers between separately compiled artifacts:

compiled = tt.compile_modules(
    [encoder, projector, decoder],
    example_inputs=(x,),
    names=["encoder", "projector", "decoder"],
)

For branches, joins, structured arguments, or nested outputs, build a ModuleGraph from ModuleNode, GraphInput, and NodeOutput. Invalid names, forward references, and output paths are rejected before export.

Opt-in training

Compilation is inference-only by default. Set allow_training=True to use the same heterogeneous schedule with autograd:

config = tt.CompileConfig(allow_training=True)
compiled = tt.compile(model, example_inputs=(x,), config=config)

optimizer = torch.optim.Adam(compiled.parameters())
compiled.train()
optimizer.zero_grad()
loss = compiled(x).sum()
loss.backward()
optimizer.step()
compiled.eval()

Training cannot currently be combined with NVMe parameter streaming, activation spill budgets, or process workers. See the full product scope for intentional limits.

Does it actually work?

On a single device TensorTorrent reaches eager parity at scale — matching or beating PyTorch on large MLPs and transformers. Eligible resident single-region graphs use the direct path by default. Measured resident CPU+accelerator branch plans can use the same low-overhead path after synchronized timing beats both schedule execution and full fusion (prefer_direct_path; override with TT_DIRECT_PATH=0/1). The product focus beyond that is multi-device placement, parameter streaming, and activation spill.

Measured tables, the same-device harness pin, and open roadmap items live in Benchmarks.

Resource budgets and guardrails

Every memory limit, CPU count, and disk quota flows through a single resolver that reads cgroup v2/v1 limits, live OS availability, and explicit config values — in that precedence order. Containers automatically see their cgroup limits, not host totals. The resolver provenance is shown by tensortorrent doctor.

See Resource budgets and guardrails for the full precedence chain, spill lifecycle, stall watchdog, and worked examples.

Development

make sync                 # create the environment and build the native extension
make check                # lint, types, Rust tests, Python tests, doctor
make audit                # cargo-audit (Rust) + pip-audit (Python)
make coverage             # run tests with coverage gate (Python 3.12)
make native-gate          # native extension smoke and execution checks
make hardware-test        # explicit: may consume most available VRAM or spill space

On a machine with a GPU, run everything that needs real hardware in one go:

bash tools/run_everything.sh     # tests + hardware suite + all benchmarks

It writes logs, JSON, and a SUMMARY.md to bench-results/<timestamp>/. Install the benchmark baselines first with uv sync --extra bench so the ONNX Runtime and Accelerate comparisons run instead of reporting as missing.

CI: PRs and pushes to main (Python 3.10 + 3.13, x86-64/ARM64). Hardware tests are opt-in. See CONTRIBUTING.md.

Repository map

python/tensortorrent/   Python control plane, public API, and serving
crates/tt-*/            Rust IR, runtime, memory, storage, backends, and FFI
tests/                  Unit, integration, end-to-end, property, and hardware tests
docs/                   Product, architecture, deployment, and reference guides
examples/               Small public API programs
bench/                  Runtime and planner comparisons
tools/                  Local quality and native-extension gates
deploy/                 Docker Compose and Kubernetes examples
Dockerfile              CPU-only production container
Dockerfile.cuda         CUDA GPU production container (validate on GPU host before use)

Documentation

Versions and releases

Versions follow Semantic Versioning; tags are vMAJOR.MINOR.PATCH. A tag builds wheels, a GitHub Release, and a PyPI publish — see docs/RELEASING.md.

License

Apache-2.0. See LICENSE.

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

tensortorrent-0.2.3.tar.gz (362.5 kB view details)

Uploaded Source

Built Distributions

If you're not sure about the file name format, learn more about wheel file names.

tensortorrent-0.2.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.4 MB view details)

Uploaded CPython 3.13manylinux: glibc 2.17+ x86-64

tensortorrent-0.2.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl (1.3 MB view details)

Uploaded CPython 3.13manylinux: glibc 2.17+ ARM64

tensortorrent-0.2.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.4 MB view details)

Uploaded CPython 3.12manylinux: glibc 2.17+ x86-64

tensortorrent-0.2.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl (1.3 MB view details)

Uploaded CPython 3.12manylinux: glibc 2.17+ ARM64

tensortorrent-0.2.3-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.4 MB view details)

Uploaded CPython 3.11manylinux: glibc 2.17+ x86-64

tensortorrent-0.2.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.4 MB view details)

Uploaded CPython 3.10manylinux: glibc 2.17+ x86-64

File details

Details for the file tensortorrent-0.2.3.tar.gz.

File metadata

  • Download URL: tensortorrent-0.2.3.tar.gz
  • Upload date:
  • Size: 362.5 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for tensortorrent-0.2.3.tar.gz
Algorithm Hash digest
SHA256 b18a4bbe42ce85607d121875b7b4b1a378a9128672ca13789ca22c8035783929
MD5 c8afd8bab2f171df403929eb9e3f5edc
BLAKE2b-256 27e33b23e1a1e05771dbbc2154dbe734e183caea58b884e17804bdeb9ef5509b

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.3.tar.gz:

Publisher: release.yml on alhussein-jamil/TensorTorrent

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tensortorrent-0.2.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 27aaf40f3a481b2afd3882f24c4e7c5dc0a93c18926c7be6d77978ce5f26a48d
MD5 db47ccb9e1c050aab3b94b3334281d6d
BLAKE2b-256 2f45ddf25110a3da22c4453d59c8f0fe3e0bbd8af470c3ac728592b46e9965ed

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl:

Publisher: release.yml on alhussein-jamil/TensorTorrent

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tensortorrent-0.2.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
Algorithm Hash digest
SHA256 62f85d3c17adb254f9d5e52bb73916cded983a9338a8448814dcf4acb10594ed
MD5 471ca056e5f8f6e1069033ead2417a35
BLAKE2b-256 8a171a4b680c55ee2c54941643c5f16fa9043bac3fc1d579263c3249c2cdb13a

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl:

Publisher: release.yml on alhussein-jamil/TensorTorrent

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tensortorrent-0.2.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 882759368ebbdcb0f4092ab44e0009e9b3b2eeab198c729678acfda26b2541fc
MD5 21cd30bd493cb4b8aabdf8a879bc617e
BLAKE2b-256 8e2e4a8c419680a62496b3c88f90ee74b1180004336ae44c939d271e8dcf5975

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl:

Publisher: release.yml on alhussein-jamil/TensorTorrent

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tensortorrent-0.2.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
Algorithm Hash digest
SHA256 c76b75680d263cc70000943afd945cce0a4cbde8e8f6aace42a6fbfb0537e411
MD5 46921a010de6512c1f02ec6896220e6c
BLAKE2b-256 72e9e5e8089d1b557dfc69cc0bea0042c86a2d45362eb8356e2428feee250037

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl:

Publisher: release.yml on alhussein-jamil/TensorTorrent

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tensortorrent-0.2.3-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.3-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 37b448e53a1a54a5b7207dbb2115c3b301fb2db29a3b122276aab5a971c1a1ea
MD5 6865f15ce3ca0fa0c5f5badcba71e812
BLAKE2b-256 83885f7591728076175c1ca50094ba8c9cad573c405545a1f5b2388cc339a377

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.3-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl:

Publisher: release.yml on alhussein-jamil/TensorTorrent

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tensortorrent-0.2.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 e9a5dac2c0a6046b52d24d0f68ef8e446517febc5b3e91ab8defad14b386aef5
MD5 a03ecb8a7e6486278ce9f5affb74adeb
BLAKE2b-256 c9ea645d1c0332362004134886d2cb368fb509453f61d1225910c02a552df610

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl:

Publisher: release.yml on alhussein-jamil/TensorTorrent

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page