Optimizing Inter-Process Communication in PyTorch Data Input Pipelines

How to Identify and Solve IPC Bottlenecks in AI/ML Applications

Photo by Kin Shing Lai on Unsplash

Multi-process data loading is a common PyTorch optimization technique — often the first you will encounter in any PyTorch performance profiling tutorial (e.g., see here). In the default single-process data loading mode, data preparation and model execution are run in the same process. When the model runs on the GPU, this mode can lead to severe GPU starvation — a condition where the expensive GPU resources lay idle while they wait for the — often lengthy — data preparation phase to complete. The natural solution to this performance bottleneck is to prepare data input batches on the CPU in parallel to running the model on the GPU. This is accomplished via multi-process data loading.

In multi-process data loading, one or more dedicated data-preparation worker processes run in parallel to the primary training (or inference) loop. The worker processes push the ready data batches to a shared multi-processing queue which are, in turn, consumed by the PyTorch model running in the main process. In many cases, simply setting the num_workers argument of the DataLoader API to a positive integer can result in remarkable speedups.

It is important to note that although enabling this feature is extremely (almost deceivingly) easy, moving from a single-process Python application to a multi-process Python application introduces a great deal of complexity including: increased memory overhead, serialization challenges, and synchronization hurdles. The PyTorch software stack does an exceptional job at hiding these complexities from the AI/ML developer. In many cases the transition is seamless. However, there are occasions where users may suddenly encounter out-of-memory faults, pickling errors, or file descriptor shortages (“Too many open files”) requiring adaptations to the PyTorch script.

In an ideal world, we would be able to achieve the parallelization between data loading and model execution using multi-threading — without needing to resort to the use of multiple processes and their associated costs. But, this is prevented by the Python Global Interpreter Lock (GIL). Note, that this may change in the not-so-distant future as PyTorch increases its support for free-threaded, no-GIL Python.

In this post, we focus on the overhead introduced by the inter-process communication (IPC) between the multiple data loading workers and the main process. Input data batches can be extremely large — especially in the case of computer vision models. Ordinarily, transmitting such large payloads over IPC would be prohibitively slow. Fortunately, PyTorch optimizes IPC of PyTorch tensors through the use of shared memory. However, a careless definition of the dataset could easily undermine PyTorch’s IPC optimization and result in an unforced performance bottleneck on the data-input pipeline.

Previous posts, (e.g., here), have covered the topic of performance bottlenecks in the data input pipeline. Some of the common causes we have seen are:

1. Slow retrieval of raw data
2. CPU resource exhaustion due to heavy preprocessing tasks
3. Host-to-device transfer bottlenecks

In this post we show how the overhead of the IPC between the data-loader workers and the main model execution process can introduce a performance bottleneck on the data-input pipeline. In this post we will:

1. Demonstrate PyTorch’s “zero copy” IPC in action
2. Show how a naive definition of the data batch can undermine PyTorch’s IPC optimization and result in a significant unforced performance bottleneck.
3. Revisit some of the tools we have covered in previous posts for analyzing the performance of PyTorch workloads and apply them to the problem of IPC bottlenecks.
4. Share a simple example of IPC optimization.

To facilitate our discussion, we will define some toy examples explicitly designed to include an IPC bottleneck. We will run our experiments on an Amazon EC2 g6e.2xlarge instance with an NVIDIA L40S GPU running an AWS Deep Learning (Ubuntu 24.04) AMI with PyTorch (2.14). Importantly, the code we will share is intended for demonstrative purposes; Please do not rely on its correctness or optimality. Please do not interpret our mention of any library, platform, or API as an endorsement of its use.

The Overhead of IPC

In a multi-process application each process maintains its own independent and isolated virtual memory space. Communication between processes is conducted via a shared kernel socket. It typically involves the following stages:

1. Process A serializes the communication payload into a raw byte stream
2. Process A passes the byte stream into a temporary Linux kernel socket buffer
3. Process B reads the byte stream from kernel memory and deserializes it

IPC requires both CPU compute resources — for payload serialization and deserialization, as well as CPU memory resources — for creating the serialized buffer and copying it into and out of kernel memory. If the payload is large, as is often the case in AI/ML workloads, IPC can result in significant performance bottlenecks.

In the following block we demonstrate the overhead of serializing and un-serializing a 12 MB image in numpy:

import time
import numpy as np
# Multiprocessing IPC uses ForkingPickler to serialize objects sent between
# processes.
from multiprocessing.reduction import ForkingPickler

IMAGE_SHAPE = (3, 1024, 1024)
REPETITIONS = 20

image = np.arange(
    np.prod(IMAGE_SHAPE), dtype=np.float32
).reshape(IMAGE_SHAPE)

# Warm up any one-time pickle and allocation costs before measurement.
payload = bytes(ForkingPickler.dumps(image))
restored = ForkingPickler.loads(payload)

