Replacing SubsetRandomSampler by RandomSampler in BATCH_SAMPLER

[NohTow]

Summary

This PR replaces the SubsetRandomSampler by a RandomSampler in the BACH_SAMPLER implementation.
It allows to drop memory usage by a large amount without losing performance or changing the current behavior.

Motivations

• SubsetRandomSampler creates an explicit list of indices that is stored in memory. This can become very large for large datasets, and even more so when using multiple datasets.
• RandomSampler actually uses lower level PyTorch tensors which are less costly and sample on the fly instead of keeping it in memory.
• When used to sample from the whole dataset, RandomSampler (without replacement, the default behavior) is equivalent to SubsetRandomSampler.

Results & tests

• I ran a PyLate training on the whole unsupervised Nomic dataset with 8 GPUs and 8 workers and the RAM usage dropped from 1T6GB to 620GB (this is fairly extreme, but it highlights how much we can gain). The results looked pretty sane
• @tomaarsen ran a small bench and the RandomSampler is both faster and more memory efficient for sampling a 1M dataset (80.0MB, 1.8s-1.9s vs 136.0MB, 2.5s-2.7s)
• We also discussed with Tom as well and I am pretty sure this is equivalent and just a free lunch as we are not sampling from a subset.

Changes

• Simply replace SubsetRandomSampler with RandomSampler in the BATCH_SAMPLER implementation.

Related Issue

I did not actually open an issue about it, but I can if needed.

Shout-out to @arn4 for finding the reason why simply sampling ids was actually very slow and costly in a resource intensive setup!

[tomaarsen]

I indeed ran some tests with this, and it seems like the RandomSampler is indeed faster and less memory-intensive! Thanks a bunch for this.