Maximizing Parallelism in Distributed Training for Huge Neural Networks

Extracted source text · arXiv 2105.14450

\documentclass{article}

% if you need to pass options to natbib, use, e.g.:
%     \PassOptionsToPackage{numbers, compress}{natbib}
% before loading neurips_2021

% ready for submission
\PassOptionsToPackage{square, numbers}{natbib}
\usepackage[preprint]{neurips_2021}

% to compile a preprint version, e.g., for submission to arXiv, add add the
% [preprint] option:
%     \usepackage[preprint]{neurips_2021}

% to compile a camera-ready version, add the [final] option, e.g.:
%     \usepackage[final]{neurips_2021}

% to avoid loading the natbib package, add option nonatbib:
%    \usepackage[nonatbib]{neurips_2021}

\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc}    % use 8-bit T1 fonts
\usepackage{hyperref}       % hyperlinks
\usepackage{url}            % simple URL typesetting
\usepackage{booktabs}       % professional-quality tables
\usepackage{amsfonts}       % blackboard math symbols
\usepackage{nicefrac}       % compact symbols for 1/2, etc.
\usepackage{microtype}      % microtypography
\usepackage{xcolor}         % colors

\usepackage{caption}
\usepackage{subcaption}
\usepackage{graphicx}
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{bm}
\usepackage{subfiles}
\usepackage{multirow}
\usepackage{multicol}
\usepackage{makecell}
\usepackage{tabularx}
\usepackage{footmisc}
\usepackage{listings}
\usepackage{algorithm}
\usepackage{algorithmic}
\renewcommand{\algorithmicrequire}{\textbf{Input:}}
\renewcommand{\algorithmicensure}{\textbf{Output:}}

\title{Maximizing Parallelism in Distributed Training for Huge Neural Networks}

% The \author macro works with any number of authors. There are two commands
% used to separate the names and addresses of multiple authors: \And and \AND.
%
% Using \And between authors leaves it to LaTeX to determine where to break the
% lines. Using \AND forces a line break at that point. So, if LaTeX puts 3 of 4
% authors names on the first line, and the last on the second line, try using
% \AND instead of \And before the third author name.

\author{%
%   David S.~Hippocampus\thanks{Use footnote for providing further information
%     about author (webpage, alternative address)---\emph{not} for acknowledging
%     funding agencies.} \\
%   Department of Computer Science\\
%   Cranberry-Lemon University\\
%   Pittsburgh, PA 15213 \\
%   \texttt{hippo@cs.cranberry-lemon.edu} \\
    Zhengda Bian\thanks{These authors have equal contributions.}\\
    National University of Singapore\\
    Singapore\\
    \texttt{zbian@comp.nus.edu.sg}\\
    \And
    Qifan Xu\footnotemark[\value{footnote}]\\
    University of California, Los Angeles\\
    Los Angeles, California, USA\\
    \texttt{QifanXu@mednet.ucla.edu}\\
    \AND
    Boxiang Wang\\
    National University of Singapore\\
    Singapore\\
    \texttt{boxiangwang059@gmail.com}\\
    \And
    Yang You\\
    National University of Singapore\\
    Singapore\\
    \texttt{youy@comp.nus.edu.sg}\\
  % examples of more authors
  % \And
  % Coauthor \\
  % Affiliation \\
  % Address \\
  % \texttt{email} \\
  % \AND
  % Coauthor \\
  % Affiliation \\
  % Address \\
  % \texttt{email} \\
  % \And
  % Coauthor \\
  % Affiliation \\
  % Address \\
  % \texttt{email} \\
  % \And
  % Coauthor \\
  % Affiliation \\
  % Address \\
  % \texttt{email} \\
}


\begin{document}

\maketitle

\begin{abstract}
%   The abstract paragraph should be indented \nicefrac{1}{2}~inch (3~picas) on
%   both the left- and right-hand margins. Use 10~point type, with a vertical
%   spacing (leading) of 11~points.  The word \textbf{Abstract} must be centered,
%   bold, and in point size 12. Two line spaces precede the abstract. The abstract
%   must be limited to one paragraph.
% As user data explodes, companies have to frequently update their models to capture the preference of users and create profiles for newly generated items to provide up-to-date service for searching, recommending and broadcasting.
% The magic power of huge language models has been brought to us by academia, so that 
% nowadays people have to think twice if the pictures and texts in front of them are taken from camera or written by human.
% So it’s no wonder that 
% Hinton declared that life, universe and everything were just 4.398 trillion parameters short after GPT-3 was published.
The recent Natural Language Processing techniques have been refreshing the state-of-the-art performance at an incredible speed.
Training huge language models is therefore an imperative demand in both industry and academy.
However, the huge models impose challenges to both hardware and software.
Graphical processing units (GPUs) are iterated frequently to meet the exploding demand, and a variety of ASICs like TPUs are spawned.
However, there is still a tension between the fast growth of the extremely huge models and fact that Moore’s law is approaching the end.
To this end, many model parallelism techniques are proposed to distribute the model parameters to multiple devices, so as to alleviate the tension on both memory and computation.
Our work is the first to introduce a 3-dimensional model parallelism for expediting huge language models.
By reaching a perfect load balance, our approach presents smaller memory and communication cost than existing state-of-the-art 1-D and 2-D model parallelism.
Our experiments on 64 TACC's V100 GPUs show that our 3-D parallelism outperforms the 1-D and 2-D parallelism with 2.32X and 1.57X speedup, respectively.

\end{abstract}

\subfile{section/introduction.tex}

\subfile{section/background.tex}

\subfile{section/design.tex}

% \subfile{section/implementation.tex}

\subfile{section/evaluation.tex}

\subfile{section/conclusion.tex}


\bibliographystyle{plainnat}
\bibliography{reference}

% \appendix

% \section{Appendix}

% Optionally include extra information (complete proofs, additional experiments and plots) in the appendix.
% This section will often be part of the supplemental material.

% \subfile{section/instruction.tex}
\iffalse

\section*{Checklist}

\begin{enumerate}

\item For all authors...
\begin{enumerate}
  \item Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
    \answerYes{}
  \item Did you describe the limitations of your work?
    \answerYes{See Section~\ref{sec:conclusion}}
  \item Did you discuss any potential negative societal impacts of your work?
    \answerNA{}
  \item Have you read the ethics review guidelines and ensured that your paper conforms to them?
    \answerYes{}
\end{enumerate}

\item If you are including theoretical results...
\begin{enumerate}
  \item Did you state the full set of assumptions of all theoretical results?
    \answerYes{See Section~\ref{sec:design}}
	\item Did you include complete proofs of all theoretical results?
    \answerYes{See Section~\ref{sec:design}}
\end{enumerate}

\item If you ran experiments...
\begin{enumerate}
  \item Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
    \answerNo{No, currently our codes are not publicly available.}
  \item Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
    \answerYes{See Section~\ref{sec:eval}.}
	\item Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
    \answerNo{No, I will work on this later due to the time limit.}
	\item Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
    \answerYes{See Section~\ref{sec:eval}.}
\end{enumerate}

\item If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
\begin{enumerate}
  \item If your work uses existing assets, did you cite the creators?
    \answerYes{See Section~\ref{sec:background}.}
  \item Did you mention the license of the assets?
    \answerNA{}
  \item Did you include any new assets either in the supplemental material or as a URL?
    \answerNA{}
  \item Did you discuss whether and how consent was obtained from people whose data you're using/curating?
    \answerNA{}
  \item Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
    \answerNA{}
\end{enumerate}

\item If you used crowdsourcing or conducted research with human subjects...
\begin{enumerate}
  \item Did you include the full text of instructions given to participants and screenshots, if applicable?
    \answerNA{}
  \item Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
    \answerNA{}
  \item Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
    \answerNA{}
\end{enumerate}

\end{enumerate}

\fi

\end{document}
\section{Background}
\label{sec:background}

\subsection{Transformer Language Models}
There has been a recent trend to use pretrained language models to help people train specific natural language processing (NLP) tasks, as pretrained models leverage the understanding of very large corpus (e.g. Wikipedia) to alleviate the efforts of training entire language models.
Transformer is one of the current dominant choices of pretraining language models, achieving the state-of-the-art compute efficiency and prediction accuracy.
The Transformer language model uses deep attention modules to transform a given sequential input into another sequential output.
Like LSTM, a complete Transformer model includes two parts, an Encoder and a Decoder, but each consists of attention layers rather than recurrent layers.
However, recent state-of-the-art Transformer applications, such as BERT \cite{BERT}, GPT-2 \cite{gpt2}, and GPT-3 \cite{gpt-3}, adopt more concise structures that only use necessary Encoders or Decoders on demand.

Nevertheless, training Transformer language models can be extremely computational expensive.
The recent language models advance their accuracy with a rising number of parameters (e.g. OpenAI GPT: 110M \cite{gpt1}; BERT: 340M \cite{BERT}; GPT-2: 1.5B \cite{gpt2}; GPT-3: 175B \cite{gpt-3}; Switch Transformers: 1.6T \cite{switch-tranformer}).
In the meantime, to deal with huge language corpus, there has been a growing interest in using large minibatches to reduce training time.
Therefore, efficient systems to expedite Transformers become an urgent demand for training large language models.

\subsection{Data and Model Parallelism}
Data parallelism is the most common paradigm to parallel the computation of deep neural networks.
This approach distributes the entire minibatch across multiple workers, which execute a single replica of the model and communicate with each other for synchronization at the end of each training step.
It is easy to use more devices to train larger minibatches.
By taking advantage of recent optimization techniques \cite{smith2017don,you2017large,you2018imagenet,you2020large,bytescheduler}, training speed reaches almost linear scaling with the number of devices.
However, due to the issue of generalization gap \cite{keskar2016large,hoffer2017train}, using very large minibatches requires extra efforts to guarantee the convergence performance \cite{goyal2017accurate,you2017large,you2018imagenet,you2020large}.
Besides, a practical limitation is that large Transformer models that have more than billions of parameters are basically not able to be accommodated into the memory of a single GPU device.

Model parallelism can successfully remove the above memory limitation.
There are two model parallel paradigms.
First, the layer-wise pipelined parallelism splits the entire model by layer, and execute the layers in the pipeline fashion.
Some approaches \cite{pipedream,gpipe} use their algorithms to tackle the inconsistency issue among different pipelines, so that the computation and communication time can be overlapped.
However, since each training step requires both forward and backward processing, there will be an inevitable bubble overhead, so that the compute resources cannot be fully utilized.

The second model parallel paradigm is the intra-layer parallelism, which is orthogonal to the pipelined parallelism.
It distributes the operations in each layer such as matrix multiplications and activations across multiple workers.
For example, Mesh-Tensorflow \cite{meshTensorflow}, a framework proposed by Google, provides a convenient way for users to allocate the partition of tensor dimensions. 
Then the matrix multiplications can be automatically partitioned in a single dimension and distributed to workers, without any user elaboration.
% However, memory redundancy is severe, and the originally nearly-free reshape operations would result in complex communication because of layout change. 
Similarly, Megatron-LM \cite{megatron} adopts a 1-D matrix partition strategy to implement the Transformer model.
It splits matrices along rows or columns, and gathers results with the all-reduce operation.
% Although parameters and optimizer states are fully distributed, there is still memory redundancy for activations, as each device has to accommodate the whole input and output activations for every multi-layer perceptrons (MLPs) and self-attention modules.
Moreover, Optimus uses the Scalable Universal Matrix Multiplication Algorithm (SUMMA) \cite{SUMMA}, a 2-D parallel strategy for matrix multiplications, which helps to reduce memory and communication costs required for executing the Transfer model.
Our approach aims to take a step further and use the 3-D parallelism to improve the performance by optimally balancing the computation, memory and communication load.

\subsection{3-D Parallel Matrix Multiplication}
In this section, we provide an outline of the 3-D parallel matrix multiplication algorithm \cite{agarwal1995three}, which is adopted in our approach.

\begin{figure}[h]
    \centering
    \includegraphics[width=0.5\linewidth]{figure/cube.pdf}
    \caption{The structure of a 3-D processing cube with $P=p^3$ processors. The colored block $(i,j,k)$ represents an example of a single processor.}
    \label{fig:cube}
\end{figure}

This algorithm computes the multiplication between two dense matrices concurrently on $P=p^3$ processors, as illustrated in \autoref{fig:cube}, where $P$ processors can be stacked into a $p^3$ cube.
In this cube, let $A, B, C$ denote the planes, $x, y, z$ denote the directions for the planes, and $i, j, l$ denote the different indices along the directions.
We also use $A, B, C$ to represent the matrices involved in the multiplication, so that $A_{ij}, 0 \le i,j < p$ is supposed to represent a submatrix of $A$.

For simplicity, we consider an example of the 3-D matrix multiplication $C=AB$ on a $2\times2\times2$ processing cube.
We split $A$ and $B$ of size $(M,N)$ and $(N,K)$ into $2\times2$ partitions as follows, so that the sizes of each partition $A_{il}$ and $B_{lj}$ is $(M/2, N/2)$ and $(N/2, K/2)$ respectively, for $0 \le i,j < 2$.
\begin{align}
    A = \left[
        \begin{matrix}
            A_{00} & A_{01} \\
            A_{10} & A_{11}
        \end{matrix}
        \right], \quad
    B = \left[
        \begin{matrix}
            B_{00} & B_{01} \\
            B_{10} & B_{11}
        \end{matrix}
        \right].
\end{align}

\begin{figure}[H]
    \centering
    \begin{subfigure}[t]{0.25\linewidth}
        \centering
        \includegraphics[width=\linewidth]{figure/3d_a.pdf}
        \caption{Broadcast $A_{il}$}
        \label{fig:3d-a}
    \end{subfigure}
    \begin{subfigure}[t]{0.25\linewidth}
        \centering
        \includegraphics[width=\linewidth]{figure/3d_b.pdf}
        \caption{Broadcast $B_{lj}$}
        \label{fig:3d-b}
    \end{subfigure}
    \begin{subfigure}[t]{0.25\linewidth}
        \centering
        \includegraphics[width=\linewidth]{figure/3d_c.pdf}
        \caption{Reduce $C_{ij}$}
        \label{fig:3d-c}
    \end{subfigure}
    \caption{An example of the 3-D parallel matrix multiplication of size $2\times2\times2$.}
    \label{fig:3d}
\end{figure}

Each partition $A_{il}$ is stored in the processor $(i,0,l)$, while $B_{lj}$ is stored in $(0,j,l)$.
Before multiplication, we first broadcast each $A_{il}$ from $(i,0,l)$ along the $y$ direction (\autoref{fig:3d-a}), so that we have $A_il$ on each $(i,j,l)$.
Similarly, we broadcast each $B_{lj}$ along the $x$ direction (\autoref{fig:3d-b}).
Next, we can compute $C_{ijl} = A_{il}B_{lj}$ on each processor $(i,j,l)$.
Then we reduce $C_{ijl}$ to $(i,j,0)$ along the $z$ direction to derive $C_{ij}=\sum_l A_{il}B_{lj}$ (\autoref{fig:3d-c}).
Thus, we finally get $C=AB$ as
\begin{align}
    C = \left[
        \begin{matrix}
            A_{00}B_{00}+A_{01}B_{10} & A_{00}B_{01}+A_{01}B_{11} \\
            A_{10}B_{00}+A_{11}B_{10} & A_{10}B_{01}+A_{11}B_{11}
        \end{matrix}
        \right].
\end{align}

\section{Conclusion}
\label{sec:conclusion}
In this work, we introduce a 3-D intra-layer model parallelism algorithm for training huge neural models.
We propose a load balanced design to store and execute linear layers with minimum memory and communication cost.
By leveraging the 3-D model parallelism, we implement a 3-D parallel Transformer model and evaluate it on up to 64 GPUs.
Compared with the existing 1-D and 2-D parallelism, our Transformer model achieves a significant speedup.
However, it is still interesting to leverage our approach to train larger models on more processors.
We expect to see promising advances of our work with larger-scale compute resources in the future.
\section{3-D Parallel Transformers}
\label{sec:design}

We now describe our main design and implementation of the 3-D parallel Transformer language model.

\subsection{3-D Operations}

As demonstrated in \autoref{fig:transformer}, the Transformer language model computes its hidden states from top to bottom, through multiple Transformer blocks.
In each block, the major computation in both the Self-Attention and Multi-Layer Perception (MLP) layer involves linear operations and activations.
While activation operations can be independently executed in parallel, to accelerate linear operations is the key to reduce execution time of the Transformer model.
We firstly tackle the load balancing issue of the original 3-D matrix multiplication, and then present the designs of both 3-D parallel matrix-matrix and matrix-vector operations that are used in the Transformer model.

\subsubsection{Load Balancing}
We consider the matrix multiplication $C=AB$ between the input matrix $A$ of size $(M, N)$ and the weight matrix $B$ of size $(N,K)$ on $P=p^3$ processors, which requires $MNK$ calculations in total.
An intuitive way to store the matrices is to partition each dimension of the matrices by $p$, and hold each partition $A_{il}$ on the processor $(i,0,l)$, $B_{lj}$ on $(0,j,l)$, and $C_{ij}$ on $(i,j,0)$ respectively.
In this way, we will need to broadcast each partition $A_{il}$ and $B_{lj}$ across at least $p$ processors before multiplication, and then reduce the partitions $C_{ij}$ after the multiplication, so as to guarantee the consistency among them.
Therefore, an obvious issue is that the storage is imbalanced, so that we waste a great amount of redundant memory in the processor $(i,j,l)$ if $i \neq 0$ or $j \neq 0$ or $l \neq 0$.
In the meantime, the imbalanced storage will result in imbalanced activation or element-wise operations, causing the inefficiency that the computation of such operations is not evenly distributed to all processors.