start = time.perf_counter()
payloads = [
    bytes(ForkingPickler.dumps(image)) for _ in range(REPETITIONS)
]
serialize_seconds = time.perf_counter() - start

start = time.perf_counter()
for payload in payloads:
    restored = ForkingPickler.loads(payload)
deserialize_seconds = time.perf_counter() - start

print(f"Forking pickle:       {len(payload) / 1024**2:.1f} MiB")
print(
    "Serialization:        "
    f"{serialize_seconds * 1000 / REPETITIONS:.2f} ms per image"
)
print(
    "Deserialization:      "
    f"{deserialize_seconds * 1000 / REPETITIONS:.2f} ms per image"
)

On our Amazon EC2 g6e.2xlarge instance this results in a payload size of 12 MB, a serialization time of 9.71 ms and a deserialization time of 1.04 ms, per image. In a high performance AI/ML setting, those kinds of numbers can be prohibitive.

To assess the full IPC overhead, we create a simple two-process producer-consumer example:

import multiprocessing as mp

def send_images(result_queue, ready, start):
    image = np.arange(
        np.prod(IMAGE_SHAPE), dtype=np.float32
    ).reshape(IMAGE_SHAPE)
    ready.set()
    start.wait()
    for _ in range(REPETITIONS + 1):
        result_queue.put(image)


if __name__ == "__main__":
    context = mp.get_context("spawn")
    result_queue = context.Queue(maxsize=1)
    ready = context.Event()
    start = context.Event()
    producer = context.Process(
        target=send_images,
        args=(result_queue, ready, start),
    )
    producer.start()
    ready.wait()
    start.set()

    warmup_image = result_queue.get()

    start = time.perf_counter()
    received_images = [result_queue.get() for _ in range(REPETITIONS)]
    transfer_seconds = time.perf_counter() - start
    producer.join()
    result_queue.close()
    result_queue.join_thread()

    print("\nSpawned producer/consumer Queue transfer:")
    print(
        f"Throughput: {REPETITIONS / transfer_seconds:.2f} images/s "
    )

This results in a throughput of 50.44 images per second.

How PyTorch Optimizes IPC

To address the potential overhead of IPC, PyTorch avoids the serialization and transmission of the raw float data across the IPC pipe/queue by implementing custom reducers for PyTorch tensors using Python’s ForkingPickler. When a torch.Tensor is sent across a process boundary, PyTorch moves the underlying storage (DataPtr) into shared memory (if not already there). Then, instead of serializing large blocks (e.g., gigabytes) of raw float data, PyTorch serializes just the metadata (shape, stride, dtype, and a shared memory handle/file descriptor). The consumer deserializes the metadata and uses it to read and process the torch.Tensor in shared memory. The actual payload size across the IPC queue drops from many megabytes or gigabytes to a few hundred bytes.

In the code block below, we repeat the serialization and IPC overhead experiments, this time with a torch.Tensor instead of a numpy array. We also include a brief report of the contents of the optimized payload:

import pickletools
import time
import torch
import torch.multiprocessing as torch_mp
from multiprocessing.reduction import ForkingPickler

IMAGE_SHAPE = (3, 1024, 1024)
REPETITIONS = 20

# toggle sharing strategy between "file_descriptor" and "file_system"
SHARING_STRATEGY = "file_descriptor" 
#SHARING_STRATEGY = "file_system" 

def summarize_pickle(payload):
    strings = []
    for opcode, argument, _ in pickletools.genops(payload):
        if opcode.name in {"BINBYTES", "SHORT_BINBYTES", "BYTEARRAY8"}:
            try:
                strings.append(argument.decode("utf-8"))
            except UnicodeDecodeError:
                pass
        elif opcode.name.endswith("UNICODE"):
            strings.append(argument)
    recipes = tuple(
        dict.fromkeys(
            value
            for value in strings
            if value.startswith("rebuild_") or value == "DupFd"
        )
    )
    shared_memory_names = tuple(
        dict.fromkeys(
            value
            for value in strings
            if value.startswith("/torch_")
        )
    )
    return recipes, shared_memory_names


def describe_payload(payload, tensor):
    recipes, shared_memory_names = summarize_pickle(payload)
    storage = tensor.untyped_storage()
    try:
        shared_fd = storage._get_shared_fd()
    except (AttributeError, RuntimeError):
        shared_fd = None
    
    print("\nForking pickle payload:")
    print(
        f"  pickle stream: {len(payload)} bytes "
        f"({len(payload) / 1024:.2f} KiB)"
    )
    print("  reconstruction recipes:")
    if recipes:
        for recipe in recipes:
            print(f"    - {recipe}")
    if shared_fd is not None:
        print(
            f"  local shared-storage FD: {shared_fd}"
            "(process-local; receiver gets a duplicate)"
        )
    if shared_memory_names:
        print(
            "  shared-memory names in pickle: "
            f"{', '.join(shared_memory_names)}"
        )


