Hi @MichaelZ thanks for posting the question. It seems to me one of your issue is that optimize method on the server side is creating a new optimizer everytime you run a epoch, this will make a new state for the model params everytime when updating the param, although autograd might recording the gradients.  You should create the optimizer once and in optimize just call the step function i think. Also I am not sure what combine_model() is and can’t be sure if it’s recording gradients correctly or not.

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?

Distributed optimizer is mainly used by model parallelism as you said, but it does respect the gradients from different workers in the case of model parallelism. For DDP, the gradients are all reduced so it’s just a local optimizer.