Program hangs when node trying to wrap torch model with DDP

Minh_Le: <p>Hi everyone. So I am working on a small project, the idea is to set up a distributed environment with master node and worker node. The master node is going to partition a dataset, and distribute data chunks, along with the model, to each worker in the system. Upon receiving, the workers will wrap the model in torch DDP. This is where I got stuck, because the program hangs here and I have no idea why.</p>
<p>FYI, I am using Macbook M2 so I think I can’t setup cuda environment, which is why I use CPU as device. I got this suggestion from my advisor, but I’m still clueless about what’s going on here: “It might be number of devices on your testing machine. Let’s say, you have world_rank = 4, assume 1 master and 3 workers, and the testing environment has to have 4 devices, like 1 CPU and 3 GPUs to spawn worker on each device, otherwise the cross-node communication would be stuck here b/c the system thinks that it doesn’t receive response from other nodes.”<br>
Can anyone help me please?</p>
<p>This is my code with core functions (data partitioning, distributing, receiving)</p>
<pre><code class="lang-auto">import torch
import torch.distributed
from torch.utils.data import Dataset, DataLoader, random_split
from torch.nn import Module
import torch.distributed as dist
import io
from torch.nn.parallel import DistributedDataParallel as DDP
import logging
import hashlib

logging.basicConfig(filename='test_real_dataset_distribute_model.log', level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s')

def partition_dataset(dataset: Dataset, num_workers: int) -&gt; list[Dataset]:
    if len(dataset) == 0:
        error_message = "The dataset is empty"
        logging.error(error_message)
        raise ValueError(error_message)
    if num_workers &lt;= 0:
        error_message = "Invalid number of workers"
        logging.error(error_message)
        raise ValueError(error_message)
    partition_size = len(dataset) // num_workers
    partitions = [partition_size] * num_workers
    remainder = len(dataset) - partition_size * num_workers
    if remainder &gt; 0:
        partitions[-1] += remainder
    return random_split(dataset, partitions)

def setup(backend, init_method, rank, world_size):
    dist.init_process_group(backend=backend, init_method=init_method, rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def calculate_model_hash(model):
    buffer = io.BytesIO()
    torch.save(model.state_dict(), buffer)
    buffer.seek(0)
    hash_value = hashlib.md5(buffer.read()).hexdigest()
    return hash_value

def master_node_distribute(backend: str, init_method: str, model: Module, dataset: Dataset, world_size: int):
    logging.info("Starting master node distribution...")

    try:
        setup(backend, init_method, 0, world_size)
        logging.info(f"Initialized process group with backend {backend} and init method {init_method}")

        logging.info(f"Master node partitions dataset for {world_size - 1} worker nodes...")
        partitions = partition_dataset(dataset, world_size)
        logging.info("Partitioned dataset")

        buffer = io.BytesIO()
        torch.save(model.state_dict(), buffer)
        buffer.seek(0)
        model_tensor = torch.tensor(bytearray(buffer.read()), dtype=torch.uint8)
        model_tensor_size = torch.tensor([model_tensor.size(0)], dtype=torch.int)
        logging.info("Serialized model")

        master_model_hash = calculate_model_hash(model)
        logging.info(f"Master model hash: {master_model_hash}")

        for rank in range(1, world_size):
            partition_loader = DataLoader(partitions[rank], batch_size=len(partitions[rank]), shuffle=True)
            for data in partition_loader:
                data_tensor = data[0]
                dist.send(tensor=data_tensor, dst=rank)
                logging.info(f"Sent data partition to worker {rank}")

            dist.send(tensor=model_tensor_size, dst=rank)
            logging.info(f"Sent model size to worker {rank}")
            dist.send(tensor=model_tensor, dst=rank)
            logging.info(f"Sent model to worker {rank}")
        dist.barrier()
    except Exception as e:
        cleanup()
        logging.error(f"An error occurred: {e}")
    finally:
        cleanup()
        logging.info("Clean up process")

def worker_node_receive(backend: str, init_method: str, rank: int, world_size: int, model: Module, dataset: Dataset):
    def calculate_data_size(dataset: Dataset, rank: int, world_size: int) -&gt; int:
        num_samples = len(dataset)
        partition_size = num_samples // world_size
        remainder = num_samples % world_size

        if rank &lt; remainder:
            data_size = partition_size + 1
        else:
            data_size = partition_size

        return data_size

    logging.info(f"Worker {rank} starting up...")

    try:
        setup(backend, init_method, rank, world_size)
        logging.info(f"Worker {rank} initialized process group")

        data_size = calculate_data_size(dataset, rank, world_size)
        data_tensor = torch.zeros(data_size, 3, 32, 32, dtype=torch.float32)
        dist.recv(tensor=data_tensor)
        logging.info(f"Worker {rank} received data")

        model_tensor_size = torch.tensor([0], dtype=torch.int)
        dist.recv(tensor=model_tensor_size, src=0)
        logging.info(f"Worker {rank} received model size")
        model_tensor = torch.zeros(model_tensor_size.item(), dtype=torch.uint8)
        dist.recv(tensor=model_tensor, src=0)
        logging.info(f"Worker {rank} received model")

        buffer = io.BytesIO(model_tensor.numpy().tobytes())
        buffer.seek(0)
        model.load_state_dict(torch.load(buffer))
        logging.info(f"Worker {rank} loaded model")

        worker_model_hash = calculate_model_hash(model)
        logging.info(f"Worker model hash: {worker_model_hash}")

        # # Move model to appropriate device
        device = torch.device("cpu")
        logging.info(f"Worker {rank} moving model to device {device}...")
        model.to(device)
        logging.info(f"Worker {rank} moved model to device {device}")

        # Wrap model with DDP without specifying device_ids
        logging.info(f"Worker {rank} wrapping model with DDP...")
        model = DDP(model)
        logging.info(f"Worker {rank} wrapped model with DDP")

        dist.barrier()

    except Exception as e:
        cleanup()
        logging.error(f"Worker {rank} encountered an error: {e}")

    finally:
        cleanup()
        logging.info(f"Worker {rank} cleaned up process group")
</code></pre>
<p>This is my main function to test the implementation with real dataset and real model:</p>
<pre><code class="lang-auto">import torch
import torch.multiprocessing as mp
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
from torch.nn import Module
from data_partitioning import master_node_distribute, worker_node_receive

class SimpleModel(Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.model = models.resnet50(pretrained=True)

    def forward(self, x):
        return self.model(x)

def main(rank, world_size):
    backend = 'gloo'
    init_method = 'tcp://localhost:12345'
    model = SimpleModel()

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)

    if rank == 0:
        master_node_distribute(backend, init_method, model, dataset, world_size)
    else:
        worker_node_receive(backend, init_method, rank, world_size, model, dataset)

if __name__ == "__main__":
    world_size = 4
    mp.spawn(main, args=(world_size,), nprocs=world_size, join=True)
</code></pre>
<p>This is the log file (as you can see, all workers hang when starting to wrap model)</p>
<pre><code class="lang-auto">2024-07-11 14:31:53,141 - INFO - Starting master node distribution...
2024-07-11 14:31:53,194 - INFO - Worker 2 starting up...
2024-07-11 14:31:53,200 - INFO - Worker 1 starting up...
2024-07-11 14:31:53,208 - INFO - Worker 3 starting up...
2024-07-11 14:31:53,229 - INFO - Initialized process group with backend gloo and init method tcp://localhost:12345
2024-07-11 14:31:53,229 - INFO - Worker 1 initialized process group
2024-07-11 14:31:53,229 - INFO - Master node partitions dataset for 3 worker nodes...
2024-07-11 14:31:53,230 - INFO - Worker 3 initialized process group
2024-07-11 14:31:53,230 - INFO - Worker 2 initialized process group
2024-07-11 14:31:53,349 - INFO - Partitioned dataset
2024-07-11 14:31:59,048 - INFO - Serialized model
2024-07-11 14:31:59,438 - INFO - Master model hash: 7fa5ebf909fdce4f36af0719ba9e32f8
2024-07-11 14:32:00,666 - INFO - Sent data partition to worker 1
2024-07-11 14:32:00,667 - INFO - Worker 1 received data
2024-07-11 14:32:00,667 - INFO - Sent model size to worker 1
2024-07-11 14:32:00,667 - INFO - Worker 1 received model size
2024-07-11 14:32:00,706 - INFO - Worker 1 received model
2024-07-11 14:32:00,706 - INFO - Sent model to worker 1
2024-07-11 14:32:00,746 - INFO - Worker 1 loaded model
2024-07-11 14:32:01,145 - INFO - Worker model hash: 7fa5ebf909fdce4f36af0719ba9e32f8
2024-07-11 14:32:01,145 - INFO - Worker 1 moving model to device cpu...
2024-07-11 14:32:01,147 - INFO - Worker 1 moved model to device cpu
2024-07-11 14:32:01,147 - INFO - Worker 1 wrapping model with DDP...
2024-07-11 14:32:01,842 - INFO - Sent data partition to worker 2
2024-07-11 14:32:01,842 - INFO - Worker 2 received data
2024-07-11 14:32:01,843 - INFO - Sent model size to worker 2
2024-07-11 14:32:01,843 - INFO - Worker 2 received model size
2024-07-11 14:32:01,880 - INFO - Worker 2 received model
2024-07-11 14:32:01,880 - INFO - Sent model to worker 2
2024-07-11 14:32:01,944 - INFO - Worker 2 loaded model
2024-07-11 14:32:02,346 - INFO - Worker model hash: 7fa5ebf909fdce4f36af0719ba9e32f8
2024-07-11 14:32:02,346 - INFO - Worker 2 moving model to device cpu...
2024-07-11 14:32:02,347 - INFO - Worker 2 moved model to device cpu
2024-07-11 14:32:02,347 - INFO - Worker 2 wrapping model with DDP...
2024-07-11 14:32:02,998 - INFO - Sent data partition to worker 3
2024-07-11 14:32:02,999 - INFO - Worker 3 received data
2024-07-11 14:32:02,999 - INFO - Sent model size to worker 3
2024-07-11 14:32:02,999 - INFO - Worker 3 received model size
2024-07-11 14:32:03,028 - INFO - Sent model to worker 3
2024-07-11 14:32:03,028 - INFO - Worker 3 received model
2024-07-11 14:32:03,107 - INFO - Worker 3 loaded model
2024-07-11 14:32:03,517 - INFO - Worker model hash: 7fa5ebf909fdce4f36af0719ba9e32f8
2024-07-11 14:32:03,517 - INFO - Worker 3 moving model to device cpu...
2024-07-11 14:32:03,518 - INFO - Worker 3 moved model to device cpu
2024-07-11 14:32:03,518 - INFO - Worker 3 wrapping model with DDP...
</code></pre>

Shantanu_Nath: <p><a class="mention" href="/u/minh_le">@Minh_Le</a> I have the similar issue. Did you able to solve?</p>

yf225: <p>To debug this, you might want to do a few ablation studies, by removing things one by one until it works.</p>
<p>e.g. is it still stuck if we remove the call to DDP?</p>

Shantanu_Nath: <p><a class="mention" href="/u/yf225">@yf225</a> It works while I used distributed backend <code>gloo</code> instead of <code>nccl</code>. But only 30% of the both GPU was working and it was slower than my single GPU training.</p>

Shantanu_Nath: <p>This is NCCL_INFO output:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/f/0f2ba352ae12a6321e41912cb6df39ae2b137c17.png" data-download-href="https://discuss.pytorch.org/uploads/default/0f2ba352ae12a6321e41912cb6df39ae2b137c17" title="Screenshot from 2024-07-24 16-06-28"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/f/0f2ba352ae12a6321e41912cb6df39ae2b137c17_2_690x152.png" alt="Screenshot from 2024-07-24 16-06-28" data-base62-sha1="2acEFENkZ0ZpP0ZtH0PlC6BupgP" width="690" height="152" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/f/0f2ba352ae12a6321e41912cb6df39ae2b137c17_2_690x152.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/f/0f2ba352ae12a6321e41912cb6df39ae2b137c17_2_1035x228.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/f/0f2ba352ae12a6321e41912cb6df39ae2b137c17_2_1380x304.png 2x" data-dominant-color="3F1C34"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2024-07-24 16-06-28</span><span class="informations">1908×422 123 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>