def send_images(result_queue, ready, start, consumer_done):
    torch_mp.set_sharing_strategy(SHARING_STRATEGY)
    image = torch.arange(
        torch.tensor(IMAGE_SHAPE).prod(), dtype=torch.float32
    ).reshape(IMAGE_SHAPE)
    ready.set()
    start.wait()
    for _ in range(REPETITIONS + 1):
        result_queue.put(image)
    consumer_done.wait()


if __name__ == "__main__":
    torch_mp.set_sharing_strategy(SHARING_STRATEGY)
    image = torch.arange(
        torch.tensor(IMAGE_SHAPE).prod(), dtype=torch.float32
    ).reshape(IMAGE_SHAPE)

    payload = bytes(ForkingPickler.dumps(image))
    restored = ForkingPickler.loads(payload)

    start = time.perf_counter()
    payloads = [
        bytes(ForkingPickler.dumps(image)) for _ in range(REPETITIONS)
    ]
    serialize_seconds = time.perf_counter() - start

    start = time.perf_counter()
    for payload in payloads:
        restored = ForkingPickler.loads(payload)
    deserialize_seconds = time.perf_counter() - start

    print(f"Forking pickle:       {len(payload) / 1024:.1f} KiB")
    print(
        "Serialization:        "
        f"{serialize_seconds * 1000 / REPETITIONS:.2f} ms per image"
    )
    print(
        "Deserialization:      "
        f"{deserialize_seconds * 1000 / REPETITIONS:.2f} ms per image"
    )

    describe_payload(payload, image)

    context = torch_mp.get_context("spawn")
    result_queue = context.Queue(maxsize=1)
    ready = context.Event()
    start = context.Event()
    consumer_done = context.Event()
    producer = context.Process(
        target=send_images,
        args=(result_queue, ready, start, consumer_done),
    )
    producer.start()
    ready.wait()
    start.set()

    warmup_image = result_queue.get()
    start = time.perf_counter()
    received_images = [result_queue.get() for _ in range(REPETITIONS)]
    transfer_seconds = time.perf_counter() - start
    consumer_done.set()
    producer.join()
    result_queue.close()
    result_queue.join_thread()

    print(
        f"Throughput: {REPETITIONS / transfer_seconds:.2f} images/s "
    )

This results in the following output:

Forking pickle:       0.3 KiB
Serialization:        0.05 ms per image
Deserialization:      0.41 ms per image

Forking pickle payload:
  pickle stream: 353 bytes (0.34 KiB)
  reconstruction recipes:
    - rebuild_tensor
    - rebuild_typed_storage
    - rebuild_storage_fd
    - DupFd
  local shared-storage FD: 3 (process-local; receiver gets a duplicate)

Throughput: 2600.86 images/s 

The size of the IPC payload drops all the way down to 0.34 KiB and the throughput increases to 2600.86 images per second — a 51.5X improvement over the numpy experiment. The pickle contents include a recipe for mapping to the tensor in shared memory, in this case using a file descriptor.

Note, that a similar zero-copy SHM-based IPC optimization could be applied to numpy arrays, as well. See the multiprocessing.shared_memory API for details.

PyTorch Memory Sharing Strategies

PyTorch offers two different strategies for passing the shared memory locations across processes. The first uses dedicated file descriptors and the second uses explicit file names. Please consult the PyTorch documentation for details on the tradeoffs between the two options.

For the sake of completion, we repeat the experiment using the “file_system” strategy:

Forking pickle payload:
  pickle stream: 330 bytes (0.32 KiB)
  reconstruction recipes:
    - rebuild_tensor
    - rebuild_typed_storage
    - rebuild_storage_filename
  shared-memory names in pickle: /torch_2806_3240771400_0

Note, the inclusion of explicit file names in place of the file descriptor metadata.

Performance Pitfalls of PyTorch IPC

In the previous section, we saw how effective PyTorch’s zero-copy custom reducer for PyTorch tensors is for minimizing the overhead of IPC. In this section, we demonstrate how easily a naive data-loading implementation can completely decimate the effects of this optimization causing significant IPC overhead and introducing unforced data-input bottlenecks.

Example 1: Inclusion of Python Lists, Dicts, and Other Objects

A common mistake when designing PyTorch datasets is to include Python objects that are serialized and streamed over the IPC channel. We demonstrate this via a simple dataset of fixed-length token sequences.

import os
import time
import torch
from torch.utils.data import DataLoader, Dataset

BATCH_SIZE = 64
MEASURED_BATCHES = 200
NUM_WORKERS = os.cpu_count()
SEQUENCE_LENGTH = 2048
SAMPLE_COUNT = BATCH_SIZE * (MEASURED_BATCHES + 1)

