Repository: pytorch
Issue: [RFC] TorchStore - A Shared-Memory Tensor Store
Labels: feature,triaged

## 🚀 Feature

TorchStore is a key-value store that holds ATen tensors in shared memory so that they can be accessed across process boundaries without any expensive copy operations.

## Problem

With the trajectory of ever-increasing model sizes the conventional handling of inference loads by multiple worker processes has started to become a memory bottleneck. Since each worker requires its own instance of the model (or models), the memory use proportionally increases with the number of workers. Considering that today’s multi-core systems can easily handle a large number of workers, this means a big portion of physical memory potentially goes wasted with redundant copies of the same model. Some users, as described [here](https://medium.com/ibm-data-ai/how-to-load-pytorch-models-340-times-faster-with-ray-8be751a6944c), have successfully leveraged Apache Arrow’s [Plasma Object Store](https://arrow.apache.org/docs/python/plasma.html) to circumvent the problem, but the overall user experience is less than ideal and a production-ready solution requires quite a bit of manual work. 

Similarly training tasks in various machine learning fields such as graph neural networks or knowledge distillation happen to use read-only models during training. The same problem as in inference exists here as well. Many redundant copies of the same model waste the precious memory resources of the host machine. The problem is even more amplified in the training case since the memory is typically already under pressure due to the optimizer and backpropagation states. [This issue](https://github.com/PyTorchLightning/pytorch-lightning/issues/8230) in the PyTorch Lightning repository describes a use case that suffers from this problem.

Yet another negative implication is the load times of such replicated models. Some large models can take up to several minutes to load and during that time they typically saturate the I/O bandwidth of the host machine. Attempting to load the same model concurrently by many worker processes will only prolong the load times, most of the time significantly. The same [blog post](https://medium.com/ibm-data-ai/how-to-load-pytorch-models-340-times-faster-with-ray-8be751a6944c) that was mentioned earlier in the first problem statement describes this problem in more depth.

**Why not `Tensor.share_memory_()`?**

Since PyTorch already provides the [share_memory_()](https://pytorch.org/docs/stable/generated/torch.Tensor.share_memory_.html) method as part of its Tensor API, a legitimate question is why we need a separate solution. There are two major reasons:

* In training tasks, in order to leverage `share_memory_()`, one has to control both the parent process and the worker processes. The typical implementation loads a tensor or a module in the parent process and then shares it with the spawned worker processes via `torch.multiprocessing`. However with TorchElastic, which is now part of PyTorch, and the upcoming `torchrun` launcher script, the task of spawning worker processes is handled by PyTorch itself. This means user-code has no chance to influence the initialization logic.
* For inference workloads the lifetime of worker processes are not bound by a parent process. They are independently initiated and terminated, usually by a 3rd party request coordinator. In such setting `share_memory_()` is of no use.
* ~With `share_memory_()` a tensor can only be shared from the parent process to its child processes. If a child process wishes to share a tensor with the rest of its peers, there is no mechanism for it.~ This turned out to be not true. See Edward's feedback below.

## What is TorchStore?

TorchStore is a key-value store that holds ATen tensors in shared memory. The tensors can be accessed across process boundaries (e.g. between workers) without any expensive serialization or de-serialization. In other words sharing tensors between processes via TorchStore is a zero-copy operation.

### Starting the Store

The store can be run either as a daemon or as an in-process service. In order to run it as a daemon the `torchstored` program can be used:

```bash
$ torchstored --name my-store

# Or alternatively run in foreground. This allows easy monitoring in terminal.
$ torchstored --name my-store --no-fork
```

The daemon mode is useful if a persistent store is desired that is not bound to any particular task. For instance an inference server can start the daemon and bootstrap it by loading models and tensors (e.g. large embedding tables). The ephemeral worker processes can simply connect to the store and get access to its entries.

In contrast the in-process mode is more appropriate if the entries held by the store are specific to a particular task. For instance if the store holds a teacher model for a knowledge distillation task that is run on a multi-tenant cluster, there is no need to preserve it after the training is finished. The in-process service can be started via a simple API call.

```python
import torchstore

store = torchstore.start(name="my-store")
```

### Connecting to the Store

To connect to a store, simply call `connect()`  using the name of the store. In order to support use cases where the store is started by one of the worker processes (i.e. in-process mode) while its peers attempt to connect to it, `connect()` has built-in retry logic and can seamlessly handle connection attempts made to a store that has not started yet, but is about to start.

```python
import torchstore

store = torchstore.connect(name="my-store")
```

### Saving to the Store

The store supports both CPU and CUDA tensors. A store entry can hold a tensor, a list of tensors, or a map of string-tensor pairs. Moreover saving multiple tensors as part of an entry (i.e. a list or a map) preserves their view relationships, meaning a shared “storage” is only saved once. One limitation though is that, similar to `Tensor.share_memory_()`, a tensor in store cannot be resized.

The `save()` API is fairly straightforward:

```python
import torch
import torchstore

cpu_tensor = torch.ones([10, 10])

gpu_tensor = torch.ones([10, 10], device="cuda:0")

list_of_tensors = [torch.rand([10, 10]) for _ in range(4)]

dict_of_tensors = {
    "tensor1": torch.rand([10, 10]),
    "tensor2": torch.rand([10, 10]),
}

store = torchstore.connect(name="my-store")

store.save("entry1", cpu_tensor)
store.save("entry2", gpu_tensor)
store.save("entry3", list_of_tensors)
store.save("entry4", dist_of_tensors)
```

Although using the `save()` API as demonstrated above is perfectly fine, there is one issue with that example. Each tensor is first initialized locally and then copied over to the store. This means, albeit temporarily, twice as much memory is required to save a tensor. The `storage()` API demonstrated below mitigates this problem. It temporarily redirects PyTorch’s CPU and CUDA allocators to their store versions so that later `save()` calls won’t require any copy operation.

```python
import torch
import torchstore

store = torchstore.connect(name="my-store"):

with store.storage():
    # The storage for this tensor is allocated from the store.
    tensor_alloc_in_store = torch.ones([10, 10])
    
# This is now a zero-copy operation.
store.save("my-tensor", tensor_alloc_in_store)
```

If a tensor gets allocated from the store as demonstrated above, but does not get saved into the store via a `save()` call, its storage will be deallocated once it goes out of scope. This means there is no harm in allocating from the store, but no gain either. The general advice is to save tensors allocated in the scope of `storage()` to the store.

Another difficulty with a regular `save()` call is the management of an entry's lifetime. In particular in publish/subscribe problems where an entry has no specific owner, determining the right time to delete it from the store poses a problem. To mitigate that problem, the `save()` API accepts a time-to-live argument that specifies how long an entry should be kept in store once its external reference count drops to zero. *Thanks to @VoVAllen for bringing up this problem.*

```python
import torch
import torchstore

store = torchstore.connect(name="my-store")

# Delete 'my-tensor' from the store after 30 seconds unless a client
# retrieves it in the meantime.
store.save("my-tensor", torch.ones([10, 10]), ttl=30)
```

### Saving a Model to the Store

Since a model’s `state_dict` is a map of string-tensor pairs, saving a model to the store resembles very much saving it to a file.

```python
import torch
import torchstore

class MyModel(torch.nn.Module):
    ...
        
store = torchstore.connect(name="my-store")

with store.storage():
    m = MyModel()

# Save the state dictionary of the model to the store.
store.save("my-model", m.state_dict())
```

### Loading from the Store

In order to retrieve a store entry, the `load()` API should be called. Its simplest use is demonstrated below.

```python
import torchstore

store = torchstore.connect(name="my-store")

tensor = store.load("my-tensor")
```

Beyond its basic usage, `load()` also accepts a `timeout` argument. If the specified timeout is greater than zero, the call will wait until either the entry is in the store or the request times out. This mechanism helps peer processes (i.e. workers) to coordinate; one process can add a new entry while the other waits for the entry to appear in the store. A timeout value of zero (default) means to fail immediately if the entry cannot be found. A value of -1 means to wait indefinitely.

```python
import torchstore

store = torchstore.connect(name="my-store")

# Wait up to 30 seconds for 'my-tensor' to show up in the store.
tensor = store.load("my-tensor", timeout=30)
```

Besides a single entry, the `load()` API can also be used to load a set of entries at the same time. This reduces the amount of IPC communication required and is also helpful to specify a single timeout for a group of logically related entries. *Thanks to @yifuwang for this feature suggestion.*

```python
import torchstore

store = torchstore.connect(name="my-store")

# Load the entries 'entry1', 'entry2', and 'entry3' at the same time. Wait up
# to 30 seconds for all three entries to show up in the store.
list_of_tensors = store.load(["entry1", "entry2", "entry3"], timeout=30)
```

Closely related to `load()` is the `contains()` API that, as its name suggests, simply checks whether the store contains a specified entry.

```python
import torchstore

store = torchstore.connect(name="my-store")

if store.contains("my-tensor"):
    ...
```

### Loading a Model from the Store

A model can be loaded from the store similar to how its loaded from a file. A state_dict that was saved earlier via `save()` can be retrieved and passed to the model.

```python
import torch
import torchstore

store = torchstore.connect(name="my-store")

class MyModel(torch.nn.Module):
    ...
    
m = MyModel()

m.load_state_dict(store.load("my-model"))
```

### Deleting from the Store

The store internally uses a reference counting mechanism. Removing an entry from the store decrements its reference count and marks the entry as removed so that it cannot be retrieved anymore. However entries that are already in use (i.e. loaded in processes) will be kept alive until they go out of scope.

```python
import torchstore

store = torchstore.connect(name="my-store")

tensor = store.load("my-tensor")

# The store will mark 'my-tensor' as removed and release its internal reference; however
# the outstanding references will keep the entry alive.
store.remove("my-tensor")

# Once the reference count of `tensor` drops to zero with the next statement, the
# store will completely discard and deallocate the entry.
tensor = None
```
## Shortcoming

One shortcoming of PyTorch is the lack of immutable tensors which requires users of TorchStore to be extra cautious. Ignoring the slight nuance between a read-only versus immutable tensor; if a tensor gets accidentally modified, the changes will be immediately reflected in other processes. There is unfortunately no mechanism to prevent this kind of behavior. Having said that, this “feature” can be helpful in certain scenarios where one writer process is allowed to modify the shared tensor, while others treat it as read-only (this use case also clarifies the difference between read-only and immutable).

## Alternatives

For training tasks one alternative is to keep using `Tensor.share_memory_()` and extending TorchElastic with some form of pre-processing hook. The user will provide a function that performs the necessary initialization in the launcher process and TorchElastic will forward the return value of the function to the worker processes via `torch.multiprocessing`. Unfortunately this alternative won’t address the problem for inference tasks or for frequently-run tasks that want to store models or tensors persistently.

Another alternative is to leverage Apache Arrow’s Plasma Object Store. The advantage is that we won’t need to implement our own store and will simply wrap Plasma. However this alternative also has its disadvantages: 1) no support for the `storage()` API, saving a tensor to the store will always require double-allocation, 2) no support for in-process store meaning the user has to manually start a separate process for the store regardless of the use case, 3) dependency on the Apache Arrow library which has a large dependency closure.

## Design Considerations

* TorchStore will be available only for Linux and macOS systems. Since the implementation is highly OS-dependent, there are no plans to support non-POSIX systems (e.g. Windows).
* There will be Python and C++ APIs that follow the conventions, idioms, and guidelines of the PyTorch code base.
* The project will be hosted in its own repository for easier iteration and faster release cadence. Similar to the evolvement of the TorchElastic project, this decision can be re-considered once the project reaches certain maturity.

cc @VitalyFedyunin @mrshenli @pritamdamania @rohan-varma @kiukchung @bowangbj @aivanou @SciPioneer @H-Huang @tchaton @rusty1s @BarclayII @VoVAllen @houseroad @anj-s @divchenko