PDB 调试 DDP

概述

PDB（Python Debugger）是 Python 的内置调试工具，而 DDP（Distributed Data Parallel）是 PyTorch 提供的分布式数据并行模块，用于在多 GPU 环境中加速模型训练。在使用 DDP 进行分布式训练时，可能会遇到各种问题，如数据加载错误、模型同步问题等。此时，PDB 可以帮助我们逐步调试代码，定位问题。

使用 PDB 调试 DDP 的方法

1. 在代码中插入断点

在需要调试的代码位置插入 breakpoint() 或 import pdb; pdb.set_trace()。例如：

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# 初始化 DDP
dist.init_process_group(backend='nccl', init_method='env://')
rank = dist.get_rank()

# 创建模型
model = torch.nn.Linear(10, 1).to(rank)
ddp_model = DDP(model, device_ids=[rank])

# 插入断点
import pdb; pdb.set_trace()

# 训练模型
inputs = torch.randn(20, 10).to(rank)
labels = torch.randn(20, 1).to(rank)
outputs = ddp_model(inputs)
loss = (outputs - labels).pow(2).sum()
loss.backward()

2. 启动调试

运行代码后，程序会在断点处暂停，进入 PDB 调试模式。此时可以使用以下命令进行调试：

n（next）：执行下一行代码。
c（continue）：继续执行代码，直到遇到下一个断点。
q（quit）：退出调试模式。
p（print）：打印变量的值。例如，p inputs 会打印 inputs 的值。

3. 查看变量和状态

在 PDB 调试模式下，可以查看变量的值和状态。例如：

(pdb) p inputs
tensor([[ 0.1234,  0.5678,  0.9012,  0.3456,  0.7890,  0.2345,  0.6789,  0.1234,
          0.5678,  0.9012],
        [ 0.3456,  0.7890,  0.2345,  0.6789,  0.1234,  0.5678,  0.9012,  0.3456,
          0.7890,  0.2345]], device='cuda:0')
(pdb) p ddp_model
DDP(
  (module): Linear(in_features=10, out_features=1, bias=True)
)

4. 单步调试

使用 n 命令可以逐行执行代码，观察每一步的执行结果。例如：

(pdb) n
> /path/to/your/script.py(15)<module>()
-> outputs = ddp_model(inputs)
(pdb) n
> /path/to/your/script.py(16)<module>()
-> loss = (outputs - labels).pow(2).sum()
(pdb) p outputs
tensor([[ 0.1234],
        [ 0.5678]], device='cuda:0', grad_fn=<AddmmBackward>)

5. 跟踪错误

如果程序抛出异常，PDB 会自动进入调试模式，显示错误信息。此时可以使用 tb（traceback）命令查看调用栈信息，定位问题的根源。例如：

(pdb) tb
  File "/path/to/your/script.py", line 16, in <module>
    loss = (outputs - labels).pow(2).sum()
RuntimeError: CUDA error: device-side assert triggered

6. 调试多进程代码

在 DDP 中，每个进程都有自己的 Python 解释器实例。如果需要调试多个进程，可以在每个进程中插入断点，并分别启动调试。例如：

if rank == 0:
    import pdb; pdb.set_trace()

7. 使用 `ipdb`

ipdb 是一个增强版的 PDB，提供了更友好的交互界面和更多的调试功能。可以使用以下命令安装 ipdb：

pip install ipdb

然后在代码中使用 ipdb 代替 pdb：

import ipdb; ipdb.set_trace()

注意事项

1. 调试性能问题

在调试 DDP 代码时，可能会遇到性能问题，如训练速度变慢。这是因为 PDB 会暂停程序的执行，导致其他进程等待。可以通过减少调试点的数量或使用异步调试工具来缓解这个问题。

2. 调试多 GPU 代码

在多 GPU 环境中调试代码时，需要注意每个进程的设备（GPU）分配。确保在正确的设备上执行代码，避免设备不匹配的问题。

3. 调试分布式通信问题

如果遇到分布式通信问题，如进程间同步失败，可以检查 dist.init_process_group 的参数设置是否正确，以及网络环境是否正常。

示例代码

以下是一个完整的示例代码，展示了如何使用 PDB 调试 DDP：

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# 初始化 DDP
dist.init_process_group(backend='nccl', init_method='env://')
rank = dist.get_rank()

# 创建模型
model = torch.nn.Linear(10, 1).to(rank)
ddp_model = DDP(model, device_ids=[rank])

# 插入断点
import pdb; pdb.set_trace()

# 训练模型
inputs = torch.randn(20, 10).to(rank)
labels = torch.randn(20, 1).to(rank)
outputs = ddp_model(inputs)
loss = (outputs - labels).pow(2).sum()
loss.backward()

# 清理 DDP
dist.destroy_process_group()

Ref

https://blog.csdn.net/flyingluohaipeng/article/details/128626943 https://blog.csdn.net/qianbin3200896/article/details/108182504 https://zhuanlan.zhihu.com/p/37294138

https://github.com/ionelmc/python-remote-pdb https://github.com/ionelmc/python-remote-pdb/blob/master/src/remote_pdb.py