class BaselineDataset(Dataset):

    def __init__(self, sample_count=SAMPLE_COUNT):
        self.sample_count = sample_count
        self.input_ids = torch.arange(SEQUENCE_LENGTH, dtype=torch.int64)

    def __len__(self):
        return self.sample_count

    def __getitem__(self, index):
        return {"input_ids": self.input_ids}

loader = DataLoader(
    BaselineDataset(),
    batch_size=BATCH_SIZE,
    num_workers=NUM_WORKERS,
)
iterator = iter(loader)
next(iterator)  # Exclude worker startup.

start = time.perf_counter()
for _ in range(MEASURED_BATCHES):
    next(iterator)
elapsed = time.perf_counter() - start
batches_per_second = MEASURED_BATCHES / elapsed

print(f"Throughput: {batches_per_second:.1f} batches/s")

The baseline throughput is 1465.1 batches/s. The baseline dataset is comprised of PyTorch tensors and utilizes the zero-copy reducer. The trouble starts when we start adding non-tensor metadata such as the raw strings from which the token sequences were calculated:

RAW_TEXT_BYTES = 8 * 1024

class RawTextDataset(Dataset):

    def __init__(self, sample_count=SAMPLE_COUNT):
        self.sample_count = sample_count
        self.input_ids = torch.arange(SEQUENCE_LENGTH, dtype=torch.int64)
        self.raw_texts = [
            f"sample-{sample_index:05d} " + "x" * (RAW_TEXT_BYTES - 13)
            for sample_index in range(self.sample_count)
        ]

    def __len__(self):
        return self.sample_count

    def __getitem__(self, index):
        return {
            "input_ids": self.input_ids,
            "raw_text": self.raw_texts[index],
        }

The inclusion of the raw text strings results in a throughput of 806.9 batches/s, a performance reduction of nearly 2X. Importantly, while the IPC overhead is a significant reason for this slowdown, it is not the necessarily the sole contributor. For example, the creation and collation of the metadata are certainly likely to slow the batch rate. Later on, when we discuss how to profile IPC bottlenecks, we will attempt to quantify how much of the degradation can be attributed to the IPC.

Note, that our example is simply iterating over batches without actually running any model. In practice, as with other bottlenecks on the data input pipeline, the time that it takes to run the model might hide the IPC overhead. However, even if it does not noticeably slow your overall throughput, the metadata inclusion increases the strain on the system resources and increases the likelihood of GPU starvation. In a high performance runtime environment, even the slightest misallocation of system resources can make a difference.

Bottom line — reduce the inclusion of metadata whenever possible. If absolutely required, use PyTorch tensors rather than non-tensor Python objects to benefit from the zero-copy reducer.

Example 2: Naive Custom Collation

Some datasets require the definition and use of a customized function for batching individual samples. Often, such customizations will omit optimizations that are applied by PyTorch’s default_collate function.

Get Chaim Rand’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

In the following code block we define a simple image dataset along with a naive custom batching function using torch.cat. Note, that the resultant tensor is not in shared memory. Consequently, before the worker process passes the batch to the main process, it needs to copy the batch to shared memory.

import time
import os
import torch
from torch.utils.data import DataLoader, Dataset

BATCH_SIZE = 32
MEASURED_BATCHES = 100
NUM_WORKERS = os.cpu_count()
IMAGE_SHAPE = (3, 1024, 1024)
SAMPLE_COUNT = BATCH_SIZE * (MEASURED_BATCHES + 1)

class ImageDataset(Dataset):

    def __init__(self):
        self.image = torch.arange(
            torch.tensor(IMAGE_SHAPE).prod(),
            dtype=torch.float32,
        ).reshape(IMAGE_SHAPE)

    def __len__(self):
        return SAMPLE_COUNT

    def __getitem__(self, index):
        return {"image": self.image}

def manual_collate(samples):
    images = torch.cat(
        [sample["image"].unsqueeze(0) for sample in samples],
        dim=0,
    )
    assert not images.is_shared()
    return {"image": images}

def measure_throughput(label, collate_fn):
    loader = DataLoader(
        ImageDataset(),
        batch_size=BATCH_SIZE,
        num_workers=NUM_WORKERS,
        collate_fn=collate_fn,
    )
    iterator = iter(loader)
    warmup_batch = next(iterator)

    start = time.perf_counter()
    for _ in range(MEASURED_BATCHES):
        batch = next(iterator)
    elapsed = time.perf_counter() - start
    batches_per_second = MEASURED_BATCHES / elapsed
    print(f"{label:<22} {batches_per_second:.1f} batches/s")


measure_throughput("default collate", None)
measure_throughput("manual private concat", manual_collate)

Using the default collate function results in a throughput of 14.6 batches/s, whereas the custom version reduces performance by 26.7%, to a throughput of 10.7 batches/s.

