PyTorch多进程推理中的进程间权重共享
内容提要
本文介绍如何在PyTorch多进程推理中通过CUDA IPC共享模型权重，避免GPU显存重复占用。示例展示了使用`model.share_memory()`共享权重，以及AOTInductor通过`load_constants`绑定父进程权重的方法。文中还讨论了配置要点、TorchScript的局限及程序依赖的重复问题，旨在提升单GPU多进程推理效率。
延伸解读
显存节省的关键机制
示例中通过`model.share_memory()`将模型权重标记为可共享，使子进程通过CUDA IPC直接引用父进程的显存地址，避免复制权重。AOTInductor模式下，通过设置`package_constants_in_so`为False，防止权重打包进.so文件，再在子进程中用`load_constants`绑定父进程的state_dict，从而复用同一份显存。实测中，若不共享，24个worker将超出16GB显存，而共享后设备显存占用接近单份模型大小。
AOTInductor的配置要点
构建AOTInductor包时，需将`aot_inductor.package_constants_in_so`设为False，避免权重嵌入.so文件导致每个进程加载时重复分配显存。同时，`load_constants`的`user_managed`参数设为False，表示权重生命周期由用户管理，runner仅持有指针，避免额外拷贝。若设为True，runner会管理引用计数，但不会复制权重，跨进程引用计数有效，确保权重在所有引用释放前不被释放。
TorchScript的局限与程序依赖开销
TorchScript虽可类似共享权重，但需先在CPU实例化模型再替换权重，导致CPU内存膨胀，除非顺序处理，因此示例未包含。此外，即使权重共享，每个进程仍会加载独立的程序依赖（如共享库），但操作系统会通过mmap共享物理内存页，减少重复开销。不过，每个进程仍有一定的主机和GPU内存开销，具体取决于依赖大小。
Q&A
为什么在PyTorch多进程推理中需要共享模型权重？
因为Python的GIL限制了多线程性能，所以常采用多进程推理。但模型权重无法在进程间原生共享，导致每个进程都会在GPU显存中复制一份权重，容易耗尽显存。通过CUDA IPC共享权重可以避免重复占用。
如何在PyTorch多进程推理中共享模型权重？
使用PyTorch的multiprocessing模块，在父进程中调用model.share_memory()，然后通过进程间传递模型对象，即可让子进程共享同一份CUDA权重，避免显存重复。
AOTInductor中如何实现权重共享？
在编译AOTInductor包时设置aot_inductor.package_constants_in_so为False，避免将权重打包进.so文件。然后在每个worker进程中用torch._inductor.aoti_load_package加载包，再调用runner.load_constants(state_dict, check_full_update=True, user_managed=False)将父进程的state_dict绑定到runner上，从而共享权重。
runner.load_constants中的user_managed参数有什么作用？
user_managed=False表示AOTInductor runner不管理权重的生命周期，只持有指向父进程权重的指针，由用户负责管理。若设为True，runner会管理权重，但不会复制权重，而是通过跨进程引用计数确保权重在所有引用释放前不被释放。
为什么在示例中要为每个worker复制一份AOTInductor包？
为了演示通过父进程state_dict共享权重。如果所有worker加载同一个包文件，OS会共享mmap页面，权重自动共享，无法体现显式绑定。复制包文件可以打破这种共享，从而展示通过load_constants绑定权重的效果。
TorchScript能否实现权重共享？有什么问题？
可以，但问题较多。需要在所有进程中先在CPU上实例化TorchScript模型，然后用父进程的共享权重替换，这会导致CPU内存占用显著增加，除非顺序处理。因此示例中未包含TorchScript实现。
多进程推理中程序依赖（如共享库）会重复吗？
每个进程会加载自己的程序依赖副本，但OS会共享相同物理内存页（如果依赖来自同一文件），从而减少内存开销。不过仍会有一定的内存开销，具体取决于依赖大小和进程数量。
