Where does the execution happen in the RPC-based parameter server tutorial?

aguirguis: <p>Hello,</p>
<p>In this nice <a href="https://pytorch.org/tutorials/intermediate/rpc_param_server_tutorial.html" rel="nofollow noopener">tutorial</a>, it is described how to implement a parameter server (PS) deployment with RPC.<br>
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?<br>
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?<br>
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? if yes, would you please explain where? if no, then what’s the strategy to incorporate all trainers’ work in the global model.<br>
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?</p>
<p>Thank you very much!</p>

mrshenli: <p>Hey <a class="mention" href="/u/aguirguis">@aguirguis</a></p>
<blockquote>
<p>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?</p>
</blockquote>
<p>The executions scattered on both trainer and PS:</p>
<ul>
<li>Forward pass: input (on trainer) → parameter (on PS) → output (on trainer) → loss (on trainer)</li>
</ul>
<p>So, for each trainer, there are three pieces of autograd graphs scattered on the trainer and the PS, and those are connected by RPC.</p>
<blockquote>
<p>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?</p>
</blockquote>
<p>Yes. Both <code>TrainerNet</code> and <code>ParameterServer</code> have a <code>forward</code> function as they are both <code>nn.Module</code> subclasses. <code>TrainerNet</code> forward calls <code>ParameterServer</code> forward. An analogy would be if you have <code>nn.Sequential(nn.Linear(10, 10))</code>, both <code>Sequential</code> and <code>Linear</code> implement their own <code>forward</code> and <code>Sequential</code>’s forward calls <code>Linear</code>’s forward. This is also why the autograd graph scatters on both the trainer and the parameter server.</p>
<blockquote>
<p>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?</p>
</blockquote>
<p>No, the aggregation is not part of the tutorial. Instead, it is doing sth similar to <a href="https://people.eecs.berkeley.edu/~brecht/papers/hogwildTR.pdf" rel="noopener nofollow ugc">Hogwild!</a> 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 <code>torch.distributed.rpc</code>, but it won’t scale to large workload on large clusters, as any straggler will kill the performance.</p>
<blockquote>
<p>if no, then what’s the strategy to incorporate all trainers’ work in the global model.</p>
</blockquote>
<p>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 <a href="https://pytorch.org/tutorials/intermediate/rpc_tutorial.html#distributed-reinforcement-learning-using-rpc-and-rref" rel="noopener nofollow ugc">RL tutorial</a> where the agent tells observers what to do and the agent owns the policy model.</p>
<blockquote>
<p>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?</p>
</blockquote>
<p>If adopting the RL tutorial above, then the PS can collect the gradients from different trainers in a list and compute the medium accordingly.</p>
<p>cc the author of the PS tutorial <a class="mention" href="/u/rvarm1">@rvarm1</a></p>

aguirguis: <p>Thanks, <a class="mention" href="/u/mrshenli">@mrshenli</a> for your detailed answers. I have a few follow-up questions:</p>
<aside class="quote no-group quote-modified" data-username="mrshenli" data-post="2" data-topic="84931">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mrshenli/48/12220_2.png" class="avatar"> mrshenli:</div>
<blockquote>
<p>input (on trainer) → parameter (on PS) → output (on trainer)</p>
</blockquote>
</aside>
<p>Does this mean the input is sent to the PS first to propagate through the network parameters in the PS and then at the end, the output is sent back to the trainer?<br>
What confuses me here is that it seems that the <code>forward</code> function of <code>TrainerNet</code> is just dummy and all what it does is calling that of <code>ParameterServer</code>. As far as I understand, in the PS architecture, data never leaves the trainer machine and that the whole gradient computation process should be done entirely locally on the trainer machine.<br>
If you can describe all the communication that happens in one training iteration, that would be great. For instance, assume that we have one PS machine and two trainer machines. PS has the model and each trainer has a few data samples. What is sent to whom?<br>
Hogwild! assumes shared memory so, the setup is inherently different from that of the PS, right? I cannot entirely digest how/why do you blend these two setups. Would you please clarify?</p>
<p>Thanks a lot.</p>

mrshenli: <aside class="quote no-group" data-username="aguirguis" data-post="3" data-topic="84931">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aguirguis/48/21870_2.png" class="avatar"> aguirguis:</div>
<blockquote>
<p>Does this mean the input is sent to the PS first to propagate through the network parameters in the PS and then at the end, the output is sent back to the trainer?</p>
</blockquote>
</aside>
<p>Yes.</p>
<blockquote>
<p>What confuses me here is that it seems that the  <code>forward</code>  function of  <code>TrainerNet</code>  is just dummy and all what it does is calling that of  <code>ParameterServer</code> .</p>
</blockquote>
<p>Yes, you are right. In this specific case, as <code>x</code> does not require grad, there is no need to link it to the distributed autograd graph. So there are only two pieces of autograd graph on PS and the trainer. (I was wrong when saying there three pieces in previous comments.)</p>
<blockquote>
<p>As far as I understand, in the PS architecture, data never leaves the trainer machine and that the whole gradient computation process should be done entirely locally on the trainer machine.</p>
</blockquote>
<p>There are different ways to implement this. Imagine there is a super large embedding table and the trainer only holds a several lookup indices in each iteration. One solution is to do training all on trainer, but then the application will need to implement update functions that converts indices and gradients from the trainer back to the embedding table gradients. Another option is let autograd engine taking care of this, and simply calling <code>loss.backward()</code> on trainer will be sufficient to update embedding table on ps.</p>
<blockquote>
<p>If you can describe all the communication that happens in one training iteration, that would be great. For instance, assume that we have one PS machine and two trainer machines. PS has the model and each trainer has a few data samples. What is sent to whom?If you can describe all the communication that happens in one training iteration, that would be great. For instance, assume that we have one PS machine and two trainer machines. PS has the model and each trainer has a few data samples. What is sent to whom?</p>
</blockquote>
<p>Sure. Since trainers are independent in that tutorial IIUC, I will only describe what happens between a PS-trainer pair.</p>
<p>In forward pass, there are two comms: 1) trainer → ps to send input sample 2) ps → trainer to send the output<br>
In the backward pass, there is one comm: trainer → ps to send the gradients for the model outputs, which will then trigger local autograd engine on the ps to compute gradients on the model.<br>
In the optimizer step pass, there is one comm: trainer → ps tell the local optimizer on ps to update model parameters. (It is possible to pack this into the comm during the backward pass using hooks.)</p>
<p>Since there are two trainers accessing the same model, instead of storing the grads in <code>param.grad</code>, the ps will put those grads in dedicated contexts associated with each distributed autograd context, and those grads will later be consumed by the distributed optimizer.</p>
<p>More details about dist autograd can be found here: <a href="https://pytorch.org/docs/stable/rpc/distributed_autograd.html#distributed-autograd-design" class="inline-onebox" rel="noopener nofollow ugc">Distributed Autograd Design — PyTorch 2.1 documentation</a></p>
<blockquote>
<p>Hogwild! assumes shared memory so, the setup is inherently different from that of the PS, right? I cannot entirely digest how/why do you blend these two setups. Would you please clarify?</p>
</blockquote>
<p>Right, the original paper was mainly focusing on shm. But the lock-free spirit can be apply to distributed training as well. This is especially useful for training using large dataset with large embedding tables.</p>

aguirguis: <p>Thanks <a class="mention" href="/u/mrshenli">@mrshenli</a> for your detailed answers. Now, everything is clear to me.<br>
Probably a better design for my use case is to use local optimizer and autograd as in the RL tutorial you referred to before.</p>