Split batch un-evenly across multi-gpus

username: Diego999
cooked: <p>Hi,</p>
<p>Is there a way to not split batches evenly between gpus ?</p>
<p>Let’s give an example. Imagine we have a dataset where each sample has one document with one or more query vectors and the position of the related answer.</p>
<p>Example:</p>
<pre><code class="lang-auto">(Doc1, (q11, p11), (q12, p12)),
(Doc2, (q21, p21), (q22, p22), (q23, p23), (q24, p24))
(Doc3, (q31, p31))
(Doc4, (q41, p41))
</code></pre>
<p>Our batches could be:</p>
<pre><code class="lang-auto">batch = {
    'doc':     [doc1_tokens, doc2_tokens, doc3_tokens, doc4_tokens],
    'query':   [q11, q12, q21, q22, q23, q24, q31, q41],
    'y_pos':   [p11, p12, p21, p22, p23, p24, p31, p41],
}

repetition_vector = [2, 4, 1, 1]
</code></pre>
<p>For efficiency, a first module could compute the representation of the documents (RNN) and then a second one could repeat the embedding and use the query with.</p>
<pre><code class="lang-auto">doc_embds = model1(batch['docs'])
doc_embds = torch.repeat_interleave(doc_embds, repetition_vector, dim=1)

y_hat_pos = model2(doc_embds, batch['query'])

loss = criterion(y_hat_pos, batch['y_pos'])
</code></pre>
<p><strong>PROBLEM</strong><br>
Let’s imagine we have 2 gpus. doc_embds is gonna be split evenly into [doc1, doc2] and [doc3, doc4]; it will generate final embeddings such as [doc1, doc1, doc2, doc2, doc2, doc2] and [doc3, doc4] while queries and y_pos gonna be [q11, q12, q21, q22] with  [q23, q24, q31, q41] and  [0, 6, 3, 7] and [4, 5, 5, 9].</p>
<p>In this case, tensors will not match on each GPUs are the query vectors have been split evenly (which shouldn’t be the case).</p>

username: Diego999
cooked: <p>I don’t what solution for this problem is available:</p>
<ul>
<li>Write a sub-class of DataParallel for this ?</li>
<li>Should I reorganize my batches in another way such that using padding (would like to avoid this approach as it takes GRAM) ?</li>
<li>Forget the repeat_interleave and re-do the computation for each doc_embs (heavy computation).</li>
</ul>
<p>The gain of multi-gpu would quitte consequent for this use-case.</p>
<p>Thank you for your help !</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="Diego999" data-post="1" data-topic="64361">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/d/ecc23a/48.png" class="avatar"> Diego999:</div>
<blockquote>
<p>For efficiency, a first module could compute the representation of the documents (RNN) and then a second one could repeat the embedding and use the query with.</p>
</blockquote>
</aside>
<p>I’m not sure I understand your use case completely, but this description sounds more like model parallel / sharding, where the model itself is split to multiple devices, not like data parallel.</p>

username: Diego999
cooked: <p>I’m working with documents where given a query; I have to find the position of the “answer”. <span class="hashtag">#queries</span> varies across documents.</p>
<p>I have a document encoder and a query encoder. Normally, we could have</p>
<pre><code class="lang-auto">batch = {
    'doc':     [d1,  d1,  d2,  d2,  d2,  d2,  d3,  d4],
    'query':   [q11, q12, q21, q22, q23, q24, q31, q41],
    'y_pos':   [p11, p12, p21, p22, p23, p24, p31, p41],
}
</code></pre>
<p>As computing document embedding takes a while with RNN, I thought of doing:</p>
<pre><code class="lang-auto">batch = {
    'doc':     [d1,       d2,                 d3,  d4],
    'query':   [q11, q12, q21, q22, q23, q24, q31, q41],
    'y_pos':   [p11, p12, p21, p22, p23, p24, p31, p41],
    'nb_q':    [2,        4,                  1,   1  ],
}
</code></pre>
<p>So one model could compute each document embeddings with only 4 computations (d1, d2, d3,d4) and then use torch.repeat_interleave instead having 8 computations (2<em>d1, 4</em>d2, 1<em>d3, 1</em>d4).</p>
<p>Is it more clear ?</p>
<p>For your answer, I could split the model on 2 gpus but they aren’t big, both fit within a single gpu. Would be the gain worth ?</p>