WSL2下PyTorch多进程传递CUDA张量出现未知错误求助

算力阿峰

2026-8-10

解决WSL2中PyTorch多进程传递CUDA张量的报错问题

问题根源

代码逻辑错误：创建Process时错误地将CUDA张量作为参数传入，而非Pipe的连接端，导致子进程启动时需要反序列化CUDA张量。
WSL2的CUDA IPC（进程间通信）实现与原生Linux存在差异，spawn模式下直接传递CUDA张量时，子进程无法正确获取共享的CUDA内存上下文，触发驱动报错。

解决方案

方案一：修正参数传递并显式初始化子进程CUDA上下文

修正代码逻辑，传递Pipe连接端而非直接传CUDA张量，同时在子进程中显式初始化CUDA上下文以适配WSL2环境：

import torch
from torch.multiprocessing import set_start_method, Pipe, Process

def func(conn):
    # WSL2环境下显式初始化CUDA上下文
    torch.cuda.init()
    data = conn.recv()
    print(data)

if __name__ == "__main__":
    set_start_method('spawn')
    parent_conn, child_conn = Pipe()
    data = torch.tensor([1, 2, 3], device='cuda')
    # 传递Pipe的子连接端给子进程
    proc = Process(target=func, args=(child_conn,))
    proc.start()
    parent_conn.send(data)
    proc.join()

方案二：通过CPU中转传递张量

如果WSL2的CUDA IPC仍存在兼容性问题，可以先将CUDA张量转为CPU张量传递，子进程接收后再转回GPU：

import torch
from torch.multiprocessing import set_start_method, Pipe, Process

def func(conn):
    torch.cuda.init()
    data_cpu = conn.recv()
    data = data_cpu.to('cuda')
    print(data)

if __name__ == "__main__":
    set_start_method('spawn')
    parent_conn, child_conn = Pipe()
    data = torch.tensor([1, 2, 3], device='cuda')
    proc = Process(target=func, args=(child_conn,))
    proc.start()
    # 传递CPU张量
    parent_conn.send(data.cpu())
    proc.join()