huggingface / huggingface/datasets
Extreme memory bandwidth.
- 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
Assessment
This issue has not been assessed yet.