To eliminate the redundancy caused by the gap between the number of computation and memory dimensions, our insight is to distribute the matrices evenly to the cube.
More specifically, we define $m=M/p^2, n=N/p^2, k=K/p^2$.
Let each processor $(i,j,l)$ hold the partitions $A_{ijl}=A[imp+jm: imp+jm+m-1, lnp: lnp+np-1]$, $B_{lji}=B[lnp: lnp+np-1, jkp+ik: jkp+ik+k-1]$, as shown in \autoref{fig:3d2-a}.
Before the multiplication, instead of the original broadcast operation, we use an all-gather operation along the $y$ direction to copy $A_{il}=A[imp: imp+mp-1,lnp: lnp+np-1]$ across processors $(i, j, l)$ for $0 \le j < p$ (\autoref{fig:3d2-b}).
Similarly, we execute an all-gather for $B_{lj}=B[lnp: np-1, jkp: jkp+kp-1]$ (\autoref{fig:3d2-c}).
Then we compute $A_{il}B_{lj}$ on each processor, and use a reduce-scatter operation in the $z$ direction to get $C_{ij}=\sum_l A_{il}B_{lj}$ across processors $(i, j, l)$ for $0 \le l < p$.
Thus, we will finally get $C_{ilj}=C_{ij}[lm: lm+m-1, :] = C[imp+lm: imp+lm+m-1, jkp: jkp+kp-1]$ in each processor $(i,j,l)$, as shown in \autoref{fig:3d2-d}.
In overall, the memory cost per processor of our approach is supposed to be $(M/p^2)*(N/p)+(N/p)*(K/p^2)+(M/p^2)*(K/p) \approx O(1/p^3)=O(1/P)$.

\begin{figure}[t]
\begin{minipage}{0.4\linewidth}
    \centering
    \includegraphics[width=\linewidth]{figure/transformer.pdf}
    \caption{Architecture of the Transformer model that consists of multiple Transformer layers.}
    \label{fig:transformer}
\end{minipage}
% \end{figure}
% \begin{figure}[t]
\begin{minipage}{0.6\linewidth}
        \centering
        \begin{subfigure}[t]{0.4\linewidth}
            \centering
            \includegraphics[width=\linewidth]{figure/3d2_a.pdf}
            \caption{The input and weight submatrices on each processor.}
            \label{fig:3d2-a}
        \end{subfigure}
        \begin{subfigure}[t]{0.4\linewidth}
            \centering
            \includegraphics[width=\linewidth]{figure/3d2_b.pdf}
            \caption{All-gather $A_{il}$ in the $y$ direction.}
            \label{fig:3d2-b}
        \end{subfigure}
        \begin{subfigure}[t]{0.4\linewidth}
            \centering
            \includegraphics[width=\linewidth]{figure/3d2_c.pdf}
            \caption{All-gather $B_{lj}$ in the $x$ direction.}
            \label{fig:3d2-c}
        \end{subfigure}
        \begin{subfigure}[t]{0.4\linewidth}
            \centering
            \includegraphics[width=\linewidth]{figure/3d2_d.pdf}
            \caption{Compute each $C_{ijl}$ and reduce-scatter them in the $z$ direction.}
            \label{fig:3d2-d}
        \end{subfigure}
        \caption{Load-balanced 3-D matrix-matrix multiplication.}
        \label{fig:3d2}
\end{minipage}
\end{figure}

For vector parameters that are used in matrix-vector operations, we store the vector $b$ of size $N$ diagonally on the B plane, i.e. only the processors $(i,j,l)$ holds $b_{ji} = b[jnp+in: jnp+in+n-1]$ for $j = l, 0 \le i,j,l < p$, as shown by the colored blocks in \autoref{fig:vector}.
To execute a matrix-vector operation, e.g. $A+b$, we need to broadcast $b_{li}$ in the $y$ direction, and then get $b_l$ by all-gathering them in the $x$ direction.
In contrast, to execute $C+b$, we need to broadcast $b_{ji}$ in the $z$ direction before all-gathering them.
Then we can get $A_{il}+b_l$ or $C_{ij}+b_j$ on each processor.

\begin{figure}[h]
    \centering
    \includegraphics[width=0.5\linewidth]{figure/vector.pdf}
    \caption{Balancing vector parameters for matrix-vector operations.}
    \label{fig:vector}
\end{figure}

We need to notice that the communication directions of the input matrix $A$ and output matrix $C$ are different.
To execute the Transformer model, since the output states of each linear layer are basically the input states of the next layer, we exchange the communication directions $y$ and $z$ for the input and output states after the 3-D parallel matrix multiplication, while keep the direction $x$ for the network weights.
The matrix-vector operations are supposed to have no affect on the communication directions.

\subsubsection{Matrix-matrix Operations}
In the Transformer model, three forms of matrix-matrix multiplication, including $C=AB$, $C=AB^T$ and $C=A^TB$, are involved w.r.t the following differentiation formulas.
\begin{align}
    C=AB, \dot{A}=\dot{C}B^T, \dot{B}=A^T\dot{C}; \\
    C=AB^T, \dot{A}=\dot{C}B, \dot{B}=\dot{C}^TA; \\
    C=A^TB, \dot{A}=B\dot{C}^T, \dot{B}=A\dot{C},
\end{align}
where $\dot{X}$ denotes the gradient of the parameter $X$.
We present the algorithms for the three forms of multiplication with the optimal 3-D execution and communication efficiency as shown in the following pseudo-codes (Algorithm \autoref{alg:parallel_ab_forward}-\autoref{alg:parallel_atb_backward}).
Note that for each multiplication algorithm, other than the input matrices, we also need to specify the directions for them as well as the output matrix.

\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Forward $C=AB$}
        \label{alg:parallel_ab_forward}
        \begin{algorithmic}[1]
            \REQUIRE $A_{ijl}$, $B_{lji}$, \\
                    directions $y,x,z$ for $A,B,C$ 
            \ENSURE $C_{ilj}$
            \STATE All-gather $A_{il}$ in $y$ \\
            \STATE All-gather $B_{lj}$ in $x$ \\
            \STATE $C_{ilj} \gets A_{il}B_{lj}$ \\
            \STATE Reduce-scatter $C_{ilj}$ in $z$ \\
            \RETURN $C_{ilj}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}
\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Backward $C=AB$}
        \label{alg:parallel_ab_backward}
        \begin{algorithmic}[1]
            \REQUIRE $\dot{C}_{ilj}$, $A_{ijl}$, $B_{lji}$\\
                    directions $y,x,z$ for $A,B,C$
            \ENSURE $\dot{A}_{ijl},\dot{B}_{lji}$ \\ [0.3\baselineskip]
            \STATE $\dot{A}_{ijl} \gets \dot{C}_{ilj}B_{lji}^T$ in $z,x,y$ \\ [0.3\baselineskip]
            \STATE $\dot{B}_{lji} \gets A_{ijl}^T\dot{C}_{ilj}$ in $y,z,x$\\ [0.3\baselineskip]
            \RETURN $\dot{A}_{ijl},\dot{B}_{lji}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}

\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Forward $C=AB^T$}
        \label{alg:parallel_abt_forward}
        \begin{algorithmic}[1]
            \REQUIRE $A_{ijl}$, $B_{jli}$, \\
                    directions $y,x,z$ for $A,B,C$
            \ENSURE $C_{ilj}$
            \STATE All-gather $A_{il}$ in $y$ \\
            \STATE All-gather $B_{jl}$ in $x$ \\
            \STATE $C_{ilj} \gets A_{il}B_{jl}^T$ \\
            \STATE Reduce-scatter $C_{ilj}$ in $z$\\
            \RETURN $C_{ilj}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}
\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Backward $C=AB^T$}
        \label{alg:parallel_abt_backward}
        \begin{algorithmic}[1]
            \REQUIRE $\dot{C}_{ilj}$, $A_{ijl}$, $B_{jli}$\\
                    directions $y,x,z$ for $A,B,C$
            \ENSURE $\dot{A}_{ijl},\dot{B}_{jli}$ \\ [0.4\baselineskip]
            \STATE $\dot{A}_{ijl} \gets \dot{C}_{ilj}B_{jli}$ in $z,x,y$ \\ [0.4\baselineskip]
            \STATE $\dot{B}_{jli} \gets \dot{C}_{ilj}^TA_{ijl}$ in $z,y,x$\\ [0.4\baselineskip]
            \RETURN $\dot{A}_{ijl},\dot{B}_{jli}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}

\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Forward $C=A^TB$}
        \label{alg:parallel_atb_forward}
        \begin{algorithmic}[1]
            \REQUIRE $A_{ilj}$, $B_{ijl}$, \\
                    directions $y,x,z$ for $A,B,C$
            \ENSURE $C_{jli}$
            \STATE All-gather $A_{ij}$ in $y$ \\
            \STATE All-gather $B_{il}$ in $x$ \\
            \STATE $C_{jl} \gets A_{ij}^TB_{il}$ \\
            \STATE Reduce-scatter $C_{jli}$ in $z$\\
            \RETURN $C_{jli}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}
\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Backward $C=AB^T$}
        \label{alg:parallel_atb_backward}
        \begin{algorithmic}[1]
            \REQUIRE $\dot{C}_{jli}$, $A_{ilj}$, $B_{ijl}$\\
                    directions $y,x,z$ for $A,B,C$
            \ENSURE $\dot{A}_{ilj},\dot{B}_{ijl}$ \\ [0.35\baselineskip]
            \STATE $\dot{A}_{ilj} \gets B_{jli}\dot{C}_{ilj}^T$ in $x,z,y$ \\ [0.35\baselineskip]
            \STATE $\dot{B}_{ijl} \gets A_{ijl}\dot{C}_{ilj}$ in $y,z,x$\\ [0.35\baselineskip]
            \RETURN $\dot{A}_{ilj},\dot{B}_{ilj}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}

Our 3-D parallel matrix multiplication can evenly distribute the computational cost to all the processors, where each processor only multiplies the submatrices of size $(M/p, N/p)$ and $(N/p, K/p)$.
If we fix the problem size, the execution cost of our approach is supposed to be $(M/p)*(N/p)*(K/p) \approx O(1/p^3)=O(1/P)$.
Furthermore, compared with 1-D and 2-D approaches, each communication operation in our 3-D approach moves data across a smaller number of processors.
Each forward algorithm uses all-gather and reduce-scatter operations to move $(MN+NK+MK)/p^3$ across $p$ processors in total.
Therefore, the bandwidth cost of our approach is supposed to be $O(1/p^2)=O(P^{-2/3})$, while the latency cost is supposed to be $O(log(p))=O(log(P^{1/3}))$.

\subsubsection{Matrix-vector Operations}
The logic to execute matrix-vector operations such $C=A+b$ and $C=A*b$ is similar.
We present an example of the add operation as shown in Algorithm \autoref{alg:parallel_add_forward}-\autoref{alg:parallel_add_backward}.

For the matrix-vector multiplication, the only changes from the above algorithms are to return $C_{ijl}=A_{ijl}*b_{l}$ for the forward pass, as well as to return $\dot{A}_{ijl}=\dot{C}_{ijl}*{b}_{ijl}$ and $\dot{b}_{ijl}=\sum_{ij} \dot{C}_{ijl}*{A}_{ijl}$ for the backward pass.

\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Forward $C=A+b$}
        \label{alg:parallel_add_forward}
        \begin{algorithmic}[1]
            \REQUIRE $A_{ijl}$, $b_{ijl}$, \\
                    directions $y,x,z$ for $A,B,C$ \\ [\baselineskip]
            \ENSURE $C_{ijl}$ \\ [\baselineskip]
            \STATE Broadcast $b_{il}$ in $y$ \\ [\baselineskip]
            \STATE All-gather $b_{l}$ in $x$ \\ [\baselineskip]
            \STATE $C_{ijl} \gets A_{ijl}+b_{l}$ \\ [\baselineskip]
            \RETURN $C_{ijl}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}
\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Backward $C=A+b$}
        \label{alg:parallel_add_backward}
        \begin{algorithmic}[1]
            \REQUIRE $\dot{C}_{ijl}$,\\
                    directions $y,x,z$ for $A,B,C$
            \ENSURE $\dot{A}_{ijl},\dot{b}_{ijl}$ \\
            \STATE $\dot{A}_{ijl} \gets \dot{C}_{ijl}$ \\
            \STATE $\dot{b}_{ijl} \gets \sum_{ij} \dot{C}_{ijl}$
            \IF{$j=l$}
                \STATE Reduce-scatter $\dot{b}_{ijl}$ in $x$\\
            \ELSE
                \STATE $\dot{b}_{ijl} \gets$ null
            \ENDIF
            \RETURN $\dot{A}_{ijl},\dot{b}_{ijl}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}

By taking advantage of balanced parameter storage, we can evenly distribute the computation cost of not only the matrix-vector operations but also the activations to all the processors.
Their computation cost is supposed to be the same scale as the memory cost $O(1/P)$.


\subsection{Parallel Transformer layers}

\begin{figure}[h]
    \centering
    \begin{subfigure}[t]{\linewidth}
        \centering
        \includegraphics[width=\linewidth]{figure/attention.pdf}
        \caption{The Self-Attention block}
        \label{fig:attention}
    \end{subfigure}
    \begin{subfigure}[t]{0.7\linewidth}
        \centering
        \includegraphics[width=\linewidth]{figure/mlp.pdf}
        \caption{The Multi-Layer Perception block}
        \label{fig:mlp}
    \end{subfigure}
    \caption{Outline of 3-D Parallel Transformer layers. Each $m,n,k$ in $[m,n,k]$ represents the size of the parameter stored on each processor. The arrows present the order of computation.}
    \label{fig:layer}
\end{figure}

In this section, we describe the major implementation details of the Transformer model that uses our 3-D parallel operations.

To leverage 3-D model parallelism, we implement our new linear layer module and layer normalization module based on the 3-D parallel operations.
The 3-D linear layer flattens an input $X$ of size $[b/p,s/p,h/p]$ into size $[b*s/p^2, h/p]$ first, and then uses Algorithm \autoref{alg:parallel_ab_forward} and \autoref{alg:parallel_add_forward} to perform $Y=XW+b$ with parameters $W$ and $b$.
Thus, after a 3-D linear layer, the input and output directions should be exchanged.
The 3-D layer normalization will not affect the directions, because it only applies matrix-vector adds and multiplications with the parameters $\gamma$ and $\beta$.

We define $b,s,n,h$ as the batch size, sequence length, number of attention heads and hidden size of the Transformer model on a $p^3$ processing cube.
For ease of presentation, we do not discuss the embedding and output layers of the language model, as they usually depend on the downstream tasks.
Then our implementation can basically divided into two modules: Self-Attention blocks and MLP blocks, as illustrated in \autoref{fig:attention} and \autoref{fig:mlp}.
Both the Self-Attention and MLP block include two 3-D linear layers, so that communication directions of their inputs and outputs are supposed to be the same.
Thus, the input and output of each Transformer layer are also supposed to have the same communication direction.

Besides, in practice, it is therefore not necessary to track all the directions of 3-D parallel operations as shown in Algorithm \autoref{alg:parallel_ab_forward}-\autoref{alg:parallel_add_backward}, because we only need to exchange the input and output direction after the first linear layer of both Self-Attention and MLP blocks.
For generalizability, we index the initial input, weight and output group as 0, 1, 2.
We provide an input group index $y$ for the first linear layer, so that the output group index is supposed to be $z=1-y$.
After the layer, we take $z$ as the new input group index, and return it to $y=1-z$ after each Self-Attention or MLP block is completed.

\section{Experiments}
\label{sec:eval}

\subsection{Setup}
We evaluate our 3-D parallel Transformer model on TACC's Longhorn supercomputers \cite{longhorn}.
Our testbed cluster consists of up to 16 GPU servers connected by Mellanox EDR Infiniband network.
Each server has 2 20-core IBM Power 9 CPUs, 256GB memory and 4 NVIDIA V100 GPUs, so there are 64 GPUs in total.
We store the training data, model checkpoints in a Hadoop Distributed File System (HDFS) via 1Gbps Ethernet.

We compare the performance of our approach against the existing state-of-the-art 1-D \cite{megatron} and 2-D \cite{optimus} model parallel approaches.
For simplicity, we only evaluate the performance of the Transformer model itself, i.e. the consecutive Transformer layers, to highlight the efficiency of different parallelism.
We evaluate the 1-D parallelism on 8,16,36 and 64 GPUs, the 2-D parallelism on a subset of 16,36 and 64 GPUs, as well as our 3-D approach on a subset of 8 and 64 GPUs.

\subsection{Performance}

\subsubsection{Weak Scaling}

We first evaluate the weak scaling performance, where we fix the network size on each processor, and scale up the number of processors.
\autoref{tab:weak} presents the results of the weak scaling experiments, including the execution time for both forward and backward pass of the Transformer model.
We also provide the average step time of each sequence, which is computed as 
\begin{align}
    \text{Average step time} = \frac{\text{forward time} + \text{backward time}}{\text{batch size}}.
\end{align}

