An idea about better balance strategy num_workers in DataLoader

username: zhenxz0121
cooked: <p>In my project, I need to load video frames into memory. Through analysis, I observed that the <code>model.forward()</code> operation takes approximately 1 second, while loading the data consumes about 18 seconds (without multi-processing). It is evident that I require (additional) workers to independently read the data.<br>
I have 2 RTX 4090 GPUs with a total of 70 GB memory for Distributed Data Parallel (DDP) training. The batch size is set to 128, and the dataset used is MSRVTT. The frame rate is only 1, and the maximum frames per video is 12. During my testing, I observed that setting <code>num_worker</code> greater than 16 for each GPU leads to some workers being killed during training due to excessive memory usage, resulting in the termination of the training process.<br>
Upon monitoring the training process, I noticed that workers employ a greedy strategy for data loading, prioritizing efficiency without considering memory constraints. Theoretically, we can schedule the workers more evenly, ensuring that the batch data in memory remains within the range of (1x batch_size, 2x batch_size). This not only prevents fluctuations in CPU work caused by a sudden consumption of workers’ efforts by GPUs(4090 is very fast) but also optimizes memory utilization more reasonably. This approach aims to avoid a peak utilization scenario that could lead to worker termination, while maintaining the original greedy strategy’s GPU utilization (as the time for each step is generally fixed, theoretically allowing for a balanced throughput by evenly arranging workers)."<br>
It’s happy to discuss about this idea.</p>

username: zhenxz0121
cooked: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/4/04f76352461b816cbf6a70c4e45886c2444d4742.png" data-download-href="https://discuss.pytorch.org/uploads/default/04f76352461b816cbf6a70c4e45886c2444d4742" title="1700753640678"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/4/04f76352461b816cbf6a70c4e45886c2444d4742_2_339x500.png" alt="1700753640678" data-base62-sha1="HVW2DZnC0IkJYGsucHxnPQYIue" width="339" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/4/04f76352461b816cbf6a70c4e45886c2444d4742_2_339x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/4/04f76352461b816cbf6a70c4e45886c2444d4742_2_508x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/4/04f76352461b816cbf6a70c4e45886c2444d4742_2_678x1000.png 2x" data-dominant-color="232222"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">1700753640678</span><span class="informations">847×1249 68.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
This interval is caused by non-even executation strategy of workers. It’s true that i can get more workers, but they will cause the memory peak.</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="zhenxz0121" data-post="1" data-topic="192478">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zhenxz0121/48/65718_2.png" class="avatar"> zhenxz0121:</div>
<blockquote>
<p>Theoretically, we can schedule the workers more evenly, ensuring that the batch data in memory remains within the range of (1x batch_size, 2x batch_size).</p>
</blockquote>
</aside>
<p>I don’t fully understand this idea as workers will already use the specified <code>prefetch_factor</code> which would thus limit the overall memory usage. Reduce it, if needed.</p>

username: zhenxz0121
cooked: <p>Thanks for your response. Let me put it simply: the <code>prefetch_factor</code> only constrains the memory usage for batches that are multiples of integer n, n equals to `num_workers’. If I’m concerned about overall memory usage, I can reduce it to 1. Then, I can set the number of workers to be equal to the time of loading the data divided by the average time for a GPU to run one step. This way, I can maintain high GPU utilization.</p>
<p>However, I think it would be better if we could limit the length of the entire prefetch batch queue. On top of that, we could add a small amount of redundant data batch (see the <code>prefetch_factor</code> = 1 and the number of workers are set like last paragraph as the non-redundant data batch) like 3~5 batches, instead of whole <code>num_workers</code> batches. I hope the limitation on prefetching data can shift from a discrete value to a continuous range, specified in terms of batches</p>

username: ptrblck
cooked: <p>Right now the data loading order should be deterministic by adding batches to the queue from each worker sequentially. If you now want to add more flexibility by limiting this queue you would still need to guarantee the data loading order as determinism is important for a lot of use cases.<br>
Of course it’s possible, but I don’t know how much complexity this would add into the <code>DataLoader</code> multiprocessing backend. You could create a feature request and discuss it with the code owners on GitHub.</p>

username: zhenxz0121
cooked: <p>Right, complexity is an important aspect to consider. Thanks for your advice and discussion <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>