PyTorch Multi-Process Inference Weight Sharing Via Inter-Process Communication

Introduction

PyTorch is probably the most popular deep learning framework for model training. Ideally, without any additional effort, PyTorch model could be used for inference seamlessly in Python. However, in practice, there are some concerns when using PyTorch for inference in Python, primarily because of the Global Interpreter Lock (GIL) in Python. The GIL prevents multiple threads from executing Python bytecodes at once, which can lead to performance bottlenecks in high-throughput applications. To overcome this limitation, PyTorch models can be run in multiple processes, where each process runs its own Python interpreter and has its own GIL. However, because model weights cannot be natively shared across processes in Python, multi-process inference on single-GPU can lead to duplicated model weights in GPU VRAM, which can quickly exhaust the available GPU memory. To mitigate this issue, model weights can be shared across processes using CUDA Inter-Process Communication (IPC).

It turns out that PyTorch multiprocessing module has abstracted away the details of CUDA IPC, and provides a simple API to share CUDA tensors across processes. In this blog post, I will demonstrate how to use PyTorch multiprocessing module to share model weights across multiple processes for inference, and avoid weight duplication in GPU VRAM.

PyTorch Multi-Process Inference Weight Sharing

In the following example, I will demonstrate how to share model weights via CUDA IPC across multiple processes for PyTorch and AOTInductor inference using PyTorch multiprocessing module.

pytorch_inference_ipc.py

import argparse
import io
import os
import tempfile
import time

import torch
import torch.nn as nn
import torch.multiprocessing as mp
import torch._inductor


# =====================================================================
# 1. MODEL DEFINITION
# =====================================================================
class SharedWeightInferenceEngine(nn.Module):

    def __init__(self,
                 in_features: int = 16384,
                 hidden_dim: int = 16384,
                 out_features: int = 12800):
        super().__init__()
        self.encoder = nn.Linear(in_features, hidden_dim)
        self.head = nn.Linear(hidden_dim, out_features)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = torch.relu(self.encoder(x))
        return self.head(x)


# =====================================================================
# 2. HELPER FUNCTIONS
# =====================================================================
def get_device_vram_usage_mb(device_id: int = 0) -> float:
    """Return device-wide GPU memory usage in MB."""
    torch.cuda.set_device(device_id)
    free_bytes, total_bytes = torch.cuda.mem_get_info(device_id)
    return (total_bytes - free_bytes) / (1024**2)


def load_aoti_runner(package_path: str, state_dict: dict, device_id: int):
    """Load an AOTI runner and bind it to the supplied model weights."""
    runner = torch._inductor.aoti_load_package(package_path,
                                               device_index=device_id)

    # To enable IPC weight sharing using the parent state dict,
    # uncomment the following lines:
    expected_keys = set(state_dict.keys())
    if hasattr(runner, "get_constant_fqns"):
        assert set(runner.get_constant_fqns()) == expected_keys
    # Now the weight tensor pointers or references are pointing to the shared weights in the parent process.
    runner.load_constants(state_dict,
                          check_full_update=True,
                          user_managed=False)

    torch.cuda.synchronize()
    time.sleep(0.2)
    return runner


# =====================================================================
# 3. WORKER PROCESS INFERENCE LOOPS
# =====================================================================
def pytorch_inference_worker(rank: int, model, input_queue: mp.Queue,
                             output_queue: mp.Queue):
    pid = os.getpid()
    torch.cuda.set_device(0)
    model.eval()

    encoder_ptr = model.encoder.weight.data_ptr()
    head_ptr = model.head.weight.data_ptr()

    print(f"\n[Worker {rank} | PID {pid}] Initialized successfully!")
    print(f"  ├── Shared 'encoder.weight' CUDA Address : {hex(encoder_ptr)}")
    print(f"  └── Shared 'head.weight' CUDA Address    : {hex(head_ptr)}")

    while True:
        task = input_queue.get()
        if task is None:
            print(
                f"[Worker {rank} | PID {pid}] Received shutdown signal. Exiting."
            )
            break

        task_id, batch_cpu = task
        start_time = time.perf_counter()

        with torch.no_grad():
            batch_gpu = batch_cpu.to('cuda', non_blocking=True)
            logits = model(batch_gpu)
            preds = torch.argmax(logits, dim=-1)

        elapsed_ms = (time.perf_counter() - start_time) * 1000
        print(
            f"[Worker {rank} | PID {pid}] Processed Task #{task_id} in {elapsed_ms:.2f}ms"
        )
        output_queue.put((task_id, preds.cpu(), rank))


