Torch distributed and num_workers>0 pickle error

username: farakiko
cooked: <p>I think the big picture question here is “how can I serialize only the data, and not the dataloader object?”.</p>
<p>I have written down sample code here <a href="https://github.com/farakiko/particleflow/blob/pyg_ddp_numworkers_fix/mlpf/ddp_error.py" rel="noopener nofollow ugc">https://github.com/farakiko/particleflow/blob/pyg_ddp_numworkers_fix/mlpf/ddp_error.py</a><br>
and I have been debugging for some time with no solution <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"></p>
<p>The catch is</p>
<ol>
<li>it works for single gpu when num_workers is None or &gt;0</li>
<li>it works for multigpu <strong>only</strong> when num_workers is None (else i get the following error)<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/4/444387c9b061bad22a5fed1868d3a87ed4069013.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/444387c9b061bad22a5fed1868d3a87ed4069013" title="zzz 7"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/444387c9b061bad22a5fed1868d3a87ed4069013_2_640x500.jpeg" alt="zzz 7" data-base62-sha1="9JT9npjszbAAK3kRBh6eqMLWBfZ" width="640" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/444387c9b061bad22a5fed1868d3a87ed4069013_2_640x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/444387c9b061bad22a5fed1868d3a87ed4069013_2_960x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/444387c9b061bad22a5fed1868d3a87ed4069013_2_1280x1000.jpeg 2x" data-dominant-color="111111"><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">zzz 7</span><span class="informations">1920×1500 269 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></li>
</ol>
<p>The data_dir can be found here <a href="https://pfvol.nrp-nautilus.io/tensorflow_datasets_small/cms_pf_ttbar/1.6.0/" rel="noopener nofollow ugc">https://pfvol.nrp-nautilus.io/tensorflow_datasets_small/cms_pf_ttbar/1.6.0/</a></p>
<p>Please any help would be appreciated!</p>

username: H-Huang
cooked: <p>Are your Dataset and collate_fn picklable? <a href="https://pytorch.org/docs/stable/data.html#platform-specific-behaviors" class="inline-onebox" rel="noopener nofollow ugc">torch.utils.data — PyTorch 2.1 documentation</a></p>
<p>You could also try setting multiprocess start method to “fork” <a href="https://docs.python.org/3/library/multiprocessing.html#multiprocessing.set_start_method" class="inline-onebox" rel="noopener nofollow ugc">multiprocessing — Process-based parallelism — Python 3.12.0 documentation</a></p>

username: farakiko
cooked: <p>thanks, “fork” works for the dataloader part but then i can’t move the tensors to CUDA because fork doesn’t work with CUDA</p>

username: farakiko
cooked: <p>i assume since the code runs fine with num_workers&gt;0 for both cpu and single gpu, then my collate function is serializable. is that not always the case?</p>

username: farakiko
cooked: <p>ok it seems that my datasource cannot be serialized. is there a way around that? and do you happen to know why the code works fine with torch distributed expect when num_workers&gt;0</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/6/062f7dff44e5dc1f271507d8d062e7279827ee4e.png" data-download-href="https://discuss.pytorch.org/uploads/default/062f7dff44e5dc1f271507d8d062e7279827ee4e" title="zzz 10"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/6/062f7dff44e5dc1f271507d8d062e7279827ee4e_2_690x182.png" alt="zzz 10" data-base62-sha1="SICddnlJtq2huY7vN9Ojl8kjRI" width="690" height="182" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/6/062f7dff44e5dc1f271507d8d062e7279827ee4e_2_690x182.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/6/062f7dff44e5dc1f271507d8d062e7279827ee4e_2_1035x273.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/6/062f7dff44e5dc1f271507d8d062e7279827ee4e_2_1380x364.png 2x" data-dominant-color="F5F4F3"><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">zzz 10</span><span class="informations">1752×464 62 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></p>

username: H-Huang
cooked: <p>This is more of a mechanism of the dataloader rather than any distributed code, see the mulitprocess section of dataloader <a href="https://pytorch.org/docs/stable/data.html#multi-process-data-loading" class="inline-onebox" rel="noopener nofollow ugc">torch.utils.data — PyTorch 2.1 documentation</a> or you can try asking for help in that topic of the forum <a href="https://discuss.pytorch.org/c/data/37" class="inline-onebox">data - PyTorch Forums</a>. My understanding is that when <code>num_workers &gt; 0</code> each worker process will pickle the dataset when it has to create an instance of the dataloader, my guess is this is to prevent multiple instatiations on the dataset across multiple processes.</p>
<p>If you are dataset cannot be pickled, you can try modifying the dataset to make it “pickle-able”, see <a href="https://docs.python.org/3/library/pickle.html#what-can-be-pickled-and-unpickled" class="inline-onebox" rel="noopener nofollow ugc">pickle — Python object serialization — Python 3.12.0 documentation</a></p>

username: farakiko
cooked: <p>Thanks for the help! Will check those references <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

username: farakiko
cooked: <p>thanks alot for the help! your references helped alot!!</p>
<p>I managed to solve it by defining this part to avoid pickling the dataset_info class and just pickle a SimpleNamespace that holds the content the dataloader would need from dataset_info<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/e/4ed857a70c77f9bfb2d2540f8940dd685da4e829.png" data-download-href="https://discuss.pytorch.org/uploads/default/4ed857a70c77f9bfb2d2540f8940dd685da4e829" title="zzz 18"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/e/4ed857a70c77f9bfb2d2540f8940dd685da4e829_2_690x184.png" alt="zzz 18" data-base62-sha1="bfuKKTbcmxKy0HtngQT8lx1gV2p" width="690" height="184" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/e/4ed857a70c77f9bfb2d2540f8940dd685da4e829_2_690x184.png, https://discuss.pytorch.org/uploads/default/original/3X/4/e/4ed857a70c77f9bfb2d2540f8940dd685da4e829.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/4/e/4ed857a70c77f9bfb2d2540f8940dd685da4e829.png 2x" data-dominant-color="F5F1DB"><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">zzz 18</span><span class="informations">1004×268 24.5 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></p>
<p>will leave it here in case it helps anyone!</p>