Ability to train on a sample of a training set for every epoch

With datasets the size of MLS it makes less sense to iterate over every train set once for every epoch. There's questions of data balance and just pure scarceness of compute resources. One reasonable approach is sampling a fixed number of samples/batches from each training set on each epoch, similar to the setup described in the GPT-3 paper. A training run could sample 10% of a huge dataset on each epoch and 200% of a smaller one.