Seq-to-Seq Encoder Decoder Models with Reinforcement Learning - CUDA memory consumption debugging

username: jabhinav
cooked: <p>Hi,</p>
<p>Problem summary: I have implemented a seq-to-seq encoder decoder model that uses Reinforcement Learning for model training. I have encountered following error while training (exactly when attention is being computed i.e.<br>
intermediate = vector.matmul(self._w_matrix).unsqueeze(1) + matrix.matmul(self._u_matrix)):</p>
<p>‘’‘RuntimeError: CUDA out of memory. Tried to allocate 944.00 MiB (GPU 0; 11.17 GiB total capacity; 9.86 GiB already allocated; 310.81 MiB free; 10.58 GiB reserved in total by PyTorch)’’’</p>
<p>Model Details:</p>
<ol>
<li>The encoder model encodes an input sequence using Pytorch’s LSTM cell (hidden_size=512).</li>
<li>The decoder model’s LSTM cell is initialised using encoder model’s output. And at every time-step(t), decoder model’s input is a concatenation of (a) previous time step’s prediction (t-1) and (b) attended input produced by attending to encoder outputs.</li>
<li>Training objective: We beam sample top-k predictions from the decoder model and generate the reward for each decoding. I am back-propagating loss = log probabilities of top-k decodings * their respective rewards (REINFORCE trick).</li>
</ol>
<p>The attention mechanism to compute weights is the following:<br>
Attention is computed between a vector, x (in our case h_t) and a matrix, y (outputs of encoder) using an additive attention function. The function has two matrices <code>W</code>, <code>U</code> and a vector <code>V</code>. The similarity between the vector <code>x</code> and the matrix <code>y</code> is computed as <code>V tanh(Wx + Uy)</code>.</p>
<p>Parameters: I am sampling 64 top decodings, batch-size: 16, hidden size of 512 for every LSTM cell (encoder and decoder both), max sequence length of 50 when decoding using beam search.</p>
<p>I know this might be confusing without exact code. Let me know if any part of the code needs to be shared for more clarity. I want to debug what is consuming most of the memory because the same error appears on GPU with larger memory of 32GBs (tried on k80 and v100).</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="jabhinav" data-post="1" data-topic="79324">
<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/jabhinav/48/23628_2.png" class="avatar"> jabhinav:</div>
<blockquote>
<p>I want to debug what is consuming most of the memory because the same error appears on GPU with larger memory of 32GBs (tried on k80 and v100).</p>
</blockquote>
</aside>
<p>I would recommend to add <code>print(torch.cuda.memory_allocated())</code> statements in your code and check each operation or layer sequentially.<br>
This would give you the overview where most of the memory is used.</p>

username: jabhinav
cooked: <p>Thanks, will debug and report if I am not able to resolve it.</p>

username: jabhinav
cooked: <p>Hi, <a class="mention" href="/u/ptrblck">@ptrblck</a> I tried debugging by checking each operation sequentially. Figured out that decoder is essentially the bottleneck because using encoder outputs, I am decoding new predictions at each time-step using previous step prediction and attended inputs (Using LSTM cell). Saving new predictions at each time step in a list is also saving its computation graph which is essentially increasing the memory consumed.</p>
<p>What I want to understand is for the following memory consumption during sequence decoding:</p>
<blockquote>
<p><em><strong>TRAINING ITERATION: 0</strong></em><br>
<code>Decoding start</code><br>
<code>allocated: 58M, max allocated: 82M, cached: 102M, max cached: 102M</code><br>
<code>-- Before Beam Sampling</code><br>
<code>allocated: 58M, max allocated: 82M, cached: 102M, max cached: 102M</code><br>
– <strong>TimeStep_0</strong> –<br>
<code>-- _prepare_output_projections | Start</code><br>
<code>allocated: 58M, max allocated: 82M, cached: 102M, max cached: 102M</code><br>
<code>-- _prepare_output_projections | New Decoder state computed</code><br>
<code>allocated: 60M, max allocated: 82M, cached: 102M, max cached: 102M</code><br>
– <strong>TimeStep_1</strong> –<br>
<code>--  _prepare_output_projections | Start</code><br>
<code>allocated: 94M, max allocated: 119M, cached: 156M, max cached: 156M</code><br>
<code>--  _prepare_output_projections | New Decoder state computed</code><br>
<code>allocated: 162M, max allocated: 207M, cached: 264M, max cached: 264M</code><br>
– <strong>TimeStep_2</strong> –<br>
<code>--  _prepare_output_projections | Start</code><br>
<code>allocated: 231M, max allocated: 256M, cached: 266M, max cached: 266M</code><br>
<code>--  _prepare_output_projections | New Decoder state computed</code><br>
<code>allocated: 299M, max allocated: 344M, cached: 374M, max cached: 374M</code><br>
…<br>
…<br>
– <strong>TimeStep_19</strong> –<br>
<code>-- 1) _prepare_output_projections | Start</code><br>
<code>allocated: 2395M, max allocated: 2420M, cached: 2442M, max cached: 2442M</code><br>
<code>-- 4) _prepare_output_projections | New Decoder state computed</code><br>
<code>allocated: 2464M, max allocated: 2508M, cached: 2548M, max cached: 2548M</code><br>
<code>-- After Beam Sampling</code><br>
<code>allocated: 2529M, max allocated: 2530M, cached: 2550M, max cached: 2550M</code></p>
</blockquote>
<blockquote>
<p><em><strong>TRAINING ITERATION: 1</strong></em><br>
<code>Decoding start</code><br>
<code>allocated: 214M, max allocated: 2600M, cached: 2664M, max cached: 2664M</code><br>
<code>-- Before Beam Sampling</code><br>
<code>allocated: 214M, max allocated: 2600M, cached: 2664M, max cached: 2664M</code><br>
– <strong>TimeStep_0</strong> –<br>
<code>-- _prepare_output_projections | Start</code><br>
<code>allocated: 214M, max allocated: 2600M, cached: 2664M, max cached: 2664M</code><br>
<code>-- _prepare_output_projections | New Decoder state computed</code><br>
<code>allocated: 218M, max allocated: 2600M, cached: 2664M, max cached: 2664M</code><br>
– <strong>TimeStep_1</strong> –<br>
<code>--  _prepare_output_projections | Start</code><br>
<code>allocated: 307M, max allocated: 2600M, cached: 2786M, max cached: 2786M</code><br>
<code>--  _prepare_output_projections | New Decoder state computed</code><br>
<code>allocated: 429M, max allocated: 2600M, cached: 2850M, max cached: 2850M</code><br>
– <strong>TimeStep_2</strong> –<br>
<code>--  _prepare_output_projections | Start</code><br>
<code>allocated: 605M, max allocated: 2600M, cached: 2974M, max cached: 2974M</code><br>
<code>--  _prepare_output_projections | New Decoder state computed</code><br>
<code>allocated: 727M, max allocated: 2600M, cached: 3036M, max cached: 3036M</code><br>
…<br>
…<br>
– <strong>TimeStep_33</strong> –<br>
<code>-- 1) _prepare_output_projections | Start</code><br>
<code>allocated: 9839M, max allocated: 9899M, cached: 9958M, max cached: 9958M</code><br>
<code>-- 4) _prepare_output_projections | New Decoder state computed</code><br>
<code>allocated: 9962M, max allocated: 10006M, cached: 10074M, max cached: 10074M</code><br>
<code>-- After Beam Sampling</code><br>
<code>allocated: 10136M, max allocated: 10137M, cached: 10196M, max cached: 10196M</code></p>
</blockquote>
<blockquote>
<p>_prepare_output_projections is the function that computes new decoder state (h_t, c_t) using previous step prediction and attended input. The function also returns new predictions by projecting h_t, output of LSTM cell into the vocabulary space.</p>
</blockquote>
<p>At each time-step, I am predicting for batch_size x beam_size sequences. That is, if <code>batch_size = 8</code> and <code>beam_size = 100</code>, I am making 8x100 = 800 predictions at each time step. Above stats are shown for <code>batch_size = 2</code> and <code>beam_size = 32</code>. <code>Max_time_steps_allowed = 50</code></p>
<p>I can see that with increase in number of predictions, allocated memory increases. Is there something wrong, or is this is the expected behaviour?  I want to experiment with <code>batch_size =16</code> for stable training. Should I go with gradient accumulation? Please, advise.</p>

