CUDA allocation lifetime for inputs to distributed.all_reduce

cbcase: <p>I am debugging a memory usage issue for tensor-parallel training, and I observe that runs with tensor-parallelism enabled that OOM have large amounts of “reserved but not unallocated” memory (pointing to fragmentation). If I <em>remove</em> the all_reduce calls in the TP region (but change nothing else), the max batch size that can fit before OOM goes up significantly (up to 2x), and there is no sign of fragmentation.</p>
<p>Digging in, I have noticed the following odd behavior when you perform an all_reduce as part of model execution (as you do in tensor-parallel training). Consider this toy program that interleaves matmuls with all_reduce operations:</p>
<pre><code class="lang-auto">import os

import torch
import torch.nn.functional as F


class Model(torch.nn.Module):
    def __init__(self, dim, nlayers):
        super().__init__()
        self.layers = torch.nn.ModuleList()
        for _ in range(nlayers):
            self.layers.append(torch.nn.Linear(dim, dim, bias=False))

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
            torch.distributed.all_reduce(x)
        return x


def main():
    rank = int(os.environ["RANK"])
    assert rank &gt;= 0
    torch.distributed.init_process_group(backend="nccl")
    torch.cuda.set_device(f"cuda:{rank}")

    dim = 8192
    nlayers = 8
    batch = 4096
    torch.cuda.memory._record_memory_history()

    model = Model(dim, nlayers)
    model.bfloat16()
    model.cuda()
    x = torch.randn((batch, dim), dtype=torch.bfloat16, device="cuda")

    with torch.no_grad():
        for _ in range(5):
            loss = model(x)

    torch.cuda.memory._dump_snapshot(f"snapshot-ar-{rank}.pkl")


if __name__ == "__main__":
    main()
</code></pre>
<p>I have run this on 2x GPUS with <code>torchrun --standalone --nproc_per_node=2 memtest.py</code>.</p>
<p>Here is what the memory visualizer shows in Active Memory for that code:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/3/2366bbd600cc75fa1e723a14a8a1819896c3ecf5.png" data-download-href="https://discuss.pytorch.org/uploads/default/2366bbd600cc75fa1e723a14a8a1819896c3ecf5" title="Screenshot 2023-11-09 at 11.29.13 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/2366bbd600cc75fa1e723a14a8a1819896c3ecf5_2_690x141.png" alt="Screenshot 2023-11-09 at 11.29.13 AM" data-base62-sha1="53aOGnraLweHz2edzzjH5RQ7Ipf" width="690" height="141" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/2366bbd600cc75fa1e723a14a8a1819896c3ecf5_2_690x141.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/2366bbd600cc75fa1e723a14a8a1819896c3ecf5_2_1035x211.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/3/2366bbd600cc75fa1e723a14a8a1819896c3ecf5_2_1380x282.png 2x" data-dominant-color="DACEC5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2023-11-09 at 11.29.13 AM</span><span class="informations">4254×874 165 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>In contrast, if I remove the call to <code>torch.distributed.all_reduce</code> (but change nothing else), the visualizer shows:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/4/342e66f086a2096a03db0bd8a121ee6154e69bb9.png" data-download-href="https://discuss.pytorch.org/uploads/default/342e66f086a2096a03db0bd8a121ee6154e69bb9" title="Screenshot 2023-11-09 at 11.29.41 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/4/342e66f086a2096a03db0bd8a121ee6154e69bb9_2_690x93.png" alt="Screenshot 2023-11-09 at 11.29.41 AM" data-base62-sha1="7rCf9W4oQSXkhu6hUIUYk6IPYql" width="690" height="93" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/4/342e66f086a2096a03db0bd8a121ee6154e69bb9_2_690x93.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/4/342e66f086a2096a03db0bd8a121ee6154e69bb9_2_1035x139.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/3/4/342e66f086a2096a03db0bd8a121ee6154e69bb9_2_1380x186.png 2x" data-dominant-color="DDD2C9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2023-11-09 at 11.29.41 AM</span><span class="informations">3598×490 19.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The non-all_reduce picture makes sense: you can see each of the five iterations of the model, and within each iteration you can see the allocations per-layer for each matmul.</p>
<p>But the run with all_reduce enabled has long-lasting allocations that live both across layers and across loops of the model. Each of those long-lived allocations are the <em>inputs</em> to the all_reduce call. It makes sense that they should live a little longer (since all_reduce occurs in place), but that they persist across model iterations is perplexing.</p>
<p>My goal is to make the top picture look like the bottom one, so I’d love suggestions about what to try and / or any deeper insight into what the allocator is doing here. Thanks.</p>

