如何通过共享内存实现多进程模型权重共享：大幅降低 App 内存占用

在端侧推理或高并发 Web 服务场景中，为了提升吞吐量，我们常会启动多个进程并行处理推理请求。然而，如果每个进程都独立加载一份模型（例如一个 2GB 的 BERT 模型），启动 4 个进程就会占用 8GB 内存，这极易导致 OOM（内存溢出）。

本文将教你如何利用 PyTorch 的共享内存机制，实现「一份权重，多个进程使用」，从而大幅压低系统内存占用。

1. 核心原理

在 Python 的 multiprocessing 模块中，不同进程的内存空间是隔离的。但 PyTorch 扩展了这一能力，通过 share_memory_() 方法，可以将 Tensor 移动到一段可被多个进程同时访问的共享内存区域（Shared Memory）。

当模型调用 .share_memory_() 后，其权重数据不会在创建子进程时被复制（Copy-on-Write），而是直接在原内存地址上进行读取。

2. 实战代码示例

下面的示例展示了如何加载一个模型，将其设置为共享模式，并在两个不同的子进程中并发执行推理。

import torch
import torch.multiprocessing as mp
from transformers import AutoModel, AutoTokenizer

# 1. 定义推理函数
def inference_worker(model, input_data, process_id):
    with torch.no_grad():
        output = model(**input_data)
        print(f\"[Process {process_id}] Output shape: {output.last_hidden_state.shape}\")

if __name__ == \"__main__\":
    # 设置多进程启动方法为 'spawn' (在 Linux/macOS 上推荐)
    try:
        mp.set_start_method('spawn')
    except RuntimeError:
        pass

    # 2. 在主进程中加载模型
    model_name = \"bert-base-uncased\"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name)

    # 3. 核心步骤：将模型权重移动到共享内存
    model.share_memory_()
    model.eval()

    # 准备测试数据
    text = \"Sharing memory saves a lot of RAM!\"
    inputs = tokenizer(text, return_tensors=\"pt\")

    # 4. 创建并启动多个子进程
    processes = []
    for i in range(3):
        p = mp.Process(target=inference_worker, args=(model, inputs, i))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

    print(\"所有进程推理完成。\")

3. 关键注意事项

1. 仅限 CPU Tensor：share_memory_() 主要针对 CPU 上的 Tensor。对于 GPU，PyTorch 默认使用 CUDA IPC（Inter-Process Communication）句柄来共享显存，原理类似但底层机制不同。
2. 只读安全：在共享模式下，建议将模型设置为 .eval() 模式并使用 torch.no_grad()。如果多个进程尝试同时修改共享的权重（例如在线学习场景），会导致竞态条件（Race Condition）。
3. 启动方式：建议使用 mp.set_start_method(‘spawn’)。虽然 fork 在某些系统上更快，但它会复制整个父进程的地址空间，在复杂的 AI 应用中容易引发死锁。

4. 收益总结

通过这种方式，无论你开启多少个推理 worker，物理内存中始终只维护一套模型权重。这对于内存受限的端侧设备（如树莓派、国产 AI 开发板）或是需要高并发处理的生产环境具有极高的实操价值。