username: ptrblck
cooked: <p>Do you need to call <code>backward</code> at some point and want Autograd to use all computation graphs?<br>
I’m not familiar with your use case, but for RNNs you could detach the last state(s) so that the backward pass only calculates the gradients for the current step.</p>

username: jabhinav
cooked: <p>I am calling backward on computed reward which is calculated in the following fashion:<br>
For each training sample in the batch, I will have to first decode <code>n</code> complete sequences (<code>n = beam_size</code>), evaluate them based on a metric to calculate reward(or loss) for back-propagation.</p>
<blockquote>
<p>For example, for each sample in the batch, decode n sequences so that output of decoder is of size: (batch_size, n, max_seq_len). The metric can evaluate sequences only after they are complete (i.e. only after end_token is sampled). Reward is computed for a total of batch_size*n sequences.</p>
</blockquote>
<blockquote>
<p>If using <code>batch_size = 8</code>, and <code>beam size = 32</code>. I am sampling <code>8*32</code> sequences from the decoder, calculating their (1) log probability and (2) reward, using Reinforce Trick to calculate loss which is <code>reward*log_probability</code> before calling backward. Here is the loss -</p>
</blockquote>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/b/6beb4da2c8b6c9d8b39e5f0baa086657b7681b2a.png" alt="image" data-base62-sha1="foHdU35PIMaSRnlPFCC4cnaO3AS" width="668" height="74"></p>
<p>The loss is the expectation of reward for decoded sequences, r(.) is the reward for complete sequences, <code>y_t</code> are hidden states sample from the decoder and <code>pi</code> is decoder model.</p>
<p>In this training regime, I have no GT to match predictions with at each time-step. I can only back-propagate after I have sampled complete sequence. Thus saving predictions for all sequences for all time-steps is saving their computation graphs which I think is incrementally increasing GPU memory consumption. Right?</p>
<p>Solutions tried: I have tried gradient accumulation where I am splitting batch_size into sub-batches, calling loss.backward for each of them before calling optimiser.step().</p>
<p>Question: Is there anything that can be done to minimise such huge memory consumption. Let me know if you need specific code excerpt to debug.</p>
<p>Paper for reference <a href="http://openreview.net/pdf?id=H1Xw62kRZ" rel="noopener nofollow ugc">https://openreview.net/pdf?id=H1Xw62kRZ</a></p>

username: ptrblck
cooked: <p>Thanks for the update.</p>
<aside class="quote no-group" data-username="jabhinav" data-post="6" data-topic="79324">
<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/jabhinav/48/23628_2.png" class="avatar"> jabhinav:</div>
<blockquote>
<p>Solutions tried: I have tried gradient accumulation where I am splitting batch_size into sub-batches, calling loss.backward for each of them before calling optimiser.step().</p>
</blockquote>
</aside>
<p>If your model is not depending on the batch size, e.g. via batchnorm layers, this should be a valid approach to trade compute for memory.<br>
Did you still run out of memory using this approach?</p>