Help Needed from vLLM team on profiling pytorch cuda memory

youkaichao1: <p>Hi, I’m working on <a href="https://github.com/vllm-project/vllm" class="inline-onebox" rel="noopener nofollow ugc">GitHub - vllm-project/vllm: A high-throughput and memory-efficient inference and serving engine for LLMs</a> , and the recently release of pytorch 2.2.0 caused some trouble to me. I came here for help in profiling cuda memory usage.</p>
<p>The basic story: vLLM tries to allocate as much memory as possible for KV Cache to accelerate LLM inference. In order to do so, it first profiles the memory usage, guess the maximum size of memory available for KV Cache, and also leave some for storing activation during inference.</p>
<p>What’s more, vLLM uses cuda graph to reduce Python overhead.</p>
<p>When PyTorch upgrades from 2.1.2 to 2.2.0 , there seems to be some internal change of memory allocator, and the memory that can be used is decreased. It can cause OOM error during inference.</p>
<p>Here are the diagnoses data (produced by <code>torch.cuda.memory_stats</code> and <code>torch.cuda.memory._dump_snapshot</code>), collected from a server with 2 L4 GPUs:</p>
<aside class="onebox googledrive" data-onebox-src="https://drive.google.com/file/d/1RzL5uDyTHJucgZXJIjh_ITKLvwqYpmuu/view?usp=sharing">
  <header class="source">

      <a href="https://drive.google.com/file/d/1RzL5uDyTHJucgZXJIjh_ITKLvwqYpmuu/view?usp=sharing" target="_blank" rel="noopener nofollow ugc">drive.google.com</a>
  </header>

  <article class="onebox-body">
      <a href="https://drive.google.com/file/d/1RzL5uDyTHJucgZXJIjh_ITKLvwqYpmuu/view?usp=sharing" target="_blank" rel="noopener nofollow ugc"><span class="googledocs-onebox-logo g-drive-logo"></span></a>



<h3><a href="https://drive.google.com/file/d/1RzL5uDyTHJucgZXJIjh_ITKLvwqYpmuu/view?usp=sharing" target="_blank" rel="noopener nofollow ugc">memory-stat.zip</a></h3>

<p>Google Drive file.</p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>PyTorch 2.1.2 works with or without cuda graph; PyTorch 2.2.0 works only without cuda graph, but OOM with cuda graph.</p>
<p>So I have three questions:</p>
<ul>
<li>Is there anything that indeed changed in cuda graph from pytorch 2.1.2 to 2.2.0 , leading to the change of memory allocation with cuda graph?</li>
<li>If yes, are there any knobs that can be used to control the behavior?</li>
<li>At worst, if pytorch 2.2.0 indeed uses more internal memory or causes more memory fragment, what is the reliable way to calculate the maximum memory for vLLM’s usecase to allocate KV Cache?</li>
</ul>
<p>Any help is appreciated!</p>

ptrblck: <p>Did you try to bisect the issue using nightly binaries?<br>
Also, is the issue still observable in the latest nightly?</p>

youkaichao1: <p>How to biset all the nightly binaries? I only find <code>2.2.0.dev20231010+cu121</code> available, other nightly binaries for <code>2.2.0</code> do not exist any more in <code>pip</code>.</p>

youkaichao1: <p>The last resort is to biset the pytorch github repo <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"></p>

albanD: <p>You can look for the index-url in the install instructions to see the nightlies hosted by us (pypi has strict limits on storage so we have very little there).<br>
You can see <a href="https://download.pytorch.org/whl/nightly/torch/">https://download.pytorch.org/whl/nightly/torch/</a> for example for the latest cuda enabled ones.</p>
<p>You can see there the <code>+cpu</code> and <code>+cu118</code> builds corresponding to cpu and cuda11.8 and the other ones are the default cuda 12.1.</p>

ptrblck: <p>Answered <a href="https://discuss.pytorch.org/t/install-particular-pytorch-nightly/186853/5">here</a>.</p>

youkaichao1: <p>Thank you, I can install previous nightly build now.</p>
<p>Another question is how to associate nightly build date with github commit. Say if I find <code>2.2.0.dev20231210+cu121</code> works but <code>2.2.0.dev20231211+cu121</code> does not work, then I need to find the commit between December 10 and December 11.</p>
<p>Would looking into the commit history by day work? e.g. looking for the dates in<br>
<a href="https://github.com/pytorch/pytorch/commits/main/" class="inline-onebox" rel="noopener nofollow ugc">Commits · pytorch/pytorch · GitHub</a> .</p>

ptrblck: <p>Yes, the commit history would reflect the changes between the nightly builds. I don’t know when exactly the nightly commit is cloned (maybe midnight Pacific Time), but narrowing down the actual commit should be easy based on the dates of the nightlies in any case.</p>