PyTorch CUDA IPC 引用计数机制深度解析共享 CUDA 张量的跨进程生命周期管理【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读PyTorch 通过torch.multiprocessing支持在不同进程间共享 CUDA 张量底层基于 CUDA IPC 与共享内存而共享显存的生命周期管理则是整个机制正确性的基石。本文以 torch/multiprocessing/cuda_multiprocessing.md 为骨架结合 torch/csrc/CudaIPCTypes.h、torch/csrc/CudaIPCTypes.cpp、torch/csrc/StorageSharing.cpp 与 torch/multiprocessing/reductions.py 的源码实现完整剖析 PyTorch 的 CUDA IPC 引用计数Refcounting实现原理。读完本文你将理解为什么发送共享 CUDA 张量后生产者不能立即释放显存、CudaIPCSentData/CudaIPCReceivedData/CudaIPCSentDataLimbo/CudaIPCRefCountersFile四个核心结构如何协作、引用计数文件与同步事件如何在进程间传递以及torch.cuda.ipc_collect()的触发时机与适用场景。一、问题背景谁拥有共享显存谁来负责释放CUDA 显存与 CPU 共享内存/dev/shm的关键差异在于共享的 CUDA 内存块归属于生产者producer进程。消费者consumer进程通过cudaIpcOpenMemHandle拿到的是对同一块显存的映射但该块内存的分配与释放仍然由生产者进程的 CUDA 缓存分配器CUDACachingAllocator管理。因此必须采取特殊措施确保这块共享显存在共享张量的整个生命周期内始终处于已分配状态绝不能因为生产者进程中的原张量对象离开作用域、被垃圾回收就立即释放——否则消费者进程可能仍在异步读写这块内存造成 use-after-free 与数据损坏。文档给出了一个朴素但存在缺陷的“手动同步”方案作为对照# Producer生产者 queue.put(tensor) event.wait() # Consumer消费者 tensor queue.get() safe_to_use_tensor tensor.clone() event.set()该方案的问题在文档中明确点出阻塞了生产者进程——生产者必须等待消费者完成克隆才能继续执行多消费者场景过于复杂——每个消费者都需要一套独立的同步协议竞态条件难以处理——事件等待/设置的先后顺序、异常路径中的信号丢失等都会引入难以调试的 bug。PyTorch 的正式实现选择了更优雅的路线为共享 CUDA及 HIP张量实现跨进程引用计数cross-process reference counting让生产者进程内存的分配状态自动跟随所有消费者进程的使用状态从而在张量的完整生命周期内保证显存不提前释放。二、核心数据结构四个关键组件整个引用计数机制由 torch/csrc/CudaIPCTypes.h 中定义的四个结构协作完成。下图概括了它们在发送、接收、销毁三条路径中的角色结构所属进程职责CudaIPCSentData生产者包装待发送张量的DataPtr销毁时检查引用计数并决定是否真正释放显存CudaIPCRefCountersFile生产者一个共享内存文件内含多个 64 位引用计数器槽位供多个张量共享CudaIPCSentDataLimbo生产者“临时收容所”保存已被生产者释放但仍被消费者引用的数据块并周期性回收引用计数归零的块CudaIPCReceivedData消费者包装收到的数据销毁时负责递减对应引用计数2.1 CudaIPCSentData生产者侧的数据包装在发送张量的时刻PyTorch 将张量的DataPtr包装进CudaIPCSentData。它仍然指向同一块显存但销毁行为被重写。其成员见 CudaIPCTypes.h包括handle_共享内存引用计数文件的句柄字符串offset_本张量在该计数文件中的槽位偏移counter_ptr_指向共享内存中引用计数器的指针original_ptr_原始的内存分配at::DataPtr真正的显存块持有者event_与event_sync_required_用于进程间同步的 CUDA 事件device_张量所在设备。构造函数CudaIPCTypes.cpp还承担了一个关键职责——建立发送前的流同步在 CUDA 平台上只要全局使用的“同步事件数”未达到上限CUDA_IPC_MAXIMUM_EVENTS_TO_USE 1000就创建一个cudaEventDisableTiming | cudaEventInterprocess | cudaEventBlockingSync事件并记录到当前流上event_sync_required_ true一旦事件数量达到上限退化为直接stream_synchronize当前流event_sync_required_ false。之所以设 1000 这个上限是因为经验测出 CUDAv10.1 及以下对已记录的阻塞式跨进程事件数量有约 22,000 个的非官方限制预留 1000 个余量足够日常共享使用在 HIPROCm平台上由于不支持cuIpcGetEventHandle一律退化为stream_synchronize。2.2 CudaIPCRefCountersFile共享内存中的计数器文件CudaIPCRefCountersFile管理一个共享内存文件通过at::RefcountedMapAllocator映射实际位于/dev/shm。每个文件默认包含CUDA_IPC_REF_COUNTER_FILE_SIZE 10000个计数器槽位每个 8 字节见 CudaIPCTypes.cpp 中的sizeof(int64_t) * 10000。当前实现采用“顺序分配”策略通过next_offset_不断递增偏移依次把下一个可用计数器分配给新的待发送张量。相关方法见 CudaIPCTypes.hcounter_ptr()返回当前偏移处计数器的指针set_counter(value)把当前槽位的计数器初始化为给定值have_offsets()是否还有未分配的槽位next_offset_ size_rotate_offset()分配一个槽位next_offset_used_slots_return_offset(offset)归还一个槽位used_slots_--offsets_in_use()当前是否有槽位仍在使用。当文件内 10000 个槽位用尽have_offsets()为假时next_available_ref_counters_file_会被重置下次发送张量时再新建一个计数文件CudaIPCTypes.cpp。2.3 CudaIPCSentDataLimbo生产者侧的“临终收容所”CudaIPCSentDataLimbo是全局单例cuda_ipc_global_entities的一部分见 CudaIPCTypes.cpp。它保存生产者已不再使用张量已离开作用域、但消费者仍然在使用或将要使用的数据块。其关键方法是collect()CudaIPCTypes.cpp加锁遍历shared_blocks_引用计数counter_value() 0仍大于零的块保留引用计数归零的块移出列表返回freed_memory true在临界区外真正reset()这些块释放原始显存避免持锁时触发销毁逻辑造成死锁。同时add()会统计列表大小一旦超过CUDA_IPC_WARN_AFTER_X_BLOCKS_IN_LIMBO 1000个未回收块就打出一条告警提示“生产者进程尝试释放超过 1000 个被消费者引用的内存块释放可能被显著拖慢”CudaIPCTypes.cpp。2.4 CudaIPCReceivedData消费者侧的包装消费者侧收到数据后将其包装进CudaIPCReceivedData见 CudaIPCTypes.h内部是一个std::shared_ptrvoid持有cudaIpcOpenMemHandle得到的基址。它的析构函数负责递减收到的张量对应的引用计数。三、发送路径从张量到 IPC 描述元组发送路径的入口在 Python 层 torch/multiprocessing/reductions.py 的reduce_tensor()。当 CUDA 张量经过mp.Queue等机制被 pickle 时reduce_tensor调用storage._share_cuda_()即 C 层的THPStorage_shareCuda见 StorageSharing.cpp。THPStorage_shareCuda的核心流程获取 IPC 句柄通过c10::cuda::CUDACachingAllocator::shareIpcHandle()取得cudaIpcMemHandle_t标识生产者进程中的整个cudaMalloc分配块而非仅 storage 本身以及 storage 在该块中的字节偏移包装 DataPtr调用torch::GetNewRefCountedSentData()生成新的带引用计数的at::DataPtr其 deleter 被替换为CudaIPCSentDataDelete原 DataPtr 存入sent_data-set_original_ptr()分配计数器槽位GetNewRefCountedSentData在全局状态中取出next_available_ref_counters_file_把该槽位计数器初始化为 1创建CudaIPCSentData并rotate_offset()占用槽位生成同步事件若event_sync_required_为真则调用cudaIpcGetEventHandle导出事件句柄返回 8 元组(device, handle, size_bytes, offset_bytes, ref_counter, ref_counter_offset, event_handle, event_sync_required)。其中offset_bytes被特别要求必须使用字节而非元素个数因为(storage_handle, offset)会作为shared_cache的键来唯一标识一个 storage见 StorageSharing.cpp 与 reductions.py。GetNewRefCountedSentDataCudaIPCTypes.cpp的完整逻辑是if (!next_available_ref_counters_file_) { // 新建共享内存计数文件10000 个 int64 槽位 std::string ref_counter_handle at::NewProcessWideShmHandle(); at::DataPtr sptr at::RefcountedMapAllocator::makeDataPtr( ref_counter_handle.c_str(), ALLOCATOR_MAPPED_SHAREDMEM | ALLOCATOR_MAPPED_EXCLUSIVE, sizeof(int64_t) * CUDA_IPC_REF_COUNTER_FILE_SIZE, nullptr); // 注册到全局 ref_counters_files_ 映射并设为当前文件 } // 当前槽位计数初始化为 1 next_available_ref_counters_file_-set_counter(1); // 构造 CudaIPCSentData占用槽位 auto sent_data new CudaIPCSentData(handle, offset, counter_ptr, device); next_available_ref_counters_file_-rotate_offset(); if (!next_available_ref_counters_file_-have_offsets()) { next_available_ref_counters_file_.reset(); // 槽位用尽下次新建文件 } return at::DataPtr(data, sent_data, CudaIPCSentDataDelete, device);四、接收路径重建 storage 并挂钩消费者引用计数消费者进程在反序列化时调用rebuild_cuda_tensorreductions.py其流程为查询共享缓存以(storage_handle, storage_offset_bytes)为键查询shared_cache若已存在对应的StorageWeakRef则直接复用 storage并调用storage_cls._release_ipc_counter(ref_counter_handle, ref_counter_offset)释放生产者这次发送带来的新引用计数首次接收调用storage_cls._new_shared_cuda(...)C 层THPStorage_newSharedCuda其内部若event_sync_required为真用收到的cudaIpcEventHandle_t构造at::cuda::CUDAEvent并在当前流上block确保生产者已写完所有数据调用CUDACachingAllocator::getIpcDevPtr(handle)打开 IPC 句柄得到basePtr通过devPtr basePtr storage_offset_bytes偏移出真实 storage 的起始地址构造一个自定义 deleter 的DataPtr销毁时递减ref_counter_offset处的引用计数并stream_synchronize当前流保证相关内核全部结束否则其他进程可能复用内存造成数据损坏。代码注释StorageSharing.cpp指出理想方案是发送“未触发”的 CUDA 事件作为释放判据但 CUDA当时 10.1不支持创建未触发事件且上千个共享事件的性能影响未知因此最终采用“同步 计数递减”的组合方案设置received_cuda(true)禁止该 storage 被 resize——否则消费者可能越界写入缓存分配块中不属于自己的区域见 reductions.py 中的Note [CUDA IPC and the caching allocator]。消费者侧的CudaIPCReceivedData正是持有basePtr的那个shared_ptr其生命周期决定了cudaIpcOpenMemHandle打开的映射何时关闭。五、销毁路径引用计数如何保护共享显存5.1 生产者销毁CudaIPCSentDataDelete当生产者进程中的张量离开作用域时其DataPtr的 deleterCudaIPCSentDataDelete被调用CudaIPCTypes.cppvoid CudaIPCSentDataDelete(void* ptr) { std::unique_ptrCudaIPCSentData sent_data(static_castCudaIPCSentData*(ptr)); if (!CudaIPCGlobalEntities::alive) return; if (sent_data-counter_value() 0) { // 仍有消费者引用 → 进入 Limbo暂不释放显存 cuda_ipc_global_entities.CudaIPCSentDataLimbo_.add(std::move(sent_data)); } cuda_ipc_global_entities.CudaIPCSentDataLimbo_.collect(); }即生产者销毁时并不直接释放显存而是检查共享计数器——只要还有消费者引用就把数据块移入CudaIPCSentDataLimbo由其持有original_ptr_保证显存存活同时立刻触发一次collect()尝试回收引用计数已归零的块。5.2 消费者销毁CudaIPCReceivedData消费者进程中的张量销毁时其 deleterStorageSharing.cpp 中的 lambda执行received_data.shared_ptr_.reset()关闭 IPC 映射stream_synchronize当前流确保所有依赖该存储的异步操作完成重新以ALLOCATOR_MAPPED_NOCREATE打开引用计数共享内存文件执行*(counter_ptr ref_counter_offset) - 1递减计数best-effort生产者若已退出则捕获c10::Error静默忽略。Python 层与之对应的THPStorage_releaseIPCCounterStorageSharing.cpp同样执行递减逻辑用于rebuild_cuda_tensor命中shared_cache的场景。5.3 Limbo 的回收触发时机CudaIPCSentDataLimbo::collect()会在以下事件发生时被调用见 cuda_multiprocessing.md 与 CudaIPCTypes.cpp每次CudaIPCSentDataDelete之后即每次生产者销毁共享张量时CUDA 缓存分配器找不到合适的块而进行下一次分配时——这是通过REGISTER_FREE_MEMORY_CALLBACK(cuda_ipc_collect, CudaIPCCollectCallback)注册的FreeMemoryCallback钩子实现的见 CudaIPCTypes.h任何共享块的删除尝试时显式调用torch.cuda.ipc_collect()时见下文第六节。六、手动回收torch.cuda.ipc_collect()Python 层提供了手动触发回收的 APItorch.cuda.ipc_collect()torch/cuda/init.py它调用torch._C._cuda_ipc_collect()最终落到 C 的CudaIPCCollect()bool CudaIPCCollect() { if (!CudaIPCGlobalEntities::alive) return true; bool freed_memory cuda_ipc_global_entities.CudaIPCSentDataLimbo_.collect(); if (cuda_ipc_global_entities.CudaIPCSentDataLimbo_.size() 0) { cuda_ipc_global_entities.safe_clean_current_file(); } return freed_memory; }其官方文档语义为“在 CUDA IPC 释放显存后强制回收 GPU 内存检查是否有已发送的 CUDA 张量可以从内存中清理若没有活跃计数器则强制关闭用于引用计数的共享内存文件。当生产者进程停止主动发送张量并希望释放未使用的内存时非常有用。”需要注意两点safe_clean_current_file()只在next_available_ref_counters_file_的offsets_in_use() 0没有任何槽位被占用时才从ref_counters_files_映射中删除并重置该文件CudaIPCTypes.cpp整个流程是best-effort的如果消费者进程异常退出、生产者先于消费者释放数据引用计数可能永远无法归零。此时CudaIPCSentDataLimbo的析构函数与CudaIPCGlobalEntities的析构函数会发出告警warnProducerTerminatedBeforeSharedTensorsReleased()CudaIPCTypes.cpp提示“生产者进程在全部共享 CUDA 张量被释放之前已终止”。七、全局状态与进程生命周期引用计数机制依赖一个进程内的全局单例cuda_ipc_global_entitiesCudaIPCGlobalEntities见 CudaIPCTypes.cpp包含ref_counters_mutex_保护计数文件映射的互斥量sync_events_used_已使用的同步事件计数原子变量用于触发事件 → 流同步的退化策略ref_counters_files_handle - CudaIPCRefCountersFile的映射next_available_ref_counters_file_当前正在顺序分配槽位的计数文件CudaIPCSentDataLimbo_临终收容所实例。该类使用静态布尔量alive追踪自身生命周期避免在静态析构阶段被再次访问导致段错误。进程退出时析构函数它会先collect()一次 Limbo再尝试安全清理当前计数文件若清理后仍存在未释放的共享块或计数文件即触发上述“生产者先终止”告警。八、设计权衡与边界情况从源码实现可以归纳出该设计几个值得注意的权衡事件预算 vs 性能优先使用跨进程阻塞事件做同步不阻塞生产者只记录事件超过 1000 个事件上限后退化为流同步保证在 CUDA 事件数量限制下系统仍可用CudaIPCTypes.cpp。引用计数文件顺序分配CudaIPCRefCountersFile顺序递增偏移分配槽位实现简单、无碎片化搜索代价是每个文件固定 10000 个槽位发送量极大时会创建多个计数文件CudaIPCTypes.cpp。Limbo 线性扫描collect()每次全量遍历shared_blocks_代码注释明确标注“TODO: 可改为 FIFO 以避免每次 collect 全量遍历”CudaIPCTypes.h。超过 1000 个块时会告警提示性能下降。共享缓存复用消费者侧通过shared_cacheStorageWeakRef集合见 reductions.py缓存已打开的 IPC 存储。命中缓存时不再重复打开句柄而是直接递减本次发送新增的引用计数避免同一cudaIpcMemHandle被反复 open/close 导致地址漂移。显式禁止 resize通过 IPC 重建的 storage 被设置为不可扩展set_resizable(false)防止消费者越界写入缓存分配块破坏其他张量的数据StorageSharing.cpp。九、相关阅读路径若要继续深入建议按以下顺序阅读仓库源码设计文档torch/multiprocessing/cuda_multiprocessing.md核心数据结构声明与常量 torch/csrc/CudaIPCTypes.h引用计数、Limbo 与全局状态实现torch/csrc/CudaIPCTypes.cppstorage 发送/接收的 C 实现_share_cuda_/_new_shared_cuda/_release_ipc_counter_cudatorch/csrc/StorageSharing.cppPython 层 reducer 与缓存torch/multiprocessing/reductions.py手动回收 API 文档torch/cuda/init.py进程与队列封装torch/multiprocessing/pool.py、torch/multiprocessing/queue.py、torch/multiprocessing/spawn.py【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考