\begin{table}[h]
    \begin{subtable}[h]{\linewidth}
        \centering
        \begin{tabular}{c|c|c|c|c|c|c}
            \toprule
            \multirow{6}{*}{1-D \cite{megatron}} & \multirow{2}{*}{\# GPUs} & \multirow{2}{*}{\makecell[c]{Batch\\size}} & \multirow{2}{*}{\makecell[c]{Hidden\\size}} & \multirow{2}{*}{\makecell[c]{Forward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Backward \\ time (s)}} & \multirow{2}{*}{\makecell[c]{Average\\ step time (s)}}\\
            & & & & & & \\
            \cline{2-7}
            & 8 & 60 & 2048 & 4.759 & 15.676 & 0.341 \\
            \cline{2-7}
            & 16 & 60 & 4096 & 12.488 & 30.894 & 0.723 \\
            \cline{2-7}
            & 36 & 40 & 6120 & 13.515 & 31.822 & 1.133 \\
            \cline{2-7}
            & 64 & 30 & 8192 & 13.915 & 32.890 & 1.560 \\
            \bottomrule
        \end{tabular}
        % \caption{1-D Parallelism \cite{megatron}}
        \label{tab:weak-1d}
    \end{subtable}
    
    \begin{subtable}[h]{\linewidth}
        \centering
        \begin{tabular}{c|c|c|c|c|c|c}
            \toprule
            \multirow{5}{*}{2-D \cite{optimus}} & \multirow{2}{*}{\# GPUs} & \multirow{2}{*}{\makecell[c]{Batch\\size}} & \multirow{2}{*}{\makecell[c]{Hidden\\size}} & \multirow{2}{*}{\makecell[c]{Forward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Backward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Average\\step time (s)}}\\
            & & & & & & \\
            \cline{2-7}
            & 16 & 192 & 4096 & 33.860 & 101.981 & 0.708 \\
            \cline{2-7}
            & 36 & 288 & 6120 & 54.760 & 165.850 & 0.766 \\
            \cline{2-7}
            & 64 & 384 & 8192 & 99.419 & 304.707 & 1.052 \\
            \bottomrule
        \end{tabular}
        % \caption{2-D Parallelism \cite{optimus}}
        \label{tab:weak-2d}
    \end{subtable}
    \begin{subtable}[h]{\linewidth}
        \centering
        \begin{tabular}{c|c|c|c|c|c|c}
            \toprule
            \multirow{4}{*}{~~~ 3-D ~~~} & \multirow{2}{*}{\# GPUs} & \multirow{2}{*}{\makecell[c]{Batch\\size}} & \multirow{2}{*}{\makecell[c]{Hidden\\size}}  &  \multirow{2}{*}{\makecell[c]{Forward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Backward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Average\\step time (s)}}\\
            & & & & & & \\
            \cline{2-7}
            & 8 & 192 & 2048 & 30.096 & 81.212 & 0.580 \\
            \cline{2-7}
            & 64 & 384 & 8192 & 79.349 & 125.037 & \textbf{0.672} \\
            \bottomrule
        \end{tabular}
        % \caption{3-D Parallelism}
        \label{tab:weak-3d}
    \end{subtable}
    \caption{Comparison of weak scaling results. The number of processors increases from 8 to 64, while the number of parameters is fixed for each approach. For each of presentation, we mainly adjust the batch size and hidden size based on the number of processors, while fix the sequence length to 512. The bolded number represents the best result.}
    \label{tab:weak}
\end{table}

We see that the 3-D parallelism has the slowest rising speed in the average step time, reaching the smallest value at the largest compute scale.
The result implies that our approach is efficient to reduce the overhead of model parallelism, as it reaches the minimum communication cost.


\subsubsection{Strong Scaling}
Next, we evaluate the strong scaling performance of different model parallelism.
This experiment aims to examine how much speedup each approach can achieve by using an increasing number of processors to execute the fix-sized problem.

The problem size and execution results are shown in \autoref{tab:strong}.
We see that our 3-D parallelism outperforms other approaches by achieving the smallest average step time on 64 GPUs, with 2.32X and 1.57X speedup over 1-D and 2-D approach, respectively.

\begin{table}[h]
    \begin{subtable}[h]{\linewidth}
        \centering
        \begin{tabular}{c|c|c|c|c|c|c}
            \toprule
            \multirow{6}{*}{1-D \cite{megatron}} & \multirow{2}{*}{\# GPUs} & \multirow{2}{*}{\makecell[c]{Batch\\size}} & \multirow{2}{*}{\makecell[c]{Hidden\\size}} & \multirow{2}{*}{\makecell[c]{Forward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Backward \\ time (s)}} & \multirow{2}{*}{\makecell[c]{Average\\ step time (s)}}\\
            & & & & & & \\
            \cline{2-7}
            & 8 & 12 & 3072 & 1.470 & 5.699 & 0.597 \\
            \cline{2-7}
            & 16 & 12 & 3072 & 1.371 & 5.152 & 0.544 \\
            \cline{2-7}
            & 36 & 12 & 3072 & 1.455 & 5.414 & 0.572 \\
            \cline{2-7}
            & 64 & 12 & 3072 & 1.433 & 5.167 & 0.550 \\
            \bottomrule
        \end{tabular}
        % \caption{1-D Parallelism \cite{megatron}}
        \label{tab:strong-1d}
    \end{subtable}
    
    \begin{subtable}[h]{\linewidth}
        \centering
        \begin{tabular}{c|c|c|c|c|c|c}
            \toprule
            \multirow{5}{*}{2-D \cite{optimus}} & \multirow{2}{*}{\# GPUs} & \multirow{2}{*}{\makecell[c]{Batch\\size}} & \multirow{2}{*}{\makecell[c]{Hidden\\size}} & \multirow{2}{*}{\makecell[c]{Forward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Backward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Average\\step time (s)}}\\
            & & & & & & \\
            \cline{2-7}
            & 16 & 24 & 3072 & 4.680 & 13.698 & 0.766 \\
            \cline{2-7}
            & 36 & 24 & 3072 & 3.900 & 11.433 & 0.639 \\
            \cline{2-7}
            & 64 & 24 & 3072 & 3.007 & 8.920 & 0.497 \\
            \bottomrule
        \end{tabular}
        % \caption{2-D Parallelism \cite{optimus}}
        \label{tab:strong-2d}
    \end{subtable}
    \begin{subtable}[h]{\linewidth}
        \centering
        \begin{tabular}{c|c|c|c|c|c|c}
            \toprule
            \multirow{4}{*}{~~~ 3-D ~~~} & \multirow{2}{*}{\# GPUs} & \multirow{2}{*}{\makecell[c]{Batch\\size}} & \multirow{2}{*}{\makecell[c]{Hidden\\size}}  &  \multirow{2}{*}{\makecell[c]{Forward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Backward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Average\\step time (s)}}\\
            & & & & & & \\
            \cline{2-7}
            & 8 & 24 & 3072 & 3.249 & 9.120 & 0.515 \\
            \cline{2-7}
            & 64 & 24 & 3072 & 2.494 & 6.129 & \textbf{0.359} \\
            \bottomrule
        \end{tabular}
        % \caption{3-D Parallelism}
        \label{tab:strong-3d}
    \end{subtable}
    \caption{Comparison of strong scaling results. The problem size is fixed while the number of processors increases from 8 to 64. The bolded number represents the best result.}
    \label{tab:strong}
\end{table}

\section{Implementation}
To do
\section{Submission instructions}

\textcolor{blue}{For checking submission format.}

\subsection{Style}

Papers to be submitted to NeurIPS 2021 must be prepared according to the
instructions presented here. Papers may only be up to {\bf nine} pages long,
including figures. Additional pages \emph{containing only acknowledgments and
references} are allowed. Papers that exceed the page limit will not be
reviewed, or in any other way considered for presentation at the conference.

The margins in 2021 are the same as those in 2007, which allow for $\sim$$15\%$
more words in the paper compared to earlier years.

Authors are required to use the NeurIPS \LaTeX{} style files obtainable at the
NeurIPS website as indicated below. Please make sure you use the current files
and not previous versions. Tweaking the style files may be grounds for
rejection.

\subsection{Retrieval of style files}

The style files for NeurIPS and other conference information are available on
the World Wide Web at
\begin{center}
  \url{http://www.neurips.cc/}
\end{center}
The file \verb+neurips_2021.pdf+ contains these instructions and illustrates the
various formatting requirements your NeurIPS paper must satisfy.

The only supported style file for NeurIPS 2021 is \verb+neurips_2021.sty+,
rewritten for \LaTeXe{}.  \textbf{Previous style files for \LaTeX{} 2.09,
  Microsoft Word, and RTF are no longer supported!}

The \LaTeX{} style file contains three optional arguments: \verb+final+, which
creates a camera-ready copy, \verb+preprint+, which creates a preprint for
submission to, e.g., arXiv, and \verb+nonatbib+, which will not load the
\verb+natbib+ package for you in case of package clash.

\paragraph{Preprint option}
If you wish to post a preprint of your work online, e.g., on arXiv, using the
NeurIPS style, please use the \verb+preprint+ option. This will create a
nonanonymized version of your work with the text ``Preprint. Work in progress.''
in the footer. This version may be distributed as you see fit. Please \textbf{do
  not} use the \verb+final+ option, which should \textbf{only} be used for
papers accepted to NeurIPS.

At submission time, please omit the \verb+final+ and \verb+preprint+
options. This will anonymize your submission and add line numbers to aid
review. Please do \emph{not} refer to these line numbers in your paper as they
will be removed during generation of camera-ready copies.

The file \verb+neurips_2021.tex+ may be used as a ``shell'' for writing your
paper. All you have to do is replace the author, title, abstract, and text of
the paper with your own.

The formatting instructions contained in these style files are summarized in
Sections \ref{gen_inst}, \ref{headings}, and \ref{others} below.

\section{General formatting instructions}
\label{gen_inst}

The text must be confined within a rectangle 5.5~inches (33~picas) wide and
9~inches (54~picas) long. The left margin is 1.5~inch (9~picas).  Use 10~point
type with a vertical spacing (leading) of 11~points.  Times New Roman is the
preferred typeface throughout, and will be selected for you by default.
Paragraphs are separated by \nicefrac{1}{2}~line space (5.5 points), with no
indentation.

The paper title should be 17~point, initial caps/lower case, bold, centered
between two horizontal rules. The top rule should be 4~points thick and the
bottom rule should be 1~point thick. Allow \nicefrac{1}{4}~inch space above and
below the title to rules. All pages should start at 1~inch (6~picas) from the
top of the page.

For the final version, authors' names are set in boldface, and each name is
centered above the corresponding address. The lead author's name is to be listed
first (left-most), and the co-authors' names (if different address) are set to
follow. If there is only one co-author, list both author and co-author side by
side.

Please pay special attention to the instructions in Section \ref{others}
regarding figures, tables, acknowledgments, and references.

\section{Headings: first level}
\label{headings}

All headings should be lower case (except for first word and proper nouns),
flush left, and bold.

First-level headings should be in 12-point type.

\subsection{Headings: second level}

Second-level headings should be in 10-point type.

\subsubsection{Headings: third level}

Third-level headings should be in 10-point type.

\paragraph{Paragraphs}

There is also a \verb+\paragraph+ command available, which sets the heading in
bold, flush left, and inline with the text, with the heading followed by 1\,em
of space.

\section{Citations, figures, tables, references}
\label{others}

These instructions apply to everyone.

\subsection{Citations within the text}

The \verb+natbib+ package will be loaded for you by default.  Citations may be
author/year or numeric, as long as you maintain internal consistency.  As to the
format of the references themselves, any style is acceptable as long as it is
used consistently.

The documentation for \verb+natbib+ may be found at
\begin{center}
  \url{http://mirrors.ctan.org/macros/latex/contrib/natbib/natnotes.pdf}
\end{center}
Of note is the command \verb+\citet+, which produces citations appropriate for
use in inline text.  For example,
\begin{verbatim}
   \citet{hasselmo} investigated\dots
\end{verbatim}
produces
\begin{quote}
  Hasselmo, et al.\ (1995) investigated\dots
\end{quote}

If you wish to load the \verb+natbib+ package with options, you may add the
following before loading the \verb+neurips_2021+ package:
\begin{verbatim}
   \PassOptionsToPackage{options}{natbib}
\end{verbatim}

If \verb+natbib+ clashes with another package you load, you can add the optional
argument \verb+nonatbib+ when loading the style file:
\begin{verbatim}
   \usepackage[nonatbib]{neurips_2021}
\end{verbatim}

As submission is double blind, refer to your own published work in the third
person. That is, use ``In the previous work of Jones et al.\ [4],'' not ``In our
previous work [4].'' If you cite your other papers that are not widely available
(e.g., a journal paper under review), use anonymous author names in the
citation, e.g., an author of the form ``A.\ Anonymous.''

\subsection{Footnotes}

Footnotes should be used sparingly.  If you do require a footnote, indicate
footnotes with a number\footnote{Sample of the first footnote.} in the
text. Place the footnotes at the bottom of the page on which they appear.
Precede the footnote with a horizontal rule of 2~inches (12~picas).

Note that footnotes are properly typeset \emph{after} punctuation
marks.\footnote{As in this example.}

\subsection{Figures}

\begin{figure}
  \centering
  \fbox{\rule[-.5cm]{0cm}{4cm} \rule[-.5cm]{4cm}{0cm}}
  \caption{Sample figure caption.}
\end{figure}

All artwork must be neat, clean, and legible. Lines should be dark enough for
purposes of reproduction. The figure number and caption always appear after the
figure. Place one line space before the figure caption and one line space after
the figure. The figure caption should be lower case (except for first word and
proper nouns); figures are numbered consecutively.

You may use color figures.  However, it is best for the figure captions and the
paper body to be legible if the paper is printed in either black/white or in
color.

\subsection{Tables}

All tables must be centered, neat, clean and legible.  The table number and
title always appear before the table.  See Table~\ref{sample-table}.

Place one line space before the table title, one line space after the
table title, and one line space after the table. The table title must
be lower case (except for first word and proper nouns); tables are
numbered consecutively.

Note that publication-quality tables \emph{do not contain vertical rules.} We
strongly suggest the use of the \verb+booktabs+ package, which allows for
typesetting high-quality, professional tables:
\begin{center}
  \url{https://www.ctan.org/pkg/booktabs}
\end{center}
This package was used to typeset Table~\ref{sample-table}.

\begin{table}
  \caption{Sample table title}
  \label{sample-table}
  \centering
  \begin{tabular}{lll}
    \toprule
    \multicolumn{2}{c}{Part}                   \\
    \cmidrule(r){1-2}
    Name     & Description     & Size ($\mu$m) \\
    \midrule
    Dendrite & Input terminal  & $\sim$100     \\
    Axon     & Output terminal & $\sim$10      \\
    Soma     & Cell body       & up to $10^6$  \\
    \bottomrule
  \end{tabular}
\end{table}

\section{Final instructions}

Do not change any aspects of the formatting parameters in the style files.  In
particular, do not modify the width or length of the rectangle the text should
fit into, and do not change font sizes (except perhaps in the
\textbf{References} section; see below). Please note that pages should be
numbered.

\section{Preparing PDF files}

Please prepare submission files with paper size ``US Letter,'' and not, for
example, ``A4.''

Fonts were the main cause of problems in the past years. Your PDF file must only
contain Type 1 or Embedded TrueType fonts. Here are a few instructions to
achieve this.

\begin{itemize}

\item You should directly generate PDF files using \verb+pdflatex+.

\item You can check which fonts a PDF files uses.  In Acrobat Reader, select the
  menu Files$>$Document Properties$>$Fonts and select Show All Fonts. You can
  also use the program \verb+pdffonts+ which comes with \verb+xpdf+ and is
  available out-of-the-box on most Linux machines.

\item The IEEE has recommendations for generating PDF files whose fonts are also
  acceptable for NeurIPS. Please see
  \url{http://www.emfield.org/icuwb2010/downloads/IEEE-PDF-SpecV32.pdf}

\item \verb+xfig+ "patterned" shapes are implemented with bitmap fonts.  Use
  "solid" shapes instead.

\item The \verb+\bbold+ package almost always uses bitmap fonts.  You should use
  the equivalent AMS Fonts:
\begin{verbatim}
   \usepackage{amsfonts}
\end{verbatim}
followed by, e.g., \verb+\mathbb{R}+, \verb+\mathbb{N}+, or \verb+\mathbb{C}+
for $\mathbb{R}$, $\mathbb{N}$ or $\mathbb{C}$.  You can also use the following
workaround for reals, natural and complex:
\begin{verbatim}
   \newcommand{\RR}{I\!\!R} %real numbers
   \newcommand{\Nat}{I\!\!N} %natural numbers
   \newcommand{\CC}{I\!\!\!\!C} %complex numbers
\end{verbatim}
Note that \verb+amsfonts+ is automatically loaded by the \verb+amssymb+ package.

\end{itemize}

If your file contains type 3 fonts or non embedded TrueType fonts, we will ask
you to fix it.

\subsection{Margins in \LaTeX{}}

Most of the margin problems come from figures positioned by hand using
\verb+\special+ or other commands. We suggest using the command
\verb+\includegraphics+ from the \verb+graphicx+ package. Always specify the
figure width as a multiple of the line width as in the example below:
\begin{verbatim}
   \usepackage[pdftex]{graphicx} ...
   \includegraphics[width=0.8\linewidth]{myfile.pdf}
\end{verbatim}
See Section 4.4 in the graphics bundle documentation
(\url{http://mirrors.ctan.org/macros/latex/required/graphics/grfguide.pdf})

A number of width problems arise when \LaTeX{} cannot properly hyphenate a
line. Please give LaTeX hyphenation hints using the \verb+\-+ command when
necessary.

\begin{ack}
Use unnumbered first level headings for the acknowledgments. All acknowledgments
go at the end of the paper before the list of references. Moreover, you are required to declare
funding (financial activities supporting the submitted work) and competing interests (related financial activities outside the submitted work).
More information about this disclosure can be found at: \url{https://neurips.cc/Conferences/2021/PaperInformation/FundingDisclosure}.

Do {\bf not} include this section in the anonymized submission, only in the final paper. You can use the \texttt{ack} environment provided in the style file to autmoatically hide this section in the anonymized submission.
\end{ack}

\section*{References}

References follow the acknowledgments. Use unnumbered first-level heading for
the references. Any choice of citation style is acceptable as long as you are
consistent. It is permissible to reduce the font size to \verb+small+ (9 point)
when listing the references.
Note that the Reference section does not count towards the page limit.
\medskip

{
\small

[1] Alexander, J.A.\ \& Mozer, M.C.\ (1995) Template-based algorithms for
connectionist rule extraction. In G.\ Tesauro, D.S.\ Touretzky and T.K.\ Leen
(eds.), {\it Advances in Neural Information Processing Systems 7},
pp.\ 609--616. Cambridge, MA: MIT Press.

[2] Bower, J.M.\ \& Beeman, D.\ (1995) {\it The Book of GENESIS: Exploring
  Realistic Neural Models with the GEneral NEural SImulation System.}  New York:
TELOS/Springer--Verlag.

[3] Hasselmo, M.E., Schnell, E.\ \& Barkai, E.\ (1995) Dynamics of learning and
recall at excitatory recurrent synapses and cholinergic modulation in rat
hippocampal region CA3. {\it Journal of Neuroscience} {\bf 15}(7):5249-5262.
}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section*{Checklist}

%%% BEGIN INSTRUCTIONS %%%
The checklist follows the references.  Please
read the checklist guidelines carefully for information on how to answer these
questions.  For each question, change the default \answerTODO{} to \answerYes{},
\answerNo{}, or \answerNA{}.  You are strongly encouraged to include a {\bf
justification to your answer}, either by referencing the appropriate section of
your paper or providing a brief inline description.  For example:
\begin{itemize}
  \item Did you include the license to the code and datasets? \answerYes{See Section~\ref{gen_inst}.}
  \item Did you include the license to the code and datasets? \answerNo{The code and the data are proprietary.}
  \item Did you include the license to the code and datasets? \answerNA{}
\end{itemize}
Please do not modify the questions and only use the provided macros for your
answers.  Note that the Checklist section does not count towards the page
limit.  In your paper, please delete this instructions block and only keep the
Checklist section heading above along with the questions/answers below.
%%% END INSTRUCTIONS %%%

\begin{enumerate}

\item For all authors...
\begin{enumerate}
  \item Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
    \answerTODO{}
  \item Did you describe the limitations of your work?
    \answerTODO{}
  \item Did you discuss any potential negative societal impacts of your work?
    \answerTODO{}
  \item Have you read the ethics review guidelines and ensured that your paper conforms to them?
    \answerTODO{}
\end{enumerate}

\item If you are including theoretical results...
\begin{enumerate}
  \item Did you state the full set of assumptions of all theoretical results?
    \answerTODO{}
	\item Did you include complete proofs of all theoretical results?
    \answerTODO{}
\end{enumerate}

\item If you ran experiments...
\begin{enumerate}
  \item Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
    \answerTODO{}
  \item Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
    \answerTODO{}
	\item Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
    \answerTODO{}
	\item Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
    \answerTODO{}
\end{enumerate}

\item If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
\begin{enumerate}
  \item If your work uses existing assets, did you cite the creators?
    \answerTODO{}
  \item Did you mention the license of the assets?
    \answerTODO{}
  \item Did you include any new assets either in the supplemental material or as a URL?
    \answerTODO{}
  \item Did you discuss whether and how consent was obtained from people whose data you're using/curating?
    \answerTODO{}
  \item Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
    \answerTODO{}
\end{enumerate}

\item If you used crowdsourcing or conducted research with human subjects...
\begin{enumerate}
  \item Did you include the full text of instructions given to participants and screenshots, if applicable?
    \answerTODO{}
  \item Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
    \answerTODO{}
  \item Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
    \answerTODO{}
\end{enumerate}

\end{enumerate}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\section{Introduction}
% In 1989, Yann LeCun proposed to train convolutional neural networks (CNNs) using back-propagation, which makes learning fully automatic, instead of manually designing convolutional kernels. However, the current prevailance of deep learning started decades after it, when AlexNet took the crown in ImageNet Large Scale Visual Recognition Challenge 2012. It owed to the development of computation capability. Since then, 
To deal with the growing amount of data, nowadays people are squeezing the computation from GPUs or FPGAs to get larger and more powerful neural networks. 
Consequently, records are being broken in a variety of Natural Language Processing (NLP) tasks ranging from classification, question answering to translation as the language models are becoming much deeper.
In the meantime, starting from the emergence of Transformer models \cite{transformer} like BERT \cite{BERT}, the trend to increase model sizes becomes even more radical.
Since then, the Moore's Law can never reach the growth of language models.
Past several years have witnessed the creation of gigantic models, like GPT-2 \cite{gpt2} and GPT-3 \cite{gpt-3}, which are able to compose texts and synthesize images.
The amazing performance is also attracting attention from industry.
% These models are now all around in our daily life, integrated in search engines, recommendation and broadcasting systems.
% They facilitate our daily life whenever we search an item in Google, browse Twitter, or shop in Amazon.
Google recently announced its Multitask Unified Model (MUM) for language understanding, LaMDA for conversation, and TPUv4 chips for model training.
% All these suggest the great demand on large language models which would last for years.

% The gap between the gradually slowed-down Moore's Law and the drastically increased model size and computational demand mainly reflects in two aspects: memory and computation.
The surging model size brings challenges in terms of both memory and computation.
Apart from developing chips such as GPU and TPU with higher computation capability and memory limit, activation checkpointing techniques are developed to squeeze a large model into a single device that would otherwise not be able to host such a big model. 
Moreover, ZeRO-infinity techniques \cite{zero,zero,rajbhandari2021zero} are proposed to utilize the memory of CPU or other memory device external to GPU without sacrificing the data transmission time.
% To fully utilize the concurrency of HPC clusters,
Data parallelism is a dominant practice to fully utilize available HPC resources.
It distributes a large minibatch to multiple devices, where each device holds an identical model replica, and finally gathers the gradients for synchronous parameter update.
With recent optimization techniques \cite{goyal2017accurate,you2017large,you2018imagenet,you2020large}, it is now able to train very large minibatches on thousands of GPU devices.
% However, studies show that large batch size would result in unstable training and degraded generalization performance.
% To address this problem, You et al propose novel optimizers facilitating the large batch training of visual and language models.


Another branch is model parallelism, the main idea of which is to divide both computation and memory of a single neural network to multiple devices. 
Pipelined parallelism \cite{pipedream,gpipe} is proposed to split the whole model by layer, which will be executed in a pipeline fashion.
Mixed precision training is utilized to reduce both memory and computation cost.
% Activation checkpointing trades computation for memory, which makes the study of huge neural networks accessible to more researchers.
Megatron-LM \cite{megatron} integrates both model parallelism and data parallelism by splitting the parameter tensor among model parallel groups, achieving the state-of-the-art language model.
Recently, a 2-dimensional matrix-matrix multiplication algorithm has been introduced into language model training \cite{optimus}, thus bridging the gap between traditional high-performance computing and machine learning.
In this work, we stride a step forward and propose a 3-D model parallelism technique to further harness the capability of GPU clusters.
We list our main contributions as follows.
\begin{itemize}
    \item We propose a 3-D model parallelism algorithm for linear operations, which can reach a perfect load balance, so as to provide the optimal efficiency;
    \item We use our 3-D model parallelism to implement an efficient Transformer model based on the existing PyTorch Transformer implementation and distributed communication package;
    \item We demonstrate the effectiveness of our 3-D model parallelism on 64 GPUs by comparing it with the existing 1-D and 2-D approaches. The results present the superiority of our 3-D parallelism with 2.32X and 1.57X speedup, respectively.
\end{itemize}
% Pipeline parallelism (PipeDream and GPipe) aims to partition a model by layer (vertically), and execute the training like a pipeline. The benefit is that, communication can be largely overlapped by computation. However, both the methods suffer from their own weaknesses. Although PipeDream elaborately arranges the forward and backward computations so that it fully utilizes all time bins without idleness, it is essentiallyan asynchronous approach, so suffers from gradient staleness. GPipe is a synchronous approach, but "bubbles" is inevitable. For extremely large networks, more GPUs should be used to host the model, and the idleness becomes more severe.

% Unlike model parallelism, ZeRO proposes to having each device host its own partition of parameters and optimizer states, which are updated locally. Parameters are gathered whenever necessary during forward and backward propagation.

% Another trend of model parallelism is to partition a model by operation (horizontally).
% Mesh-Tensorflow is a framework proposed by Google, which provides an elegant way for users to annotate the partition of tensor dimensions. Matrix multiplications can be automatically derived from these annotations, without user elaboration. However, memory redundancy is severe, and the originally nearly-free reshape operations would result in complex communication because of layout change. Megatron-LM adopts a 1-dimensional matrix partition strategy, and alternatively splits parameter tensors along rows or columns. An elegant conjunction is given to address the gradient propagation by simply wrapping the all-reduce operation. Although parameters and optimizer states are fully distributed, there is still memory redundancy for activations, as each device has to accommodate the whole input and output activations for every multi-layer perceptrons (MLPs) and self-attention modules. Optimus takes the first step addressing this issue, utilizing a 2D matrix-matrix multiplication strategy. Here, our method takes a step further, and uses a 3D approach, while also eliminating memory redundancy.

% Besides fitting existing models into parallel and distributed versions, other approaches adapt models to better capture the concurrency of modern HPC clusters. Mixture-of-experts (MoE) utilizes sparsity to discretely send each hidden state to distinct dense layers, or "experts". In this way, the amount of parameters could be easily scaled without significantly adding to computation burden. The rationale behind behind this is that more parameters result in higher performance. A lot of works achieve excellent performance based on this method.
\documentclass{article}

% if you need to pass options to natbib, use, e.g.:
%     \PassOptionsToPackage{numbers, compress}{natbib}
% before loading neurips_2021

% ready for submission
\PassOptionsToPackage{square, numbers}{natbib}
\usepackage[preprint]{neurips_2021}

% to compile a preprint version, e.g., for submission to arXiv, add add the
% [preprint] option:
%     \usepackage[preprint]{neurips_2021}

% to compile a camera-ready version, add the [final] option, e.g.:
%     \usepackage[final]{neurips_2021}

% to avoid loading the natbib package, add option nonatbib:
%    \usepackage[nonatbib]{neurips_2021}

\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc}    % use 8-bit T1 fonts
\usepackage{hyperref}       % hyperlinks
\usepackage{url}            % simple URL typesetting
\usepackage{booktabs}       % professional-quality tables
\usepackage{amsfonts}       % blackboard math symbols
\usepackage{nicefrac}       % compact symbols for 1/2, etc.
\usepackage{microtype}      % microtypography
\usepackage{xcolor}         % colors

\usepackage{caption}
\usepackage{subcaption}
\usepackage{graphicx}
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{bm}
\usepackage{subfiles}
\usepackage{multirow}
\usepackage{multicol}
\usepackage{makecell}
\usepackage{tabularx}
\usepackage{footmisc}
\usepackage{listings}
\usepackage{algorithm}
\usepackage{algorithmic}
\renewcommand{\algorithmicrequire}{\textbf{Input:}}
\renewcommand{\algorithmicensure}{\textbf{Output:}}

\title{Maximizing Parallelism in Distributed Training for Huge Neural Networks}

% The \author macro works with any number of authors. There are two commands
% used to separate the names and addresses of multiple authors: \And and \AND.
%
% Using \And between authors leaves it to LaTeX to determine where to break the
% lines. Using \AND forces a line break at that point. So, if LaTeX puts 3 of 4
% authors names on the first line, and the last on the second line, try using
% \AND instead of \And before the third author name.

\author{%
%   David S.~Hippocampus\thanks{Use footnote for providing further information
%     about author (webpage, alternative address)---\emph{not} for acknowledging
%     funding agencies.} \\
%   Department of Computer Science\\
%   Cranberry-Lemon University\\
%   Pittsburgh, PA 15213 \\
%   \texttt{hippo@cs.cranberry-lemon.edu} \\
    Zhengda Bian\thanks{These authors have equal contributions.}\\
    National University of Singapore\\
    Singapore\\
    \texttt{zbian@comp.nus.edu.sg}\\
    \And
    Qifan Xu\footnotemark[\value{footnote}]\\
    University of California, Los Angeles\\
    Los Angeles, California, USA\\
    \texttt{QifanXu@mednet.ucla.edu}\\
    \AND
    Boxiang Wang\\
    National University of Singapore\\
    Singapore\\
    \texttt{boxiangwang059@gmail.com}\\
    \And
    Yang You\\
    National University of Singapore\\
    Singapore\\
    \texttt{youy@comp.nus.edu.sg}\\
  % examples of more authors
  % \And
  % Coauthor \\
  % Affiliation \\
  % Address \\
  % \texttt{email} \\
  % \AND
  % Coauthor \\
  % Affiliation \\
  % Address \\
  % \texttt{email} \\
  % \And
  % Coauthor \\
  % Affiliation \\
  % Address \\
  % \texttt{email} \\
  % \And
  % Coauthor \\
  % Affiliation \\
  % Address \\
  % \texttt{email} \\
}


\begin{document}

\maketitle

\begin{abstract}
%   The abstract paragraph should be indented \nicefrac{1}{2}~inch (3~picas) on
%   both the left- and right-hand margins. Use 10~point type, with a vertical
%   spacing (leading) of 11~points.  The word \textbf{Abstract} must be centered,
%   bold, and in point size 12. Two line spaces precede the abstract. The abstract
%   must be limited to one paragraph.
% As user data explodes, companies have to frequently update their models to capture the preference of users and create profiles for newly generated items to provide up-to-date service for searching, recommending and broadcasting.
% The magic power of huge language models has been brought to us by academia, so that 
% nowadays people have to think twice if the pictures and texts in front of them are taken from camera or written by human.
% So it’s no wonder that 
% Hinton declared that life, universe and everything were just 4.398 trillion parameters short after GPT-3 was published.
The recent Natural Language Processing techniques have been refreshing the state-of-the-art performance at an incredible speed.
Training huge language models is therefore an imperative demand in both industry and academy.
However, the huge models impose challenges to both hardware and software.
Graphical processing units (GPUs) are iterated frequently to meet the exploding demand, and a variety of ASICs like TPUs are spawned.
However, there is still a tension between the fast growth of the extremely huge models and fact that Moore’s law is approaching the end.
To this end, many model parallelism techniques are proposed to distribute the model parameters to multiple devices, so as to alleviate the tension on both memory and computation.
Our work is the first to introduce a 3-dimensional model parallelism for expediting huge language models.
By reaching a perfect load balance, our approach presents smaller memory and communication cost than existing state-of-the-art 1-D and 2-D model parallelism.
Our experiments on 64 TACC's V100 GPUs show that our 3-D parallelism outperforms the 1-D and 2-D parallelism with 2.32X and 1.57X speedup, respectively.

\end{abstract}

\subfile{section/introduction.tex}

\subfile{section/background.tex}

\subfile{section/design.tex}

% \subfile{section/implementation.tex}

\subfile{section/evaluation.tex}

\subfile{section/conclusion.tex}


\bibliographystyle{plainnat}
\bibliography{reference}

% \appendix

% \section{Appendix}

% Optionally include extra information (complete proofs, additional experiments and plots) in the appendix.
% This section will often be part of the supplemental material.

% \subfile{section/instruction.tex}
\iffalse

\section*{Checklist}

\begin{enumerate}

\item For all authors...
\begin{enumerate}
  \item Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
    \answerYes{}
  \item Did you describe the limitations of your work?
    \answerYes{See Section~\ref{sec:conclusion}}
  \item Did you discuss any potential negative societal impacts of your work?
    \answerNA{}
  \item Have you read the ethics review guidelines and ensured that your paper conforms to them?
    \answerYes{}
\end{enumerate}

\item If you are including theoretical results...
\begin{enumerate}
  \item Did you state the full set of assumptions of all theoretical results?
    \answerYes{See Section~\ref{sec:design}}
	\item Did you include complete proofs of all theoretical results?
    \answerYes{See Section~\ref{sec:design}}
\end{enumerate}

\item If you ran experiments...
\begin{enumerate}
  \item Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
    \answerNo{No, currently our codes are not publicly available.}
  \item Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
    \answerYes{See Section~\ref{sec:eval}.}
	\item Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
    \answerNo{No, I will work on this later due to the time limit.}
	\item Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
    \answerYes{See Section~\ref{sec:eval}.}
\end{enumerate}

\item If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
\begin{enumerate}
  \item If your work uses existing assets, did you cite the creators?
    \answerYes{See Section~\ref{sec:background}.}
  \item Did you mention the license of the assets?
    \answerNA{}
  \item Did you include any new assets either in the supplemental material or as a URL?
    \answerNA{}
  \item Did you discuss whether and how consent was obtained from people whose data you're using/curating?
    \answerNA{}
  \item Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
    \answerNA{}
\end{enumerate}

\item If you used crowdsourcing or conducted research with human subjects...
\begin{enumerate}
  \item Did you include the full text of instructions given to participants and screenshots, if applicable?
    \answerNA{}
  \item Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
    \answerNA{}
  \item Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
    \answerNA{}
\end{enumerate}

\end{enumerate}

\fi

\end{document}
\section{Background}
\label{sec:background}

\subsection{Transformer Language Models}
There has been a recent trend to use pretrained language models to help people train specific natural language processing (NLP) tasks, as pretrained models leverage the understanding of very large corpus (e.g. Wikipedia) to alleviate the efforts of training entire language models.
Transformer is one of the current dominant choices of pretraining language models, achieving the state-of-the-art compute efficiency and prediction accuracy.
The Transformer language model uses deep attention modules to transform a given sequential input into another sequential output.
Like LSTM, a complete Transformer model includes two parts, an Encoder and a Decoder, but each consists of attention layers rather than recurrent layers.
However, recent state-of-the-art Transformer applications, such as BERT \cite{BERT}, GPT-2 \cite{gpt2}, and GPT-3 \cite{gpt-3}, adopt more concise structures that only use necessary Encoders or Decoders on demand.

Nevertheless, training Transformer language models can be extremely computational expensive.
The recent language models advance their accuracy with a rising number of parameters (e.g. OpenAI GPT: 110M \cite{gpt1}; BERT: 340M \cite{BERT}; GPT-2: 1.5B \cite{gpt2}; GPT-3: 175B \cite{gpt-3}; Switch Transformers: 1.6T \cite{switch-tranformer}).
In the meantime, to deal with huge language corpus, there has been a growing interest in using large minibatches to reduce training time.
Therefore, efficient systems to expedite Transformers become an urgent demand for training large language models.

\subsection{Data and Model Parallelism}
Data parallelism is the most common paradigm to parallel the computation of deep neural networks.
This approach distributes the entire minibatch across multiple workers, which execute a single replica of the model and communicate with each other for synchronization at the end of each training step.
It is easy to use more devices to train larger minibatches.
By taking advantage of recent optimization techniques \cite{smith2017don,you2017large,you2018imagenet,you2020large,bytescheduler}, training speed reaches almost linear scaling with the number of devices.
However, due to the issue of generalization gap \cite{keskar2016large,hoffer2017train}, using very large minibatches requires extra efforts to guarantee the convergence performance \cite{goyal2017accurate,you2017large,you2018imagenet,you2020large}.
Besides, a practical limitation is that large Transformer models that have more than billions of parameters are basically not able to be accommodated into the memory of a single GPU device.

Model parallelism can successfully remove the above memory limitation.
There are two model parallel paradigms.
First, the layer-wise pipelined parallelism splits the entire model by layer, and execute the layers in the pipeline fashion.
Some approaches \cite{pipedream,gpipe} use their algorithms to tackle the inconsistency issue among different pipelines, so that the computation and communication time can be overlapped.
However, since each training step requires both forward and backward processing, there will be an inevitable bubble overhead, so that the compute resources cannot be fully utilized.

The second model parallel paradigm is the intra-layer parallelism, which is orthogonal to the pipelined parallelism.
It distributes the operations in each layer such as matrix multiplications and activations across multiple workers.
For example, Mesh-Tensorflow \cite{meshTensorflow}, a framework proposed by Google, provides a convenient way for users to allocate the partition of tensor dimensions. 
Then the matrix multiplications can be automatically partitioned in a single dimension and distributed to workers, without any user elaboration.
% However, memory redundancy is severe, and the originally nearly-free reshape operations would result in complex communication because of layout change. 
Similarly, Megatron-LM \cite{megatron} adopts a 1-D matrix partition strategy to implement the Transformer model.
It splits matrices along rows or columns, and gathers results with the all-reduce operation.
% Although parameters and optimizer states are fully distributed, there is still memory redundancy for activations, as each device has to accommodate the whole input and output activations for every multi-layer perceptrons (MLPs) and self-attention modules.
Moreover, Optimus uses the Scalable Universal Matrix Multiplication Algorithm (SUMMA) \cite{SUMMA}, a 2-D parallel strategy for matrix multiplications, which helps to reduce memory and communication costs required for executing the Transfer model.
Our approach aims to take a step further and use the 3-D parallelism to improve the performance by optimally balancing the computation, memory and communication load.

\subsection{3-D Parallel Matrix Multiplication}
In this section, we provide an outline of the 3-D parallel matrix multiplication algorithm \cite{agarwal1995three}, which is adopted in our approach.

\begin{figure}[h]
    \centering
    \includegraphics[width=0.5\linewidth]{figure/cube.pdf}
    \caption{The structure of a 3-D processing cube with $P=p^3$ processors. The colored block $(i,j,k)$ represents an example of a single processor.}
    \label{fig:cube}
\end{figure}

This algorithm computes the multiplication between two dense matrices concurrently on $P=p^3$ processors, as illustrated in \autoref{fig:cube}, where $P$ processors can be stacked into a $p^3$ cube.
In this cube, let $A, B, C$ denote the planes, $x, y, z$ denote the directions for the planes, and $i, j, l$ denote the different indices along the directions.
We also use $A, B, C$ to represent the matrices involved in the multiplication, so that $A_{ij}, 0 \le i,j < p$ is supposed to represent a submatrix of $A$.

For simplicity, we consider an example of the 3-D matrix multiplication $C=AB$ on a $2\times2\times2$ processing cube.
We split $A$ and $B$ of size $(M,N)$ and $(N,K)$ into $2\times2$ partitions as follows, so that the sizes of each partition $A_{il}$ and $B_{lj}$ is $(M/2, N/2)$ and $(N/2, K/2)$ respectively, for $0 \le i,j < 2$.
\begin{align}
    A = \left[
        \begin{matrix}
            A_{00} & A_{01} \\
            A_{10} & A_{11}
        \end{matrix}
        \right], \quad
    B = \left[
        \begin{matrix}
            B_{00} & B_{01} \\
            B_{10} & B_{11}
        \end{matrix}
        \right].
\end{align}

\begin{figure}[H]
    \centering
    \begin{subfigure}[t]{0.25\linewidth}
        \centering
        \includegraphics[width=\linewidth]{figure/3d_a.pdf}
        \caption{Broadcast $A_{il}$}
        \label{fig:3d-a}
    \end{subfigure}
    \begin{subfigure}[t]{0.25\linewidth}
        \centering
        \includegraphics[width=\linewidth]{figure/3d_b.pdf}
        \caption{Broadcast $B_{lj}$}
        \label{fig:3d-b}
    \end{subfigure}
    \begin{subfigure}[t]{0.25\linewidth}
        \centering
        \includegraphics[width=\linewidth]{figure/3d_c.pdf}
        \caption{Reduce $C_{ij}$}
        \label{fig:3d-c}
    \end{subfigure}
    \caption{An example of the 3-D parallel matrix multiplication of size $2\times2\times2$.}
    \label{fig:3d}
\end{figure}

Each partition $A_{il}$ is stored in the processor $(i,0,l)$, while $B_{lj}$ is stored in $(0,j,l)$.
Before multiplication, we first broadcast each $A_{il}$ from $(i,0,l)$ along the $y$ direction (\autoref{fig:3d-a}), so that we have $A_il$ on each $(i,j,l)$.
Similarly, we broadcast each $B_{lj}$ along the $x$ direction (\autoref{fig:3d-b}).
Next, we can compute $C_{ijl} = A_{il}B_{lj}$ on each processor $(i,j,l)$.
Then we reduce $C_{ijl}$ to $(i,j,0)$ along the $z$ direction to derive $C_{ij}=\sum_l A_{il}B_{lj}$ (\autoref{fig:3d-c}).
Thus, we finally get $C=AB$ as
\begin{align}
    C = \left[
        \begin{matrix}
            A_{00}B_{00}+A_{01}B_{10} & A_{00}B_{01}+A_{01}B_{11} \\
            A_{10}B_{00}+A_{11}B_{10} & A_{10}B_{01}+A_{11}B_{11}
        \end{matrix}
        \right].
\end{align}

\section{Conclusion}
\label{sec:conclusion}
In this work, we introduce a 3-D intra-layer model parallelism algorithm for training huge neural models.
We propose a load balanced design to store and execute linear layers with minimum memory and communication cost.
By leveraging the 3-D model parallelism, we implement a 3-D parallel Transformer model and evaluate it on up to 64 GPUs.
Compared with the existing 1-D and 2-D parallelism, our Transformer model achieves a significant speedup.
However, it is still interesting to leverage our approach to train larger models on more processors.
We expect to see promising advances of our work with larger-scale compute resources in the future.
\section{3-D Parallel Transformers}
\label{sec:design}

We now describe our main design and implementation of the 3-D parallel Transformer language model.

\subsection{3-D Operations}

As demonstrated in \autoref{fig:transformer}, the Transformer language model computes its hidden states from top to bottom, through multiple Transformer blocks.
In each block, the major computation in both the Self-Attention and Multi-Layer Perception (MLP) layer involves linear operations and activations.
While activation operations can be independently executed in parallel, to accelerate linear operations is the key to reduce execution time of the Transformer model.
We firstly tackle the load balancing issue of the original 3-D matrix multiplication, and then present the designs of both 3-D parallel matrix-matrix and matrix-vector operations that are used in the Transformer model.

\subsubsection{Load Balancing}
We consider the matrix multiplication $C=AB$ between the input matrix $A$ of size $(M, N)$ and the weight matrix $B$ of size $(N,K)$ on $P=p^3$ processors, which requires $MNK$ calculations in total.
An intuitive way to store the matrices is to partition each dimension of the matrices by $p$, and hold each partition $A_{il}$ on the processor $(i,0,l)$, $B_{lj}$ on $(0,j,l)$, and $C_{ij}$ on $(i,j,0)$ respectively.
In this way, we will need to broadcast each partition $A_{il}$ and $B_{lj}$ across at least $p$ processors before multiplication, and then reduce the partitions $C_{ij}$ after the multiplication, so as to guarantee the consistency among them.
Therefore, an obvious issue is that the storage is imbalanced, so that we waste a great amount of redundant memory in the processor $(i,j,l)$ if $i \neq 0$ or $j \neq 0$ or $l \neq 0$.
In the meantime, the imbalanced storage will result in imbalanced activation or element-wise operations, causing the inefficiency that the computation of such operations is not evenly distributed to all processors.

To eliminate the redundancy caused by the gap between the number of computation and memory dimensions, our insight is to distribute the matrices evenly to the cube.
More specifically, we define $m=M/p^2, n=N/p^2, k=K/p^2$.
Let each processor $(i,j,l)$ hold the partitions $A_{ijl}=A[imp+jm: imp+jm+m-1, lnp: lnp+np-1]$, $B_{lji}=B[lnp: lnp+np-1, jkp+ik: jkp+ik+k-1]$, as shown in \autoref{fig:3d2-a}.
Before the multiplication, instead of the original broadcast operation, we use an all-gather operation along the $y$ direction to copy $A_{il}=A[imp: imp+mp-1,lnp: lnp+np-1]$ across processors $(i, j, l)$ for $0 \le j < p$ (\autoref{fig:3d2-b}).
Similarly, we execute an all-gather for $B_{lj}=B[lnp: np-1, jkp: jkp+kp-1]$ (\autoref{fig:3d2-c}).
Then we compute $A_{il}B_{lj}$ on each processor, and use a reduce-scatter operation in the $z$ direction to get $C_{ij}=\sum_l A_{il}B_{lj}$ across processors $(i, j, l)$ for $0 \le l < p$.
Thus, we will finally get $C_{ilj}=C_{ij}[lm: lm+m-1, :] = C[imp+lm: imp+lm+m-1, jkp: jkp+kp-1]$ in each processor $(i,j,l)$, as shown in \autoref{fig:3d2-d}.
In overall, the memory cost per processor of our approach is supposed to be $(M/p^2)*(N/p)+(N/p)*(K/p^2)+(M/p^2)*(K/p) \approx O(1/p^3)=O(1/P)$.

\begin{figure}[t]
\begin{minipage}{0.4\linewidth}
    \centering
    \includegraphics[width=\linewidth]{figure/transformer.pdf}
    \caption{Architecture of the Transformer model that consists of multiple Transformer layers.}
    \label{fig:transformer}
\end{minipage}
% \end{figure}
% \begin{figure}[t]
\begin{minipage}{0.6\linewidth}
        \centering
        \begin{subfigure}[t]{0.4\linewidth}
            \centering
            \includegraphics[width=\linewidth]{figure/3d2_a.pdf}
            \caption{The input and weight submatrices on each processor.}
            \label{fig:3d2-a}
        \end{subfigure}
        \begin{subfigure}[t]{0.4\linewidth}
            \centering
            \includegraphics[width=\linewidth]{figure/3d2_b.pdf}
            \caption{All-gather $A_{il}$ in the $y$ direction.}
            \label{fig:3d2-b}
        \end{subfigure}
        \begin{subfigure}[t]{0.4\linewidth}
            \centering
            \includegraphics[width=\linewidth]{figure/3d2_c.pdf}
            \caption{All-gather $B_{lj}$ in the $x$ direction.}
            \label{fig:3d2-c}
        \end{subfigure}
        \begin{subfigure}[t]{0.4\linewidth}
            \centering
            \includegraphics[width=\linewidth]{figure/3d2_d.pdf}
            \caption{Compute each $C_{ijl}$ and reduce-scatter them in the $z$ direction.}
            \label{fig:3d2-d}
        \end{subfigure}
        \caption{Load-balanced 3-D matrix-matrix multiplication.}
        \label{fig:3d2}
\end{minipage}
\end{figure}

For vector parameters that are used in matrix-vector operations, we store the vector $b$ of size $N$ diagonally on the B plane, i.e. only the processors $(i,j,l)$ holds $b_{ji} = b[jnp+in: jnp+in+n-1]$ for $j = l, 0 \le i,j,l < p$, as shown by the colored blocks in \autoref{fig:vector}.
To execute a matrix-vector operation, e.g. $A+b$, we need to broadcast $b_{li}$ in the $y$ direction, and then get $b_l$ by all-gathering them in the $x$ direction.
In contrast, to execute $C+b$, we need to broadcast $b_{ji}$ in the $z$ direction before all-gathering them.
Then we can get $A_{il}+b_l$ or $C_{ij}+b_j$ on each processor.

\begin{figure}[h]
    \centering
    \includegraphics[width=0.5\linewidth]{figure/vector.pdf}
    \caption{Balancing vector parameters for matrix-vector operations.}
    \label{fig:vector}
\end{figure}

We need to notice that the communication directions of the input matrix $A$ and output matrix $C$ are different.
To execute the Transformer model, since the output states of each linear layer are basically the input states of the next layer, we exchange the communication directions $y$ and $z$ for the input and output states after the 3-D parallel matrix multiplication, while keep the direction $x$ for the network weights.
The matrix-vector operations are supposed to have no affect on the communication directions.

\subsubsection{Matrix-matrix Operations}
In the Transformer model, three forms of matrix-matrix multiplication, including $C=AB$, $C=AB^T$ and $C=A^TB$, are involved w.r.t the following differentiation formulas.
\begin{align}
    C=AB, \dot{A}=\dot{C}B^T, \dot{B}=A^T\dot{C}; \\
    C=AB^T, \dot{A}=\dot{C}B, \dot{B}=\dot{C}^TA; \\
    C=A^TB, \dot{A}=B\dot{C}^T, \dot{B}=A\dot{C},
\end{align}
where $\dot{X}$ denotes the gradient of the parameter $X$.
We present the algorithms for the three forms of multiplication with the optimal 3-D execution and communication efficiency as shown in the following pseudo-codes (Algorithm \autoref{alg:parallel_ab_forward}-\autoref{alg:parallel_atb_backward}).
Note that for each multiplication algorithm, other than the input matrices, we also need to specify the directions for them as well as the output matrix.

\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Forward $C=AB$}
        \label{alg:parallel_ab_forward}
        \begin{algorithmic}[1]
            \REQUIRE $A_{ijl}$, $B_{lji}$, \\
                    directions $y,x,z$ for $A,B,C$ 
            \ENSURE $C_{ilj}$
            \STATE All-gather $A_{il}$ in $y$ \\
            \STATE All-gather $B_{lj}$ in $x$ \\
            \STATE $C_{ilj} \gets A_{il}B_{lj}$ \\
            \STATE Reduce-scatter $C_{ilj}$ in $z$ \\
            \RETURN $C_{ilj}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}
\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Backward $C=AB$}
        \label{alg:parallel_ab_backward}
        \begin{algorithmic}[1]
            \REQUIRE $\dot{C}_{ilj}$, $A_{ijl}$, $B_{lji}$\\
                    directions $y,x,z$ for $A,B,C$
            \ENSURE $\dot{A}_{ijl},\dot{B}_{lji}$ \\ [0.3\baselineskip]
            \STATE $\dot{A}_{ijl} \gets \dot{C}_{ilj}B_{lji}^T$ in $z,x,y$ \\ [0.3\baselineskip]
            \STATE $\dot{B}_{lji} \gets A_{ijl}^T\dot{C}_{ilj}$ in $y,z,x$\\ [0.3\baselineskip]
            \RETURN $\dot{A}_{ijl},\dot{B}_{lji}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}

\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Forward $C=AB^T$}
        \label{alg:parallel_abt_forward}
        \begin{algorithmic}[1]
            \REQUIRE $A_{ijl}$, $B_{jli}$, \\
                    directions $y,x,z$ for $A,B,C$
            \ENSURE $C_{ilj}$
            \STATE All-gather $A_{il}$ in $y$ \\
            \STATE All-gather $B_{jl}$ in $x$ \\
            \STATE $C_{ilj} \gets A_{il}B_{jl}^T$ \\
            \STATE Reduce-scatter $C_{ilj}$ in $z$\\
            \RETURN $C_{ilj}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}
\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Backward $C=AB^T$}
        \label{alg:parallel_abt_backward}
        \begin{algorithmic}[1]
            \REQUIRE $\dot{C}_{ilj}$, $A_{ijl}$, $B_{jli}$\\
                    directions $y,x,z$ for $A,B,C$
            \ENSURE $\dot{A}_{ijl},\dot{B}_{jli}$ \\ [0.4\baselineskip]
            \STATE $\dot{A}_{ijl} \gets \dot{C}_{ilj}B_{jli}$ in $z,x,y$ \\ [0.4\baselineskip]
            \STATE $\dot{B}_{jli} \gets \dot{C}_{ilj}^TA_{ijl}$ in $z,y,x$\\ [0.4\baselineskip]
            \RETURN $\dot{A}_{ijl},\dot{B}_{jli}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}

\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Forward $C=A^TB$}
        \label{alg:parallel_atb_forward}
        \begin{algorithmic}[1]
            \REQUIRE $A_{ilj}$, $B_{ijl}$, \\
                    directions $y,x,z$ for $A,B,C$
            \ENSURE $C_{jli}$
            \STATE All-gather $A_{ij}$ in $y$ \\
            \STATE All-gather $B_{il}$ in $x$ \\
            \STATE $C_{jl} \gets A_{ij}^TB_{il}$ \\
            \STATE Reduce-scatter $C_{jli}$ in $z$\\
            \RETURN $C_{jli}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}
\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Backward $C=AB^T$}
        \label{alg:parallel_atb_backward}
        \begin{algorithmic}[1]
            \REQUIRE $\dot{C}_{jli}$, $A_{ilj}$, $B_{ijl}$\\
                    directions $y,x,z$ for $A,B,C$
            \ENSURE $\dot{A}_{ilj},\dot{B}_{ijl}$ \\ [0.35\baselineskip]
            \STATE $\dot{A}_{ilj} \gets B_{jli}\dot{C}_{ilj}^T$ in $x,z,y$ \\ [0.35\baselineskip]
            \STATE $\dot{B}_{ijl} \gets A_{ijl}\dot{C}_{ilj}$ in $y,z,x$\\ [0.35\baselineskip]
            \RETURN $\dot{A}_{ilj},\dot{B}_{ilj}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}