cbcase: <p>Bumping this. <a class="mention" href="/u/ptrblck">@ptrblck</a> for the beg for pointers to internals / things I could look at to start debugging.</p>
<p>-Carl</p>

ptrblck: <p>Hi Carl!</p>
<p>Thanks for reporting the issue and the great memory allocation plots.<br>
You might be running into <a href="https://github.com/pytorch/pytorch/pull/76861">this issue</a>. In <a href="https://github.com/pytorch/pytorch/pull/89880">this PR</a> we exposed an env variable to avoid the stream recording via <code>TORCH_NCCL_AVOID_RECORD_STREAMS=1</code>. Could you use it and let me know if the memory allocation behavior changes?</p>

cbcase: <p>Oh my gosh, I’m floored. Cannot imagine I’d ever find that issue on my own (of course I’d forgotten even that PyTorch puts NCCL kernels on its own internal stream). Setting that env var and running the same code <em>with</em> all_reduce enabled gives me this beautiful picture:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/2/72f5be673273e9a387cf7808e449ee44a178b989.png" data-download-href="https://discuss.pytorch.org/uploads/default/72f5be673273e9a387cf7808e449ee44a178b989" title="Screenshot 2023-11-17 at 10.35.48 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/72f5be673273e9a387cf7808e449ee44a178b989_2_690x128.png" alt="Screenshot 2023-11-17 at 10.35.48 AM" data-base62-sha1="goYVYXxFYfu2U1ch3qnLjB8zVnb" width="690" height="128" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/72f5be673273e9a387cf7808e449ee44a178b989_2_690x128.png, https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/72f5be673273e9a387cf7808e449ee44a178b989_2_1035x192.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/72f5be673273e9a387cf7808e449ee44a178b989_2_1380x256.png 2x" data-dominant-color="D4C7BB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2023-11-17 at 10.35.48 AM</span><span class="informations">3982×740 102 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Ie, it’s ~identical to the no-all_reduce allocator behavior. I will confirm that this pays off as better perf in real runs, but I am optimistic.</p>
<p>Thanks so much, Piotr! (Love that this is a Carilli fix, tell him I say hi next time you guys cross paths!)</p>

ngimel: <p><code>record_stream</code> calls are generally considered harmful, and there is a long post in dev-discuss describing their dangers, you can read if you are interested <a href="https://dev-discuss.pytorch.org/t/fsdp-cudacachingallocator-an-outsider-newb-perspective/1486" class="inline-onebox" rel="noopener nofollow ugc">FSDP &amp; CUDACachingAllocator: an outsider newb perspective - distributed - PyTorch Dev Discussions</a></p>

kwen2501: <p>Just FYI: we started rolling out avoid_record_stream as the default behavior for synchronous comm ops.<br>
Currently covering all_gather, reduce_scatter, scatter and broadcast. See these two PRs:</p><aside class="onebox githubpullrequest" data-onebox-src="https://github.com/pytorch/pytorch/pull/111431">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/pull/111431" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">



    <div class="github-icon-container" title="Pull Request">
      <svg width="60" height="60" class="github-icon" viewBox="0 0 12 16" aria-hidden="true"><path fill-rule="evenodd" d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
    </div>

  <div class="github-info-container">



      <h4>
        <a href="https://github.com/pytorch/pytorch/pull/111431" target="_blank" rel="noopener nofollow ugc">[ProcessGroupNCCL] Avoid recording stream for synchronous ops</a>
      </h4>

    <div class="branches">
      <code>pytorch:main</code> ← <code>pytorch:avoid_record_ag_rs</code>
    </div>

      <div class="github-info">
        <div class="date">
          opened <span class="discourse-local-date" data-format="ll" data-date="2023-10-17" data-time="19:30:59" data-timezone="UTC">07:30PM - 17 Oct 23 UTC</span>
        </div>

        <div class="user">
          <a href="https://github.com/kwen2501" target="_blank" rel="noopener nofollow ugc">
            <img alt="kwen2501" src="https://avatars.githubusercontent.com/u/6676466?v=4" class="onebox-avatar-inline" width="20" height="20">
            kwen2501
          </a>
        </div>

        <div class="lines" title="4 commits changed 8 files with 75 additions and 28 deletions">
          <a href="https://github.com/pytorch/pytorch/pull/111431/files" target="_blank" rel="noopener nofollow ugc">
            <span class="added">+75</span>
            <span class="removed">-28</span>
          </a>
        </div>
      </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">For synchronous ops (i.e. `asyncOp = False`), we don't want to record streams be<span class="show-more-container"><a href="https://github.com/pytorch/pytorch/pull/111431" target="_blank" rel="noopener nofollow ugc" class="show-more">…</a></span><span class="excerpt hidden">cause we know that the NCCL stream will join back to the "current" stream right after this op. So we might just as well keep the stream ownership of the input/output tensors unchanged (i.e. not telling the caching allocator there is a stream-stream hand off). The benefit would be that the allocation/free of the tensors would look deterministic to the "current" stream so that the caching allocator can reuse memory pool for this stream in a clever way. 

