TorchStore

Distributed, asynchronous tensor storage for PyTorch, built on Monarch.

TorchStore makes it easy to share tensors and model weights across distributed processes. Store and retrieve PyTorch tensors (including DTensors and arbitrary Python objects), exchange state_dicts between actors for workflows like reinforcement learning weight sync, and reshard data across different device meshes — all with automatic transport selection that picks the fastest available path based on topology and hardware availability.

Key Features:

• Async put/get API with batch operations (put_batch/get_batch) and key management (delete, exists, keys)
• DTensor-aware storage with tensor-slice retrieval and resharding across different layouts
• state_dict exchange for checkpoint-style weight sync between actors
• Direct RDMA weight sync for zero-copy GPU-to-GPU model weight transfer (one-hop, no intermediate storage for cases like synchronous RL)
• Automatic transport selection — POSIX shared memory for same-host, RDMA when available, with Gloo and Monarch RPC fallbacks
• Configurable storage strategies — LocalRankStrategy (per-rank volumes), HostStrategy (per-host), extensible for your specific use case

Note: TorchStore supports both Monarch-native launches and SPMD entry points (torchrun / torchx). See SPMD Usage below for the latter.

⚠️ Early Development Warning TorchStore is currently in an experimental stage. You should expect bugs, incomplete features, and APIs that may change in future versions. The project welcomes bugfixes, but to make sure things are well coordinated you should discuss any significant change before starting the work. It's recommended that you signal your intention to contribute in the issue tracker, either by filing a new issue or by claiming an existing one.

Installation

uv (Recommended)

uv venv --python 3.12
source .venv/bin/activate

git clone git@github.com:meta-pytorch/torchstore.git
cd torchstore

# Core development install.
uv sync --extra dev

# CUDA 12.8 + torchcomms 0.2.0 from the PyTorch cu128 index.
uv sync --extra dev --extra cu128

# CUDA 13.0 + torchcomms 0.2.0 from the PyTorch cu130 index.
uv sync --extra dev --extra cu130

pip

If you prefer pip, use the matching PyTorch CUDA index explicitly:

conda create -n torchstore python=3.12
conda activate torchstore

# Base install
pip install -e .

# Development install with CUDA 12.8 wheels
pip install -e '.[dev,cu128]' --extra-index-url https://download.pytorch.org/whl/cu128

# Development install with CUDA 13.0 wheels
pip install -e '.[dev,cu130]' --extra-index-url https://download.pytorch.org/whl/cu130

Regular Installation

pip install 'torchstore[cu128] @ git+https://github.com/meta-pytorch/torchstore.git' \
  --extra-index-url https://download.pytorch.org/whl/cu128

Performance: For the best transfer speeds, install with a CUDA extra (cu128 or cu130) to include torchcomms.

Once installed, you can import it in your Python code:

import torchstore

Usage

Reusing keys

Ordinary tensors can overwrite an existing key with a different shape or dtype. Shared-memory transfers reuse compatible allocations and retire obsolete cached handles when a replacement succeeds or a reader attaches to a new handle. This also bounds cached allocations when keys cycle through delete/put operations.

Callers must finish consuming a key before deleting or recycling it. The shared-memory lock serializes PUT/GET operations sharing one client cache; it does not coordinate independent clients or DELETE. Overwrites do not provide atomic snapshots or batch rollback. Changing a DTensor's shard shape, dtype, or layout under an existing key is unsupported; delete the old key before storing a different layout.

TorchStore APIs are called from within Monarch actors. Each actor interacts with the store through the module-level ts.* functions.

import asyncio

import torch

from monarch.actor import Actor, current_rank, endpoint

import torchstore as ts
from torchstore.utils import spawn_actors


WORLD_SIZE = 4


# In monarch, Actors are the way we represent multi-process/node applications. For additional details, see:
# https://github.com/meta-pytorch/monarch
class ExampleActor(Actor):
    def __init__(self, world_size=WORLD_SIZE):
        self.rank = current_rank().rank
        self.world_size = world_size

    @endpoint
    async def store_tensor(self):
        t = torch.tensor([self.rank])
        await ts.put(f"{self.rank}_tensor", t)

    @endpoint
    async def print_tensor(self):
        other_rank = (self.rank + 1) % self.world_size
        t = await ts.get(f"{other_rank}_tensor")
        print(f"Rank=[{self.rank}] Fetched {t} from {other_rank=}")


