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] Supported target: Linux with Python 3.10–3.13 and PyTorch 2.4 or newer. Validate every deployment machine with tensortorrent validate-hardware before serving production traffic.

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
Native placement planner (Rayon subset + beam search) crates/tt-planner
Discrete-event schedule simulation (batch finalist ranking) crates/tt-runtime/simulator
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 / IR / profile]
    E --> N[Native Rust planner]
    N --> F[Top-K finalist schedules]
    F --> D[Parallel Rust DES]
    D --> W[Compile winner only]
    W --> R[Rust runtime]

TensorTorrent profiles the hardware, generates candidate heterogeneous execution strategies, and uses a native Rust planner to search a large placement space with measured compute/transfer characteristics. It shortlists strong distinct plans, builds real executable schedules, simulates those finalists with a Rust discrete-event simulator, and selects the best feasible candidate for the requested objective — then compiles only the winner and executes across CPUs/GPUs/storage. Planner parallelism is automatic and falls back to serial execution when the search is too small to benefit. Not every combinatorial plan is exhaustively simulated.

The Python control plane owns export, normalization, partitioning, region compilation, public APIs, and diagnostics. The Rust data plane owns placement search, discrete-event simulation, the artifact, schedule, workers, residency, transfers, storage, cancellation, and telemetry. Torch compute regions may call back into Python; planning, 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 and the same-device harness pin 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.7.tar.gz (382.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.7-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.5 MB view details)

Uploaded CPython 3.13manylinux: glibc 2.17+ x86-64

tensortorrent-0.2.7-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl (1.5 MB view details)

Uploaded CPython 3.13manylinux: glibc 2.17+ ARM64

tensortorrent-0.2.7-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.5 MB view details)

Uploaded CPython 3.12manylinux: glibc 2.17+ x86-64

tensortorrent-0.2.7-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl (1.5 MB view details)

Uploaded CPython 3.12manylinux: glibc 2.17+ ARM64

tensortorrent-0.2.7-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.5 MB view details)

Uploaded CPython 3.11manylinux: glibc 2.17+ x86-64

tensortorrent-0.2.7-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.5 MB view details)

Uploaded CPython 3.10manylinux: glibc 2.17+ x86-64

File details

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

File metadata

  • Download URL: tensortorrent-0.2.7.tar.gz
  • Upload date:
  • Size: 382.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.7.tar.gz
Algorithm Hash digest
SHA256 c4d01198657078930fa07edbb9c8d2f09d1062b8bd919516097be1ae14b92dab
MD5 10b9842956e6621d9c747d64ef9b9b78
BLAKE2b-256 e6d8b0d8186401950a3dcde4af9f4310b2e576fa030fe1c67c6a7a6850fe81d7

See more details on using hashes here.

Provenance

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

File metadata

File hashes

Hashes for tensortorrent-0.2.7-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 53cc75b21ade138b0b265d0edef9de8431dbc72a6cf8a56082afd7f927bcfbaa
MD5 cbb139ade8484e728d3d5f7e8966b307
BLAKE2b-256 095f92d1f0b3337e01d276898155a6003ee610ae328fcf74d766a6d156360ef6

See more details on using hashes here.

Provenance

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

File metadata

File hashes

Hashes for tensortorrent-0.2.7-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
Algorithm Hash digest
SHA256 23ad79a11cbd16a1e724d595e9e302bd1b1c2f145df2154e1daec17744ec36b2
MD5 5c6a7c2ebb3d1588d7660f0d33c5bfa3
BLAKE2b-256 967a5ef49949275e1d384741b20a70a8aa292310e0e89ef0f38ff0dda4e936b5

See more details on using hashes here.

Provenance

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

File metadata

File hashes

Hashes for tensortorrent-0.2.7-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 441356185f3c9da5a3f8d536fdcc59a26e02cc6ebae85e1f97f294e10c21aada
MD5 15bd4a82b655a76546460a6cf0b009be
BLAKE2b-256 e77d56bd301ef925da3631be5fb4870bbdfaf84f08649273ea66affa312b113b

See more details on using hashes here.

Provenance

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

File metadata

File hashes

Hashes for tensortorrent-0.2.7-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
Algorithm Hash digest
SHA256 3a810b0acb6e2b2520af83517b0a0717a203f8a52cd04cb7813b67584afb67ce
MD5 3833b08b5b72d49a0c029cdb0b9894ed
BLAKE2b-256 324f5bf40e8c2c25e0bbf38792fccf2ee78106dfb8ea613f529c60eedc0f7efd

See more details on using hashes here.

Provenance

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

File metadata

File hashes

Hashes for tensortorrent-0.2.7-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 860298469762d6895c846c9af7365ec06b7e0374b4a1ac9b5dae5e44aedfa2f0
MD5 faa9eb967a1e418c21943df8ac5a4542
BLAKE2b-256 3bb7d7b93fe78399252ed99884f331fd9a8c97867a1963633025620a20f6643b

See more details on using hashes here.

Provenance

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

File metadata

File hashes

Hashes for tensortorrent-0.2.7-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 6cc2df8c9c04b2d7695d3e6f24570595fb6888cd3c4bf13acd48b73250e72333
MD5 8a6a82c08dad576a76f3d58a6f110629
BLAKE2b-256 38b4cba524e7aa422fee6b7dd0589c6a113723bd875673568a6fa0a085597ea4

See more details on using hashes here.

Provenance

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