Parallelizing a loss function over CPU cores

jonathanking: <p>Hello everyone,</p>
<p>I’m training a sequence-based model on a single machine with 1 GPU and 16 CPU cores. My loss function is computationally expensive and performs best on the CPU. The reason it performs well on a CPU is that computing this particular loss is a sequential process and, while it can’t be parallelized/GPU optimized for a single data point, it can be parallelized across a batch of samples. The rest of the model works well on a GPU.</p>
<p><strong>I’d like to know how to parallelize the loss of my model over a batch of samples and CPU cores while keeping the rest of my model on the GPU.</strong> I know this will involve transferring the model’s predictions from GPU - &gt; CPU, but I believe the performance increase by using CPUs for the loss will outweigh this fact.</p>
<p>I’ve tried <code>torch.multiprocessing</code> as follows (where <code>single_item_loss</code> is the loss function I described above, that takes a single model prediction from the batch of predictions and returns a single Torch float tensor representing the loss):</p>
<pre><code class="lang-python">    with multiprocessing.Pool(multiprocessing.cpu_count()) as p:
        results = p.map(single_item_loss, predictions)
</code></pre>
<p>but this yields the error:</p>
<pre><code class="lang-auto">RuntimeError: Cowardly refusing to serialize non-leaf tensor which requires_grad, since autograd does not support crossing process boundaries.  If you just want to transfer the data, call detach() on the tensor before serializing (e.g., putting it on the queue).
</code></pre>
<p>I’ve also tried using the <code>joblib</code> library:</p>
<pre><code class="lang-auto">results = Parallel(n_jobs=multiprocessing.cpu_count(), 
            backend="loky")(delayed(single_item_loss)(pred) for pred in predictions)
</code></pre>
<p>but this must sever each of the idividual losses from the computation graph, because calling <code>torch.mean(torch.stack(results)).backward()</code> (and later <code>optimizer.step()</code>) has no effect and the model does not train.</p>
<p>I’ve also tried calling <code>loss.backward()</code> before returning <code>loss</code> from my <code>single_item_loss</code> function, but this also does not work.</p>
<p>I’m eager to hear any feedback and work to solve this problem. Evidently this is not a very common situation as GPUs are almost always preferred, but in this case, I believe the custom loss function warrants this special treatment. Thank you very much for your help!</p>

tom: <p>So currently this isn’t supported.</p>
<aside class="quote no-group" data-username="jonathanking" data-post="1" data-topic="60097">
<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/jonathanking/48/16256_2.png" class="avatar"> jonathanking:</div>
<blockquote>
<p>RuntimeError: Cowardly refusing to serialize non-leaf tensor which requires_grad, since autograd does not support crossing process boundaries. If you just want to transfer the data, call detach() on the tensor before serializing (e.g., putting it on the queue).</p>
</blockquote>
</aside>
<p>You could implement forward and backward separately and stick it in a autograd.Function. Then inside the forward, you don’t have grad-requiring things.<br>
For a loss function in particular, you could compute the backward inside the forward as well (you can turn also autograd back on inside the forward with <code>torch.enable_grad</code>).<br>
Then in the backward you just hand back the pre-computed result. (This is what CuDNN’s CTC Loss implementation does.)</p>
<p>It might be more efficient in terms of development time as well as runtime to drop into C++ and use PyTorch’s <code>parallel_for</code>, though.</p>
<p>Note that PyTorch tries to use intra-op parallelism, you would not want to “overbook” your cores.</p>
<p>Best regards</p>
<p>Thomas</p>

kako_lelouch: <p>Hi,</p>
<p>Did you find an answer to make the following code work? I am also looking for a way to backward the total loss which is construct from the result from multiprocessing.</p>
<pre><code class="lang-auto">results = Parallel(n_jobs=multiprocessing.cpu_count(), 
            backend="loky")(delayed(single_item_loss)(pred) for pred in predictions)
torch.mean(torch.stack(results)).backward()
optimizer.step()
</code></pre>