Our 3-D parallel matrix multiplication can evenly distribute the computational cost to all the processors, where each processor only multiplies the submatrices of size $(M/p, N/p)$ and $(N/p, K/p)$.
If we fix the problem size, the execution cost of our approach is supposed to be $(M/p)*(N/p)*(K/p) \approx O(1/p^3)=O(1/P)$.
Furthermore, compared with 1-D and 2-D approaches, each communication operation in our 3-D approach moves data across a smaller number of processors.
Each forward algorithm uses all-gather and reduce-scatter operations to move $(MN+NK+MK)/p^3$ across $p$ processors in total.
Therefore, the bandwidth cost of our approach is supposed to be $O(1/p^2)=O(P^{-2/3})$, while the latency cost is supposed to be $O(log(p))=O(log(P^{1/3}))$.

\subsubsection{Matrix-vector Operations}
The logic to execute matrix-vector operations such $C=A+b$ and $C=A*b$ is similar.
We present an example of the add operation as shown in Algorithm \autoref{alg:parallel_add_forward}-\autoref{alg:parallel_add_backward}.

For the matrix-vector multiplication, the only changes from the above algorithms are to return $C_{ijl}=A_{ijl}*b_{l}$ for the forward pass, as well as to return $\dot{A}_{ijl}=\dot{C}_{ijl}*{b}_{ijl}$ and $\dot{b}_{ijl}=\sum_{ij} \dot{C}_{ijl}*{A}_{ijl}$ for the backward pass.

