Torch.cuda.memory_snapshot() shows an unexpected 8519680B memory allocation during matrix multiplication

HermitSun: <p>I met an unexpected memory allocation in <code>torch.cuda.memory_snapshot()</code> when I ran the following code:</p>
<pre><code class="lang-py">import torch

W = torch.rand(1000, 10).cuda()
b = torch.rand(10).cuda()
X = torch.rand(1000).cuda()

y = X @ W + b
print(torch.cuda.memory_snapshot())
</code></pre>
<p>Since <code>W</code> requests 1000 * 10 * 4B = 40000B, <code>b</code> requests 10 * 4B = 40B, <code>X</code> requests 1000 * 4B = 4000B, <code>y</code> requests 10 * 4B = 40B, and there is a matrix alignment mechanism, if everything is as expected, the output should look something like this (and I did get the similar result in another env, this output is from torch 1.10.0):</p>
<pre><code class="lang-auto">[
  {
    device: 0,
    address: 68780294144,
    total_size: 2097152,
    allocated_size: 45568,
    active_size: 45568,
    segment_type: "small",
    blocks: [
      { size: 40448, state: "active_allocated" },
      { size: 512, state: "active_allocated" },
      { size: 4096, state: "active_allocated" },
      { size: 512, state: "inactive" },
      { size: 512, state: "active_allocated" },
      { size: 2051072, state: "inactive" },
    ],
  },
];
</code></pre>
<p>But this time, besides the output above, I got some additional output like this:</p>
<pre><code class="lang-auto">{
  device: 0,
  address: 23007248515072,
  total_size: 20971520,
  allocated_size: 8519680,
  active_size: 8519680,
  requested_size: 8519680,
  stream: 0,
  segment_type: "large",
  blocks: [
    { size: 8519680, requested_size: 8519680, state: "active_allocated" },
    { size: 12451840, requested_size: 0, state: "inactive" },
  ],
}
</code></pre>
<p>It always occupies this fixed size: <strong>8519680B</strong>, and this only happens when I do matrix multiplications or use a <code>nn.Linear</code>. If I use <code>nn.Conv2d</code> or <code>nn.RNN</code>, this magic number will disappear, just as expected.</p>
<p>The environment in which I met this problem: PyTorch 2.0.0 + CUDA 11.4 + Ubuntu 20.04.</p>
<p>The environment in which I didn’t meet this problem: PyTorch 1.10.0 + CUDA 11.6 + Windows 10.</p>
<p>Any information will be helpful.</p>

ptrblck: <p>This behavior looks expected as we are allocating a cuBLAS workspace.<br>
You could delete it via <code>torch._C._cuda_clearCublasWorkspaces()</code> and rerun the memory snapshot again.<br>
Also <a class="mention" href="/u/eqy">@eqy</a> pointed out that e.g. removing the cuBLAS workspace via:</p>
<pre><code class="lang-python">import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":0:0"
</code></pre>
<p>before creating the handle should avoid the memory allocation.</p>

HermitSun: <p>That’s right, thank you!</p>
<p>Could you tell me something about what’s under the hood of the 8519680B memory? I want to know more about why it is this value.</p>

ptrblck: <p>The default cuBLAS workspace size is defined <a href="https://github.com/pytorch/pytorch/blob/31f311a816c026bbfca622d6121d6a7fab44260d/aten/src/ATen/cuda/CublasHandlePool.cpp#L47">here</a> as:</p>
<pre><code class="lang-auto">const size_t default_size = sm90 ? 4096 * 8 * 1024 : 4096 * 1024 * 2 + 16 * 1024 * 8;
</code></pre>
<p>Assuming you are not using a Hopper device (<code>sm_90</code>) you would thus use:</p>
<pre><code class="lang-auto">4096 * 1024 * 2 + 16 * 1024 * 8 = 8519680 bytes
</code></pre>
<p>which is split into different internal workspaces.</p>

HermitSun: <p>Thank you for your excellent reply <img src="https://discuss.pytorch.org/images/emoji/apple/smiling_face_with_three_hearts.png?v=12" title=":smiling_face_with_three_hearts:" class="emoji" alt=":smiling_face_with_three_hearts:" loading="lazy" width="20" height="20"></p>

HermitSun: <p>Thank you for your patience.</p>
<p>I have another question that why this magic number(i.e. 8519680B, if not using <code>sm_90</code> devices) be chosen as the default? Is it for performance consideration, or for kernel scheduling correctness, or some other reason?</p>

ptrblck: <p>For performance reasons a larger workspace is used by default on Hopper GPUs (H100 in this case).<br>
Since these devices also come with a large memory pool increasing the cuBLAS workspace will most likely not be noticed.</p>

HermitSun: <p>That’s interesting. So the workspace’s size is an empirical value observed in practice, or is there some theoretical basis behind it?</p>

ptrblck: <p>From the <a href="https://docs.nvidia.com/cuda/cublas/index.html#cublassetworkspace">docs</a>:</p>
<blockquote>
<p>Too small <code>workspaceSizeInBytes</code> may cause some routines to fail with <code>CUBLAS_STATUS_ALLOC_FAILED</code> error returned or cause large regressions in performance. Workspace size equal to or larger than 16KiB is enough to prevent <code>CUBLAS_STATUS_ALLOC_FAILED</code> error, while a larger workspace can provide performance benefits for some routines. Recommended size of user-provided workspace is at least 4MiB (to match cuBLAS’ default workspace pool).</p>
</blockquote>
<p>The workspace is used for routines running in parallel streams which use this workspace to store intermediates.</p>

HermitSun: <p>Thanks again for your kindness.<br>
As the documentation says, the recommended size is at least 4MiB. I’m just wondering what will happen if I use 4MiB (rather than the current <code>default_size</code>) as the default workspace size. Will it cause performance degradation or non-determinism problems?</p>