def aoti_inference_worker(rank: int, package_path: str,
                          parent_state_dict: dict, input_queue: mp.Queue,
                          output_queue: mp.Queue):
    pid = os.getpid()
    device_id = 0
    torch.cuda.set_device(device_id)

    optimized_runner = load_aoti_runner(package_path, parent_state_dict,
                                        device_id)

    encoder_ptr = parent_state_dict["encoder.weight"].data_ptr()
    head_ptr = parent_state_dict["head.weight"].data_ptr()

    print(
        f"\n[Worker {rank} | PID {pid}] Loaded AOTInductor package from {os.path.basename(package_path)}!"
    )
    print(
        f"  ├── Parent-process 'encoder.weight' CUDA Address : {hex(encoder_ptr)}"
    )
    print(
        f"  └── Parent-process 'head.weight' CUDA Address    : {hex(head_ptr)}"
    )

    while True:
        task = input_queue.get()
        if task is None:
            print(
                f"[Worker {rank} | PID {pid}] Received shutdown signal. Exiting."
            )
            break

        task_id, batch_cpu = task
        start_time = time.perf_counter()

        with torch.no_grad():
            batch_gpu = batch_cpu.to('cuda', non_blocking=True)
            logits = optimized_runner(batch_gpu)
            preds = torch.argmax(logits, dim=-1)

        elapsed_ms = (time.perf_counter() - start_time) * 1000
        print(
            f"[Worker {rank} | PID {pid}] Processed Task #{task_id} in {elapsed_ms:.2f}ms"
        )
        output_queue.put((task_id, preds.cpu(), rank))


# =====================================================================
# 4. MAIN ORCHESTRATOR
# =====================================================================
def parse_args(argv=None):
    parser = argparse.ArgumentParser(
        description=
        "Run a CUDA IPC demo using standard PyTorch or AOTInductor execution"
    )
    parser.add_argument("--mode",
                        choices=["pytorch", "aoti"],
                        default="pytorch")
    parser.add_argument("--workers", type=int, default=3)
    parser.add_argument("--batches", type=int, default=6)

    parser.add_argument("--in-dim",
                        type=int,
                        default=16384,
                        help="Input feature size")
    parser.add_argument("--hidden-dim",
                        type=int,
                        default=16384,
                        help="Encoder output / hidden layer size")
    parser.add_argument("--out-dim",
                        type=int,
                        default=12800,
                        help="Output classification size")
    return parser.parse_args(argv)


