huggingface / huggingface/datasets

Extreme memory bandwidth.

Open
#7,498 0 comments 0 reactions 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

When I use hf datasets on 4 GPU with 40 workers I get some extreme memory bandwidth of constant ~3GB/s.

However, if I wrap the dataset in `IterableDataset`, this issue is gone and the data also loads way faster (4x faster training on 1 worker).

It seems like the workers don't share memory and basically duplicate the data 4x40.

### Steps to reproduce the bug

Trainer arguments:
```
dataloader_pin_memory=True,
dataloader_num_workers=40,
dataloader_prefetch_factor=2,
dataloader_persistent_workers=True,
```
Call trainer:
```
trainer = Trainer(
model=model,
args=train_args,
train_dataset=load_from_disk('..').with_fromat('torch'),
)
```
The dataset has 600GB and consists of 1225 files.

### Expected behavior

The optimal bandwidth should be 100MB/s to keep up with GPU.

### Environment info

Linux
Python 3.11
datasets==3.2.0

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.