I attempted to create my own parameter server implementation where the PS only combines gradients, does the optimization step and some synchronization.
my code is something like this:
worker:
def run_training_loop(rank, num_gpus, train_loader, test_loader, param_server_rref):
    
    model = resnet18(pretrained=True)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, NUMBER_OF_CLASSES)
    optimizer = optim.SGD(model.parameters(), lr=0.0002, momentum=0.9)

    for epoch in range(NUMBER_OF_EPOCHS):

        for i, (data, target) in enumerate(train_loader):
            
            model_output = model(data)
            loss = F.cross_entropy(model_output, target) 
            loss.backward()

            model = param_server_rref.rpc_sync().optimize(rank, model)

        get_accuracy(test_loader, model)

    print("Training complete!")

server:
def optimize(self, rank, model):
        self.barrier.wait()
  
        if rank == 1:
               self.combined_model = self.combine_model()
               optimizer = optim.SGD(combined_model.parameters(), lr=0.0002, momentum=0.9)
               optimizer.step()
               optimizer.zero_grad()

        self.barrier.wait()
        return copy.deepcopy(self.combined_model)

Unfortunately the optimization step on the paramteter server doesn’t seem to have any effect!
I suspect that autograd is not recording a graph since the optimizer is created after calling the forward/backward functions?!
How can I get this working without doing all computation on the server?
The distributed_optimizer seems to be only useful for model parallelism and not data parallelism since it does not respect gradients from different workers. Is this correct?
Thank you in advance.