Running torchscript on multiple GPUs

yrath: <p>Hi,<br>
I am trying to run an pytorch object detector using triton server. I used tracing for the model and scripting for the post-processing function.  For a single GPU the torchscript runs smoothly on the server.</p>
<p>But, in the multi-GPU case, I am not able to run the compiled script. This is because the scripted post-process function memories the GPU id (<code>cuda:0</code>) which I used to run the torchscript and expects all the tensor operations to be performed using that id. This invariably fails when the triton server passes any other cuda device.</p>
<p>Is there any workaround this?</p>

ptrblck: <p>How are you defining the multi-GPU use case? Could you explain your deployment and where it’s currently breaking?</p>

yrath: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="103797">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>How are you defining the multi-GPU use case?</p>
</blockquote>
</aside>
<p>In multi-GPU usecase I start one triton server on two GPUs and place one instance of the torchscript model on each GPU using triton config.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="103797">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Could you explain your deployment</p>
</blockquote>
</aside>
<p>I am using  <a href="https://github.com/triton-inference-server/server" rel="noopener nofollow ugc">triton server</a>. What other information is required?</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="103797">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>where it’s currently breaking?</p>
</blockquote>
</aside>
<ul>
<li>The triton server is running on <code>cuda:0</code> and <code>cuda:1</code></li>
<li>Torchscript was complied using <code>cuda:0</code></li>
</ul>
<p>The error is caused by model instance on <code>cuda:1</code>  of the server fails with the error message <code>expected device cuda:1 but got device cuda:0</code></p>

ptrblck: <p>Would it work, if you write the PyTorch model device-agnostic, i.e. use <code>cuda:0</code> as now and mask the other GPUs with <code>CUDA_VISIBLE_DEVICES</code>?</p>

koonyook: <p>As I have to load the same traced model to both cuda:0 and cuda:1 in the same script, I cannot mask any of the GPUs as you suggest.</p>
<p>The workaround I am using is to have two redundant traces on both cuda:0 and cuda:1 and load the right file to the right device. However, this double the space required for the model in deployment. And what if the third GPU is added? not very scalable.</p>
<p>It would be nice if torch.jit.load’s map_location can accept dictionary to allow me to map from cuda:0 to any cuda device I want and also keep cpu device as it is.<br>
Is there any reason that torch.jit.load’s map_location doesn’t accept dictionary like torch.load does?</p>

ptrblck: <p><code>TorchScript</code> is in maintenance mode and not ins active development anymore so I would recommend switching to the <code>torch.compile</code> stack.</p>