To prevent the input/output tensors from being recycled by python, we rely on the stashing mechanism in ProcessGroupNCCL (which can be also turned on by setting `TORCH_NCCL_AVOID_RECORD_STREAMS=1`).

This mechanism change is for libraries like FSDP which uses `all_gather_into_tensor` and `reduce_scatter_tensor` in a synchronous way and which cannot set `TORCH_NCCL_AVOID_RECORD_STREAMS=1` for their users. And therefore, this change is limited to these two collectives for now.

Cc: @awgu @janeyx99 @albanD</span></p>
  </div>

  </article>

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

  <div style="clear: both"></div>
</aside>
<aside class="onebox githubpullrequest" data-onebox-src="https://github.com/pytorch/pytorch/pull/112896">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/pull/112896" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">



    <div class="github-icon-container" title="Pull Request">
      <svg width="60" height="60" class="github-icon" viewBox="0 0 12 16" aria-hidden="true"><path fill-rule="evenodd" d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
    </div>

  <div class="github-info-container">



      <h4>
        <a href="https://github.com/pytorch/pytorch/pull/112896" target="_blank" rel="noopener nofollow ugc">[ProcessGroupNCCL] Avoid recording stream for broadcast and scatter</a>
      </h4>

    <div class="branches">
      <code>pytorch:main</code> ← <code>kwen2501:export-D50985365</code>
    </div>

      <div class="github-info">
        <div class="date">
          opened <span class="discourse-local-date" data-format="ll" data-date="2023-11-03" data-time="19:56:43" data-timezone="UTC">07:56PM - 03 Nov 23 UTC</span>
        </div>

        <div class="user">
          <a href="https://github.com/kwen2501" target="_blank" rel="noopener nofollow ugc">
            <img alt="kwen2501" src="https://avatars.githubusercontent.com/u/6676466?v=4" class="onebox-avatar-inline" width="20" height="20">
            kwen2501
          </a>
        </div>

        <div class="lines" title="1 commits changed 8 files with 40 additions and 15 deletions">
          <a href="https://github.com/pytorch/pytorch/pull/112896/files" target="_blank" rel="noopener nofollow ugc">
            <span class="added">+40</span>
            <span class="removed">-15</span>
          </a>
        </div>
      </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">Summary: Follows PR #111431, save memory for DTensor init

Test Plan: Sandcastle<span class="show-more-container"><a href="https://github.com/pytorch/pytorch/pull/112896" target="_blank" rel="noopener nofollow ugc" class="show-more">…</a></span><span class="excerpt hidden">

Reviewed By: wanchaol

Differential Revision: D50985365</span></p>
  </div>

  </article>

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

  <div style="clear: both"></div>
</aside>
<p>
We plan to cover all collectives by 2.2.0 release.</p>

cbcase: <p><a class="mention" href="/u/ngimel">@ngimel</a> – thanks for the pointer, I learned a lot reading that.</p>
<p><a class="mention" href="/u/kwen2501">@kwen2501</a> – good to know, sad that all_reduce missed the first rollout, as I would have been able to blindly go about my life knowing nothing of this <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

cbcase: <p>Follow-up question <a class="mention" href="/u/ptrblck">@ptrblck</a> : how do these settings interact with cuda graph capture? Is it largely the same, except that the allocations are going to be coming from the graph-private memory pool? If that’s the case, then there should also be benefits to setting the AVOID_RECORD_STREAMS flag for graph-capture, as it will allow the allocator to re-use the same buffers for collective input <em>during capture</em> (and therefore during replay) as opposed to (likely) needing a bunch of distinct buffers, depending on the exact timing characteristics of when allocations occur.</p>