def main(argv=None):
    args = parse_args(argv)
    mp.set_start_method('spawn', force=True)

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for this demonstration.")

    main_pid = os.getpid()
    print("=" * 70)
    print(
        f"MAIN PROCESS [PID {main_pid}] STARTING {args.mode.upper()} CUDA IPC DEMO"
    )
    print("=" * 70)

    print(
        f"\n[Step A] Loading model ({args.in_dim} -> {args.hidden_dim} -> {args.out_dim}) onto GPU VRAM..."
    )
    vram_before = torch.cuda.memory_allocated() / (1024**2)

    model = SharedWeightInferenceEngine(
        in_features=args.in_dim,
        hidden_dim=args.hidden_dim,
        out_features=args.out_dim).cuda().eval()

    model.share_memory()

    vram_after = torch.cuda.memory_allocated() / (1024**2)
    model_vram = vram_after - vram_before

    main_encoder_ptr = model.encoder.weight.data_ptr()
    main_head_ptr = model.head.weight.data_ptr()

    print(f"  ├── Model VRAM Allocation : {model_vram:.2f} MB")
    print(f"  ├── Main 'encoder' Pointer: {hex(main_encoder_ptr)}")
    print(f"  └── Main 'head' Pointer   : {hex(main_head_ptr)}")

    if args.mode == "pytorch":
        payloads = [model] * args.workers
        worker_target = pytorch_inference_worker
        label = "Standard PyTorch model"
    else:
        print(
            f"\n[Step A.1] Exporting model and writing {args.workers} separate AOTI packages..."
        )
        example_input = torch.randn(16, args.in_dim, device='cuda')
        exported_program = torch.export.export(model, (example_input, ))
        temp_dir = tempfile.mkdtemp()

        # Build initial template package
        template_package = os.path.join(temp_dir, "model_aoti_base.pt2")
        template_package = torch._inductor.aoti_compile_and_package(
            exported_program,
            package_path=template_package,
            inductor_configs={
                "always_keep_tensor_constants": True,
                # If it sets to True, the so file will contain model constant weights, which could be large.
                # In our example, because we will load an AOTI package for each process,
                # If it sets to True, it is possible that the GPU already runs out of memory before we set the weights to shared weights.
                # If it sets to False, and later we did not load the weights to the AOTI package, there will be undefined behaviors.
                "aot_inductor.package_constants_in_so": False,
                # Save a copy of the model constant weights to disk in pickle format.
                # It is useless for this example, because the weights were loaded from the PyTorch model state dict.
                "aot_inductor.package_constants_on_disk_format":
                "pickle_weights",
            })

        # Read template binary bytes
        with open(template_package, "rb") as f:
            package_bytes = f.read()

        # Save unique file copies for each worker to break OS mmap page sharing
        # If every process loads the same package file, the OS will share the same mmap pages for all processes.
        # Consequently, the model weights will be automatically shared across processes,
        # even without explicitly sharing the weights via the parent state dict.
        payloads = []
        for rank in range(args.workers):
            worker_pkg_path = os.path.join(temp_dir,
                                           f"model_aoti_worker_{rank}.pt2")
            with open(worker_pkg_path, "wb") as f:
                f.write(package_bytes)
            payloads.append(worker_pkg_path)

        worker_target = aoti_inference_worker
        label = "AOTInductor packages (unique copies per worker)"

    print(f"\n[Step B] Creating Inter-Process Communication queues...")
    input_queue = mp.Queue()
    output_queue = mp.Queue()

    print(
        f"\n[Step C] Spawning {args.workers} worker processes for {label}...")
    workers = []
    for rank in range(args.workers):
        payload = payloads[rank]
        if args.mode == "aoti":
            parent_state_dict = model.state_dict()
            process_args = (rank, payload, parent_state_dict, input_queue,
                            output_queue)
        else:
            process_args = (rank, payload, input_queue, output_queue)

        p = mp.Process(target=worker_target, args=process_args)
        p.start()
        workers.append(p)

    time.sleep(1.0)

    main_process_vram = torch.cuda.memory_allocated() / (1024**2)
    device_vram_used = get_device_vram_usage_mb()
    print("\n" + "-" * 70)
    print(f"VRAM VERIFICATION AFTER SPAWNING {args.workers} WORKERS:")
    print(
        f"  ├── Expected VRAM if duplicated ({args.workers + 1} copies) : ~{model_vram * (args.workers + 1):.2f} MB"
    )
    print(
        f"  ├── Main-process allocated VRAM                             : ~{main_process_vram:.2f} MB"
    )
    print(
        f"  └── Device-wide VRAM in use                                 : ~{device_vram_used:.2f} MB"
    )
    print("-" * 70 + "\n")

    print(f"\n[Step D] Dispatching {args.batches} inference tasks...")
    for task_id in range(args.batches):
        batch = torch.randn(16, args.in_dim)
        input_queue.put((task_id, batch))

    print(f"\n[Step E] Receiving predictions from workers...")
    for _ in range(args.batches):
        task_id, predictions, worker_rank = output_queue.get()
        print(
            f"  └── Result Task #{task_id} from Worker {worker_rank} | Predictions: {predictions[:4].tolist()}..."
        )

    print("\n[Step F] Terminating workers cleanly...")
    for _ in range(args.workers):
        input_queue.put(None)

    for p in workers:
        p.join()

    print("\n" + "=" * 70)
    print("ALL PROCESSES FINISHED SUCCESSFULLY")
    print("=" * 70)


