Shared variable in GPU

[Yi_Zhang]

Hi,

I want to have two parallel processes in one GPU, one for training (calculating) and the other for communicating parameter updates with other GPUs. And both processes can modify a shared variable, (like a buffer to store the most updated parameters). Anyone knows is it possible to do this?
I checked this documentation: https://pytorch.org/docs/stable/notes/multiprocessing.html, and it mentions “multiprocessing.Queue”, not sure is it suitable in my case? Or any good examples?

[mrshenli (Shen Li)]

This should be doable, see the example code in this post:

`Exception: process 0 terminated with exit code 1` error when using `torch.multiprocessing.spawn` to parallelize over multiple GPUs

You can use share_momery_() and torch.multiprocessing.SimpleQueue to implement IPC. E.g.: import numpy as np import torch import torch.multiprocessing as mp def func(rank, x, p2c, c2p): x_power = x.to(rank) ** rank c2p.put(x_power) # citing multiprocessing doc: Unlike CPU tensors, the # sending process is required to keep the original tensor # as long as the receiving process retains a copy of # the tensor. The refcounting is implemented under the # hood but re…