The following collator attempts to correct this by stacking the individual samples directly into shared memory:

def manual_shared_collate(samples):
    source_images = [sample["image"] for sample in samples]
    first_image = source_images[0]
    storage = first_image._typed_storage()._new_shared(
        sum(image.numel() for image in source_images),
        device=first_image.device,
    )
    shared_images = first_image.new(storage).resize_(
        len(source_images),
        *list(first_image.size()),
    )
    torch.stack(source_images, dim=0, out=shared_images)
    return {"image": shared_images}

This results in a throughput of 14.4 batches/s—much closer to the performance of the default collation but still slightly lower. This indicates that there may be additional optimizations we have neglected.

Bottom line — use the default collator whenever possible.

Example 3: Explicit Call of Pickling Function

Sometimes, developers will include explicit calls to the serialization and deserialization functions, as in the code block below.

import pickle

def pickle_collate(samples):
    batch = torch.utils.data.default_collate(samples)
    return pickle.dumps(batch)

loader = DataLoader(
    ImageDataset(),
    batch_size=BATCH_SIZE,
    num_workers=NUM_WORKERS,
    collate_fn=pickle_collate,
)
iterator = iter(loader)
warmup_batch = pickle.loads(next(iterator))

start = time.perf_counter()
for _ in range(MEASURED_BATCHES):
    batch = pickle.loads(next(iterator))
elapsed = time.perf_counter() - start
batches_per_second = MEASURED_BATCHES / elapsed
print(f"Explicit pickle: {batches_per_second:.1f} batches/s")

Unfortunately, an immediate byproduct of this is that it prevents the use of the PyTorch zero-copy reducer. In our case, the resultant throughput is 0.8 batches/s, nearly 20 times slower than the baseline result.

Bottom line — rely on the default PyTorch serialization rather than calling pickle explicitly (or use multiprocessing.reduction.ForkingPickler instead).

Identifying Performance Bottlenecks in IPC

Having demonstrated the overhead of IPC in AI/ML workloads and its potential impact on throughput, let us discuss ways of identifying performance bottlenecks resulting from suboptimal IPC.

Performance Analysis with PyTorch Profiler

In the code block below we return to the token dataset example from the previous section (example 1), this time with a simple PyTorch model and a PyTorch profiler-enabled training loop:

import os
import time
import torch
from torch import nn
from torch.nn import functional as F
from torch.profiler import ProfilerActivity

BATCH_SIZE = 256
VOCAB_SIZE = 5632
NUM_WORKERS = os.cpu_count()
TOTAL_STEPS = 200
PROFILE_START_STEP = 40
PROFILE_STOP_STEP = 50
TIMER_START_STEP = 100
TIMER_STOP_STEP = 200
SAMPLE_COUNT = BATCH_SIZE * (TOTAL_STEPS + 1)
TEST_TRIALS = 5
TEST_WARMUP_STEPS = 20
TEST_MEASURED_STEPS = 200


class NextTokenModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(SEQUENCE_LENGTH, VOCAB_SIZE)

    def forward(self, input_ids):
        return self.linear(input_ids.float())


def train_step(model, optimizer, batch, device):
    input_ids = batch["input_ids"].to(device, non_blocking=True)
    targets = (input_ids[:, -1] + 1) % VOCAB_SIZE
    optimizer.zero_grad(set_to_none=True)
    logits = model(input_ids)
    loss = F.cross_entropy(logits, targets)
    loss.backward()
    optimizer.step()


def profile_dataset(label, dataset_type, device):
    dataset = dataset_type()
    loader = DataLoader(
        dataset,
        batch_size=BATCH_SIZE,
        num_workers=NUM_WORKERS,
        pin_memory=True,
    )
    iterator = iter(loader)
    model = NextTokenModel().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)



    trace_path = f"raw_text_ipc_{label}.json"
    schedule = torch.profiler.schedule(
        wait=PROFILE_START_STEP,
        warmup=0,
        active=PROFILE_STOP_STEP - PROFILE_START_STEP,
        repeat=1,
    )

    def export_trace(profiler):
        profiler.export_chrome_trace(trace_path)


    with torch.profiler.profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        schedule=schedule,
        on_trace_ready=export_trace,
        with_stack=True,
    ) as profiler:
        for step in range(TOTAL_STEPS):
            if step == TIMER_START_STEP:
                torch.cuda.synchronize(device)
                timer_start = time.perf_counter()
            with torch.profiler.record_function("data_loader_next"):
                batch = next(iterator)
            with torch.profiler.record_function("gpu_train_step"):
                train_step(model, optimizer, batch, device)
            if step + 1 == TIMER_STOP_STEP:
                torch.cuda.synchronize(device)
                timer_elapsed = time.perf_counter() - timer_start
            profiler.step()

    torch.cuda.synchronize(device)
    iterator._shutdown_workers()
    timer_steps = TIMER_STOP_STEP - TIMER_START_STEP
    batches_per_second = timer_steps / timer_elapsed
    samples_per_second = BATCH_SIZE * batches_per_second

    print(f"{samples_per_second:.0f} samples/s")

