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.4.tar.gz (364.4 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.4-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.4-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.4-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.4-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.4-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.4-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.4.tar.gz.

File metadata

  • Download URL: tensortorrent-0.2.4.tar.gz
  • Upload date:
  • Size: 364.4 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.4.tar.gz
Algorithm Hash digest
SHA256 9be8d9536bd6bf3919b2424c08a3bf6466ca1c7db91bf43b164fbbfae0de7c93
MD5 5f75f5d11fa940d03e11b9652c9eb821
BLAKE2b-256 5613443f117ae9ec9b980f3bd49ce3b9fa5b3cad2fb4694fdcddc97162ca641a

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.4.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.4-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.4-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 ff4f94e5554f6577935531e7c6e6f9fa9134190a14ca93b82b36c5589fd89158
MD5 474552e6a317c5d2c94544ef129842e6
BLAKE2b-256 c2dca5a5852de0290c1b5b5fcbebb136ddbd058e3db7241f5faaeba7d8a0a193

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.4-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.4-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.4-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
Algorithm Hash digest
SHA256 451a5f117154a783b156dc0b4be8ee27d9b608c5ef65df2afdf5e34ff02f926b
MD5 45552187b24ee673ba6eb5e7cab28ff8
BLAKE2b-256 6fec13b66866746ae4162583ca00d8c88fb2c042bcf9515d777a3565d750376f

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.4-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.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 252f3b11c95ef922b78d137a43a8ad625cec9374e0f11497b8d9421c90c5b0e0
MD5 aab288d91230d48ff5cb7ef9785b8c9a
BLAKE2b-256 8130d05910996dd115bc5dd96f98d0330585c9d59972370354c055099719ffcd

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.4-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.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
Algorithm Hash digest
SHA256 66c2b7e8d666681a191a1261bf30d9fd00da2f9a1fabad6be4a08519a07f1ff9
MD5 47bddcdddda0c6e30fa4d8aa681c05ab
BLAKE2b-256 40312b94b3cb75fd9a8913979360c3607b5ae5bc503118e06acfee9af9188791

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.4-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.4-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.4-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 450fa2c4e95e1f38ba87879ff825c7d1c26b9a516289a82fd60e592490d9485b
MD5 d4f095d3f5342e127ea6827244029075
BLAKE2b-256 dc6d3340d53e25c5c5ca95dd2ece25922cee4ac96dae5e1567c69eb07cfef1a7

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.4-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.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for tensortorrent-0.2.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 3404bd2b1cc22ec9983a16b0fb073a76cfef39e139b4524032db673c568f0e96
MD5 4dc69d9cdf72a4643978a6efcc308af4
BLAKE2b-256 37b5fbab8a9a340a95433cdbfec475140df174360494004bca07d489d58e399d

See more details on using hashes here.

Provenance

The following attestation bundles were made for tensortorrent-0.2.4-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