What’s the Best Way to Debug FSDP When Hitting a C++ Backend Error?

username: nariaki3551
cooked: <p>Hello,</p>
<p>I’m currently working on adding MPI backend support for both FSDP1 and FSDP2 as part of an effort to contribute to PyTorch’s distributed training stack. During testing, I encountered an error in FSDP2 during the backward pass. I would appreciate any advice on how to investigate this issue further — especially regarding debugging techniques for distributed code involving C++ internals.</p>
<h2><a name="p-470017-environment-and-reproduction-1" class="anchor" href="#p-470017-environment-and-reproduction-1"></a>Environment and Reproduction</h2>
<p>Here is the code I used:</p>
<details>
<summary> code </summary>
<pre data-code-wrap="python"><code class="lang-python">import os
import time
import argparse
import dataclasses

import torch
from torch import nn
import torch.distributed as dist
from torch.distributed.fsdp import fully_shard
import torch.optim as optim


class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.sequential0 = nn.Linear(512, 1024)
        self.sequential1 = nn.Linear(1024, 1024)
        self.sequential2 = nn.Linear(1024, 512)
        self.last = nn.Linear(512, 10)

    def forward(self, x):
        x = torch.relu(self.sequential0(x))
        x = torch.relu(self.sequential1(x))
        x = torch.relu(self.sequential2(x))
        return self.last(x)


def fsdp_training(local_rank: int):
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device = torch.device('cuda', local_rank)

    model = SimpleModel().to(device)
    fully_shard(model)

    if rank == 0:
        print(model)

    optimizer = optim.AdamW(model.parameters(), lr=0.01)
    input_data = torch.randn(2, 512).to(device)
    target = torch.randint(0, 10, (2,)).to(device)

    print(f"Rank {rank}/{world_size}: Start training")
    model.train()
    num_epochs = 10
    for epoch in range(num_epochs):
        optimizer.zero_grad()
        output = model(input_data)
        loss = nn.CrossEntropyLoss()(output, target)
        loss.backward()
        optimizer.step()



if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='FSDP Training')
    parser.add_argument(
        '--backend', type=str, choices=['mpi', 'nccl', 'ucc', 'gloo'], default='mpi',
        help='Backend to use for training'
    )
    args = parser.parse_args()

    local_rank = int(os.getenv("LOCAL_RANK", 0))
    if args.backend == 'mpi':
        local_rank = int(os.getenv("OMPI_COMM_WORLD_RANK", 0))
    torch.cuda.set_device(local_rank)

    print("local_rank", local_rank, "backend", args.backend)
    if args.backend == 'mpi':
        dist.init_process_group(backend="mpi")
    else: # nccl, ucc, gloo
        dist.init_process_group(backend=args.backend, init_method="env://")


    fsdp_training(local_rank)
    
    dist.destroy_process_group()