async def main():

    # Create a store instance
    await ts.initialize()

    actors = await spawn_actors(WORLD_SIZE, ExampleActor, "example_actors")

    # Calls "store_tensor" on each actor instance
    await actors.store_tensor.call()
    await actors.print_tensor.call()

    await ts.shutdown()


if __name__ == "__main__":
    asyncio.run(main())

# Expected output
# [0] [2] Rank=[2] Fetched tensor([3]) from other_rank=3
# [0] [0] Rank=[0] Fetched tensor([1]) from other_rank=1
# [0] [3] Rank=[3] Fetched tensor([0]) from other_rank=0
# [0] [1] Rank=[1] Fetched tensor([2]) from other_rank=2

SPMD Usage (torchrun / torchx)

If your job is launched with torchrun or torchx rather than from inside a Monarch runtime, use ts.initialize_spmd() instead of ts.initialize(). It reads the standard RANK / WORLD_SIZE / MASTER_ADDR / MASTER_PORT env vars (via ts.spmd.SPMDEnv.from_env()) and bootstraps a fresh Monarch context for you. Cross-rank coordination is the caller's responsibility, typically through torch.distributed.

import asyncio
import os

import torch
import torch.distributed as dist
import torchstore as ts


async def main() -> None:
    rank = int(os.environ["RANK"])
    dist.init_process_group("gloo")

    await ts.initialize_spmd(strategy=ts.LocalRankStrategy())
    try:
        if rank == 0:
            await ts.put("demo", torch.tensor([123], dtype=torch.int64))
            print(f"[rank=0] wrote", flush=True)

        dist.barrier()  # wait for rank 0's write before reading
        value = await ts.get("demo")
        print(f"[rank={rank}] got {value}", flush=True)
    finally:
        dist.barrier()  # wait for readers before tearing down storage
        await ts.shutdown()
        dist.destroy_process_group()


if __name__ == "__main__":
    asyncio.run(main())

Launch with:

torchrun --standalone --nnodes=1 --nproc-per-node=2 your_script.py

To pass an explicit SPMD config instead of reading env vars, construct ts.spmd.SPMDEnv(...) directly and pass it via env=. For a complete runnable example, see example/torchstore_spmd.py.

Use configure_job= to enable Monarch Job API components before the rank 0 client attaches to the store-backed host mesh. The callback runs only on global rank 0:

def configure_job(job):
    job.enable_telemetry()
    job.remote_mount(
        "/path/to/local/source",
        mntpoint="/path/on/workers",
        python_exe=None,
    )


await ts.initialize_spmd(
    strategy=ts.LocalRankStrategy(),
    configure_job=configure_job,
)

API Overview

Beyond put/get, TorchStore exposes batch operations and key-management helpers:

# Batch put/get for efficient multi-key transfers
await ts.put_batch({"key1": tensor1, "key2": tensor2})
results = await ts.get_batch(["key1", "key2"])

# Key management
await ts.exists("key1")        # True
keys = await ts.keys("key")    # ["key1", "key2"] — prefix search
await ts.delete("key1")

# In-place retrieval — writes directly into a pre-allocated tensor
pre_allocated = torch.empty(100, 100)
await ts.get("my_tensor", inplace_tensor=pre_allocated)

Resharding Support with DTensor

TorchStore makes it easy to fetch arbitrary slices of a distributed tensor and to reshard between different meshes. For a full DTensor example, see example/dtensor.py. For end-to-end resharding coverage, see tests/test_resharding_basic.py and tests/test_resharding_ext.py.

