NCCL ParameterServer

Osama_Abu_Hamdan: <p>Hello,<br>
I wrote this ParameterServer code, it’s naive but I don’t know why it doesn’t work.</p>
<pre><code class="lang-auto">import os

import torch
import torch.distributed as dist
from torch.optim import SGD
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


class Worker:
    def __init__(self, rank, device, model, data_loader):
        self.rank = rank
        self.device = device
        self.model = model
        self.data_loader = data_loader
        self.criterion = torch.nn.CrossEntropyLoss()

    def receive_weights(self):
        for param in self.model.parameters():
            dist.recv(param.data, src=0)
        print(f"Worker {self.rank} received weights")

    def train_one_step(self, data, target):
        output = self.model(data)
        loss = self.criterion(output, target)
        loss.backward()

    def send_gradients(self):
        for param in self.model.parameters():
            dist.send(param.grad.data, dst=0)
        print(f"Worker {self.rank} sent gradients")

    def train_one_epoch(self):
        self.model.train()
        for batch_idx, (data, target) in enumerate(self.data_loader):
            print(f"Worker {self.rank} - Step: {batch_idx}")
            data, target = data.to(self.device), target.to(self.device)
            self.receive_weights()
            self.train_one_step(data, target)
            self.send_gradients()
            dist.barrier()
            print(f"Worker {self.rank} finished step {batch_idx}")

    def run(self, epochs):
        for epoch in range(epochs):
            print(f"Worker {self.rank} - Epoch: {epoch}")
            self.train_one_epoch()
            print(f"Worker {self.rank} finished epoch {epoch}")
        print(f"Worker {self.rank} finished training")


class ParameterServer:
    def __init__(self, device, model, data_loader):
        self.model = model
        self.device = device
        self.optimizer = SGD(self.model.parameters(), lr=0.01)
        self.data_loader = data_loader
        self.initialize_weights()

    def initialize_weights(self):
        for param in self.model.parameters():
            torch.nn.init.uniform_(param.data, -0.1, 0.1)
        print("Parameter Server initialized weights")

    def send_initial_weights(self):
        for param in self.model.parameters():
            dist.broadcast(param.data, src=0)

    def gather_and_update(self):
        for param in self.model.parameters():
            grad = torch.zeros_like(param.data)
            dist.reduce(grad, dst=0, op=dist.ReduceOp.AVG)
            param.grad = grad
        self.optimizer.zero_grad()
        self.optimizer.step()
        print("Parameter Server gathered and updated weights")

    def run(self, epochs):
        for epoch in range(epochs):
            print(f"Parameter Server - Epoch: {epoch}")
            for batch_idx, (data, target) in enumerate(self.data_loader):
                print(f"Parameter Server - Step: {batch_idx}")
                self.send_initial_weights()
                self.gather_and_update()
                dist.barrier()
                print(f"Parameter Server finished step {batch_idx}")
            print(f"Parameter Server finished epoch {epoch}")
        print("Parameter Server finished training")


def run(rank):
    torch.manual_seed(42)
    model = torch.nn.Linear(784, 10)
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
    dataset = datasets.MNIST('./data', train=True, download=False, transform=transform)
    data_loader = DataLoader(dataset, batch_size=64, shuffle=True)
    epochs = 10
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model.to(device)
    if rank == 0:
        ps = ParameterServer(device, model, data_loader)
        ps.run(epochs)
    else:
        worker = Worker(rank, device, model, data_loader)
        worker.run(epochs)


if __name__ == "__main__":
    rank = int(os.environ["RANK"])
    world_size = int(os.environ['WORLD_SIZE'])
    torch.distributed.init_process_group(
        backend='nccl',
        init_method='env://',
        rank=rank,
        world_size=world_size
    )
    run(rank)
    dist.destroy_process_group()
</code></pre>
<p>I run it with SLURM with this command</p>
<pre><code class="lang-auto">srun torchrun \
    --nnodes $NUM_NODES \
    --nproc_per_node 1 \
    --rdzv_id $RANDOM \
    --rdzv_endpoint $MASTER_ADDR:$MASTER_PORT \
    --rdzv_backend c10d \
    ps_pytorch.py
</code></pre>
<p>It’s stuck at</p>
<pre><code class="lang-auto">Worker 2 - Epoch: 0
Worker 1 - Epoch: 0
Worker 1 - Step: 0
Worker 2 - Step: 0
Parameter Server initialized weights
Parameter Server - Epoch: 0
Parameter Server - Step: 0
</code></pre>