<h3 dir="auto">Feature request</h3>
<p dir="auto">For IterableDataset, it seems there's no way to handle the uneven input tensors in trainer.py. (Please correct me if I misunderstant this).<br>
Pytorch document suggests to use join: <a href="https://pytorch.org/tutorials/advanced/generic_join.html#what-is-join" rel="nofollow">https://pytorch.org/tutorials/advanced/generic_join.html#what-is-join</a></p>
<p dir="auto">Secondly, the dataloader for IterableDataset doesn't have sampler, which may be an issue when num workers &gt; 0.<br>
<a href="https://github.com/huggingface/transformers/blob/v4.18.0/src/transformers/trainer.py#L672">https://github.com/huggingface/transformers/blob/v4.18.0/src/transformers/trainer.py#L672</a></p>
<h3 dir="auto">Motivation</h3>
<p dir="auto">Currently in trainer.py, the IterableDataset is wrapped into IterableDatasetShard in distributed training. It seems that it requires every process to have the same whole dataset and distribute the samples in <a href="https://github.com/huggingface/transformers/blob/v4.18.0/src/transformers/trainer_pt_utils.py#L678">IterableDatasetShard</a>, which makes every batch has the same amount of data by <a href="https://github.com/huggingface/transformers/blob/v4.18.0/src/transformers/trainer_pt_utils.py#L770">pad the first batch data</a>.</p>
<p dir="auto">When in training, since the data is guranteed to be equal in every process, there is no need to use <a href="https://pytorch.org/tutorials/advanced/generic_join.html#what-is-join" rel="nofollow">Join</a> to process the uneven input datasets.</p>
<p dir="auto">Here's my case. I have a very large audo dataset of 1 million audio files and each file requires the processing step just like <a href="https://github.com/huggingface/transformers/blob/main/examples/pytorch/speech-recognition/run_speech_recognition_ctc.py">this example</a>. If I do the same step as this example, the large training datasets not only requires a lot of time processing but also a lot of disk storage to store the features.<br>
Therefore I decide to use IterableDataset to do on-the-fly preprocessing. However, I need to maintain a buffer to do some shuffling then the CPU memory increases greatly and every process is doing many duplicate work that the other process has already done. For example, let's say we have 2 processes to process 11 samples. Every process has to process the 11 samples instead of 5 + 6.</p>
<p dir="auto">So I decide to give up the  IterableDatasetShard and use  IterableDataset directly. Every process has its own part of data to process. For example, process0 has 5 samples and process 1 has 6 samples. However this also means some process will finish training earlier. But I see there's no join used in trainer.py. So I was wondering what's the best practice to do this training and inference in this use case.</p>
<p dir="auto">Thanks!</p>
<h3 dir="auto">Your contribution</h3>
<p dir="auto">Not sure</p>
<h2 dir="auto"><g-emoji class="g-emoji" alias="rocket" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f680.png">🚀</g-emoji> Feature</h2>

<p dir="auto">I think Albert is popular enough to not say anything more. The link to the paper is below.<br>
<a href="https://arxiv.org/pdf/1909.11942v1.pdf" rel="nofollow">https://arxiv.org/pdf/1909.11942v1.pdf</a></p>
<h2 dir="auto">Motivation</h2>

<h2 dir="auto">Additional context</h2>