Use joblib to train an ensemble of small models on the same GPU in parallel

SunHaozhe: <p>My use case is to train multiple small models to form an parallel ensemble (for example, a bagging ensemble which can be trained in parallel), <a href="https://github.com/TorchEnsemble-Community/Ensemble-Pytorch/blob/master/torchensemble/bagging.py" rel="noopener nofollow ugc">an example code</a> can be found in the <a href="https://github.com/TorchEnsemble-Community/Ensemble-Pytorch" rel="noopener nofollow ugc">TorchEnsemble</a> library (which is part of PyTorch ecosystem).</p>
<p>This example code uses <code>joblib</code> <a href="https://joblib.readthedocs.io/en/latest/generated/joblib.Parallel.html" rel="noopener nofollow ugc">library</a> to train multiple small models in parallel on the same GPU. The core part of the parallel training logic is here:</p>
<pre><code class="lang-python">from joblib import Parallel, delayed

# Maintain a pool of workers
with Parallel(n_jobs=self.n_jobs) as parallel:
    # Training loop
    for epoch in range(epochs):
        rets = parallel(delayed(_parallel_fit_per_epoch)(...))
</code></pre>
<p>where <code>_parallel_fit_per_epoch </code> function is responsible to train each base model for only one epoch, its core logic is like this:</p>
<pre><code class="lang-python">def _parallel_fit_per_epoch(...):
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
</code></pre>
<p>Suppose that we want to train 10 base models to form our ensemble, it seems to me that setting <code>n_jobs</code> &gt; 1 does not provide faster training speed. If I set <code>n_jobs=2</code> to train 10 base models, I would expect each job trains 5 base models and the 2 jobs are run in parallel, so in theory this should result in 1/2 of the training time compared to using <code>n_jobs=1</code> to train 10 base models, which means a factor of 2 of training speedup. However, in practice this did not give any speedup but slower down the training: <code>n_jobs=2</code> takes slightly longer time than <code>n_jobs=1</code> (to train 10 base models). I would wonder why this happens?</p>
<p>My hypothesis is that when joblib is combined with PyTorch, what truly happens is <a href="https://qr.ae/pvMEjG" rel="noopener nofollow ugc">Time Sharing</a> (as suggested by this <a href="https://qr.ae/pvMEjG" rel="noopener nofollow ugc">Quora answer</a>). At any given time, only one single job is using the GPU (cuda). By design, the Time Sharing mechanism does not save the total training time and it also adds a context switch overhead cost, which explains why the total training time with joblib multiprocessing is even longer for model training with GPU. Is this correct?</p>
<p>Anyway, is there any way I could <a href="https://qr.ae/pvMEjG" rel="noopener nofollow ugc">divide the cores of a single GPU into multiple groups and divide the GPU memory accordingly</a> to train a parallel ensemble of small models with PyTorch? Is this even possible?</p>
<p>P.S. I do not know how to write cuda code and I am not an expert for multiprocessing.</p>
<p>As we suppose that each base model in the ensemble is a small model, we assume that the GPU memory of this single GPU is sufficient to host all these models at the same time (for forward pass + backward pass).</p>

ptrblck: <aside class="quote no-group" data-username="SunHaozhe" data-post="1" data-topic="157831">
<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/sunhaozhe/48/14036_2.png" class="avatar"> SunHaozhe:</div>
<blockquote>
<p>If I set <code>n_jobs=2</code> to train 10 base models, I would expect each job trains 5 base models and the 2 jobs are run in parallel, so in theory this should result in 1/2 of the training time compared to using <code>n_jobs=1</code> to train 10 base models, which means a factor of 2 of training speedup.</p>
</blockquote>
</aside>
<p>That’s not necessarily the case as you would need to have free compute resources on your GPU and would also need to use different streams while still providing enough resources for a parallel execution. If a single stream already saturates the compute resources from your GPU, the stream launches will be serialized as they cannot run in parallel.</p>

SunHaozhe: <p>Thanks a lot for your answer, do I need to write the <a href="https://pytorch.org/docs/stable/notes/cuda.html#cuda-streams" rel="noopener nofollow ugc">cuda stream code</a> myself to train several tiny/small models in parallel on a single GPU or <code>joblib</code> <a href="https://joblib.readthedocs.io/en/latest/generated/joblib.Parallel.html" rel="noopener nofollow ugc">library</a> combined with PyTorch automatically handles this under the hood?</p>
<p>This <a href="https://www.quora.com/Can-I-run-multiple-deep-learning-models-on-the-same-GPU/answer/Luming-Wang-2?ch=10&amp;oid=41184475&amp;share=330d5a63&amp;srid=uiFeQ&amp;target_type=answer" rel="noopener nofollow ugc">Quora answer</a> says that there are two possible mechanisms to break one GPU for several tasks running at the same time: Dividing the cores of GPU into multiple groups and dividing the GPU memory accordingly, or Time Sharing. Which mechanism among these two is happening when I write the <a href="https://pytorch.org/docs/stable/notes/cuda.html#cuda-streams" rel="noopener nofollow ugc">cuda stream code</a>?</p>

ptrblck: <p>“Time sharing” is not a programming approach and the user describes it as a way for e.g. cloud providers to swap between jobs of different users. Of course PyTorch does not implement such a behavior for your workstation.</p>
<blockquote>
<p>Divide the cores of GPU into multiple groups and divide the GPU memory accordingly.</p>
</blockquote>
<p>I don’t fully understand this claim and guess the user points out that each kernel should not occupy the entire GPU resources to allow concurrent kernels to be executed.</p>
<p>Take a look at <a href="https://discuss.pytorch.org/t/cuda-stream-not-working/155749/4">this topic</a> and the linked GTC talk which gives you a good overview how the GPU works.</p>