## 4 Task Allocation Prediction Network
### Joint Training of Policy Gradient  
The policy gradient is jointly trained with the task partition and allocation network in the prediction network which can obtain a task allocation plan that minimizes the total training time of a single training batch. The total time of training a batch of DNN models under task allocation plan, \(\rho\), is defined as \(t(\rho)\), including one forward calculation, one backward propagation, and one parameter update for which the unit of time is seconds. In order to avoid the variance of the running time measurement, for each predicted task allocation plan, we run the DNN model on the computing node according to the plan. Then we measure it several times to calculate the average value to get the reward function of the current plan. Finally, after several iterations, the task allocation plan with the shortest actual running time is obtained. In the experiment, each task allocation plan was run 20 times. And the average value of the last 10 running results was taken as the running time measurement value. In the actual running process, sampling the running time as the reward function directly will cause poor training effect. On one hand, the task allocation scheme which was proposed in the early stage of training may lead to a longer running time with producing inappropriate learning signals for the convergence of the prediction network. On the other hand, as the prediction network converges, the predicted task allocation schemes are more similar, which will result in smaller differences in the corresponding execution time. Therefore, the reward function for a task allocation plan, \(\rho\), is defined as:  
\[R_{\rho}=-\sqrt{t(\rho)} \tag{1}\]  
The prediction network should minimize the expected value of \(R_{\rho}\) to achieve a better training effect. In summary, the loss function [31], \(J\left(\theta_{g},\theta_{d}\right)\), is defined as:  
\[ \begin{array}{l}J\left(\theta_{g},\theta_{d}\right)=\mathrm{E}_{\mathrm{P} \left(\rho;\theta_{g},\theta_{d}\right)}[R_{\rho}]\\ =\sum\limits_{g\sim\tau_{g}}\sum\limits_{d\sim\tau_{d}}p(g;\theta_{g})p(d|g; \theta_{d})\,R_{\rho}\end{array}  \tag{2}\]  
where \(\theta_{g}\) and \(\theta_{d}\), respectively, represent the parameters of the task partition network and the task allocation network; \(p(g;\theta_{g})\) represents the probability of the task partition network generating the grouping result, \(g\); and \(p(d|g;\theta_{d})\) represents the probability of the task allocation network generating the task allocation result, \(d\), according to the grouping result, \(g\). The task partition network independently predicts the partition of the DNN model. Then the task allocation network allocates computing nodes to each group according to the task partition. The changes in the grouping results make the predicted task allocation plan correspond to a large difference in running time. This paper introduces exponential moving average (EMA) to better reflect the actual execution time of the predicted allocation plan. Combining EMA to obtain the partial derivative of Equation (2), the gradients of parameters \(\theta_{g}\) and \(\theta_{d}\)[31] of the prediction network are as follows:  
\[ \begin{array}{l}\
abla_{\theta_{g}}J(\theta_{g},\theta_{d})=\sum\limits_{g \sim\tau_{g}}\
abla_{\theta_{g}}\log p(g,\theta_{g})\sum\limits_{d\sim\tau_{ d}}(R_{\rho}-Base)\\ =\sum\limits_{g\sim\tau_{g}}\
abla_{\theta_{g}}p(g;\theta_{g})\sum\limits_{d \sim\tau_{d}}p(d|g;\theta_{d})\,(R_{\rho}-Base)\end{array}  \tag{3}\]\[ \begin{array}{l}\
abla_{\theta_{d}}I(\theta_{g},\theta_{d})=\sum\limits_{d \sim\pi_{d}}\sum\limits_{g\sim\pi_{g}}p(g;\theta_{g})\
abla_{\theta_{d}}p(d|g; \theta_{d})(R_{\rho}-Base)\\ =\sum\limits_{d\sim\pi_{d}}(\sum\limits_{g\sim\pi_{g}}\
abla_{\theta_{d}}\log p (d|g;\theta_{d}))(R_{\rho}-Base)\end{array}  \tag{4}\]  
Among them, \(Base\) is the EMA that rewards \(R_{\rho}\). The gradients of parameters, \(\theta_{g}\) and \(\theta_{d}\), related to the prediction network can be calculated by Equations (3) and (4).