Quickstart

---

As the main aim of managed is to work with multiple GPUs, we will assume that you have access to a machine with multiple GPUs. Suppose you have a tessalated surface and want to solve a PDE on it. Now, the PDE can be shown to be equivalent to a linear system of equations. Consider the simple single GPU code below:

import torch
from concurrent.futures import ThreadPoolExecutor

# CPU Bound
def get_neighbors(x):
    """
    Returns the neighbors of each element in x.
    """
    ....
    return neighbors

# CPU Bound
def construct_pde_matrix(x, neighbors):
    """
    Constructs the PDE matrix for the given x.
    """
    ...
    return pde_matrix

# Can be put on GPU!!
def solve_pde(x):
    """
    Solves the PDE for the given x.
    """
    neighbors = get_neighbors(x)
    # Goes on the default cuda device
    pde_matrix = construct_pde_matrix(x, neighbors).cuda()
    # Same - default cuda device
    x = x.cuda() 
    solution = torch.linalg.solve(pde_matrix, x)
    ...
    return solution

# CPU Bound
executor = ThreadPoolExecutor(max_workers=4)
for _ in range(ITERATIONS):
    futures = []
    for point in tessalation:
        futures.append(executor.submit(solve_pde, point))
    for future in futures:
        ret = future.result()
        # Do something with ret

Here there are many cpu bound operations like finding the neighbors and creating the system. However, the main bottleneck is solving the PDE. To use multiple GPUs, you will have to add extra logic on where the allocation must happen. Some systems might take more time and might keep the GPU busy for longer. A simple polling based approach will only give you a marginal improvement. In addition, if there is a spill over at some point, it will have to be dealt separately.

With managed, you can write the same code as:

from managed import ManagedTensor as mt

# Can be put on GPU!!
def solve_pde(x):
    """
    Solves the PDE for the given x.
    """
    neighbors = get_neighbors(x)
    # Goes to GPU device with least memory usage
    pde_matrix = construct_pde_matrix(x, neighbors).as_subclass(mt).cuda(disperse=True)
    # Same device as pde_matrix
    x = x.as_subclass(mt).cuda(pde_matrix.device)
    # That's it! No need to change anything else
    solution = torch.linalg.solve(pde_matrix, x)
    ...
    return solution

There are two things to note here - the disperse flag and that we give a device for x. The reason we give a device to x is because it is going to be shifted to the GPU with PDE matrices. If we had not done this, there was a chance that x would be allocated on a different device and then we would have to again move it. The disperse flag tells managed to disperse the tensors across the GPUs. This is done by the DeviceManager object. The DeviceManager object is a singleton object and is shared across all the ManagedTensor objects. It is responsible for allocating tensors on the GPUs and moving tensors between the CPU and the GPUs. It also moves tensors between devices (CPU/GPU) according to memory availability.

---

Now let us consider the case where we want to train multiple small models. However, the convergence of the models is not uniform. Some models might converge faster than others. In such a case, again we will face a similar issue as above. We will have to add extra logic to allocate the models to the GPUs. With managed, you can write the code as:

import torch
from managed import ManagedModule as mm
from concurrent.futures import ThreadPoolExecutor

# Inherit ManagedModule instead of torch.nn.Module
class MyModel(mm):
    """
    A simple model that takes in a tensor and returns the sum of its elements.
    """
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(10, 1)

    def forward(self, x):
        return self.linear(x).sum()

def run_experiment(trainloader):
    """
    Runs an experiment on the given trainloader.
    """
    model = MyModel().cuda()
    model.pin() # Pin the model - this is important!
    optimizer = torch.optim.Adam(model.parameters())
    while True:
        net_loss = 0
        for x, y in trainloader:
            optimizer.zero_grad()
            loss = model(x) - y
            net_loss += loss.item()
            loss.backward()
            optimizer.step()
        if net_loss < 1e-3:
            break
    return model

if __name__ == "__main__":
    executor = ThreadPoolExecutor(max_workers=4)
    trainloader = torch.utils.data.DataLoader(torch.randn(100, 10), batch_size=10)
    futures = []
    for _ in range(100):
        futures.append(executor.submit(run_experiment, trainloader))
    for future in futures:
        model = future.result()
        print(model(torch.randn(10)))

Here we have used the ManagedModule class which is a wrapper around the torch.nn.Module class. In case you do not want to inherit ManagedModule, you can use the ManagedModule.from_module method to wrap a preinitialized torch.nn.Module object.

model = mm.from_module(model)

In the above code, we pin the model. Pinning is a property of ManagedTensor objects. When a tensor is pinned, it is not moved between devices. This is useful when you want to keep a tensor on a particular device. Pinning a model is same as pinning all the parameters of the model. The exact reasons why we need to pin the model are explained in the Limitations section. Note that we do not need to pin the model in the previous example as it is a simple linear model. However, there are cases like when the model has a recurrent layer we need to pin the recurrent layer not necessarily the whole model. Pinning the whole model is a safe option but it might not be the most efficient option.

Pinning is required only when training the model. It is not required when we are not relying on the autogard gradient calculation.