Memory spikes with large DataPipes

🐛 Describe the bug

I’ve noticed large “spikes” in memory usage at the start of epochs when using IterDataPipes with attributes that take a lot of memory. These can cause my training jobs to fail with out-of-memory errors.

Here’s a minimal example to reproduce:

import torch
import torchdata.datapipes as dp
from torch.utils.data import DataLoader
from tqdm import tqdm

NUM_WORKERS = 2
NUM_ITEMS = 100
ITEM_SIZE = 5_000_000


def get_item(x):
    return torch.rand(ITEM_SIZE)


def get_datapipe():
    datapipe = dp.iter.IterableWrapper(range(NUM_ITEMS))
    datapipe = datapipe.map(get_item)
    datapipe = datapipe.in_memory_cache()
    return datapipe


def main():
    datapipe = get_datapipe()
    dataloader = DataLoader(
        datapipe, batch_size=1, num_workers=NUM_WORKERS, persistent_workers=True
    )

    for epoch in range(3):
        print(f"Epoch {epoch + 1}")
        for _ in tqdm(dataloader, total=NUM_ITEMS * NUM_WORKERS):
            pass


if __name__ == "__main__":
    main()

The memory usage (logged with psutil) looks like this:

Here, start_epoch indicates the start of an epoch and first_iter corresponds to the first time each epoch we reach the pass statement in the dataloader loop. (To simplify the example code above I removed the code that logs start_epoch and first_iter. I logged the memory usage from a separate process.)

After some debugging, I can say that the memory spikes occur during the traversal of the graph that occurs in torch/utils/data/graph_settings.py::apply_random_seed() at the beginning of each epoch. Disabling the body of this function removes the memory spikes.

The spikes seem to be caused by the pickling in https://github.com/pytorch/pytorch/blob/99ded8bbcea896b02f1c0babb055329c503ca95e/torch/utils/data/graph.py#L23
The code here defines f = io.BytesIO() and pickles to f. If there are large datapipes to be pickled, it makes sense that the memory usage will blow up quickly and then fall again when f goes out of scope.

I tried replacing f = io.BytesIO() with f = open(os.devnull, "wb") (and adding f.close() at the end of the function). This didn’t eliminate the memory spikes but it did make them a bit smaller.

A few notes:

• it’s not necessary to use .in_memory_cache() to see these spikes; it seems that any datapipe that occupies a lot of memory will cause them
• I verified that the spikes do not occur with similar Dataset and IterableDataset subclasses
• the spikes do not occur if we remove the Dataloader and iterate directly over the datapipe.

Versions

I have tested the above with both

• torch 1.13.1 and torchdata 0.5.1 (my development environment)
• torch 2.0 and torchdata 0.6.0

I observed the same behavior in both cases.

Start with the minimal reproduction and inspect torch/utils/data/graph_settings.py::apply_random_seed(), then follow the pickling code referenced in torch/utils/data/graph.py around line 23. Confirm the memory spike and determine a change that preserves graph traversal while avoiding the excessive allocation; done means the reproduction no longer shows epoch-start spikes.