Safe Parallel Inference and Performance Optimization for Deep Learning Models

Safe Parallel Inference and Performance Optimization for Deep Learning Models
Photo by Amy Chen / Unsplash

Hello!

Today we address one of the questions we hear most often: "Can a single model instance safely run parallel inference?"

Safety of Parallel Inference in eval Mode

When a PyTorch model has been set to eval mode using model.eval(), it is generally safe for parallel inference.

(Here, "parallel" refers to multithreaded processing. Batch inference is discussed later.)

Here is why:

  1. Parameter immutability
    In eval mode, the model's parameters are not updated during the forward pass.
  2. Deactivation of training-specific layers
    Layers such as BatchNorm switch to a mode that uses running statistics rather than computing batch statistics.
  3. Independence of input data
    Each thread or process operates on its own input data, which resides in a separate region of memory.

Here is a basic example of safe parallel inference in eval mode:

import torch
import threading

def safe_inference(model, data):
    with torch.no_grad():
        return model(data)

model = YourModel()
model.eval()  # Important: set the model to eval mode

# Run inference on multiple threads
threads = []
for i in range(10):
    t = threading.Thread(target=safe_inference, args=(model, your_data[i]))
    threads.append(t)
    t.start()

for t in threads:
    t.join()

Cases That Require Caution

However, caution is required in situations such as the following:

  1. Custom layers
    If your model contains layers you implemented yourself, you need to carefully verify how they behave under parallel execution.
class CustomLayer(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.counter = 0  # A potential source of problems

    def forward(self, x):
        self.counter += 1  # Not thread-safe
        return x + self.counter

# Custom layers like this can cause problems under parallel execution
  1. GPU memory constraints
    When multiple threads process large amounts of data simultaneously, you may run out of GPU memory.
  2. Complex model architectures
    Some complex model structures, such as certain types of attention mechanisms, may behave unexpectedly when executed in parallel.

Using a Model Pool

In the cautionary cases above, using a pool of model instances can sometimes avoid these problems.
Here is a simple example of a model pool implementation:

import torch
from queue import Queue

class ModelPool:
    def __init__(self, model_class, num_instances):
        self.pool = Queue()
        for _ in range(num_instances):
            model = model_class().to('cuda')
            model.eval()
            self.pool.put(model)

    def get_model(self):
        return self.pool.get()

    def return_model(self, model):
        self.pool.put(model)

def safe_pooled_inference(pool, data):
    model = pool.get_model()
    try:
        with torch.no_grad():
            result = model(data)
        return result
    finally:
        pool.return_model(model)

# Usage example
pool = ModelPool(YourModel, num_instances=3)
results = [safe_pooled_inference(pool, data) for data in your_data_list]

With this approach, each inference task uses its own independent model instance, avoiding the problems that arise under parallel execution.

Batching Is the Foundation of Performance Optimization

Parallel inference offers flexibility, but its overhead can degrade performance. Here are some important tips for improving performance.

Taking Advantage of Batch Processing

Rather than running individual inferences in parallel, you can expect substantial performance gains from batch processing. GPUs excel at processing large amounts of data simultaneously, so batching makes the most of a GPU's capabilities.

1. Static batching

The simplest method is to use fixed-size batches:

def batch_inference(model, data_list, batch_size=32):
    results = []
    for i in range(0, len(data_list), batch_size):
        batch = torch.stack(data_list[i:i+batch_size])
        with torch.no_grad():
            batch_results = model(batch)
        results.extend(batch_results)
    return results

# Usage example
results = batch_inference(model, your_data_list)

Requests Never Arrive at Convenient Batch Timing

When you are building an on-demand inference service as a web service, simple parallel inference on the GPU is not enough.

That is because users do not conveniently arrive at the same time.
In fact, it is rare for requests to line up neatly into a batch.

2. Dynamic batching

To efficiently process data arriving in real time, you can use dynamic batching:

import time
from collections import deque

class DynamicBatcher:
    def __init__(self, model, max_batch_size=32, max_wait_time=0.1):
        self.model = model
        self.max_batch_size = max_batch_size
        self.max_wait_time = max_wait_time
        self.queue = deque()
        self.results = {}

    def add_item(self, item_id, data):
        self.queue.append((item_id, data))
        if len(self.queue) >= self.max_batch_size:
            self.process_batch()

    def process_batch(self):
        batch_ids, batch_data = zip(*[self.queue.popleft() for _ in range(len(self.queue))])
        batch_tensor = torch.stack(batch_data)
        with torch.no_grad():
            batch_results = self.model(batch_tensor)
        for item_id, result in zip(batch_ids, batch_results):
            self.results[item_id] = result

    def get_result(self, item_id):
        start_time = time.time()
        while item_id not in self.results:
            if time.time() - start_time > self.max_wait_time:
                self.process_batch()
            time.sleep(0.01)
        return self.results.pop(item_id)

# Usage example
batcher = DynamicBatcher(model)

def process_item(item_id, data):
    batcher.add_item(item_id, data)
    return batcher.get_result(item_id)

# Call process_item from multiple threads

With this approach, data is added to a batch as it arrives, and processing runs when the batch reaches its maximum size or the maximum wait time is exceeded.

3. Continuous batching

When data is generated continuously, continuous batching like the following is effective:

import torch
from torch.utils.data import DataLoader, IterableDataset

class ContinuousDataset(IterableDataset):
    def __iter__(self):
        while True:
            yield self.get_next_item()  # Implement your data generation logic

    def get_next_item(self):
        # Implement the actual data generation logic here
        pass

def continuous_batch_inference(model, dataset, batch_size=32):
    dataloader = DataLoader(dataset, batch_size=batch_size)
    for batch in dataloader:
        with torch.no_grad():
            yield model(batch)

# Usage example
dataset = ContinuousDataset()
for batch_results in continuous_batch_inference(model, dataset):
    process_results(batch_results)  # Process the results

This method enables efficient batch processing even when data is generated continuously.

Summary

In this article, we focused on parallelization and performance on a single GPU.
Parallel inference in eval mode is safe in most cases, but batching is essential for maximizing performance. In deep learning and LLM-based services, inference scenarios often depend on techniques such as dynamic batching and continuous batching. At Qualiteg, we have been researching dynamic and continuous batching from the beginning and apply them to LLMs, video generation, and AI character responses.
By choosing and implementing these techniques appropriately, you can substantially improve inference throughput.

Although separate from parallelization, combining additional techniques such as model quantization, TorchScript, and GPU-level optimization can deliver further performance gains.

GPUs are very expensive hardware, so the perspective of "using up" a single GPU to its fullest is extremely important, and we at Qualiteg refine these techniques every day.

For Even Larger-Scale Traffic, Consider a GPU Cluster

On the other hand, in scenarios where heavy concurrent access is expected, load balancing across multiple GPUs becomes essential. We plan to cover those techniques in a separate blog post, but the video below explains how to configure a GPU cluster for LLMs, so please take a look if you are interested.

See you next time!

Read more