device = torch.device("cuda")
profile_dataset("baseline", BaselineDataset, device=device)
profile_dataset("raw_text", RawTextDataset, device=device)

The results of this experiment are 123644 samples/s baseline dataset vs. 57523 samples/s for the raw-text dataset. The resultant traces are captured below:

PyTorch Profiler Trace for Baseline Dataset (By Author)

PyTorch Profiler Trace for Raw Text Dataset (By Author)

The raw-text trace shows a significant bottleneck in the data_loader_next call, coupled with extended periods of GPU idling. However, the trace does not point to the specific cause of this bottleneck.

Performance Analysis with NVIDIA Nsight Systems Profiler

We make another attempt at detecting the source of the bottleneck using the NVIDIA Nsight™ Systems (nsys) profiler.

wget https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2026_5/nsightsystems-linux-cli-public-2026.5.1.161-3889610.deb
sudo apt install ./nsightsystems-linux-cli-public-2026.5.1.161-3889610.deb

In the code block below we adapt the training step to use nsys. For more details on profiling with the nsys profiler, please see our previous post on the topic.

def profile_dataset(label, dataset_type, device):
    dataset = dataset_type()
    loader = DataLoader(
        dataset,
        batch_size=BATCH_SIZE,
        num_workers=NUM_WORKERS,
        pin_memory=True,
    )
    iterator = iter(loader)
    model = NextTokenModel().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

    for step in range(TOTAL_STEPS):
        if step == PROFILE_START_STEP:
            torch.cuda.profiler.start()
            torch.cuda.nvtx.range_push("profiled_training_steps")

        torch.cuda.nvtx.range_push("data_loader_next")
        batch = next(iterator)
        torch.cuda.nvtx.range_pop()
        torch.cuda.nvtx.range_push("gpu_train_step")
        train_step(model, optimizer, batch, device)
        torch.cuda.nvtx.range_pop()

        if step + 1 == PROFILE_STOP_STEP:
           torch.cuda.nvtx.range_pop()
           torch.cuda.profiler.stop()

    torch.cuda.synchronize(device)
    iterator._shutdown_workers()

To capture an nsys trace we run:

nsys profile \
  --capture-range=cudaProfilerApi \
  --trace=cuda,nvtx,osrt,syscall \
  --output=trace \
  python train.py

The resultant trace is shown below:

NVIDIA Systems Profiler Trace for Raw Text Dataset (By Author)

Once again, the raw-text trace shows a clear bottleneck alongside some increased activity in the OS runtime (OSRT) row which may hint at the IPC overhead — but there is still no clear smoking gun.

Identifying IPC Bottlenecks via Data Input Caching

In a previous post, A Caching Strategy for Identifying Bottlenecks on the Data Input Pipeline, we highlighted the difficulty of diagnosing bottlenecks on the data input pipeline in multi-worker settings. We described a multi-step solution for diagnosing such bottlenecks involving caching the input at various stages of the pipeline and measuring the impact on throughput performance. In this section, we extend the caching-based strategy to detect IPC bottlenecks.

In the following code block we extend the raw-text experiment with three caching options:

1. Cache a batch on the main process. This avoids IPC altogether and demonstrates the potential throughput absent a bottleneck on the batch creation pipeline.
2. Cache the samples returned from the [dataset.__getitems__]function as well as the batch in the collate function. This will result in the IPC execution but avoid all other sample preprocessing and batching.
3. Cache the samples returned from the [dataset.__getitems__] function. This will result in execution of the collate function and the IPC.

By comparing the throughput in each of these cases, we aim to understand the contribution of the inclusion of the metadata to the IPC overhead: If the throughput when caching a batch on the main process is significantly higher than when caching in the collate function, this may indicate a bottleneck in the IPC.

from torch.utils.data import default_collate

WARMUP_STEPS = 20
MEASURED_STEPS = 500
SAMPLE_COUNT = BATCH_SIZE * (WARMUP_STEPS + MEASURED_STEPS + 1)

class CacheableRawTextDataset(RawTextDataset):

    def __init__(self, sample_count):
        super().__init__(sample_count=sample_count)
        self.cache_getitems = False
        self._cached_items = None

    def __getitems__(self, indices):
        if not self.cache_getitems:
            return [self[index] for index in indices]
        if self._cached_items is None:
            self._cached_items = [self[index] for index in indices]
        return self._cached_items

class FullBatchCacheCollator:

    def __init__(self):
        self._cached_batch = None

    def __call__(self, samples):
        if self._cached_batch is None:
            self._cached_batch = default_collate(samples)
        return self._cached_batch

