Fix tfrecord dataset to load less files into memory

[sdtblck]

The current tfrecord dataset loads 1 tfrecord at a time into memory.
The deepspeed distributed wrapper causes the dataset to do this once, for every sample, for every GPU.
Maybe it would be best to preprocess / prefetch n samples, write them to disk, then load the correct sample from disk at train time.