Hey @aguirguis

What confuses me is that it is not clear where does the execution (forward and backward passes) happen in this example…is it on the trainer machines or the PS machine?

The executions scattered on both trainer and PS:

Forward pass: input (on trainer) -> parameter (on PS) -> output (on trainer) -> loss (on trainer)

So, for each trainer, there are three pieces of autograd graphs scattered on the trainer and the PS, and those are connected by RPC.

It makes sense that these computations happen on the trainers yet, the forward pass is defined as a part of the PS code, which is called remotely by the trainer! would you please clarify this part?

Yes. Both TrainerNet and ParameterServer have a forward function as they are both nn.Module subclasses. TrainerNet forward calls ParameterServer forward. An analogy would be if you have nn.Sequential(nn.Linear(10, 10)), both Sequential and Linear implement their own forward and Sequential's forward calls Linear's forward. This is also why the autograd graph scatters on both the trainer and the parameter server.

Also, as far as I understand, usually, the PS collects the trainers’ gradients and aggregates them (for instance by averaging) to update the global model. Does this aggregation part exist in the tutorial?

No, the aggregation is not part of the tutorial. Instead, it is doing sth similar to Hogwild! 1 training. The aggregation method you mentioned are one type of synchronized training, where all trainers need to wait for that aggregation to finish before they can proceed. This works, and it is also possible to implement this synchronous aggregation using torch.distributed.rpc, but it won’t scale to large workload on large clusters, as any straggler will kill the performance.

if no, then what’s the strategy to incorporate all trainers’ work in the global model.

One option would be letting the PS to send parameters to trainers, and then get back gradients. Then, the PS aggregates the grads as you mentioned above, and updates its parameters. The PS can keep doing this until the model converges. This is similar to this RL tutorial 3 where the agent tells observers what to do and the agent owns the policy model.

A follow-up question on that last point: how can I apply more sophisticated aggregation functions on the gradients (let’s say I want to take median instead of mean on the PS machine)…how can I do this?

If adopting the RL tutorial above, then the PS can collect the gradients from different trainers in a list and compute the medium accordingly.
cc the author of the PS tutorial @rvarm1