cuda-checkpoint --toggle hangs when checkpointing one process of torch.distributed NCCL job

[Taya666]

Hi,

I’m seeing cuda-checkpoint --toggle hang when checkpointing one process from a PyTorch torch.distributed (NCCL) job.

To Reproduce
code:

import os
import torch
import torch.distributed as dist

local_rank = int(os.environ["LOCAL_RANK"])
rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])

torch.cuda.set_device(local_rank)

dist.init_process_group(
    backend="nccl",
    init_method="env://",
    rank=rank,
    world_size=world_size,
)

x = torch.tensor([rank], device="cuda")
dist.all_reduce(x)
del x
torch.cuda.empty_cache()

print("finish all")

while True:
    pass

running command:
python -m torch.distributed.run --standalone --nproc_per_node=8 measure_cuda_nccl.py

After all 8 processes print finish all, pick one worker process and run:
./cuda-checkpoint --toggle --pid <PID>

Observed Behavior:

• cuda-checkpoint --toggle hangs indefinitely
• No error output
• Target process keeps running