Multiplying with a diagonal matrix becomes very expensive

username: legoh
cooked: <p>Hi! I have a very simple set of embeddings of size <code>|x| = [b, p, dim]</code>, where along each axis of <code>p</code>, I applying some form of scaling to the <code>dim</code> axis. In the original setting, matrix <code>x</code> is supposed to be multiplied by matrix <code>z</code>, where <code>|z| = [b, p, dim]</code>. This gives us <code>|original_result| = [b, p, p]</code>.</p>
<p>Now, I wish to multiply this by a diagonal matrix <code>y</code> for every <code>p</code>, where the diagonal is of size <code>(dim, dim)</code>. This matrix <code>y</code> also tracks gradient operations. The size of <code>y</code> is <code>|y| = [b, p, dim]</code> As such, I have been doing it this way:</p>
<pre><code class="lang-auto">x = x.unsqueeze(-2) #(b, p, 1, dim)

y = linear(x) # (b, p, dim)

y = torch.diag_embed(y)  # (b, p, dim, dim)

result = torch.matmul(x, y)  # (b, p, 1, dim)

new_result = torch.matmul(result.squeeze(-2), z)  # (b, p, n)
</code></pre>
<p>Now, from what I have noticed, this multiplication makes the GPU memory explode. Without producing <code>y</code> and doing any matmul, the GPU usage is about 5GB. However, once I introduce <code>y</code> and do the operations, the GPU usage hits 21GB.</p>
<p>I am wondering how we can make this operation smoother. Are there more efficient operations that I can take advantage of so as to reduce the computing cost?</p>

username: KFrank
cooked: <p>Hi Legoh!</p>
<aside class="quote no-group" data-username="legoh" data-post="1" data-topic="181580" data-full="true">
<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/l/82dd89/48.png" class="avatar"> legoh:</div>
<blockquote>
<p>I applying some form of scaling to the <code>dim</code> axis.</p>
</blockquote>
</aside>
<p>Based on the code fragment you posted, it appears that you are performing<br>
element-wise multiplication of two tensors in a roundabout way.</p>
<p>If so, there is no need for the (potentially large amount of) memory required<br>
to materialize the full <code>[b, p, dim, dim]</code> result of <code>diag_embed (y)</code>.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.0.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2023)
&gt;&gt;&gt;
&gt;&gt;&gt; b = 2
&gt;&gt;&gt; p = 2
&gt;&gt;&gt; dim = 3
&gt;&gt;&gt;
&gt;&gt;&gt; x = torch.randn (b, p, dim)     # (b, p, dim)
&gt;&gt;&gt; u = x                           # save a reference to original x
&gt;&gt;&gt; x = x.unsqueeze (-2)            # (b, p, 1, dim)
&gt;&gt;&gt; y = torch.randn (b, p, dim)     # (b, p, dim)
&gt;&gt;&gt; v = y                           # save a reference to original y
&gt;&gt;&gt; y = torch.diag_embed (y)        # (b, p, dim, dim)
&gt;&gt;&gt;
&gt;&gt;&gt; result = torch.matmul(x, y)     # (b, p, 1, dim)
&gt;&gt;&gt; result = result.squeeze (-2)    # (b, p, dim)
&gt;&gt;&gt;
&gt;&gt;&gt; resultB = u * v                 # element-wise multiplication of original x and y
&gt;&gt;&gt;
&gt;&gt;&gt; torch.equal (result, resultB)   # same result
True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: legoh
cooked: <p>Thank you! I completely forgot that it was the case. This seems to have solved the memory issues in a big way, thanks!</p>