Refreshing a sharded dataset in a PyTorch dataloader - pytorch

Question

I have a pretty large text dataset that needs to be preprocessed, with additional live augmentation on batches during training. I implemented a custom torch Dataset where I read my raw text file, chunk it into shards that contain ~5000 samples processed per our requirements. Each shard is a TensorDataset containing, for each sample, the tokens, token types, position ids, etc from HuggingFace tokenizers.
Since each shard is pretty large, we can only load a single shard in memory at a time. We don't really to the DDP mode, because our work requires a single GPU. Since we can't randomly access every index in the dataset due to sharding, we have to forgo the idx argument in the Dataset's __getitem__, and instead implement our own approach, where our Dataset object tracks the index inside the shard that we have retrieved so far. We can shuffle the shard contents and the loading order of the shards themselves, so that is not a concern.
Unfortunately, when we use multiple worker processes, they create replicas of the shard in the Dataset object, ergo they also create replicas of our internal index tracker. This means for 2 worker processes, we are returning 2 copies of the same entry in the shard.
I realize there is a way to directly control what indices a worker process can access inside our shard with the worker_init_fn, where we could use the size of the current shard to basically tell worker 1 to work on the first half, and worker 2 to work on the second half.
However, I am not too sure about how to refresh the shard when we exhaust samples. That is, since each worker contains a replica of the original shard, and their own counting method presumably, is there a way to share information between workers so we know both workers have retrieved, say 5000 samples, exhausting the shard. Then we can load a new shard into both workers. Ideally, the workers should share the shard as well, instead of creating another copy of the same data.
Addendum: I looked at a few examples of IterableDatasets from Pytorch at this link. It is unclear to me, for example in the following snipper, when the stream is being refreshed. If I could figure that out, I could probably synchronize a refresh of the shard across all workers.
An alternate solution for us is to use HDF5 files, but the read overhead is too large compared to out current sharding method.
def __getitem__(self, idx):
"""
Given an index, retrieve a single sample from the post-processed
dataset.
We retrieve from shards, so we need to track the current shard we are
working with, as well as which sample in the shard we are retrieving.
`self.getcount` tracks the current index, and is initialized to 0 in __init__
Shuffling is controlled with the `data_shuffle` argument during
initialization of the Dataset.
"""
self.last_idx = idx
response = self.sharded_dataset[self.getcount]
self.getcount += 1
if self.getcount == self.current_shardsize: # we have exhausted examples in this shard
self.getcount = 0
self.shard_load_index += 1 # increment the shard index that we will load
if self.shard_load_index == self.max_shard_number: # we have processed all shards.
self.shard_load_index = 0
self.sharded_dataset = self.load_shard(self.shard_load_index)
self.current_shardsize = len(self.sharded_dataset)
if self.masking:
self.sharded_dataset = self.sharded_refresh_mask_ids(self.sharded_dataset)
return response
Here, getcount is our internal tracker of the index. Each worker seems to get its own copy of getcount, so we get each sample from self.sharded_dataset twice.