Simple parallel GPU inference

aemilius65: <p>I have a trained model and 4 GPUs available. I have a dataset which I want to split in 4 (and process with the same batch size on each GPU) independently of each other and essentially add the results I get from each GPU. I did not get too much wiser from e.g. <a href="https://discuss.pytorch.org/t/how-do-i-run-inference-in-parallel/126757" class="inline-onebox">How do I run Inference in parallel?</a>. Please not that I am only doing inference with my model, and no gradient computations etc are required.</p>
<p>A minimal example of what I’m trying to do is this:</p>
<pre><code class="lang-auto">import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import os
from torch.nn.parallel import DistributedDataParallel as DDP
torch.random.manual_seed(123)

input_dim, out_dim = 10, 1
net = nn.Linear(input_dim, out_dim) # I load my model from a saved state_dict

m = torch.cuda.device_count()
n = 5 # number of data pts sent to each GPU
x = torch.rand((n,input_dim,m)) # full data set

# without parallel processing:
s0 = 0 
for i in range(m):
    s0 += net(x[:,:,i]).sum()
print('s0', s0)

# with parallel processing:    
def example(rank, world_size):
    print('rank', rank)
    # create default process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    # create local model
    model = net.to(rank)
    # construct DDP model
    ddp_model = DDP(model, device_ids=[rank])
    # forward pass
    x_local = x[:,:,rank].to(rank)
    outputs = ddp_model(x_local)
    print(outputs.sum()) # these add up to the desired value s0 
    # but how do I return these values from each GPU process and add them? 

def main():
    world_size = torch.cuda.device_count()
    mp.spawn(example,
        args=(world_size,),
        nprocs=world_size,
        join=True)

if __name__=="__main__":
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29500"
    main()
</code></pre>
<p>I’ve seen examples where people do all of the data/model preparation <em>inside</em> the example function, but I don’t know why.<br>
Also, I feel like there should be a cleaner way of loading the data using DataLoader which would work smoother with DPP, so pointers for this are also very welcome. I also posted this question to stackexchange but got no replies.</p>

H-Huang: <p>Your code example looks right. If you need to add the outputs across GPUs you need to do some sort of communication, e.g. all_reduce to sum across different ranks: <a href="https://pytorch.org/docs/main/distributed.html#torch.distributed.all_reduce" class="inline-onebox" rel="noopener nofollow ugc">Distributed communication package - torch.distributed — PyTorch main documentation</a></p>
<p>For loading the data you can use the regular pytorch dataloader, but in the sampler, specify the DistributedSampler <a href="https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler" class="inline-onebox" rel="noopener nofollow ugc">torch.utils.data — PyTorch 2.3 documentation</a></p>

aemilius65: <p>Edited post with Pipe solution:</p>
<p>Thanks for the pointers. I have found an acceptable solution using mp.Pipe instead, which seems to be needed since I want to get back the results from the individual GPUs to the main scope (do correct me if I’m wrong). Also, there is a mysterious issue (to me) when trying to send tensors to parent_conn, please see below.</p>
<pre><code class="lang-auto">import os, sys 
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP

def get_model_and_data(N=12, num_devices=2):
    torch.random.manual_seed(0)
    torch.set_default_tensor_type('torch.DoubleTensor') # for more exact comparisons
    net = nn.Linear(1, 1)
    x = torch.linspace(0, 1, N).unsqueeze(1)
    if (N % num_devices) != 0:
        sys.exit(f'Full tensor length {N} not evenly divisible over {num_devices} GPUs.')
    return net, x.reshape(num_devices, -1, 1)

def inference(rank, world_size, net, x, conn):
    torch.set_printoptions(8)
    torch.set_default_tensor_type('torch.DoubleTensor')
    dist.init_process_group('nccl', rank=rank, world_size=world_size)
    xloc = x[rank].to(rank)
    model = net.to(rank)
    model = DDP(model, device_ids=[rank])
    output = model(xloc).sum().detach() # some operation I want to do on each GPU (no grads required)
    print('Rank ', rank, 'partial sum:', output) # partial sum from each process 
    r = float(output)
    conn.send((rank, r)) # this works OK 
    #conn.send(output) # this yields warnings when sending and error when recieving
          
def main():
    torch.set_printoptions(8)
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29500"
    net, x = get_model_and_data()
    s0 = net(x).sum() # compute without parallel processing
    world_size = x.shape[0]
    
    parent_conn, child_conn = mp.Pipe()
    mp.spawn(inference,
        args=(world_size, net, x, child_conn),
        nprocs=world_size,
        join=True)
    
    res = []
    while parent_conn.poll():
        c = parent_conn.recv()
        print(f'recieving {c}')
        res.append(c)
    print(res)
        
        
if __name__=="__main__":
    main()
</code></pre>
<p>However, this only works if I cast output to e.g. float. If I do conn.send(output) then I get a warning during sending and eventually an error when receiving. The same behavior results if<br>
output is replaced with simply e.g. torch.ones(1) or torch.ones(1).to(rank). So if anyone knows how to properly send tensors this way I would like to know.</p>
<pre><code class="lang-auto">Rank  0 partial sum: tensor(3.77580110, device='cuda:0')
Rank  1 partial sum: tensor(6.85251166, device='cuda:1')
[W CudaIPCTypes.cpp:15] Producer process has been terminated before all shared CUDA tensors released. See Note [Sharing CUDA tensors]
[W CudaIPCTypes.cpp:15] Producer process has been terminated before all shared CUDA tensors released. See Note [Sharing CUDA tensors]
recieving (0, 3.7758011049971456)
Traceback (most recent call last):...
RuntimeError: CUDA error: invalid resource handle
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
</code></pre>