</code></pre>
</details>
<p>Here’s the launcher command I ran.<br>
<code>OMP_NUM_THREADS=1 mpirun -n 2 -- python3 test_fsdp2_mpi.py --backend mpi</code></p>
<p>Then, I encounter the following error on both rank 0 and 1:</p>
<pre data-code-wrap="python"><code class="lang-python">[rank0]: Traceback (most recent call last):
[rank0]:   File "/app/pytorch_tests/src/training/test_fsdp2_mpi.py", line 81, in &lt;module&gt;
[rank0]:     fsdp_training(local_rank)
[rank0]:   File "/app/pytorch_tests/src/training/test_fsdp2_mpi.py", line 52, in fsdp_training
[rank0]:     loss.backward()
[rank0]:   File "/app/pytorch/torch/_tensor.py", line 648, in backward
[rank0]:     torch.autograd.backward(
[rank0]:   File "/app/pytorch/torch/autograd/__init__.py", line 354, in backward
[rank0]:     _engine_run_backward(
[rank0]:   File "/app/pytorch/torch/autograd/graph.py", line 824, in _engine_run_backward
[rank0]:     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank0]:   File "/app/pytorch/torch/distributed/fsdp/_fully_shard/_fsdp_state.py", line 297, in _root_post_backward_final_callback
[rank0]:     fsdp_param_group.post_backward()
[rank0]:   File "/app/pytorch/torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py", line 445, in post_backward
[rank0]:     ) = foreach_reduce(
[rank0]:   File "/app/pytorch/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/app/pytorch/torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py", line 421, in foreach_reduce
[rank0]:     dist.reduce_scatter_tensor(
[rank0]:   File "/app/pytorch/torch/distributed/c10d_logger.py", line 81, in wrapper
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/app/pytorch/torch/distributed/distributed_c10d.py", line 4416, in reduce_scatter_tensor
[rank0]:     work.wait()
[rank0]: IndexError: map::at
</code></pre>
<ul>
<li>OS: Ubuntu22.04</li>
<li>Pytorch: v2.8.0a0+git6120cc8</li>
<li>MPI: Open MPI v5.0.7rc2</li>
</ul>
<h2><a name="p-470017-my-question-how-to-debug-this-effectively-2" class="anchor" href="#p-470017-my-question-how-to-debug-this-effectively-2"></a>My question: How to Debug This Effectively?</h2>
<p><strong>How do PyTorch developers or contributors typically debug such errors from internal C++ script?</strong></p>
<p>I’m trying to investigate this issue, but since FSDP is triggered from Python and the error happens in the C++ backend, it’s difficult to debug. Debugging with gdb is usually straightforward in pure C++, but at least in my case, I’m having trouble using it effectively in this hybrid case.</p>
<p>Any insight or guidance would be greatly appreciated!</p>
<p>Thanks.</p>

username: nariaki3551
cooked: <p>I was able to identify the cause of the IndexError using gdb, so I’m sharing the steps here for reference.</p>
<hr>
<h3><a name="p-470024-h-1-run-the-script-directly-without-mpirun-1" class="anchor" href="#p-470024-h-1-run-the-script-directly-without-mpirun-1"></a>1. Run the script directly without mpirun</h3>
<pre><code class="lang-auto">OMP_NUM_THREADS=1 python3 test_fsdp2_mpi.py --backend mpi
</code></pre>
<p>This produced the same <code>IndexError: map::at</code> as when using mpirun.</p>
<h3><a name="p-470024-h-2-launch-the-script-under-gdb-2" class="anchor" href="#p-470024-h-2-launch-the-script-under-gdb-2"></a>2. Launch the script under gdb</h3>
<pre><code class="lang-auto">OMP_NUM_THREADS=1 gdb --args python3 test_fsdp2_mpi.py --backend mpi
</code></pre>
<p>Gdb launched correctly, however, the IndexError thrown from the C++ was not caught immediately. It likely did not propagate back to the Python main thread in a way that gdb could catch by default, and it may have been masked or wrapped by the Python runtime or pybind11 interface.</p>
<h3><a name="p-470024-h-3-set-a-catchpoint-for-all-c-exceptions-3" class="anchor" href="#p-470024-h-3-set-a-catchpoint-for-all-c-exceptions-3"></a>3. Set a catchpoint for all C++ exceptions</h3>
<pre><code class="lang-auto">(gdb) catch throw
Catchpoint 1 (throw)
</code></pre>
<p>This instructs gdb to break whenever any C++ exception is thrown, regardless of whether it is propagated to Python.</p>
<p>As a result, the catchpoint triggered when the exception was thrown:</p>
<pre><code class="lang-auto">Thread 4 "python3" hit Catchpoint 1 (exception thrown), 0x00007f308cd914a1 in __cxa_throw()
from /lib/x86_64-linux-gnu/libstdc++.so.6
</code></pre>
<p>Then, running bt (backtrace) gave the following:</p>
<pre><code class="lang-auto">(gdb) bt
#0 __cxa_throw
#1 std::__throw_out_of_range
#2 std::map::at (c10d::ReduceOp::AVG)
#3 c10d::ProcessGroupMPI::_reduce_scatter_base(...)
#4 c10d::ProcessGroupMPI::runLoop()
...
</code></pre>
<hr>
<p>Note: I also tried using <code>import faulthandler; faulthandler.disable()</code> and <code>import signal; signal.signal(signal.SIGSEGV, signal.SIG_DFL)</code>, but it had no effect on catching the exception.</p>