class DTensorActor(Actor):
    """
    Example pseudo-code for an Actor utilizing DTensor support.

    See example/dtensor.py for the full actor definition.
    """

    @endpoint
    async def do_put(self):
        # Typical dtensor boiler-plate
        self.initialize_distributed()
        device_mesh = init_device_mesh("cpu", self.mesh_shape)
        tensor = self.original_tensor.to("cpu")
        dtensor = distribute_tensor(tensor, device_mesh, placements=self.placements)

        print(f"Calling put with {dtensor=}")
        # This will place only the local shard into TorchStore
        await ts.put(self.shared_key, dtensor)

    @endpoint
    async def do_get(self):
        # Typical dtensor boiler-plate
        self.initialize_distributed()
        device_mesh = init_device_mesh("cpu", self.mesh_shape)
        tensor = self.original_tensor.to("cpu")
        dtensor = distribute_tensor(tensor, device_mesh, placements=self.placements)

        # TorchStore will use the metadata in the local dtensor to only fetch tensor data
        # which belongs to the local shard.
        fetched_tensor = await ts.get(self.shared_key, dtensor)
        print(fetched_tensor)

State Dict Sync

TorchStore supports sharded state_dict exchange between actors, making it straightforward to synchronize model weights (e.g. learner → generator in RL workflows):

# Learner: publish weights after each training step
await ts.put_state_dict(model.state_dict(), "v0")

# Generator: pull weights into its own model
await ts.get_state_dict("v0", user_state_dict=serving_model.state_dict())

For a sample learner/generator example, see example/torchstore_rl.py.

Direct RDMA Weight Sync

The default (buffered) path already uses RDMA when available. When your use case calls for it, direct_rdma=True bypasses the intermediate StorageVolume entirely — the destination reads directly from the source's GPU memory via one-sided RDMA.

await ts.put_state_dict(model.state_dict(), "policy", direct_rdma=True)
await ts.get_state_dict("policy", user_state_dict=model.state_dict(), direct_rdma=True)

transfer_dtype can cast weights for transfer (e.g. float32 master weights transferred as bfloat16).

Transport Backends

TorchStore automatically selects the best available transport for each transfer. No configuration is needed — the selection happens at runtime:

Priority	Transport	When used
1	POSIX Shared Memory	Client and storage volume are on the same host
2	Monarch RDMA	Cross-host, monarch.rdma available
3	TorchComms RDMA	Cross-host, torchcomms installed
4	Gloo	Cross-host fallback via collective transport
5	Monarch RPC	Universal fallback, always available

Setting TORCHSTORE_TORCH_TRANSPORT_BACKEND to a backend listed by torch.distributed._transport.available_transports() (e.g. nixl) enables the torch.distributed transport, preferred after shared memory and over NIXL. Requires a PyTorch build with torch.distributed._transport. The variable is read at import time. This transport registers every tensor per request, so large state dicts are slower than with NIXL until registrations are cached.

To force a specific transport, pass default_transport_type when constructing a strategy:

from torchstore.transport import TransportType

strategy = ts.LocalRankStrategy(default_transport_type=TransportType.Gloo)
await ts.initialize(num_storage_volumes=N, strategy=strategy)

A custom transport subclasses torchstore.transport.buffers.TransportBuffer. Clients pickle each buffer to its storage volume, and the base class's __getstate__ sets storage_volume_ref to None in that copy, in the reply and in copies made with the copy module, so volume-side code reads its ctx argument and the fields the client set. A subclass hook extends super().__getstate__(), a copy of the instance dict. Where a class after TransportBuffer in the subclass's method resolution order defines its own hook, such as a mixin listed after the transport class, the base hook starts from that hook's state instead, copying a dict and passing anything else through unchanged, handle included. A subclass that keeps fields in __slots__ adds them there and restores them in its own __setstate__. Otherwise slot fields that no hook of the subclass returns are lost in copies, and a __setstate__ that unpacks the base hook's state as Python's (dict, slots) pair fails.

Testing

Pytest is used for testing.

For a quick installation smoke test, run: uv run pytest -q tests/test_store.py -k 'test_basic and MonarchRPC'

This selects the Monarch RPC parametrizations of test_basic without needing the extra transport env-var overrides.

For a more verbose test run with logs, use: TORCHSTORE_LOG_LEVEL=DEBUG uv run pytest -vs --log-cli-level=DEBUG tests/test_store.py::test_basic

License

TorchStore is BSD-3 licensed, as found in the LICENSE file.