\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Forward $C=A+b$}
        \label{alg:parallel_add_forward}
        \begin{algorithmic}[1]
            \REQUIRE $A_{ijl}$, $b_{ijl}$, \\
                    directions $y,x,z$ for $A,B,C$ \\ [\baselineskip]
            \ENSURE $C_{ijl}$ \\ [\baselineskip]
            \STATE Broadcast $b_{il}$ in $y$ \\ [\baselineskip]
            \STATE All-gather $b_{l}$ in $x$ \\ [\baselineskip]
            \STATE $C_{ijl} \gets A_{ijl}+b_{l}$ \\ [\baselineskip]
            \RETURN $C_{ijl}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}
\begin{minipage}{0.45\linewidth}
    \begin{algorithm}[H]
        \caption{Backward $C=A+b$}
        \label{alg:parallel_add_backward}
        \begin{algorithmic}[1]
            \REQUIRE $\dot{C}_{ijl}$,\\
                    directions $y,x,z$ for $A,B,C$
            \ENSURE $\dot{A}_{ijl},\dot{b}_{ijl}$ \\
            \STATE $\dot{A}_{ijl} \gets \dot{C}_{ijl}$ \\
            \STATE $\dot{b}_{ijl} \gets \sum_{ij} \dot{C}_{ijl}$
            \IF{$j=l$}
                \STATE Reduce-scatter $\dot{b}_{ijl}$ in $x$\\
            \ELSE
                \STATE $\dot{b}_{ijl} \gets$ null
            \ENDIF
            \RETURN $\dot{A}_{ijl},\dot{b}_{ijl}$
        \end{algorithmic}
    \end{algorithm}
