Excessive RAM Usage After Dataset Concatenation concatenate_datasets

[sam-hey]

Describe the bug

When loading a dataset from disk, concatenating it, and starting the training process, the RAM usage progressively increases until the kernel terminates the process due to excessive memory consumption.

#2276

Steps to reproduce the bug

rom datasets import  DatasetDict, concatenate_datasets

dataset = DatasetDict.load_from_disk("data")

...
...

combined_dataset = concatenate_datasets(
        [dataset[split] for split in dataset]
    )

#start SentenceTransformer training

Expected behavior

I would not expect RAM utilization to increase after concatenation. Removing the concatenation step resolves the issue

Environment info

sentence-transformers==3.1.1
datasets==3.2.0

python3.10

[sam-hey]

Adding a img from memray
https://gist.github.com/sam-hey/00c958f13fb0f7b54d17197fe353002f