## 3 Methods
### Training Algorithm
#### 3.1.2 Heterogeneous Language Model Training  
The proposed heterogeneous training consists of three major steps: forward propagation, backward propagation, and optimization. Each step proceeds sequentially for the divided sub-model. Forward propagation takes data batches as input and finally obtains a loss function. Backward propagation calculates the gradient for each parameter via backpropagation with the acquired loss value. Optimization uses a gradient to update parameters.  
1. Forward Propagation Forward propagation propagates data input to sequential sub-models in order and acquires the final loss value. One sub-model is forward propagated across multiple GPUs. The context for forward propagation of the sub-model in the GPU consists of parameters and inputs. Before executing the sub-model, the context on the CPU side is requested and sent to the GPU. When the context is ready, it forwards and sends the sub-model result to the CPU. Sub-model results are reused in the backward propagation process. When the result is sent to the CPU, all memory allocated to the GPU is initialized. The purpose is to repeat this process to obtain the final loss value. If the given model has \(N\) sub-models, and we define the \(i\)-th sub-model as \(S_{i}\), the output of the \(i\)-th sub-model as \(O_{i}\), and the entire parameter set corresponding to the i-th sub-model as \(W_{i}\). Then, the process of calculating the total loss value is as follows. (Here, \(O_{0}\) is the initial input to the network, and \(Y\) is the label.) \[O_{i}=S_{i}(O_{i-1};W_{i}),for\ i=1,2,\ldots,N\] \[Loss=CrossEntropy(Y,O_{N})\]
2. Backward Propagation Backward propagation obtains the gradient for each model's parameters by back-propagating the loss value in the reverse order of the previously defined sub-models. One sub-model is backpropagated on multiple GPUs. The context needed to perform backward propagation of the sub-model in the GPU consists of the gradient computed  
Figure 1: Division of the transformer-based model according to main components.  
from the previous sub-model and the parameters of the current sub-model. Using the conditions utilized earlier in the forward propagation, the gradient with respect to the parameter set \(W_{i}\) corresponding to the \(i\)-th sub-model \(S_{i}\) is computed as follows:  
\[\frac{\delta Loss}{\delta W_{i}}=\frac{\delta Loss}{\delta O_{N}}\frac{\delta O _{N}}{\delta O_{N-1}}\cdots\frac{\delta O_{i+1}}{\delta O_{i}}\frac{\delta O_{ i}}{\delta O_{W_{i}}}\]  
The sub-model requests a context on the CPU side to backpropagate the gradient and sends it to the GPU. To perform backpropagation, forward propagation must be conducted once for a given context to generate a computational graph. Therefore, in this process, when the context is transmitted to the GPU, forward propagation is performed once, and the gradient of the parameter is obtained by providing the gradient of the previous sub-model as an input to the generated computation graph. The calculated gradient is sent back to the CPU, and when the transfer is complete, all memory allocated to run the sub-model is released.
3. Optimization Optimization proceeds using the AdamW optimizer with the \(i\)-th sub-model's parameters \(W_{i}\) and the computed gradient \(\frac{\delta Loss}{\delta W_{i}}\). Optimization is performed using CPU multithreading for the AdamW operation [33; 34]. Compared to forward or backward propagation, which have a computational complexity of \(O(n^{3})\), AdamW has a significantly lower computational complexity of \(O(n^{2})\). Therefore, even in a CPU multi-thread environment, computations can be completed in a shorter time compared to propagation operations. Moreover, by conducting optimization operations on the CPU and main memory, the GPU can focus on allocating memory for computation-intensive propagation operations.  
Figure 2 summarizes the entire training process.  
Figure 2: Overview of the three main phases in the training process.