PyTorch中基于CUDA IPC的进程间Tensor共享简介

本文作者kaiyuanPyTorch中基于CUDA IPC的进程间Tensor共享简介想深耕AI Infra领域欢迎访问InfraTech库内容涵盖大模型基础、PyTorch/vLLM/SGLang框架入门、性能加速等核心方向配套50知识干货及适合初学者的notebook练习:https://github.com/CalvinXKY/InfraTech在大模型训练/推理过程中经常会遇到跨进程共享GPU Tensor的需求例如KV Cache跨卡共享、强化学习RL训练/推理间的权重传输等。在PyTorch及其上层框架中这类GPU Tensor的跨进程共享通常可以通过CUDA IPCInter-Process Communication实现从而避免不必要的数据拷贝提升吞吐并节省显存。本文将围绕IPC的基本机制介绍跨进程共享Tensor的核心原理、常见使用方式以及可运行的示例同卡与跨卡两种情形。场景一KV Cache跨卡共享在Attention计算中跨GPU共享KV Cache是一个常见需求例如构建分布式KV Cache、在多卡推理中复用历史token的缓存等。跨GPU的数据访问/传递方式通常包括通过host内存中转传递GPU间P2PPeer-to-Peer直接访问相较于host中转P2P的优势在于可以减少一次或多次数据落地与拷贝降低延迟在满足拓扑与权限条件时也能获得更高的带宽利用率。场景二RL训练/推理之间的权重传输在强化学习训练中训练进程Trainer与推理/采样进程Rollout/Gemerator之间往往需要频繁传递最新权重用于生成数据并迭代更新。当采用训推共卡co-location或多进程共享同一GPU的部署方式时理想状态是实现零冗余传输尽量避免为权重创建额外副本以提升GPU显存利用率并减少更新传播的时间开销。1 CUDA IPCCUDA IPC是CUDA提供的跨进程共享GPU资源的机制。简单来说它允许两个或多个独立进程访问同一块GPU内存或共享同一个CUDA event以支持进程间的高效数据交换与同步。进一步地如果硬件与运行环境支持P2P则GPU之间能够直接访问对方的显存并完成数据传输。将P2P与IPC结合后可以更自然地支持“跨进程跨卡”共享GPU资源的场景。更通俗地讲CUDA IPC的核心不是“拷贝数据到另一个进程”而是发送进程先在GPU上cudaMalloc一块显存发送进程把这块显存导出成 cudaIpcMemHandle_t句柄接收进程拿到句柄后在自己的进程上下文里“打开”这块显存两个进程都能访问同一块物理显存只是各自进程里的虚拟地址可能不同。因此接收进程对这块内存做原地写入发送进程随后读回时会看到变化。1.1 CUDA IPC关键接口与用途cudaIpcGetMemHandle(handle,d_ptr)发送端把cudaMalloc出来的设备指针导出为IPC句柄。cudaIpcOpenMemHandle(base_ptr,handle,cudaIpcMemLazyEnablePeerAccess)接收端根据句柄打开共享显存得到本进程可用的设备指针。cudaIpcCloseMemHandle(base_ptr)接收端使用结束后关闭映射句柄。cudaDeviceSynchronize()两端都使用确保 kernel 执行完成避免异步导致的“看起来没改成功”。下面将讨论IPC的两类典型用法同卡条件下的进程数据共享以及跨卡条件下的进程数据共享。1.2 同卡数据共享多个进程协作使用同一块GPU内存数据时可以采用“发送-接收”的模式。如下图所示有两个独立进程一个作为发送端Sender另一个作为接收端Receiver。整体流程可以概括为Sender创建并初始化数据Sender将用于访问该数据的IPC句柄传递给ReceiverReceiver基于IPC句柄映射并读写该数据C用例演示中Sender初始化一个数据Receiver修改同地址数据最后让Sender校验具体步骤Sender发送端在GPU上分配float数组用init_kernel初始化为0,1,2,...导出IPC句柄并写入文件等待receiver完成标记把数据从GPU拷回host检查是否变成100,101,102,...。Receiver接收端等待sender写好句柄文件读取句柄并cudaIpcOpenMemHandle打开共享显存运行add_kernel(...,100)原地修改这块共享显存关闭句柄并写完成标记。代码https://github.com/CalvinXKY/BasicCUDA/blob/master/common_methods/cuda_ipc/cuda_ipc_demo.cu运行方式Linux编译nvcc -stdc14 -O2 cuda_ipc_demo.cu -o cuda_ipc_demo终端 1接收端./cuda_ipc_demo --mode receiver --prefix /tmp/ipc_demo终端 2发送端./cuda_ipc_demo --mode sender --prefix /tmp/ipc_demo如何判断“确实共享了同一块显存”? 看发送端最终输出若输出类似values after receiver modify:100,101,...并显示PASS说明接收端对共享显存的原地修改被发送端直接观察到了这就证明了是同一块GPU内存被跨进程共享而不是两边各自有一份副本。输出示例1.3 跨卡IPC远程读写示例与同卡过程有所差异sender在src-device例如GPU0分配显存并导出IPC句柄receiver在access-device例如GPU1打开该句柄receiver在GPU1上发kernel直接修改GPU0上的那块共享显存远程写sender在GPU0读回结果并校验变化。代码位置[1]https://github.com/CalvinXKY/BasicCUDA/blob/master/common_methods/cuda_ipc/cuda_ipc_cross_gpu_demo.cu输出内容receiver日志会显示已在access-device完成远程写sender日志会显示读回值由i变为i1000最终输出PASS:cross-GPU IPC remote write verified即可证明跨卡IPC远程读写成功。2 Pyorch进程共享Tensor示例在了解CUDA IPC的数据共享基本过程后再看Pytorch中利用IPC共享Tensor的过程。在PyTorch中封装了CUDA的接口也就能够调用对于IPC功能。值得注意的是在CUDA C中传递数据主要关注三个信息假设数据是连续的首地址、数据长度数量、数据类型。而Tensor的包含的信息会更多所以两个进程共享Tensor需要传递Tensor类相关的基础数据如果stride、storage_offset、ref_counter_offset、ref_counter_handle等。这里介绍两种方式隐式传递和显示传递。通过构造一个示例如下图所示sender创建一个Tensor数据sender将数据地址给到recevierrecevier修改数据内容sender读取已修改数据验证是否被改变。2.1 隐式方式使用torch.multiprocessingspawn通过Queue/Pipe直接传递cudaTensor。序列化时会走与pickle相同的reduction底层对UntypedStorage调用_share_cuda_()/_new_shared_cuda()即CUDA IPC。主要代码def _receiver_queue(q_dev: mp.Queue, q_tensor: mp.Queue, q_ack: mp.Queue) - None: raw q_dev.get() if not isinstance(raw, int): raise TypeError( fexpected int CUDA device index on q_dev, got {type(raw).__name__} ) device_idx raw _cuda_startup(device_idx) t: torch.Tensor q_tensor.get() print( f[receiver] 收到 tensor: shape{tuple(t.shape)} fdevice{t.device} dtype{t.dtype} ) print(f[receiver] mean{float(t.mean()):.6f}) # 与发送方共享同一块显存发送方仍持有引用时此处原地修改对发送方可见 t.add_(1.0) q_ack.put(ok) def run_queue_demo() - None: 主进程创建 Tensor 放入 Queue子进程取出并原地加 1。 if not torch.cuda.is_available(): print(CUDA 不可用跳过 Queue 示例。, filesys.stderr) return _ensure_spawn() ctx mp.get_context(spawn) q_dev: mp.Queue ctx.Queue() q_tensor: mp.Queue ctx.Queue() q_ack: mp.Queue ctx.Queue() p ctx.Process(target_receiver_queue, args(q_dev, q_tensor, q_ack)) p.start() _cuda_startup(0) x torch.randn(4, 5, devicecuda, dtypetorch.float32).detach() dev int(x.device.index if x.device.index is not None else 0) print(f[sender] 发送前 mean{float(x.mean()):.6f} devicecuda:{dev}) q_dev.put(dev) q_tensor.put(x) ack q_ack.get() print(f[sender] 子进程回执: {ack}) print(f[sender] 接收端 add_(1) 后的 mean{float(x.mean()):.6f}) p.join()2.2 显式方式手动从TypedStorage._share_cuda_()取出元数据再在子进程里调用内部的rebuild_cuda_tensor还原Tensor与 PyTorch 多进程序列化逻辑一致便于对照到底传了哪些IPC字段。主要代码def _receiver_explicit(conn: Connection) - None: from torch.multiprocessing.reductions import rebuild_cuda_tensor _cuda_startup(0) payload conn.recv() t rebuild_cuda_tensor(*payload) print( f[receiver/explicit] 还原 tensor: shape{tuple(t.shape)} fsum{float(t.sum()):.4f} ) conn.send(done) def run_explicit_ipc_demo() - None: 主进程手动调用 storage._share_cuda_()把与 reduce_tensor 相同的参数包 发给子进程用 rebuild_cuda_tensor 还原。发送后须保持原 tensor 存活直至 子进程完成打开 IPC此处用 Pipe 同步。 if not torch.cuda.is_available(): print(CUDA 不可用跳过显式 IPC 示例。, filesys.stderr) return from torch.multiprocessing.reductions import ( StorageWeakRef, rebuild_cuda_tensor, shared_cache, ) _ensure_spawn() ctx mp.get_context(spawn) parent_conn, child_conn ctx.Pipe() p ctx.Process(target_receiver_explicit, args(child_conn,)) p.start() _cuda_startup(0) tensor torch.arange(12, dtypetorch.float32, devicecuda).view(3, 4).detach() storage tensor._typed_storage() ( device, handle, storage_size_bytes, storage_offset_bytes, ref_counter_handle, ref_counter_offset, event_handle, event_sync_required, ) storage._share_cuda_() # 与 torch.multiprocessing.reductions.reduce_tensor 中 CUDA 分支一致 shared_cache[handle] StorageWeakRef(storage) payload ( type(tensor), tensor.size(), tensor.stride(), tensor.storage_offset(), type(storage), tensor.dtype, device, handle, storage_size_bytes, storage_offset_bytes, tensor.requires_grad, ref_counter_handle, ref_counter_offset, event_handle, event_sync_required, ) parent_conn.send(payload) parent_conn.recv() # 等子进程 rebuild 完成后再退出作用域 p.join()完整代码地址[2]https://github.com/CalvinXKY/InfraTech/blob/main/deeplearning_framework/torch_process_share_tensor.ipynb执行输出用例[2]里面还包含了一个跨脚本示例演示两个进程之间数据互换过程。运行方式开两个终端先启动接收端python cuda_ipc_cross_script_demo.py –mode receiver –host 127.0.0.1 –port 29531 再启动发送端python cuda_ipc_cross_script_demo.py –mode sender –host 127.0.0.1 –port 29531 –device 0预期现象接收端收到 CUDA tensor 后执行 t.add_(1.0)发送端在收到 ack 后能看到本地 x 的均值同步变化说明两进程共享同一块显存PyTorch工程实践上常用隐式方案。其依托PyTorch与CUDA运行时的封装例如多进程启动、Tensor在进程间的传递等多数场景下开发者无需直接操作IPC句柄由框架与底层完成映射与生命周期管理上手成本更低也是日常训练/推理代码中的主流用法。在需要精细控制或与非PyTorch组件对接时可显式使用CUDA IPC相关能力但若深度绑定PyTorch内部或未公开接口可移植性与升级兼容性风险会上升因此多见于定制infra或实验性方案使用面相对较窄。参考:[1]https://github.com/CalvinXKY/BasicCUDA/blob/master/common_methods/cuda_ipc/cuda_ipc_cross_gpu_demo.cu[2]abhttps://github.com/CalvinXKY/InfraTech/blob/main/deeplearning_framework/torch_process_share_tensor.ipynbInfraTech申明未经允许不得转载