def measure_mode(label, dataset, device, reuse_first_batch=False,
                 cache_getitems=False, collate_fn=None):
    dataset.cache_getitems = cache_getitems
    dataset._cached_items = None
    loader = DataLoader(
        dataset,
        batch_size=BATCH_SIZE,
        num_workers=NUM_WORKERS,
        pin_memory=True,
        collate_fn=collate_fn,
    )

    iterator = iter(loader)
    model = NextTokenModel().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

    for _ in range(WARMUP_STEPS):
        batch = next(iterator)
        train_step(model, optimizer, batch, device)

    torch.cuda.synchronize(device)
    cached_batch = next(iterator) if reuse_first_batch else None
    start = time.perf_counter()
    for _ in range(MEASURED_STEPS):
        batch = cached_batch if reuse_first_batch else next(iterator)
        train_step(model, optimizer, batch, device)
    torch.cuda.synchronize(device)

    elapsed = time.perf_counter() - start
    iterator._shutdown_workers()
    step_ms = elapsed * 1000 / MEASURED_STEPS
    batches_per_second = MEASURED_STEPS / elapsed
    print(
        f"{label}: {batches_per_second:8.2f} batches/s"
    )
    return step_ms


device = torch.device("cuda")
dataset = CacheableRawTextDataset(sample_count=SAMPLE_COUNT)

measure_mode(
    "no cache",
    dataset,
    device,
)

measure_mode(
    "first next() batch reused",
    dataset,
    device,
    reuse_first_batch=True,
)

measure_mode(
    "__getitems__ + full collator cache",
    dataset,
    device,
    cache_getitems=True,
    collate_fn=FullBatchCacheCollator(),
)

measure_mode(
    "__getitems__ cache",
    dataset,
    device,
    cache_getitems=True,
)

The results of the experiment are summarized below:

Caching Experiment Results (By Author)

The fact that caching the batch in the main process results in a speedup of 3.69X is strong indication of an input-pipeline bottleneck. The considerable drop in performance that we see when moving the cache to before the IPC tells us that the IPC overhead is a significant contributor to this bottleneck. However, it is not the only contributor as we see from the additional drop when moving the cache to before the collate function.

IPC Optimization— A Toy Example

Our previous experiments demonstrated common pitfalls in dataset implementation that can increase IPC overhead unnecessarily. However, IPC optimization is not limited to just avoiding unforced errors. As we demonstrate in this section, sometimes, we can optimize IPC by simply restructuring how the raw samples are formatted.

In the code block below, we define a synthetic dataset of heterogenous graphs commonly used as inputs to graph neural networks (GNNs). We design each sample to consist of precisely 12 megabytes of data — matching the size of the image dataset we defined above:

import os
import time
import torch
from torch.utils.data import Dataset
from torch_geometric.data import HeteroData
from torch_geometric.loader import DataLoader as PyGDataLoader

TENSOR_COUNT = 200
DTYPE = torch.float32
TOTAL_BYTES = 12 * 1024**2
TOTAL_ELEMENTS = TOTAL_BYTES // 4
BASE_ELEMENTS, EXTRA_ELEMENTS = divmod(TOTAL_ELEMENTS, TENSOR_COUNT)
BATCH_SIZE = 32
MEASURED_BATCHES = 100
WARMUP_BATCHES = 10
NUM_WORKERS = os.cpu_count()
SAMPLE_COUNT = BATCH_SIZE * (WARMUP_BATCHES + MEASURED_BATCHES + 1)

def feature_shapes():
    shapes = {}
    for index in range(TENSOR_COUNT):
        element_count = BASE_ELEMENTS + int(index < EXTRA_ELEMENTS)
        shapes[f"feature_{index:03d}"] = (element_count,)
    return shapes

class HeterographDataset(Dataset):
    def __init__(self):
        flat = torch.arange(TOTAL_ELEMENTS, dtype=DTYPE)
        self.graph = HeteroData()
        offset = 0
        for name, shape in feature_shapes().items():
            element_count = shape[0]
            node_type = name.replace("feature_", "node_")
            self.graph[node_type].x = (
                flat[offset:offset + element_count].reshape(1, -1)
            )
            offset += element_count

    def __len__(self):
        return SAMPLE_COUNT

    def __getitem__(self, index):
        return self.graph

def measure_loader(label):
    loader = PyGDataLoader(
        HeterographDataset(),
        batch_size=BATCH_SIZE,
        num_workers=NUM_WORKERS,
    )
    iterator = iter(loader)
    for _ in range(WARMUP_BATCHES):
        next(iterator)

    start = time.perf_counter()
    for _ in range(MEASURED_BATCHES):
        next(iterator)
    elapsed = time.perf_counter() - start
    iterator._shutdown_workers()

    batches_per_second = MEASURED_BATCHES / elapsed
    print(f"{label}: {batches_per_second:8.2f} batches/s")