\end{minipage}

By taking advantage of balanced parameter storage, we can evenly distribute the computation cost of not only the matrix-vector operations but also the activations to all the processors.
Their computation cost is supposed to be the same scale as the memory cost $O(1/P)$.


\subsection{Parallel Transformer layers}

\begin{figure}[h]
    \centering
    \begin{subfigure}[t]{\linewidth}
        \centering
        \includegraphics[width=\linewidth]{figure/attention.pdf}
        \caption{The Self-Attention block}
        \label{fig:attention}
    \end{subfigure}
    \begin{subfigure}[t]{0.7\linewidth}
        \centering
        \includegraphics[width=\linewidth]{figure/mlp.pdf}
        \caption{The Multi-Layer Perception block}
        \label{fig:mlp}
    \end{subfigure}
    \caption{Outline of 3-D Parallel Transformer layers. Each $m,n,k$ in $[m,n,k]$ represents the size of the parameter stored on each processor. The arrows present the order of computation.}
    \label{fig:layer}
\end{figure}

In this section, we describe the major implementation details of the Transformer model that uses our 3-D parallel operations.

To leverage 3-D model parallelism, we implement our new linear layer module and layer normalization module based on the 3-D parallel operations.
The 3-D linear layer flattens an input $X$ of size $[b/p,s/p,h/p]$ into size $[b*s/p^2, h/p]$ first, and then uses Algorithm \autoref{alg:parallel_ab_forward} and \autoref{alg:parallel_add_forward} to perform $Y=XW+b$ with parameters $W$ and $b$.
Thus, after a 3-D linear layer, the input and output directions should be exchanged.
The 3-D layer normalization will not affect the directions, because it only applies matrix-vector adds and multiplications with the parameters $\gamma$ and $\beta$.

We define $b,s,n,h$ as the batch size, sequence length, number of attention heads and hidden size of the Transformer model on a $p^3$ processing cube.
For ease of presentation, we do not discuss the embedding and output layers of the language model, as they usually depend on the downstream tasks.
Then our implementation can basically divided into two modules: Self-Attention blocks and MLP blocks, as illustrated in \autoref{fig:attention} and \autoref{fig:mlp}.
Both the Self-Attention and MLP block include two 3-D linear layers, so that communication directions of their inputs and outputs are supposed to be the same.
Thus, the input and output of each Transformer layer are also supposed to have the same communication direction.

Besides, in practice, it is therefore not necessary to track all the directions of 3-D parallel operations as shown in Algorithm \autoref{alg:parallel_ab_forward}-\autoref{alg:parallel_add_backward}, because we only need to exchange the input and output direction after the first linear layer of both Self-Attention and MLP blocks.
For generalizability, we index the initial input, weight and output group as 0, 1, 2.
We provide an input group index $y$ for the first linear layer, so that the output group index is supposed to be $z=1-y$.
After the layer, we take $z$ as the new input group index, and return it to $y=1-z$ after each Self-Attention or MLP block is completed.

\section{Experiments}
\label{sec:eval}

\subsection{Setup}
We evaluate our 3-D parallel Transformer model on TACC's Longhorn supercomputers \cite{longhorn}.
Our testbed cluster consists of up to 16 GPU servers connected by Mellanox EDR Infiniband network.
Each server has 2 20-core IBM Power 9 CPUs, 256GB memory and 4 NVIDIA V100 GPUs, so there are 64 GPUs in total.
We store the training data, model checkpoints in a Hadoop Distributed File System (HDFS) via 1Gbps Ethernet.

We compare the performance of our approach against the existing state-of-the-art 1-D \cite{megatron} and 2-D \cite{optimus} model parallel approaches.
For simplicity, we only evaluate the performance of the Transformer model itself, i.e. the consecutive Transformer layers, to highlight the efficiency of different parallelism.
We evaluate the 1-D parallelism on 8,16,36 and 64 GPUs, the 2-D parallelism on a subset of 16,36 and 64 GPUs, as well as our 3-D approach on a subset of 8 and 64 GPUs.

\subsection{Performance}

\subsubsection{Weak Scaling}

We first evaluate the weak scaling performance, where we fix the network size on each processor, and scale up the number of processors.
\autoref{tab:weak} presents the results of the weak scaling experiments, including the execution time for both forward and backward pass of the Transformer model.
We also provide the average step time of each sequence, which is computed as 
\begin{align}
    \text{Average step time} = \frac{\text{forward time} + \text{backward time}}{\text{batch size}}.
\end{align}

