torch.distributed 分布式训练(04):调试

2 卡及以上不是不能使用断点,但普通 PyCharm 断点不能像单卡一样随意使用。

核心原因是:多个 rank 必须按相同顺序执行 barrier、all_reduce、all_gather 等集合通信。如果 rank 0 被断点暂停,而 rank 1 已经进入集合通信,rank 1 会等待 rank 0,最终表现为卡死或超时。

所以,推荐的调试顺序是:

单卡 pycharm 断点 -> 2 卡使用日志输出

2 卡使用日志输出

import os
import torch  # 仅在需要打印张量示例时使用
import torch.distributed as dist


def dlog(message: str) -> None:
    """打印带分布式进程标识的调试日志。

    每条日志自动加上 rank(进程组内编号)和 pid(操作系统进程号),
    便于在多进程 / 多 GPU 场景下定位日志来源;输出立即刷新(flush),
    避免多进程日志交错或进程崩溃时丢失内容。

    Args:
        message: 待打印的日志文本。打印变量/张量的推荐写法::

            x = 42
            dlog(f"x = {x}")                        # 普通变量:直接 f-string 拼接

            t = torch.randn(4, device="cuda")
            dlog(f"t = {t.cpu()}")                  # 整个张量:先 .cpu() 挪回内存再打印
            dlog(f"t[0] = {t[0].item():.4f}")       # 单个标量:.item(),触发一次 GPU→CPU 同步
            dlog(f"t.shape = {tuple(t.shape)}")     # 只看形状/类型:无 GPU 同步开销

            # 注意:.item() 会让 CPU 阻塞等待 GPU,训练/推理热路径中
            # 应攒批后再打印,避免每个 step 高频调用拖慢吞吐。

    Returns:
        None
    """
    # 分布式未初始化时(如单机调试)将 rank 视为 0,
    # 防止 dist.get_rank() 在未初始化状态下抛异常
    rank = dist.get_rank() if dist.is_initialized() else 0
    # flush=True:多进程并发写 stdout 时立即刷新缓冲区,
    # 防止日志滞留导致顺序错乱,崩溃时也不会丢已打印内容
    print(
        f"[rank={rank}][pid={os.getpid()}] {message}",
        flush=True,
    )

评论