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

PyTorch's multiprocessing module provides a CUDA Inter-Process Communication (IPC) API that enables sharing model weights across multiple processes for inference, avoiding duplication in GPU VRAM. The technique, demonstrated with PyTorch and AOTInductor, addresses the Global Interpreter Lock (GIL) bottleneck by running models in separate processes while sharing weights via CUDA IPC.

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. 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322 | python import argparseimport ioimport osimport tempfileimport timeimport torchimport torch.nn as nnimport torch.multiprocessing as mpimport 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. 1 | bash $ 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. 12 | bash $ 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: - 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. - 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. - 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. - 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. PyTorch Multi-Process Inference Weight Sharing Via Inter-Process Communication https://leimao.github.io/blog/PyTorch-IPC-Multi-Process-Inference-Weight-Sharing/ https://leimao.github.io/blog/PyTorch-IPC-Multi-Process-Inference-Weight-Sharing/