huggingface / huggingface/datasets

Memory leak / Large memory usage with num_workers = 0 and numerous dataset within DatasetDict

Open
#7,829 4 comments 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
22k
Forks
3.4k
Avg merge
5d 7h
Merged PRs (30d)
17

Description

### Describe the bug

Hi team, first off, I love the datasets library! 🥰

I'm encountering a potential memory leak / increasing memory usage when training a model on a very large DatasetDict.

Setup: I have a DatasetDict containing 362 distinct datasets, which sum up to ~2.8 billion rows.

Training Task: I'm performing contrastive learning with SentenceTransformer and Accelerate on a single node with 4 H100, which requires me to sample from only one dataset at a time.

Training Loop: At each training step, I sample ~16,000 examples from a single dataset, and then switch to a different dataset for the next step. I iterate through all 362 datasets this way.

Problem: The process's memory usage continuously increases over time, eventually causing a stale status where GPUs would stop working. It seems memory from previously sampled datasets isn't being released. I've set num_workers=0 for all experiments.

Chart 1: Standard DatasetDict The memory usage grows steadily until it make the training stale (RSS memory) Image

Chart 2: IterableDatasetDict I also tried to use IterableDatasetDict and IterableDataset. The memory curve is "smoother," but the result is the same: it grows indefinitely and the training become stale. Image

Any feedback or guidance on how to manage this memory would be greatly appreciated!

### Steps to reproduce the bug

WIP, I'll add some code that manage to reproduce this error, but not straightforward.

### Expected behavior

The memory usage should remain relatively constant or plateau after a few steps. Memory used for sampling one dataset should be released before or during the sampling of the next dataset.

### Environment info

Python: 3.12
Datasets: 4.3.0
SentenceTransformers: 5.1.1

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.