measure_loader("GNN baseline")

The resultant throughput is 2.48 batches/s, nearly 6 times (!!) slower than an image dataset of the same number of input bytes.

Our caching-based analysis (the details of which are omitted for brevity) points to IPC being a significant contributor towards this slowdown. Although the IPC payload has the same total number of bytes as the image dataset, the nature of the HeteroData Python object, and particularly the fact that the data is divided into 200 tensors — each requiring its own shared memory allocation and independent bookkeeping, makes the IPC much slower and potentially leads to GPU starvation.

In the code block below we attempt to address this bottleneck by reformatting the graph data input into a single packed tensor:

from torch.utils.data import DataLoader as TorchDataLoader
from torch_geometric.data import Batch

def pack_collated_graph(graph_batch):
    """Copy graph tensors directly into one preallocated shared buffer."""
    descriptor = []
    tensor_copies = []
    byte_offset = 0
    for store in graph_batch.stores:
        store_key = store._key
        for attr_key, value in store.items():
            if not isinstance(value, torch.Tensor):
                descriptor.append({
                    "store_key": store_key,
                    "attr_key": attr_key,
                    "value": value,
                })
                continue

            element_size = value.element_size()
            byte_offset += (-byte_offset) % element_size
            nbytes = value.numel() * element_size
            descriptor.append({
                "store_key": store_key,
                "attr_key": attr_key,
                "shape": list(value.shape),
                "dtype": str(value.dtype),
                "byte_offset": byte_offset,
                "nbytes": nbytes,
            })
            tensor_copies.append((value, byte_offset))
            byte_offset += nbytes

    empty_bytes = torch.empty(0, dtype=torch.uint8)
    shared_storage = empty_bytes._typed_storage()._new_shared(
        byte_offset,
        device=empty_bytes.device,
    )
    packed = empty_bytes.new(shared_storage).resize_(byte_offset)
    for value, offset in tensor_copies:
        destination = (
            packed[offset:offset + value.numel() * value.element_size()]
            .view(value.dtype)
            .reshape(value.shape)
        )
        destination.copy_(value.detach())
    return {
        "packed": packed,
        "descriptor": descriptor,
        "num_graphs": graph_batch.num_graphs,
    }


def packed_graph_collate(samples):
    graph_batch = Batch.from_data_list(samples)
    return pack_collated_graph(graph_batch)


def unpack_packed_graph(batch):
    """Rebuild a PyG HeteroDataBatch from packed stores and metadata."""
    packed = batch["packed"]
    graph = Batch(_base_cls=HeteroData)
    for entry in batch["descriptor"]:
        store_key = entry["store_key"]
        store = (
            graph._global_store
            if store_key is None
            else graph[store_key]
        )
        if "nbytes" in entry:
            start = entry["byte_offset"]
            end = start + entry["nbytes"]
            value = (
                packed[start:end]
                .view(getattr(torch, entry["dtype"].removeprefix("torch.")))
                .reshape(entry["shape"])
            )
        else:
            value = entry["value"]
        store[entry["attr_key"]] = value
    graph._num_graphs = batch["num_graphs"]
    return graph


def measure_loader(label, collate_fn, unpack_fn):
    loader = TorchDataLoader(
        HeterographDataset(),
        batch_size=BATCH_SIZE,
        num_workers=NUM_WORKERS,
        collate_fn=collate_fn,
    )
    iterator = iter(loader)
    for _ in range(WARMUP_BATCHES):
        unpack_fn(next(iterator))

    start = time.perf_counter()
    for _ in range(MEASURED_BATCHES):
        unpack_fn(next(iterator))
    elapsed = time.perf_counter() - start
    iterator._shutdown_workers()

    batches_per_second = MEASURED_BATCHES / elapsed
    print(f"{label}: {batches_per_second:8.2f} batches/s")

measure_loader(
    "GNN direct-to-SHM",
    packed_graph_collate,
    unpack_packed_graph,
)

The packing optimization results in a throughput of 7.2 batches/s, nearly 3 times faster than the unpacked experiment.

Summary

Multi-worker data preparation is a staple of modern-day AI/ML PyTorch applications. While it is essential for maximizing the utilization of our GPU (or alternative) accelerators, it is important to keep in mind that it relies on efficient communication between multiple OS processes. The details of inter-process communication (IPC) are mostly hidden from the user in PyTorch. As a result, its potential contribution to GPU starvation is sometimes overlooked. As we showed, although PyTorch optimizes IPC under the hood, these optimizations can be easily undermined by an inefficient dataset implementation — introducing a performance bottleneck, slowing down the application, and ultimately increasing the overall costs of AI/ML compute. In this post, we demonstrated how to identify and overcome such bottlenecks.

We hope you have found this post useful. Please be sure to check out our many other posts on the topic of AI/ML performance analysis and optimization.