Memory optimization for compositions of skinny matrix multiplications with pointwise operations

askorikov: <p>Let’s say I have some composition of skinny matrix multiplications and pointwise operations:</p>
<pre><code class="lang-auto">A = torch.rand(1024**2, 5)
B = torch.rand(5, 256**2)
C = torch.rand(256**2, 2)
result = torch.exp(A @ B) @ C
</code></pre>
<p>When calculated naively, the size of the matrix A @ B will be huge, and the memory utilization can be dramatically reduced by instead calculating the required matrix-vector products on-the-fly:</p>
<pre><code class="lang-auto">result = torch.zeros(1024**2, 2)
for i in range(256**2):
    result += torch.exp(A @ B[:, i]) * C[i, :]
</code></pre>
<p>Does anyone know if there is a generic way to achieve this kind of memory-targeting optimization with skinny matrix multiplications in PyTorch, perhaps either by reformulating the computation, or by using an optimizing compiler / other external tool?</p>
<p>Many thanks for any help.</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/askorikov">@askorikov</a>,</p>
<p>You could try using <code>vmap</code> from functorch (which comes packaged with the latest version of PyTorch). On the experimental namespace, there’s a ‘chunk’ version of <code>vmap</code>, which allows mini-batches of calculation to be computed in parallel (vmap example here: <a href="https://github.com/pytorch/functorch#vmap" class="inline-onebox" rel="noopener nofollow ugc">GitHub - pytorch/functorch: functorch is JAX-like composable function transforms for PyTorch.</a>). Using the chunk version of vmap will allow you to compute mini-batches (subject to memory constraints) in parallel, which should speed up the for example above.</p>

askorikov: <p>Thanks for your suggestion <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>!</p>
<p>Unfortunately, just <code>vmap</code> doesn’t seem to solve the memory problem, since the main memory optimization comes from the sequential reduction in the for loop. I guess I would need a (chunked) map-reduce operation to achieve what I want. It seems that JAX can do the required optimization using JIT compilation (<a href="https://github.com/google/jax/discussions/9505" class="inline-onebox" rel="noopener nofollow ugc">Reduce functionality for vmap · Discussion #9505 · google/jax · GitHub</a>), but if I understood correctly, <code>functorch</code> is not supported by TorchScript yet.</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/askorikov">@askorikov</a>,</p>
<p>Functorch does have a chunked vmap, it’s not on the docs (as I shared the standard vmap docs) but it’s in there. Instead of <code>functorch.vmap</code>, you use <code>functorch.experimental.chunk_vmap</code>, more info this issue: <a href="https://github.com/pytorch/functorch/issues/680" class="inline-onebox" rel="noopener nofollow ugc">vmap, jacrev should have an option to vectorize "chunks" instead of the entire batch at a time · Issue #680 · pytorch/functorch · GitHub</a></p>

askorikov: <p>The problem is that even the chunked version of <code>vmap</code> needs a large buffer to store the result of <code>A @ B</code> (<code>1024**2 x 256**2</code> in my example), since the entire output of <code>vmap</code> is stored in memory. Instead, I need to map a function to a chunk of data, reduce the intermediate result, and then go to the next chunk.</p>