How to debug pytorch distributed?

username: wconstab
cooked: <p>It can be tricky to use python debugger from a multi-rank setup.  The first thing you’d notice if you try this is that pdb may crash your program if you use it from inside a mpirun or torchrun launcher.  Fortunately, this is fixable and you can use pdb almost like usual.</p>
<p>There is a catch- it’s not too easy to attach the debugger on each rank, but it’s pretty easy to attach it to just one particular rank (and let all the other ranks pause).</p>
<p>This PR from <a class="mention" href="/u/ezyang">@ezyang</a> adds a new helper called <code>torch.distributed.breakpoint</code>.  It can be used more or less like python’s breakpoitn statement, except you’re supposed to have it called on all ranks (but always pass the same int for rank, so across all ranks one rank in particular is the one that will listen for the debugger input).</p><aside class="onebox githubpullrequest" data-onebox-src="https://github.com/pytorch/pytorch/pull/113775/files">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/pull/113775/files" 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/113775/files" target="_blank" rel="noopener nofollow ugc">Add torch.distributed.breakpoint</a>
      </h4>

    <div class="branches">
      <code>pytorch:gh/ezyang/2444/base</code> ← <code>pytorch:gh/ezyang/2444/head</code>
    </div>

      <div class="github-info">
        <div class="date">
          opened <span class="discourse-local-date" data-format="ll" data-date="2023-11-15" data-time="17:03:59" data-timezone="UTC">05:03PM - 15 Nov 23 UTC</span>
        </div>

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

        <div class="lines" title="5 commits changed 2 files with 35 additions and 0 deletions">
          <a href="https://github.com/pytorch/pytorch/pull/113775/files" target="_blank" rel="noopener nofollow ugc">
            <span class="added">+35</span>
            <span class="removed">-0</span>
          </a>
        </div>
      </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __<span class="show-more-container"><a href="https://github.com/pytorch/pytorch/pull/113775" target="_blank" rel="noopener nofollow ugc" class="show-more">…</a></span><span class="excerpt hidden">-&gt;__ #113775

I tested it works by patching

```
diff --git a/test/distributed/test_dynamo_distributed.py b/test/distributed/test_dynamo_distributed.py
index 96b3a82bdfa..dea9bac9302 100644
--- a/test/distributed/test_dynamo_distributed.py
+++ b/test/distributed/test_dynamo_distributed.py
@@ -18,6 +18,7 @@ from torch._dynamo import config
 from torch._dynamo.utils import same
 from torch._dynamo.testing import collect_results
 from torch.utils._triton import has_triton
+import torch.distributed as dist
 from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy, lambda_auto_wrap_policy
 from torch._higher_order_ops.wrap import tag_activation_checkpoint
 from torch.nn.parallel import DistributedDataParallel as DDP
@@ -398,6 +399,7 @@ class TestMultiProc(DynamoDistributedMultiProcTestCase):
     @unittest.skipIf(not has_triton(), "Inductor+gpu needs triton and recent GPU arch")
     def test_fsdp_activation_checkpointing(self):
         with _dynamo_dist_per_rank_init(self.rank, self.world_size):
+            dist.breakpoint()
             model, inputs = get_toy_model_for_activation_checkpointing(f"cuda:{self.rank}")
             is_inner = lambda module: isinstance(module, ToyInnerModel)  # noqa: E731
             wrap_policy = functools.partial(lambda_auto_wrap_policy, lambda_fn=is_inner)
```

and then running `python test/distributed/test_dynamo_distributed.py -k test_fsdp_activation_checkpointing`

It prints:

```
ATTENTION!!!

Type 'up' to get to the frame that called dist.breakpoint(rank=0)

&gt; /data/users/ezyang/c/pytorch/torch/distributed/__init__.py(71)breakpoint()
-&gt; barrier()
(Pdb) up
&gt; /data/users/ezyang/c/pytorch/test/distributed/test_dynamo_distributed.py(402)test_fsdp_activation_checkpointing()
-&gt; dist.breakpoint()
(Pdb) list
397  
398         @skip_if_lt_x_gpu(1)
399         @unittest.skipIf(not has_triton(), "Inductor+gpu needs triton and recent GPU arch")
400         def test_fsdp_activation_checkpointing(self):
401             with _dynamo_dist_per_rank_init(self.rank, self.world_size):
402  -&gt;             dist.breakpoint()
403                 model, inputs = get_toy_model_for_activation_checkpointing(f"cuda:{self.rank}")
404                 is_inner = lambda module: isinstance(module, ToyInnerModel)  # noqa: E731
405                 wrap_policy = functools.partial(lambda_auto_wrap_policy, lambda_fn=is_inner)
406                 model = apply_fsdp_with_checkpointing(model, wrap_policy, is_inner)
407                 correct_outputs = model(inputs)
```

Signed-off-by: Edward Z. Yang &lt;ezyang@meta.com&gt;</span></p>
  </div>

  </article>

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

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