if __name__ == '__main__':
    main()

To run the example, we could use the following command to start an NVIDIA PyTorch container with GPU support.

$ docker run -it --rm --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -v $(pwd):/mnt -w /mnt nvcr.io/nvidia/pytorch:26.07-py3

To run the example, we could use the following commands to run the PyTorch or AOTInductor multi-process inference with weight sharing via CUDA IPC.

$ python pytorch_inference_ipc.py --mode aoti --batches 128 --workers 24
$ python pytorch_inference_ipc.py --mode pytorch --batches 128 --workers 24

By default, each model instance will take 2 GB of GPU VRAM. My GPU is an NVIDIA RTX 5080 with 16 GB of GPU VRAM. If there is no weight sharing, I would not be able to run the example with 24 workers successfully.

There are a few key operations or caveats in the example:

1. If using PyTorch multi-process inference, the model weights are shared across processes using model.share_memory(). This allows the model weights to be shared across processes without duplicating them in GPU VRAM. Although it seems that without using model.share_memory(), the multiprocessing module will still share the model weights across processes, but it is still a good practice to explicitly call model.share_memory() to ensure that the model weights are shared across processes.
2. To avoid materializing the model weights for AOTInductor packages in each worker process, to build the AOTInductor packages, we set the aot_inductor.package_constants_in_so to False in the inductor_configs. This will prevent the model weights from being packaged into the AOTInductor packages, and when the model is loaded in each worker process using torch._inductor.aoti_load_package, the no model weights will be loaded from the package. Consequently, the GPU VRAM will not be inflated by model loading from many processes. Then, we explicitly load the model weights from the parent process state dict using runner.load_constants. This will bind the model weights in the AOTInductor package to the shared model weights in the parent process, and avoid weight duplication in GPU VRAM.
3. The runner.load_constants function has a user_managed argument, which is set to False in the example. This means that the AOTInductor runner will not manage the model weights, and the user is responsible for managing the model weights. The AOTInductor runner will just have a pointer to the model weights in the parent process, and will not manage the life cycle of model weights. If user_managed is set to True, the AOTInductor runner will manage the model weights. But instead of making a copy of the model weights in GPU VRAM, the AOTInductor runner will still have a reference to the model weights in the parent process. The reference counter is effective across processes, so the model weights will not be freed until all references to the model weights are released.
4. In this example, the AOTInductor package was replicated for each worker process to break the OS mmap page sharing, in order to demonstrate the weight sharing via the parent process state dict. If we only have one AOTInductor package file, and each process loads the same package file, the OS will still share the same mmap pages for all processes, no model weights duplication will occur on GPU VRAM, because those loaded weights are read-only and the OS knows that it can be safely shared across processes.

TorchScript Weight Sharing

If the user is still using the deprecated TorchScript, it is still possible to share model weights across processes using a similar approach. However, I encountered a lot of problems. For example, I would have to instantiate the TorchScript model on CPU in all the processes first and then replace the model weights with the shared weights from the parent process. This will inflate the CPU memory usages significantly, unless I load the TorchScript model and bind the shared weights in the parent process sequentially for each process. Consequently, I did not include the TorchScript example in the previous example implementation.

Program Dependency Duplication

Even though we saved duplicate model weights in multi-process inference, the program dependencies (e.g., shared libraries) are still duplicated in each process. This is because each process has its own address space and loads its own copy of the program dependencies. However, the OS can still share the same physical memory pages for the program dependencies across processes, if those dependencies are loaded from the same shared library files. This is a common behavior in modern operating systems, and it helps to reduce the overall memory footprint of multi-process applications. Nevertheless, it is still expected that some memory overhead will be incurred for both host memory and GPU memory for each process, and the amount of overhead really depends.

Conclusions

In this blog post, we demonstrated how to share model weights across multiple processes in PyTorch multi-process inference using CUDA IPC. We also discussed the caveats and best practices for both PyTorch and AOTInductor, as well as the implications for TorchScript and program dependency duplication.