PyTorch多进程共享CUDA张量时被初始化为零的问题咨询

问题原因 #

• Windows系统中PyTorch多进程默认采用spawn启动方式，会创建全新的Python解释器进程，而非直接复制主进程的内存空间，子进程会重新初始化对象。
• 未显式启用共享内存时，传递GPU张量（或包含GPU张量的Module）给子进程，子进程会创建新的张量实例，而非复用主进程的GPU内存，因此出现全零初始化的异常。

修正代码 #

import torch
import torch.multiprocessing as mp
from torch.nn import Embedding
import time

device = 'cuda' if torch.cuda.is_available() else 'cpu'

def worker(rank, tensor):
    if rank == 0:
        tensor.weight.data[rank, rank] = 99
    if rank == 1:
        time.sleep(1)
        print("worker", rank, tensor.weight)

if __name__ == '__main__':
    tensor = Embedding(2, 5, sparse=False)
    tensor = tensor.to(device)
    # 启用共享内存，确保子进程能访问主进程的GPU张量数据
    tensor.share_memory_()
    print("main", tensor.weight)
    
    processes = []
    for rank in range(2):
        p = mp.Process(target=worker, args=(rank, tensor))
        p.start()
        processes.append(p)
    
    for p in processes:
        p.join()

关键说明 #

• 调用share_memory_()方法：该方法会递归地将Module的所有参数张量注册到共享内存中，确保spawn模式下的子进程能正确访问主进程的GPU张量数据，而非重新初始化。
• 针对独立GPU张量：若传递的是单独的GPU张量而非Module，同样需要调用tensor.share_memory_()来启用共享内存。

内容的提问来源于stack exchange，提问作者Daniel