\begin{table}[h]
    \begin{subtable}[h]{\linewidth}
        \centering
        \begin{tabular}{c|c|c|c|c|c|c}
            \toprule
            \multirow{6}{*}{1-D \cite{megatron}} & \multirow{2}{*}{\# GPUs} & \multirow{2}{*}{\makecell[c]{Batch\\size}} & \multirow{2}{*}{\makecell[c]{Hidden\\size}} & \multirow{2}{*}{\makecell[c]{Forward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Backward \\ time (s)}} & \multirow{2}{*}{\makecell[c]{Average\\ step time (s)}}\\
            & & & & & & \\
            \cline{2-7}
            & 8 & 60 & 2048 & 4.759 & 15.676 & 0.341 \\
            \cline{2-7}
            & 16 & 60 & 4096 & 12.488 & 30.894 & 0.723 \\
            \cline{2-7}
            & 36 & 40 & 6120 & 13.515 & 31.822 & 1.133 \\
            \cline{2-7}
            & 64 & 30 & 8192 & 13.915 & 32.890 & 1.560 \\
            \bottomrule
        \end{tabular}
        % \caption{1-D Parallelism \cite{megatron}}
        \label{tab:weak-1d}
    \end{subtable}
    
    \begin{subtable}[h]{\linewidth}
        \centering
        \begin{tabular}{c|c|c|c|c|c|c}
            \toprule
            \multirow{5}{*}{2-D \cite{optimus}} & \multirow{2}{*}{\# GPUs} & \multirow{2}{*}{\makecell[c]{Batch\\size}} & \multirow{2}{*}{\makecell[c]{Hidden\\size}} & \multirow{2}{*}{\makecell[c]{Forward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Backward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Average\\step time (s)}}\\
            & & & & & & \\
            \cline{2-7}
            & 16 & 192 & 4096 & 33.860 & 101.981 & 0.708 \\
            \cline{2-7}
            & 36 & 288 & 6120 & 54.760 & 165.850 & 0.766 \\
            \cline{2-7}
            & 64 & 384 & 8192 & 99.419 & 304.707 & 1.052 \\
            \bottomrule
        \end{tabular}
        % \caption{2-D Parallelism \cite{optimus}}
        \label{tab:weak-2d}
    \end{subtable}
    \begin{subtable}[h]{\linewidth}
        \centering
        \begin{tabular}{c|c|c|c|c|c|c}
            \toprule
            \multirow{4}{*}{~~~ 3-D ~~~} & \multirow{2}{*}{\# GPUs} & \multirow{2}{*}{\makecell[c]{Batch\\size}} & \multirow{2}{*}{\makecell[c]{Hidden\\size}}  &  \multirow{2}{*}{\makecell[c]{Forward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Backward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Average\\step time (s)}}\\
            & & & & & & \\
            \cline{2-7}
            & 8 & 192 & 2048 & 30.096 & 81.212 & 0.580 \\
            \cline{2-7}
            & 64 & 384 & 8192 & 79.349 & 125.037 & \textbf{0.672} \\
            \bottomrule
        \end{tabular}
        % \caption{3-D Parallelism}
        \label{tab:weak-3d}
    \end{subtable}
    \caption{Comparison of weak scaling results. The number of processors increases from 8 to 64, while the number of parameters is fixed for each approach. For each of presentation, we mainly adjust the batch size and hidden size based on the number of processors, while fix the sequence length to 512. The bolded number represents the best result.}
    \label{tab:weak}
\end{table}

We see that the 3-D parallelism has the slowest rising speed in the average step time, reaching the smallest value at the largest compute scale.
The result implies that our approach is efficient to reduce the overhead of model parallelism, as it reaches the minimum communication cost.


\subsubsection{Strong Scaling}
Next, we evaluate the strong scaling performance of different model parallelism.
This experiment aims to examine how much speedup each approach can achieve by using an increasing number of processors to execute the fix-sized problem.

The problem size and execution results are shown in \autoref{tab:strong}.
We see that our 3-D parallelism outperforms other approaches by achieving the smallest average step time on 64 GPUs, with 2.32X and 1.57X speedup over 1-D and 2-D approach, respectively.

\begin{table}[h]
    \begin{subtable}[h]{\linewidth}
        \centering
        \begin{tabular}{c|c|c|c|c|c|c}
            \toprule
            \multirow{6}{*}{1-D \cite{megatron}} & \multirow{2}{*}{\# GPUs} & \multirow{2}{*}{\makecell[c]{Batch\\size}} & \multirow{2}{*}{\makecell[c]{Hidden\\size}} & \multirow{2}{*}{\makecell[c]{Forward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Backward \\ time (s)}} & \multirow{2}{*}{\makecell[c]{Average\\ step time (s)}}\\
            & & & & & & \\
            \cline{2-7}
            & 8 & 12 & 3072 & 1.470 & 5.699 & 0.597 \\
            \cline{2-7}
            & 16 & 12 & 3072 & 1.371 & 5.152 & 0.544 \\
            \cline{2-7}
            & 36 & 12 & 3072 & 1.455 & 5.414 & 0.572 \\
            \cline{2-7}
            & 64 & 12 & 3072 & 1.433 & 5.167 & 0.550 \\
            \bottomrule
        \end{tabular}
        % \caption{1-D Parallelism \cite{megatron}}
        \label{tab:strong-1d}
    \end{subtable}
    
    \begin{subtable}[h]{\linewidth}
        \centering
        \begin{tabular}{c|c|c|c|c|c|c}
            \toprule
            \multirow{5}{*}{2-D \cite{optimus}} & \multirow{2}{*}{\# GPUs} & \multirow{2}{*}{\makecell[c]{Batch\\size}} & \multirow{2}{*}{\makecell[c]{Hidden\\size}} & \multirow{2}{*}{\makecell[c]{Forward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Backward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Average\\step time (s)}}\\
            & & & & & & \\
            \cline{2-7}
            & 16 & 24 & 3072 & 4.680 & 13.698 & 0.766 \\
            \cline{2-7}
            & 36 & 24 & 3072 & 3.900 & 11.433 & 0.639 \\
            \cline{2-7}
            & 64 & 24 & 3072 & 3.007 & 8.920 & 0.497 \\
            \bottomrule
        \end{tabular}
        % \caption{2-D Parallelism \cite{optimus}}
        \label{tab:strong-2d}
    \end{subtable}
    \begin{subtable}[h]{\linewidth}
        \centering
        \begin{tabular}{c|c|c|c|c|c|c}
            \toprule
            \multirow{4}{*}{~~~ 3-D ~~~} & \multirow{2}{*}{\# GPUs} & \multirow{2}{*}{\makecell[c]{Batch\\size}} & \multirow{2}{*}{\makecell[c]{Hidden\\size}}  &  \multirow{2}{*}{\makecell[c]{Forward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Backward\\time (s)}} & \multirow{2}{*}{\makecell[c]{Average\\step time (s)}}\\
            & & & & & & \\
            \cline{2-7}
            & 8 & 24 & 3072 & 3.249 & 9.120 & 0.515 \\
            \cline{2-7}
            & 64 & 24 & 3072 & 2.494 & 6.129 & \textbf{0.359} \\
            \bottomrule
        \end{tabular}
        % \caption{3-D Parallelism}
        \label{tab:strong-3d}
    \end{subtable}
    \caption{Comparison of strong scaling results. The problem size is fixed while the number of processors increases from 8 to 64. The bolded number represents the best result.}
    \label{tab:strong}
\end{table}

\section{Implementation}
To do
\section{Submission instructions}

\textcolor{blue}{For checking submission format.}

\subsection{Style}

Papers to be submitted to NeurIPS 2021 must be prepared according to the
instructions presented here. Papers may only be up to {\bf nine} pages long,
including figures. Additional pages \emph{containing only acknowledgments and
references} are allowed. Papers that exceed the page limit will not be
reviewed, or in any other way considered for presentation at the conference.

The margins in 2021 are the same as those in 2007, which allow for $\sim$$15\%$
more words in the paper compared to earlier years.

Authors are required to use the NeurIPS \LaTeX{} style files obtainable at the
NeurIPS website as indicated below. Please make sure you use the current files
and not previous versions. Tweaking the style files may be grounds for
rejection.

\subsection{Retrieval of style files}

The style files for NeurIPS and other conference information are available on
the World Wide Web at
\begin{center}
  \url{http://www.neurips.cc/}
\end{center}
The file \verb+neurips_2021.pdf+ contains these instructions and illustrates the
various formatting requirements your NeurIPS paper must satisfy.

The only supported style file for NeurIPS 2021 is \verb+neurips_2021.sty+,
rewritten for \LaTeXe{}.  \textbf{Previous style files for \LaTeX{} 2.09,
  Microsoft Word, and RTF are no longer supported!}

The \LaTeX{} style file contains three optional arguments: \verb+final+, which
creates a camera-ready copy, \verb+preprint+, which creates a preprint for
submission to, e.g., arXiv, and \verb+nonatbib+, which will not load the
\verb+natbib+ package for you in case of package clash.

\paragraph{Preprint option}
If you wish to post a preprint of your work online, e.g., on arXiv, using the
NeurIPS style, please use the \verb+preprint+ option. This will create a
nonanonymized version of your work with the text ``Preprint. Work in progress.''
in the footer. This version may be distributed as you see fit. Please \textbf{do
  not} use the \verb+final+ option, which should \textbf{only} be used for
papers accepted to NeurIPS.

At submission time, please omit the \verb+final+ and \verb+preprint+
options. This will anonymize your submission and add line numbers to aid
review. Please do \emph{not} refer to these line numbers in your paper as they
will be removed during generation of camera-ready copies.

The file \verb+neurips_2021.tex+ may be used as a ``shell'' for writing your
paper. All you have to do is replace the author, title, abstract, and text of
the paper with your own.

The formatting instructions contained in these style files are summarized in
Sections \ref{gen_inst}, \ref{headings}, and \ref{others} below.

\section{General formatting instructions}
\label{gen_inst}

The text must be confined within a rectangle 5.5~inches (33~picas) wide and
9~inches (54~picas) long. The left margin is 1.5~inch (9~picas).  Use 10~point
type with a vertical spacing (leading) of 11~points.  Times New Roman is the
preferred typeface throughout, and will be selected for you by default.
Paragraphs are separated by \nicefrac{1}{2}~line space (5.5 points), with no
indentation.

The paper title should be 17~point, initial caps/lower case, bold, centered
between two horizontal rules. The top rule should be 4~points thick and the
bottom rule should be 1~point thick. Allow \nicefrac{1}{4}~inch space above and
below the title to rules. All pages should start at 1~inch (6~picas) from the
top of the page.

For the final version, authors' names are set in boldface, and each name is
centered above the corresponding address. The lead author's name is to be listed
first (left-most), and the co-authors' names (if different address) are set to
follow. If there is only one co-author, list both author and co-author side by
side.

Please pay special attention to the instructions in Section \ref{others}
regarding figures, tables, acknowledgments, and references.

\section{Headings: first level}
\label{headings}

All headings should be lower case (except for first word and proper nouns),
flush left, and bold.

First-level headings should be in 12-point type.

\subsection{Headings: second level}

Second-level headings should be in 10-point type.

\subsubsection{Headings: third level}

Third-level headings should be in 10-point type.

\paragraph{Paragraphs}

There is also a \verb+\paragraph+ command available, which sets the heading in
bold, flush left, and inline with the text, with the heading followed by 1\,em
of space.

\section{Citations, figures, tables, references}
\label{others}

These instructions apply to everyone.

\subsection{Citations within the text}

The \verb+natbib+ package will be loaded for you by default.  Citations may be
author/year or numeric, as long as you maintain internal consistency.  As to the
format of the references themselves, any style is acceptable as long as it is
used consistently.

The documentation for \verb+natbib+ may be found at
\begin{center}
  \url{http://mirrors.ctan.org/macros/latex/contrib/natbib/natnotes.pdf}
\end{center}
Of note is the command \verb+\citet+, which produces citations appropriate for
use in inline text.  For example,
\begin{verbatim}
   \citet{hasselmo} investigated\dots
\end{verbatim}
produces
\begin{quote}
  Hasselmo, et al.\ (1995) investigated\dots
\end{quote}

If you wish to load the \verb+natbib+ package with options, you may add the
following before loading the \verb+neurips_2021+ package:
\begin{verbatim}
   \PassOptionsToPackage{options}{natbib}
\end{verbatim}

If \verb+natbib+ clashes with another package you load, you can add the optional
argument \verb+nonatbib+ when loading the style file:
\begin{verbatim}
   \usepackage[nonatbib]{neurips_2021}
\end{verbatim}

As submission is double blind, refer to your own published work in the third
person. That is, use ``In the previous work of Jones et al.\ [4],'' not ``In our
previous work [4].'' If you cite your other papers that are not widely available
(e.g., a journal paper under review), use anonymous author names in the
citation, e.g., an author of the form ``A.\ Anonymous.''

\subsection{Footnotes}

Footnotes should be used sparingly.  If you do require a footnote, indicate
footnotes with a number\footnote{Sample of the first footnote.} in the
text. Place the footnotes at the bottom of the page on which they appear.
Precede the footnote with a horizontal rule of 2~inches (12~picas).

Note that footnotes are properly typeset \emph{after} punctuation
marks.\footnote{As in this example.}

\subsection{Figures}

\begin{figure}
  \centering
  \fbox{\rule[-.5cm]{0cm}{4cm} \rule[-.5cm]{4cm}{0cm}}
  \caption{Sample figure caption.}
\end{figure}

All artwork must be neat, clean, and legible. Lines should be dark enough for
purposes of reproduction. The figure number and caption always appear after the
figure. Place one line space before the figure caption and one line space after
the figure. The figure caption should be lower case (except for first word and
proper nouns); figures are numbered consecutively.

You may use color figures.  However, it is best for the figure captions and the
paper body to be legible if the paper is printed in either black/white or in
color.

\subsection{Tables}

All tables must be centered, neat, clean and legible.  The table number and
title always appear before the table.  See Table~\ref{sample-table}.

Place one line space before the table title, one line space after the
table title, and one line space after the table. The table title must
be lower case (except for first word and proper nouns); tables are
numbered consecutively.

Note that publication-quality tables \emph{do not contain vertical rules.} We
strongly suggest the use of the \verb+booktabs+ package, which allows for
typesetting high-quality, professional tables:
\begin{center}
  \url{https://www.ctan.org/pkg/booktabs}
\end{center}
This package was used to typeset Table~\ref{sample-table}.

\begin{table}
  \caption{Sample table title}
  \label{sample-table}
  \centering
  \begin{tabular}{lll}
    \toprule
    \multicolumn{2}{c}{Part}                   \\
    \cmidrule(r){1-2}
    Name     & Description     & Size ($\mu$m) \\
    \midrule
    Dendrite & Input terminal  & $\sim$100     \\
    Axon     & Output terminal & $\sim$10      \\
    Soma     & Cell body       & up to $10^6$  \\
    \bottomrule
  \end{tabular}
\end{table}

\section{Final instructions}

Do not change any aspects of the formatting parameters in the style files.  In
particular, do not modify the width or length of the rectangle the text should
fit into, and do not change font sizes (except perhaps in the
\textbf{References} section; see below). Please note that pages should be
numbered.

\section{Preparing PDF files}

Please prepare submission files with paper size ``US Letter,'' and not, for
example, ``A4.''

Fonts were the main cause of problems in the past years. Your PDF file must only
contain Type 1 or Embedded TrueType fonts. Here are a few instructions to
achieve this.

\begin{itemize}

\item You should directly generate PDF files using \verb+pdflatex+.

\item You can check which fonts a PDF files uses.  In Acrobat Reader, select the
  menu Files$>$Document Properties$>$Fonts and select Show All Fonts. You can
  also use the program \verb+pdffonts+ which comes with \verb+xpdf+ and is
  available out-of-the-box on most Linux machines.

\item The IEEE has recommendations for generating PDF files whose fonts are also
  acceptable for NeurIPS. Please see
  \url{http://www.emfield.org/icuwb2010/downloads/IEEE-PDF-SpecV32.pdf}

\item \verb+xfig+ "patterned" shapes are implemented with bitmap fonts.  Use
  "solid" shapes instead.

\item The \verb+\bbold+ package almost always uses bitmap fonts.  You should use
  the equivalent AMS Fonts:
\begin{verbatim}
   \usepackage{amsfonts}
\end{verbatim}
followed by, e.g., \verb+\mathbb{R}+, \verb+\mathbb{N}+, or \verb+\mathbb{C}+
for $\mathbb{R}$, $\mathbb{N}$ or $\mathbb{C}$.  You can also use the following
workaround for reals, natural and complex:
\begin{verbatim}
   \newcommand{\RR}{I\!\!R} %real numbers
   \newcommand{\Nat}{I\!\!N} %natural numbers
   \newcommand{\CC}{I\!\!\!\!C} %complex numbers
\end{verbatim}
Note that \verb+amsfonts+ is automatically loaded by the \verb+amssymb+ package.

\end{itemize}

If your file contains type 3 fonts or non embedded TrueType fonts, we will ask
you to fix it.

\subsection{Margins in \LaTeX{}}

Most of the margin problems come from figures positioned by hand using
\verb+\special+ or other commands. We suggest using the command
\verb+\includegraphics+ from the \verb+graphicx+ package. Always specify the
figure width as a multiple of the line width as in the example below:
\begin{verbatim}
   \usepackage[pdftex]{graphicx} ...
   \includegraphics[width=0.8\linewidth]{myfile.pdf}
\end{verbatim}
See Section 4.4 in the graphics bundle documentation
(\url{http://mirrors.ctan.org/macros/latex/required/graphics/grfguide.pdf})

A number of width problems arise when \LaTeX{} cannot properly hyphenate a
line. Please give LaTeX hyphenation hints using the \verb+\-+ command when
necessary.

\begin{ack}
Use unnumbered first level headings for the acknowledgments. All acknowledgments
go at the end of the paper before the list of references. Moreover, you are required to declare
funding (financial activities supporting the submitted work) and competing interests (related financial activities outside the submitted work).
More information about this disclosure can be found at: \url{https://neurips.cc/Conferences/2021/PaperInformation/FundingDisclosure}.

Do {\bf not} include this section in the anonymized submission, only in the final paper. You can use the \texttt{ack} environment provided in the style file to autmoatically hide this section in the anonymized submission.
\end{ack}

\section*{References}

References follow the acknowledgments. Use unnumbered first-level heading for
the references. Any choice of citation style is acceptable as long as you are
consistent. It is permissible to reduce the font size to \verb+small+ (9 point)
when listing the references.
Note that the Reference section does not count towards the page limit.
\medskip

{
\small

[1] Alexander, J.A.\ \& Mozer, M.C.\ (1995) Template-based algorithms for
connectionist rule extraction. In G.\ Tesauro, D.S.\ Touretzky and T.K.\ Leen
(eds.), {\it Advances in Neural Information Processing Systems 7},
pp.\ 609--616. Cambridge, MA: MIT Press.

[2] Bower, J.M.\ \& Beeman, D.\ (1995) {\it The Book of GENESIS: Exploring
  Realistic Neural Models with the GEneral NEural SImulation System.}  New York:
TELOS/Springer--Verlag.

[3] Hasselmo, M.E., Schnell, E.\ \& Barkai, E.\ (1995) Dynamics of learning and
recall at excitatory recurrent synapses and cholinergic modulation in rat
hippocampal region CA3. {\it Journal of Neuroscience} {\bf 15}(7):5249-5262.
}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section*{Checklist}

%%% BEGIN INSTRUCTIONS %%%
The checklist follows the references.  Please
read the checklist guidelines carefully for information on how to answer these
questions.  For each question, change the default \answerTODO{} to \answerYes{},
\answerNo{}, or \answerNA{}.  You are strongly encouraged to include a {\bf
justification to your answer}, either by referencing the appropriate section of
your paper or providing a brief inline description.  For example:
\begin{itemize}
  \item Did you include the license to the code and datasets? \answerYes{See Section~\ref{gen_inst}.}
  \item Did you include the license to the code and datasets? \answerNo{The code and the data are proprietary.}
  \item Did you include the license to the code and datasets? \answerNA{}
\end{itemize}
Please do not modify the questions and only use the provided macros for your
answers.  Note that the Checklist section does not count towards the page
limit.  In your paper, please delete this instructions block and only keep the
Checklist section heading above along with the questions/answers below.
%%% END INSTRUCTIONS %%%

\begin{enumerate}

\item For all authors...
\begin{enumerate}
  \item Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
    \answerTODO{}
  \item Did you describe the limitations of your work?
    \answerTODO{}
  \item Did you discuss any potential negative societal impacts of your work?
    \answerTODO{}
  \item Have you read the ethics review guidelines and ensured that your paper conforms to them?
    \answerTODO{}
\end{enumerate}

\item If you are including theoretical results...
\begin{enumerate}
  \item Did you state the full set of assumptions of all theoretical results?
    \answerTODO{}
	\item Did you include complete proofs of all theoretical results?
    \answerTODO{}
\end{enumerate}

\item If you ran experiments...
\begin{enumerate}
  \item Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
    \answerTODO{}
  \item Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
    \answerTODO{}
	\item Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
    \answerTODO{}
	\item Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
    \answerTODO{}
\end{enumerate}

\item If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
\begin{enumerate}
  \item If your work uses existing assets, did you cite the creators?
    \answerTODO{}
  \item Did you mention the license of the assets?
    \answerTODO{}
  \item Did you include any new assets either in the supplemental material or as a URL?
    \answerTODO{}
  \item Did you discuss whether and how consent was obtained from people whose data you're using/curating?
    \answerTODO{}
  \item Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
    \answerTODO{}
\end{enumerate}

\item If you used crowdsourcing or conducted research with human subjects...
\begin{enumerate}
  \item Did you include the full text of instructions given to participants and screenshots, if applicable?
    \answerTODO{}
  \item Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
    \answerTODO{}
  \item Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
    \answerTODO{}
\end{enumerate}

\end{enumerate}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\section{Introduction}
% In 1989, Yann LeCun proposed to train convolutional neural networks (CNNs) using back-propagation, which makes learning fully automatic, instead of manually designing convolutional kernels. However, the current prevailance of deep learning started decades after it, when AlexNet took the crown in ImageNet Large Scale Visual Recognition Challenge 2012. It owed to the development of computation capability. Since then, 
To deal with the growing amount of data, nowadays people are squeezing the computation from GPUs or FPGAs to get larger and more powerful neural networks. 
Consequently, records are being broken in a variety of Natural Language Processing (NLP) tasks ranging from classification, question answering to translation as the language models are becoming much deeper.
In the meantime, starting from the emergence of Transformer models \cite{transformer} like BERT \cite{BERT}, the trend to increase model sizes becomes even more radical.
Since then, the Moore's Law can never reach the growth of language models.
Past several years have witnessed the creation of gigantic models, like GPT-2 \cite{gpt2} and GPT-3 \cite{gpt-3}, which are able to compose texts and synthesize images.
The amazing performance is also attracting attention from industry.
% These models are now all around in our daily life, integrated in search engines, recommendation and broadcasting systems.
% They facilitate our daily life whenever we search an item in Google, browse Twitter, or shop in Amazon.
Google recently announced its Multitask Unified Model (MUM) for language understanding, LaMDA for conversation, and TPUv4 chips for model training.
% All these suggest the great demand on large language models which would last for years.

% The gap between the gradually slowed-down Moore's Law and the drastically increased model size and computational demand mainly reflects in two aspects: memory and computation.
The surging model size brings challenges in terms of both memory and computation.
Apart from developing chips such as GPU and TPU with higher computation capability and memory limit, activation checkpointing techniques are developed to squeeze a large model into a single device that would otherwise not be able to host such a big model. 
Moreover, ZeRO-infinity techniques \cite{zero,zero,rajbhandari2021zero} are proposed to utilize the memory of CPU or other memory device external to GPU without sacrificing the data transmission time.
% To fully utilize the concurrency of HPC clusters,
Data parallelism is a dominant practice to fully utilize available HPC resources.
It distributes a large minibatch to multiple devices, where each device holds an identical model replica, and finally gathers the gradients for synchronous parameter update.
With recent optimization techniques \cite{goyal2017accurate,you2017large,you2018imagenet,you2020large}, it is now able to train very large minibatches on thousands of GPU devices.
% However, studies show that large batch size would result in unstable training and degraded generalization performance.
% To address this problem, You et al propose novel optimizers facilitating the large batch training of visual and language models.


Another branch is model parallelism, the main idea of which is to divide both computation and memory of a single neural network to multiple devices. 
Pipelined parallelism \cite{pipedream,gpipe} is proposed to split the whole model by layer, which will be executed in a pipeline fashion.
Mixed precision training is utilized to reduce both memory and computation cost.
% Activation checkpointing trades computation for memory, which makes the study of huge neural networks accessible to more researchers.
Megatron-LM \cite{megatron} integrates both model parallelism and data parallelism by splitting the parameter tensor among model parallel groups, achieving the state-of-the-art language model.
Recently, a 2-dimensional matrix-matrix multiplication algorithm has been introduced into language model training \cite{optimus}, thus bridging the gap between traditional high-performance computing and machine learning.
In this work, we stride a step forward and propose a 3-D model parallelism technique to further harness the capability of GPU clusters.
We list our main contributions as follows.
\begin{itemize}
    \item We propose a 3-D model parallelism algorithm for linear operations, which can reach a perfect load balance, so as to provide the optimal efficiency;
    \item We use our 3-D model parallelism to implement an efficient Transformer model based on the existing PyTorch Transformer implementation and distributed communication package;
    \item We demonstrate the effectiveness of our 3-D model parallelism on 64 GPUs by comparing it with the existing 1-D and 2-D approaches. The results present the superiority of our 3-D parallelism with 2.32X and 1.57X speedup, respectively.
\end{itemize}
% Pipeline parallelism (PipeDream and GPipe) aims to partition a model by layer (vertically), and execute the training like a pipeline. The benefit is that, communication can be largely overlapped by computation. However, both the methods suffer from their own weaknesses. Although PipeDream elaborately arranges the forward and backward computations so that it fully utilizes all time bins without idleness, it is essentiallyan asynchronous approach, so suffers from gradient staleness. GPipe is a synchronous approach, but "bubbles" is inevitable. For extremely large networks, more GPUs should be used to host the model, and the idleness becomes more severe.

% Unlike model parallelism, ZeRO proposes to having each device host its own partition of parameters and optimizer states, which are updated locally. Parameters are gathered whenever necessary during forward and backward propagation.

% Another trend of model parallelism is to partition a model by operation (horizontally).
% Mesh-Tensorflow is a framework proposed by Google, which provides an elegant way for users to annotate the partition of tensor dimensions. Matrix multiplications can be automatically derived from these annotations, without user elaboration. However, memory redundancy is severe, and the originally nearly-free reshape operations would result in complex communication because of layout change. Megatron-LM adopts a 1-dimensional matrix partition strategy, and alternatively splits parameter tensors along rows or columns. An elegant conjunction is given to address the gradient propagation by simply wrapping the all-reduce operation. Although parameters and optimizer states are fully distributed, there is still memory redundancy for activations, as each device has to accommodate the whole input and output activations for every multi-layer perceptrons (MLPs) and self-attention modules. Optimus takes the first step addressing this issue, utilizing a 2D matrix-matrix multiplication strategy. Here, our method takes a step further, and uses a 3D approach, while also eliminating memory redundancy.

% Besides fitting existing models into parallel and distributed versions, other approaches adapt models to better capture the concurrency of modern HPC clusters. Mixture-of-experts (MoE) utilizes sparsity to discretely send each hidden state to distinct dense layers, or "experts". In this way, the amount of parameters could be easily scaled without significantly adding to computation burden. The rationale behind behind this is that more parameters result in higher performance. A lot of works achieve excellent performance based on this method.
\begin{thebibliography}{24}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
  \providecommand{\doi}[1]{doi: #1}\else
  \providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi

\bibitem[Agarwal et~al.(1995)Agarwal, Balle, Gustavson, Joshi, and
  Palkar]{agarwal1995three}
Ramesh~C Agarwal, Susanne~M Balle, Fred~G Gustavson, Mahesh Joshi, and Prasad
  Palkar.
\newblock A three-dimensional approach to parallel matrix multiplication.
\newblock \emph{IBM Journal of Research and Development}, 39\penalty0
  (5):\penalty0 575--582, 1995.

\bibitem[Brown et~al.(2020)Brown, Mann, Ryder, Subbiah, Kaplan, Dhariwal,
  Neelakantan, Shyam, Sastry, Askell, et~al.]{gpt-3}
Tom~B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla
  Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell,
  et~al.
\newblock Language models are few-shot learners.
\newblock \emph{arXiv preprint arXiv:2005.14165}, 2020.

\bibitem[Center(2021)]{longhorn}
Texas Advanced~Computing Center.
\newblock {LONGHORN - TEXAS ADVANCED COMPUTING CENTER}.
\newblock \url{https://www.tacc.utexas.edu/systems/longhorn}, 2021.

\bibitem[Devlin et~al.(2018)Devlin, Chang, Lee, and Toutanova]{BERT}
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.
\newblock Bert: Pre-training of deep bidirectional transformers for language
  understanding, 2018.

\bibitem[Fedus et~al.(2021)Fedus, Zoph, and Shazeer]{switch-tranformer}
William Fedus, Barret Zoph, and Noam Shazeer.
\newblock Switch transformers: Scaling to trillion parameter models with simple
  and efficient sparsity.
\newblock \emph{arXiv preprint arXiv:2101.03961}, 2021.

\bibitem[Goyal et~al.(2017)Goyal, Doll{\'a}r, Girshick, Noordhuis, Wesolowski,
  Kyrola, Tulloch, Jia, and He]{goyal2017accurate}
Priya Goyal, Piotr Doll{\'a}r, Ross Girshick, Pieter Noordhuis, Lukasz
  Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He.
\newblock Accurate, large minibatch sgd: Training imagenet in 1 hour.
\newblock \emph{arXiv preprint arXiv:1706.02677}, 2017.

\bibitem[Harlap et~al.(2018)Harlap, Narayanan, Phanishayee, Seshadri, Devanur,
  Ganger, and Gibbons]{pipedream}
Aaron Harlap, Deepak Narayanan, Amar Phanishayee, Vivek Seshadri, Nikhil
  Devanur, Greg Ganger, and Phil Gibbons.
\newblock Pipedream: Fast and efficient pipeline parallel dnn training, 2018.

\bibitem[Hoffer et~al.(2017)Hoffer, Hubara, and Soudry]{hoffer2017train}
Elad Hoffer, Itay Hubara, and Daniel Soudry.
\newblock Train longer, generalize better: closing the generalization gap in
  large batch training of neural networks.
\newblock In \emph{Advances in Neural Information Processing Systems}, pages
  1731--1741, 2017.

\bibitem[Huang et~al.(2019)Huang, Cheng, Bapna, Firat, Chen, Chen, Lee, Ngiam,
  Le, Wu, and Chen]{gpipe}
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia~Xu Chen, Dehao
  Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc~V. Le, Yonghui Wu, and Zhifeng Chen.
\newblock Gpipe: Efficient training of giant neural networks using pipeline
  parallelism, 2019.

\bibitem[Keskar et~al.(2016)Keskar, Mudigere, Nocedal, Smelyanskiy, and
  Tang]{keskar2016large}
Nitish~Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy,
  and Ping Tak~Peter Tang.
\newblock On large-batch training for deep learning: Generalization gap and
  sharp minima.
\newblock \emph{arXiv preprint arXiv:1609.04836}, 2016.

\bibitem[Peng et~al.(2019)Peng, Zhu, Chen, Bao, Yi, Lan, Wu, and
  Guo]{bytescheduler}
Yanghua Peng, Yibo Zhu, Yangrui Chen, Yixin Bao, Bairen Yi, Chang Lan, Chuan
  Wu, and Chuanxiong Guo.
\newblock A generic communication scheduler for distributed dnn training
  acceleration.
\newblock In \emph{Proceedings of the 27th ACM Symposium on Operating Systems
  Principles}, SOSP ’19, page 16–29, New York, NY, USA, 2019. Association
  for Computing Machinery.
\newblock ISBN 9781450368735.
\newblock \doi{10.1145/3341301.3359642}.
\newblock URL \url{https://doi.org/10.1145/3341301.3359642}.

\bibitem[Radford et~al.()Radford, Narasimhan, Salimans, and Sutskever]{gpt1}
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever.
\newblock Improving language understanding by generative pre-training.

\bibitem[Radford et~al.(2019)Radford, Wu, Child, Luan, Amodei, and
  Sutskever]{gpt2}
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya
  Sutskever.
\newblock Language models are unsupervised multitask learners.
\newblock \emph{OpenAI blog}, 1\penalty0 (8):\penalty0 9, 2019.

\bibitem[Rajbhandari et~al.(2020)Rajbhandari, Rasley, Ruwase, and He]{zero}
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He.
\newblock Zero: Memory optimizations toward training trillion parameter models,
  2020.

\bibitem[Rajbhandari et~al.(2021)Rajbhandari, Ruwase, Rasley, Smith, and
  He]{rajbhandari2021zero}
Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He.
\newblock Zero-infinity: Breaking the gpu memory wall for extreme scale deep
  learning.
\newblock \emph{arXiv preprint arXiv:2104.07857}, 2021.

\bibitem[Shazeer et~al.(2018)Shazeer, Cheng, Parmar, Tran, Vaswani,
  Koanantakool, Hawkins, Lee, Hong, Young, Sepassi, and
  Hechtman]{meshTensorflow}
Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn
  Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young,
  Ryan Sepassi, and Blake Hechtman.
\newblock Mesh-tensorflow: Deep learning for supercomputers, 2018.

\bibitem[Shoeybi et~al.(2020)Shoeybi, Patwary, Puri, LeGresley, Casper, and
  Catanzaro]{megatron}
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper,
  and Bryan Catanzaro.
\newblock Megatron-lm: Training multi-billion parameter language models using
  model parallelism, 2020.

\bibitem[Smith et~al.(2017)Smith, Kindermans, Ying, and Le]{smith2017don}
Samuel~L Smith, Pieter-Jan Kindermans, Chris Ying, and Quoc~V Le.
\newblock Don't decay the learning rate, increase the batch size.
\newblock \emph{arXiv preprint arXiv:1711.00489}, 2017.

\bibitem[van~de Geijin and Watts(1998)]{SUMMA}
Robert~A. van~de Geijin and Jerrel Watts.
\newblock Summa: scalable universal matrix multiplication algorithm.
\newblock \emph{Concurrency and Computation Practice and Experience},
  9:\penalty0 255--274, 1998.

\bibitem[Vaswani et~al.(2017)Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez,
  Kaiser, and Polosukhin]{transformer}
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones,
  Aidan~N. Gomez, Lukasz Kaiser, and Illia Polosukhin.
\newblock Attention is all you need, 2017.

\bibitem[Xu et~al.(2021)Xu, Li, Gong, and You]{optimus}
Qifan Xu, Shenggui Li, Chaoyu Gong, and Yang You.
\newblock An efficient 2d method for training super-large deep learning models,
  2021.

\bibitem[You et~al.(2017)You, Gitman, and Ginsburg]{you2017large}
Yang You, Igor Gitman, and Boris Ginsburg.
\newblock Large batch training of convolutional networks, 2017.

\bibitem[You et~al.(2018)You, Zhang, Hsieh, Demmel, and
  Keutzer]{you2018imagenet}
Yang You, Zhao Zhang, Cho-Jui Hsieh, James Demmel, and Kurt Keutzer.
\newblock Imagenet training in minutes, 2018.

\bibitem[You et~al.(2020)You, Li, Reddi, Hseu, Kumar, Bhojanapalli, Song,
  Demmel, Keutzer, and Hsieh]{you2020large}
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh
  Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh.
\newblock Large batch optimization for deep learning: Training bert in 76
  minutes, 2020.

\end{thebibliography}

\begin{thebibliography}{24}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
  \providecommand{\doi}[1]{doi: #1}\else
  \providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi

\bibitem[Agarwal et~al.(1995)Agarwal, Balle, Gustavson, Joshi, and
  Palkar]{agarwal1995three}
Ramesh~C Agarwal, Susanne~M Balle, Fred~G Gustavson, Mahesh Joshi, and Prasad
  Palkar.
\newblock A three-dimensional approach to parallel matrix multiplication.
\newblock \emph{IBM Journal of Research and Development}, 39\penalty0
  (5):\penalty0 575--582, 1995.

\bibitem[Brown et~al.(2020)Brown, Mann, Ryder, Subbiah, Kaplan, Dhariwal,
  Neelakantan, Shyam, Sastry, Askell, et~al.]{gpt-3}
Tom~B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla
  Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell,
  et~al.
\newblock Language models are few-shot learners.
\newblock \emph{arXiv preprint arXiv:2005.14165}, 2020.

\bibitem[Center(2021)]{longhorn}
Texas Advanced~Computing Center.
\newblock {LONGHORN - TEXAS ADVANCED COMPUTING CENTER}.
\newblock \url{https://www.tacc.utexas.edu/systems/longhorn}, 2021.

\bibitem[Devlin et~al.(2018)Devlin, Chang, Lee, and Toutanova]{BERT}
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.
\newblock Bert: Pre-training of deep bidirectional transformers for language
  understanding, 2018.

\bibitem[Fedus et~al.(2021)Fedus, Zoph, and Shazeer]{switch-tranformer}
William Fedus, Barret Zoph, and Noam Shazeer.
\newblock Switch transformers: Scaling to trillion parameter models with simple
  and efficient sparsity.
\newblock \emph{arXiv preprint arXiv:2101.03961}, 2021.

\bibitem[Goyal et~al.(2017)Goyal, Doll{\'a}r, Girshick, Noordhuis, Wesolowski,
  Kyrola, Tulloch, Jia, and He]{goyal2017accurate}
Priya Goyal, Piotr Doll{\'a}r, Ross Girshick, Pieter Noordhuis, Lukasz
  Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He.
\newblock Accurate, large minibatch sgd: Training imagenet in 1 hour.
\newblock \emph{arXiv preprint arXiv:1706.02677}, 2017.

\bibitem[Harlap et~al.(2018)Harlap, Narayanan, Phanishayee, Seshadri, Devanur,
  Ganger, and Gibbons]{pipedream}
Aaron Harlap, Deepak Narayanan, Amar Phanishayee, Vivek Seshadri, Nikhil
  Devanur, Greg Ganger, and Phil Gibbons.
\newblock Pipedream: Fast and efficient pipeline parallel dnn training, 2018.

\bibitem[Hoffer et~al.(2017)Hoffer, Hubara, and Soudry]{hoffer2017train}
Elad Hoffer, Itay Hubara, and Daniel Soudry.
\newblock Train longer, generalize better: closing the generalization gap in
  large batch training of neural networks.
\newblock In \emph{Advances in Neural Information Processing Systems}, pages
  1731--1741, 2017.

\bibitem[Huang et~al.(2019)Huang, Cheng, Bapna, Firat, Chen, Chen, Lee, Ngiam,
  Le, Wu, and Chen]{gpipe}
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia~Xu Chen, Dehao
  Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc~V. Le, Yonghui Wu, and Zhifeng Chen.
\newblock Gpipe: Efficient training of giant neural networks using pipeline
  parallelism, 2019.

\bibitem[Keskar et~al.(2016)Keskar, Mudigere, Nocedal, Smelyanskiy, and
  Tang]{keskar2016large}
Nitish~Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy,
  and Ping Tak~Peter Tang.
\newblock On large-batch training for deep learning: Generalization gap and
  sharp minima.
\newblock \emph{arXiv preprint arXiv:1609.04836}, 2016.

\bibitem[Peng et~al.(2019)Peng, Zhu, Chen, Bao, Yi, Lan, Wu, and
  Guo]{bytescheduler}
Yanghua Peng, Yibo Zhu, Yangrui Chen, Yixin Bao, Bairen Yi, Chang Lan, Chuan
  Wu, and Chuanxiong Guo.
\newblock A generic communication scheduler for distributed dnn training
  acceleration.
\newblock In \emph{Proceedings of the 27th ACM Symposium on Operating Systems
  Principles}, SOSP ’19, page 16–29, New York, NY, USA, 2019. Association
  for Computing Machinery.
\newblock ISBN 9781450368735.
\newblock \doi{10.1145/3341301.3359642}.
\newblock URL \url{https://doi.org/10.1145/3341301.3359642}.

\bibitem[Radford et~al.()Radford, Narasimhan, Salimans, and Sutskever]{gpt1}
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever.
\newblock Improving language understanding by generative pre-training.

\bibitem[Radford et~al.(2019)Radford, Wu, Child, Luan, Amodei, and
  Sutskever]{gpt2}
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya
  Sutskever.
\newblock Language models are unsupervised multitask learners.
\newblock \emph{OpenAI blog}, 1\penalty0 (8):\penalty0 9, 2019.

\bibitem[Rajbhandari et~al.(2020)Rajbhandari, Rasley, Ruwase, and He]{zero}
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He.
\newblock Zero: Memory optimizations toward training trillion parameter models,
  2020.

\bibitem[Rajbhandari et~al.(2021)Rajbhandari, Ruwase, Rasley, Smith, and
  He]{rajbhandari2021zero}
Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He.
\newblock Zero-infinity: Breaking the gpu memory wall for extreme scale deep
  learning.
\newblock \emph{arXiv preprint arXiv:2104.07857}, 2021.

\bibitem[Shazeer et~al.(2018)Shazeer, Cheng, Parmar, Tran, Vaswani,
  Koanantakool, Hawkins, Lee, Hong, Young, Sepassi, and
  Hechtman]{meshTensorflow}
Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn
  Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young,
  Ryan Sepassi, and Blake Hechtman.
\newblock Mesh-tensorflow: Deep learning for supercomputers, 2018.

\bibitem[Shoeybi et~al.(2020)Shoeybi, Patwary, Puri, LeGresley, Casper, and
  Catanzaro]{megatron}
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper,
  and Bryan Catanzaro.
\newblock Megatron-lm: Training multi-billion parameter language models using
  model parallelism, 2020.

\bibitem[Smith et~al.(2017)Smith, Kindermans, Ying, and Le]{smith2017don}
Samuel~L Smith, Pieter-Jan Kindermans, Chris Ying, and Quoc~V Le.
\newblock Don't decay the learning rate, increase the batch size.
\newblock \emph{arXiv preprint arXiv:1711.00489}, 2017.

\bibitem[van~de Geijin and Watts(1998)]{SUMMA}
Robert~A. van~de Geijin and Jerrel Watts.
\newblock Summa: scalable universal matrix multiplication algorithm.
\newblock \emph{Concurrency and Computation Practice and Experience},
  9:\penalty0 255--274, 1998.

\bibitem[Vaswani et~al.(2017)Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez,
  Kaiser, and Polosukhin]{transformer}
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones,
  Aidan~N. Gomez, Lukasz Kaiser, and Illia Polosukhin.
\newblock Attention is all you need, 2017.

\bibitem[Xu et~al.(2021)Xu, Li, Gong, and You]{optimus}
Qifan Xu, Shenggui Li, Chaoyu Gong, and Yang You.
\newblock An efficient 2d method for training super-large deep learning models,
  2021.

\bibitem[You et~al.(2017)You, Gitman, and Ginsburg]{you2017large}
Yang You, Igor Gitman, and Boris Ginsburg.
\newblock Large batch training of convolutional networks, 2017.

\bibitem[You et~al.(2018)You, Zhang, Hsieh, Demmel, and
  Keutzer]{you2018imagenet}
Yang You, Zhao Zhang, Cho-Jui Hsieh, James Demmel, and Kurt Keutzer.
\newblock Imagenet training in minutes, 2018.

\bibitem[You et~al.(2020)You, Li, Reddi, Hseu, Kumar, Bhojanapalli, Song,
  Demmel, Keutzer, and Hsieh]{you2020large}
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh
  Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh.
\newblock Large batch optimization for deep learning: Training bert in 76
  minutes, 2020.

\end{thebibliography}