\section{0pt}{5pt plus 4pt minus 2pt}{5pt plus 2pt minus 2pt}
\titlespacing\subsection{0pt}{5pt plus 4pt minus 2pt}{5pt plus 2pt minus 2pt}
\titlespacing\subsubsection{0pt}{5pt plus 4pt minus 2pt}{5pt plus 2pt minus 2pt}

\newcommand\pythonstyle{\lstset{
  language=Python,
  backgroundcolor=\color{white},
  basicstyle=\ttfamily\footnotesize,
  otherkeywords={self},
  keywordstyle=\footnotesize\color{deepblue},
  emph={pipeline,unroll,tile},
  emphstyle=\footnotesize\color{deepred},
  stringstyle=\color{deepgreen},
  commentstyle=\color{red},
  frame=tb,
  showstringspaces=false,
}}

\lstnewenvironment{env_python}[2]
{
\lstset{
  language=Python,
  backgroundcolor=\color{light-gray},
  otherkeywords={self},
  emph={to,tile,reuse_at,reorder,unroll,split,vectorize,quantize},
  emphstyle=\footnotesize\color{dmlred},
  belowcaptionskip=0.7\baselineskip,
  aboveskip=-2mm,
  belowskip=5mm,
  showstringspaces=false,
  columns=flexible,
  escapechar=@,
  basicstyle={\linespread{0.99}\fontencoding{T1}\footnotesize\fontfamily{lmtt}\fontseries{m}\selectfont},
  numbers={left},
  xleftmargin={2em},%
  framesep=5pt,
  framexleftmargin={1em},
  numberstyle=\tiny\color{gray},
  keywordstyle=\color{blue},
  commentstyle=\color{deepgreen},
  stringstyle=\color{mauve},
  breaklines=true,
  breakatwhitespace=true,
  tabsize=3,
  frame=tb,
  caption= #1,
  label= #2

  \vspace{\baselineskip}

{}

\lstset{
  language=Python,
  backgroundcolor=\color{aliceblue},

  emph={},
  emphstyle=\small\color{dmlred},
  belowcaptionskip=0.7\baselineskip,
  aboveskip=0mm,
  belowskip=3mm,
  showstringspaces=false,
  columns=flexible,
  basicstyle={\linespread{1.1}\fontencoding{T1}\scriptsize\fontfamily{lmtt}\fontseries{m}\selectfont},
  numbers={left},
  xleftmargin={2em},%
  numberstyle=\tiny\color{gray},
  keywordstyle=\color{blue},
  commentstyle=\color{deepgreen},
  stringstyle=\color{mauve},
  breaklines=true,
  breakatwhitespace=true,
  tabsize=3,
  morecomment=[f][\lstbg{red!20}]-,
  morecomment=[f][\lstbg{green!20}]+,
  morecomment=[f][\textit]{@@},
 }

\newcommand{\lstbg}[3][0pt]{{\fboxsep#1\colorbox{#2}{\strut #3}}}
\lstdefinelanguage{diff}{
  backgroundcolor=\color{aliceblue},
  emph={},
  emphstyle=\small\color{dmlred},
  belowcaptionskip=0.7\baselineskip,
  aboveskip=0mm,
  belowskip=3mm,
  showstringspaces=false,
  columns=flexible,
  basicstyle={\linespread{1.1}\fontencoding{T1}\scriptsize\fontfamily{lmtt}\fontseries{m}\selectfont},
  numbers={left},
  xleftmargin={2em},%
  breaklines=true,
  breakatwhitespace=true,
  tabsize=3,
  morecomment=[f][\lstbg{red!20}]-,
  morecomment=[f][\lstbg{green!20}]+,
  morecomment=[f][\textit]{@@},
}

\lstnewenvironment{env_cpp}[2]
{
\lstset{
  language=C++,
  backgroundcolor=\color{gray!10},
  otherkeywords={ap_int},
  escapeinside={<@}{@>},
  escapechar=!,
  keywordstyle=\ttb\color{deepblue},%
  emph={},
  emphstyle=\footnotesize\color{dmlred},
  belowcaptionskip=0.7\baselineskip,
  aboveskip=-2mm,
  belowskip=5mm,
  showstringspaces=false,
  columns=flexible,
  basicstyle={\linespread{1.1}\fontencoding{T1}\footnotesize\fontfamily{lmtt}\fontseries{m}\selectfont},
  numbers={left},
  xleftmargin={2em},%
  framexleftmargin={1em},
  numberstyle=\tiny\color{gray},
  keywordstyle=\color{blue},
  commentstyle=\color{deepgreen},
  stringstyle=\color{mauve},
  breaklines=true,

  showspaces = false,
  showstringspaces = false,
  tabsize=3,
  frame=tb,
  caption= #1,
  label= #2

  \vspace{\baselineskip}

{}

\lstdefinestyle{custom_python}
{%
  language=Python,%
  alsoletter={.},
  basicstyle={\small\ttm},%
  deletekeywords={type}, %
  keywordstyle=\ttb\color{deepblue},%
  commentstyle={\color{dmlgreen}},
  stringstyle=\color{deepgreen},
  morekeywords={with, as},
  emph={placeholder},
  emphstyle=\ttb\color{deepred},
  lineskip={-0.01in},%
  numbers={left},%
  numbersep={10pt},%
  numberstyle={\small},%
  xleftmargin={2em},%
  showstringspaces={false},%
  keepspaces={true},%
  escapechar={\$},

  columns={fullflexible},
  frame=tb
}%

\lstdefinestyle{custom_c}
{%
  language=C++,%
  alsoletter={.\#=},
  basicstyle={\ttfamily},%
  deletekeywords={type}, %
  keywordstyle={\bfseries\color{dmlblue}},%
  keywordstyle=[2]{\color{dmlred}},%
  keywordstyle=[3]{\color{dmlgreen}},%
  commentstyle={\color{dmlgreen}},
  stringstyle={\color{brown}},
  morecomment=[l][\color{dmlred}]{\#pragma},
  lineskip={-0.01in},%
  numbers={left},%
  numbersep={5pt},%
  numberstyle={\small},%

  showstringspaces={false},%
  keepspaces={true},%

  columns={fullflexible},
  gobble={4}
}%

\newcommand{{Slapo}\xspace}{{Slapo}\xspace}
\newcommand{\texttt{torch.fx}\xspace}{\texttt{torch.fx}\xspace}

\newcommand{yes}{yes}
\newcommand\fixme[1]{
    \ifthenelse{\equal{yes}{yes}}{\textcolor{red}{[#1]}}{\ignorespaces}
}
\newcommand\hz[1]{
    \ifthenelse{\equal{yes}{yes}}{\textcolor{red}{[hz: #1~]}}{\ignorespaces}
}
\newcommand\cody[1]{
    \ifthenelse{\equal{yes}{yes}}{\textcolor{blue}{[cody: #1~]}}{\ignorespaces}
}
\newcommand\shuai[1]{
    \ifthenelse{\equal{yes}{yes}}{\textcolor{green}{[shuai: #1~]}}{\ignorespaces}
}
\newcommand\yida[1]{
    \ifthenelse{\equal{yes}{yes}}{\textcolor{purple}{[yida: #1~]}}{\ignorespaces}
}
\newcommand\zz[1]{
    \ifthenelse{\equal{yes}{yes}}{\textcolor{violet}{[zhiru: #1]}}{\ignorespaces}
}

\begin{document}

\date{}

\title{\Large \bf Decoupled Model Schedule for Deep Learning Training}

\author{
{\rm Hongzheng Chen\thanks{Work done while interning at Amazon.}} \\
Cornell University \\
<EMAIL>
\and
{\rm Cody Hao Yu} \\
Amazon Web Services \\
<EMAIL>
\and
{\rm Shuai Zheng} \\
Amazon Web Services \\
<EMAIL>
\and
{\rm Zhen Zhang} \\
Amazon Web Services \\
<EMAIL>
\and
{\rm Zhiru Zhang} \\
Cornell University \\
<EMAIL>
\and
{\rm Yida Wang} \\
Amazon Web Services \\
<EMAIL>
}

\maketitle

\input{sections/0-sec-abstract.tex}
\input{sections/1-sec-introduction.tex}
\input{sections/2-sec-background.tex}
\input{sections/3-sec-design.tex}
\input{sections/4-sec-implementation.tex}
\input{sections/5-sec-experiments.tex}
\input{sections/7-sec-related-work.tex}
\input{sections/8-sec-conclusion.tex}
\newpage

\bibliographystyle{plain}

\subsection{Schedule Modules and Parameters}
\label{sec:sch-mod}

We first present scheduling a module and its parameters, which typically does not change the computation and thus does not require static graphs.

\subsubsection{Schedule Modules}
\label{sub:sch-mod}

For important workloads such as \texttt{attention} in Fig.~\ref{fig:schedule_mod}(a), researchers or hardware vendors may manually implement efficient kernels. These highly customized, hand-crafted kernels sometimes could outperform the ones generated by DL compilers.
With {Slapo}\xspace, we can use \texttt{.replace(new\_module)} primitive to replace a native implementation with an efficient one, where \texttt{new\_module} is the custom module to be replaced, as shown in Fig.~\ref{fig:schedule_mod}(b).

In addition to replacing a module with an efficient implementation, activation checkpointing is another important feature for large model training, as mentioned in \S~\ref{sub:training}.
{Slapo}\xspace offers \texttt{.checkpoint()} primitive that wraps a module with a checkpointed module.
This primitive allows developers to add activation checkpointing without changing the model implementation.
Consequently, compared to existing activation checkpointing implemented in DL models with fixed strategies, {Slapo}\xspace enables developers to flexibly adjust the number of checkpoints via our schedule primitive or leverage the auto-tuner for better memory and throughput trade-offs.

\subsubsection{Parameter Sharding}
\label{sub:sch-mod-sharding}

In \Circled{3} of \S~\ref{sub:example}, we introduced the steps to enable tensor parallelism -- shard parameters and aggregate the outputs. With {Slapo}\xspace, we can shard a parameter with \texttt{.shard(name, axis)} primitive, and aggregate the results using \texttt{.sync(type)} primitive.
The \texttt{type} can be ``forward'' (aggregate the forward activations), ``backward'' (aggregate the gradients), or ``both''.

With the considerations of flexibility and comprehensive schedule primitives, instead of implicitly aggregating the outputs, {Slapo}\xspace exposes \texttt{.sync()} primitive to developers to determine a proper aggregation point with an optionally customized aggregation function. Meanwhile, we leverage the verifier (\S~\ref{sec:verifier}) to check the correctness after the scheduling. In the future, we plan to develop an auto-scheduler that automatically generates these primitives.

We again use Fig.~\ref{fig:schedule_mod} to illustrate the advantage of flexible output aggregation points. In Fig.~\ref{fig:schedule_mod}(c), we shard the parameters of a \texttt{QKV} module along the axis of output features (\texttt{.shard(["weight", "bias"], 0)}).
In this case, each device computes a part of the output.
Since we also shard the weights of its subsequent linear module, \texttt{OUT}, along input features, each device only takes a part of inputs for computation.
As a result, we can defer the output aggregation after \texttt{OUT} to reduce memory footprint.

\section{Implementations}
\label{sec:impl}

We implemented {Slapo}\xspace with $\sim$3K LoC in Python on top of PyTorch~\cite{paszke2019pytorch} to benefit from its dynamic graph and usability. In this section, we highlight some implementation details.

\noindent{\textbf{Static Graph Tracing.}} Our tracer and the static graph are based on \texttt{torch.fx}\xspace~\cite{reed2022fx}, which captures PyTorch models via symbolic tracing and constructs a static graph with a 6-instruction Python-based intermediate representation (IR). Like other DL compilers, \texttt{torch.fx}\xspace tracer also has unsupported coding styles when capturing the graph, such as certain types of control flow and data structures. It also flattens the IR and discards all the model structural hierarchy.
Thus, instead of simply invoking \texttt{torch.fx}\xspace tracer from the top module, we invoke the tracer module by module and carefully maintain the hierarchy.
When a particular module cannot be traced by \texttt{torch.fx}\xspace, developers can simply disable the corresponding schedule primitives that require a static graph while the rest primitives can still be applied.

\noindent{\textbf{Framework Dialects.}} {Slapo}\xspace scheduled model is by design compatible with native PyTorch runtime and can be executed on PyTorch directly.
To integrate with the dedicated runtime of SOTA distributed systems for certain parallelism (e.g., pipeline), we also implemented two framework dialects for Megatron-LM~\cite{shoeybi2019megatron} and DeepSpeed~\cite{rasley2020deepspeed}.
These systems require the model to be in a certain form or wrapped by their custom module.
For example, DeepSpeed~\cite{rasley2020deepspeed} pipeline runtime requires a single tuple of inputs and outputs in each pipeline stage, so our DeepSpeed pipeline stage module includes the logic to (1) unpack the inputs and pack the outputs, and (2) perform liveness analysis to bypass the tensors required by subsequent stages.

\section{Evaluation}
\label{sec:exp}

In this section, we evaluate {Slapo}\xspace on different training configurations, in terms of the number of GPUs, number of instances, and parallelism strategies, to demonstrate {Slapo}\xspace is able to align or even outperform existing solutions while preserving usability.
Note that {Slapo}\xspace does not change the semantics of models and optimizers, so model convergence is not affected.
We also provide ablation studies to show the effectiveness of the schedule primitives and the auto-tuner.

\noindent\textbf{Setups.}
All experiments are conducted on Amazon EC2 p3 instances.
More specifically, we use p3.16xlarge instances with 8 NVIDIA V100 16GB GPUs for single-device and single-node evaluations, and use p3dn.24xlarge instances with 8 NVIDIA V100 32GB GPUs for multi-node evaluations.
GPUs in these instances are connected via NVLink, which provides 300 GB/s theoretical aggregated GPU interconnect bandwidth, and the inter-node bandwidth is 100 Gbps.
The software environment includes CUDA 11.7, PyTorch v1.13, Megatron-LM (git-hash 0bb597b), DeepSpeed (customization from v0.6.2 with pipeline parallelism fix to match Megatron-LM pipeline performance), and NCCL v2.14.3.

\noindent\textbf{Models and Metrics.}
We apply schedules to a set of popular PyTorch models from HuggingFace Hub~\cite{wolf2019huggingface} and torchvision~\cite{torchvision2016}, covering encoder-/decoder-based Transformer models and image classification models to demonstrate the generality of {Slapo}\xspace.
Detailed model settings are shown in Table~\ref{tab:models}.
All models are trained by AdamW optimizer~\cite{loshchilov2017decoupled} with mixed precision, and the micro-batch size (i.e., the number of samples per data parallel rank) is selected based on the memory footprint maximizing the system performance.
We use the training throughput (the number of total processed samples per second) as our evaluation metric.
For each setting, we train the models for tens of steps and take the average throughput after discarding the first few warm-up steps.

\input{figs/models}

\subsection{Evaluation on a Single GPU}
\label{sub:single_device_exp}
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{figs/exp/single_device_v100.pdf}
\caption{End-to-end throughput on a single NVIDIA V100 16GB GPU. $\times$ denotes unsupported implementations.}
\label{fig:single_v100}
\end{figure}

\begin{figure*}[t]
\centering
\includegraphics[width=\linewidth]{figs/exp/single_node_v100.pdf}
\caption{Training throughput on Amazon EC2 p3.16xlarge instance with 8 NVIDIA V100 16GB GPUs. ``{Slapo}\xspace-TP'' uses tensor parallelism to align Megatron-LM. ``{Slapo}\xspace-ZeRO3'' uses DeepSpeed ZeRO-3 as the parallelism technique.}
\label{fig:8_v100}
\end{figure*}

We first conduct end-to-end training efficiency evaluations on a single GPU to demonstrate that {Slapo}\xspace is generally applicable to various model implementations.

\noindent\textbf{Systems.}
We use PyTorch Eager mode~\cite{paszke2019pytorch} as the baseline, which provides out-of-the-box performance.
If activation checkpointing is implemented in a model, we evaluate the performance with and without activation checkpointing, and report the better one.
When applicable, we also include the performance result with TorchScript~\cite{torchscript}, a commonly-used compiler framework in PyTorch, as another baseline.
Since PyTorch 1.12, TorchScript adopts nvFuser~\cite{nvfuser} as the operator fusion mechanism, and supports training workloads.
For {Slapo}\xspace, we schedule models with efficient kernels, operator fusion, and activation checkpointing, and leverage the auto-tuner proposed in \S~\ref{sec:autotune} to search for the best configuration.

\noindent\textbf{Results.}
Fig.~\ref{fig:single_v100} shows the throughput of each configuration. In most cases, {Slapo}\xspace can consistently improve the training throughput.
Compared to Eager mode, {Slapo}\xspace achieves 1.05$\times$-2.11$\times$ speedup.
Even compared to TorchScript-optimized models, {Slapo}\xspace still achieves 1.45$\times$ speedup on average.
This is mainly attributed to {Slapo}\xspace employing the similar operator fusion but with more optimizations such as activation checkpointing and efficient custom kernels.
Note that although {Slapo}\xspace's fusion capability is limited by module boundaries for structure-preserving scheduling, most performance bottleneck subgraphs do not cross modules in training workloads.

Furthermore, {Slapo}\xspace enables more optimization opportunities than TorchScript by progressive tracing.
For example, the GPT model implementation we adopted is GPT-Neo~\cite{gao2020pile} from the HuggingFace hub. This implementation, unfortunately, cannot be optimized by TorchScript due to its coding style, thus showing no results in Fig.~\ref{fig:single_v100}.
However, {Slapo}\xspace allows us to selectively trace the FFN (Feed-Forward Network) module but not others, resulting in effective optimization.

In addition, {Slapo}\xspace tunes the ratio of activation checkpointing to maximize GPU memory utilization.
For example, after tuning the ratio of activation checkpointing in BERT, we only checkpoint 25\% of attention layers.
This brings another 1.06$\times$ speedup compared to checkpointing all layers.

\subsection{Evaluation on Multiple GPUs}
\label{sec:eval:multi-gpu}
This subsection evaluates the end-to-end training efficiency on 2, 4, and 8 NVIDIA V100 16GB GPUs in a single p3.16xlarge instance to showcase the effectiveness of {Slapo}\xspace.

\noindent\textbf{Systems.}
We select Megatron-LM v2~\cite{narayanan2021megatronv2} as the baseline, which is a SOTA system built on top of PyTorch for training large Transformer-based language models on GPUs.
Megatron-LM implements its own data loader, tokenizer, and optimizer for better training efficiency.
In addition, it implements popular Transformer models with tensor parallelism as well as efficient customized CUDA kernels.
We also choose DeepSpeed~\cite{rasley2020deepspeed} as another baseline. DeepSpeed is a SOTA framework that incorporates ZeRO-powered data parallelism (ZeRO-DP)~\cite{rajbhandari2020zero}.
Since ZeRO-DP
is applicable to arbitrary PyTorch models, DeepSpeed has been widely used to train large DL models.
As {Slapo}\xspace is agnostic to parallelism strategies, we evaluate two configurations for every model to show that {Slapo}\xspace is compatible with the existing distributed training frameworks.
Specifically, ``{Slapo}\xspace-ZeRO3'' schedules models with ZeRO-3~\cite{rajbhandari2020zero} that automatically partitions optimizer states, gradients, and parameters to enable memory-efficient data parallelism; while ``{Slapo}\xspace-TP'' schedules them to enable tensor parallelism.
For each configuration, we auto-tune the activation checkpointing along with the batch size.

\noindent\textbf{Results.}
We first compare two baselines, Megatron-LM and DeepSpeed ZeRO-3, in Fig.~\ref{fig:8_v100}.
It is worth noting that Megatron-LM officially only supports three (i.e., BERT, GPT, and T5) out of the seven models listed in Table~\ref{tab:models}.
Comparing these three models, we find that
there is no one solution that is always superior to the other, highlighting the importance of {Slapo}\xspace which enables developers to easily implement the best parallelism strategies using schedules for different models.

We now compare the results of our approach, {Slapo}\xspace, to the baselines. In Fig.~\ref{fig:8_v100}, for the models supported by Megatron-LM, {Slapo}\xspace-TP achieves 87\% to 103\% throughput on 8 GPUs. We have investigated the performance gap of up to 13\% between those three models and concluded that it is mainly due to different model implementations between Megatron-LM and HuggingFace. For instance, Megatron-LM T5 adopts a fixed embedding instead of relative position embedding in self-attention layers. This avoids the overhead of computing relative position embedding, resulting in higher throughput.

The model difference is eliminated when comparing {Slapo}\xspace-ZeRO3 against DeepSpeed since both run the same HuggingFace models.
With this premise, we can see that {Slapo}\xspace-ZeRO3 consistently outperforms DeepSpeed by a margin of 1.08$\times$-3.35$\times$.
Similar to the single GPU evaluation, the speedup mainly comes from operator fusion, efficient custom kernels, and tuned activation checkpointing ratio.

\begin{comment}
\cody{remove candidate.}
Moreover, for GPT on 2 GPUs, both DeepSpeed and {Slapo}\xspace-ZeRO3 have no results due to out-of-memory, but Megatron-LM and {Slapo}\xspace-TP have.
This is because we adopt a GPT model with 1.3 billion parameters in this evaluation. Even with ZeRO-3 that partitions parameters and optimizer states, each GPU still consumes $\sim$10~GB memory before kicking off the training process, so it is easy to run out of 16~GB device.
On the other hand, while data parallelism requires the minimal micro-batch size to be the number of GPUs (i.e., 2 in this configure), tensor parallelism does not have this requirement and is capable of supporting an even smaller micro-batch size (i.e., 1). Thus, tensor parallelism is a more desirable strategy than ZeRO-3 in this configuration.
For OPT on 8 GPUs, we observe that enabling the activation checkpointing for all attention layers in {Slapo}\xspace-ZeRO3 results in 3.03$\times$ speedup over DeepSpeed with the same configuration. However, after we use auto-tuner to realize the best checkpoint ratio (i.e., 50\%), {Slapo}\xspace-ZeRO-3 ultimately achieves 3.35$\times$ speedup.
Detailed evaluation on auto-tuning will be discussed in \S~\ref{sub:ablation}.
\end{comment}

\subsection{Evaluation on Multiple Machines}
This subsection presents the results of distributed training performance in multi-machine setups.

\noindent\textbf{Systems and Setups.} The testbed in this evaluation is p3dn.24xlarge instances with 8 NVIDIA V100 32 GB GPUs each.
We again use DeepSpeed and Megatron-LM v2, as baselines.
Following the common practices~\cite{rajbhandari2020zero, narayanan2021megatronv2}, we use ZeRO-3 for DeepSpeed and set model-parallel and pipeline-parallel size as 8 and 2, respectively, for Megatron-LM.
We consider the strong scaling efficiency, in which the global batch size is fixed.
Typically, distributed training for large models runs on hundreds to thousands of GPUs and uses global batch sizes up to thousands out of consideration for model quality~\cite{narayanan2021megatronv2, brown2020language,chowdhery2022palm}.
We fix the global batch size at 256 for our clusters with up to 64 GPUs, and tune the micro-batch size of each system for the best performance.
The evaluation uses a GPT model~\cite{radford2019gpt2} with 10 billion parameters. The input sequence length for the experiments is 1024.

\begin{figure}[t]
\centering
\includegraphics[width=0.95\linewidth]{figs/exp/multi-node.pdf}
\caption{Training throughputs of a GPT model with 10 billion parameters on up to 64 GPUs.}
\label{fig:multinode}
\end{figure}

\noindent\textbf{Results.}
First of all, Fig.~\ref{fig:multinode} again shows that there is no one parallelism strategy that always performs the best.
On the other hand, {Slapo}\xspace consistently outperforms the best available baseline.
This is because {Slapo}\xspace does not tie with a particular parallelism strategy or framework, meaning that we can pick whatever performs better and schedule additional effective optimizations as we will show in the ablation study.

\subsection{Ablation Study}
\label{sub:ablation}
We design an ablation study for applied optimizations to investigate the performance gain of the schedule.

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{figs/exp/ablation.pdf}
\caption{Ablation study with HuggingFace BERT model.
\vspace{-1em}
\label{fig:ablation}
\end{figure}

\noindent\textbf{Schedule Primitives.}
We start from a vanilla HuggingFace BERT model that is only capable of running on a single device.
As shown in Fig.~\ref{fig:ablation}, by progressively applying the schedule primitives to the model, we see a consistent increase in performance.
Kernel optimizations, such as the use of Flash Attention~\cite{dao2022flashattention} and fusing the \texttt{Bias-GeLU} kernel, have shown a 1.09$\times$ speedup on a single device.
Enabling selective checkpoints allows for a larger batch size, resulting in an additional 7\% speedup.
With the \texttt{attention} and \texttt{ffn} module sharded as discussed in Fig.~\ref{fig:schedule_mod}, we can scale out to 8 V100 GPUs and achieve 3.25$\times$ speedup.
Furthermore, the word embedding layer contains parameters regarding the vocabulary size (>3K). Sharding this parameter results in large batch sizes, leading to a final 4.02$\times$ speedup and similar scalability as Megatron-LM BERT in Fig.~\ref{fig:8_v100}.

\noindent\textbf{Auto-Tuning.}
We then investigate the efficiency of {Slapo}\xspace auto-tuner.
We exhaustively evaluate a predefined search space (including white and yellow regions in Fig.~\ref{fig:polygon}) with 91 configurations composed of various batch sizes and activation checkpointing ratios for OPT on 8 V100 GPUs.
As shown in Fig.~\ref{fig:autotune}, the optimal configuration only checkpoints 50\% of the layers with the maximum batch size below the memory threshold.
With the search space that has already pruned many inefficient configurations, as well as the coordinate descent algorithm deployed by the {Slapo}\xspace auto-tuner, only 17 configurations (19\% of the total candidates) are explored to obtain the optimal one with the best throughput.

\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{figs/exp/autotune.pdf}
\caption{Auto-tuning an OPT model. The contour line shows the training throughput of different combinations of batch size and checkpoint ratio. Throughput 0 means OOM. The purple \textcolor{purple}{$\star$} shows the explored configurations by coordinate descent, and red \textcolor{red}{$\times$} depicts the optimal one.}
\label{fig:autotune}
\end{figure}

\subsection{Schedule Computations}
\label{sec:sch-compute}

\begin{figure}[t]
    \centering
    \includegraphics[width=\linewidth]{figs/example-2.pdf}
    \caption{An example of scheduling computation of a BERT model. For illustration proposes, we do not show the entire IR of the traced \texttt{eff\_attn} module, but depict it in an identical \texttt{forward} function. The bias of the new \texttt{OUT} module is set as \texttt{None}, which has been fused into the \texttt{LN} module.}
    \label{fig:schedule_compute}
\end{figure}

The prerequisite of scheduling computations is tracing the \texttt{forward} method of the target module, and constructing a static graph in a certain intermediate representation (IR).
There are several approaches to obtaining the static graph IR. First, \textit{run-with-dummy-data}~\cite{torchscript} is an approach that directly executes the method with dummy inputs and captures all executed operators in order.
Second, \textit{AST-analysis}~\cite{torchscript} directly analyzes Python abstract-syntax-tree (AST) to obtain the static graph.
Third, \textit{just-in-time (JIT)}~\cite{suhan2021lazytensor,jax2018github} captures the execution graph in every training iteration, compiles it once, and reuses it in the rest process.
Finally, \textit{bytecode-analysis}~\cite{torchdynamo} hooks into the Python frame evaluation to construct the graph from Python bytecode.

In {Slapo}\xspace, we define \texttt{.trace(leaves, flatten)} primitive on top of all approaches. This primitive lets developers configure the granularity and the form of a traced graph. Specifically, \texttt{leaves} indicates the submodules we will not trace into, and \texttt{flatten} indicates whether to flatten a traced static graph.
By default, the predefined modules (e.g., \texttt{nn.Linear}) in DL frameworks are all leaves, and we create the static graph in a hierarchical way.
The specification is then passed to the underlying tracing engine, and the traced module and submodules become static graphs so that compiler-related primitives can be enabled.
We show a traced BERT attention module in Fig.~\ref{fig:schedule_compute}.
In the rest of this subsection, we discuss the scheduling with static graphs.

\subsubsection{Partial Computation Scheduling}
\label{sub:sch-compute-partial}

The computation latency of a module is usually dominated by a part of its computational logic, such as \texttt{attn} and \texttt{output} in Fig.~\ref{fig:schedule_compute} (highlighted in yellow and green).
As a result, it is effective if developers can offload the performance bottleneck logic to efficient kernels or DL compilers.
Since most DL models usually have repetitive layers~\cite{devlin2018bert,zagoruyko2016wide}, {Slapo}\xspace offers a \texttt{.find(regex\_or\_pattern\_fn)} primitive that performs pattern matching algorithm based on subgraph isomorphism~\cite{eppstein1995isomorphism}, with user-specified regular expression or a function with an identical subgraph.
This helps find all target subgraphs at once.
These subgraphs then can be scheduled in the same fashion with operator fusion, partial computation replacement, activation checkpointing, etc.

\noindent\textbf{Operator Fusion.} Operator fusion is an important optimization technique, as it can obviously reduce the data transfer and kernel invocation overheads to improve the latency, throughput, and memory footprint.
Existing DL compilers~\cite{xla,chen2018tvm,baghdadi2019tiramisu,torchscript} usually incorporate pattern- and rule-based fusion strategies. The former leverages a predefined pattern set to fuse matched subgraphs; while the latter traverses the dataflow graph and fuses all valid operators satisfying predefined rules together.
As shown in Fig.~\ref{fig:schedule_compute}(a), {Slapo}\xspace takes the advantage of this feature from DL compilers by defining the \texttt{.fuse(compiler, subgraph)} primitive, where \texttt{compiler} indicates the DL compiler that will be used to generate a fused kernel of the subgraph.
We currently support a pattern-based fusion strategy and use TorchScript~\cite{torchscript} as the DL compiler to improve the performance of a kernel.

\noindent\textbf{Partial Computation Replacement.} In addition to DL compilers, when the subgraph is the performance bottleneck in widely used models, developers may manually implement an efficient kernel and encapsulate it in a module. If this custom kernel achieves better performance than the one generated by a DL compiler, developers are capable of using \texttt{.replace(new\_mod, subgraph)} primitive to directly replace the corresponding computation logic with the custom one, such as Fig.~\ref{fig:schedule_compute}(b).
The decision of leveraging a DL compiler or a custom kernel can be made by developers with a one-line change. Developers can also rely on {Slapo}\xspace auto-tuner (\S~\ref{sec:autotune}) to realize a better one.

\noindent\textbf{Partial Activation Checkpointing.} While we have introduced \texttt{.checkpoint()} in \S~\ref{sec:sch-mod} that enables activation checkpointing for an entire module, {Slapo}\xspace also allows developers to use \texttt{.checkpoint(subgraph)} primitive to only checkpoint a subgraph in the computation, providing a fine-grained control for performance-memory trade-off dilemma, which has been studied by many existing works~\cite{chen2016checkpoint,marisa2021dtr,jain2020checkmate}.

\begin{comment}
Operator fusion is an important optimization technique that combines several operators into one to reduce memory access overheads.
Traditional deep learning compilers~\cite{xla,chen2018tvm} can only tackle fixed pattern fusion with explicit data dependency like Conv2d$+$BatchNorm, known as \emph{vertical fusion}; while recent research shows fusing operators in different paths, called \emph{horizontal fusion}, can also achieve performance speedup~\cite{ding2021ios,tensorrt}.
To support both fusion techniques, our system first allows users to identify the target operators to be fused using pattern matching.
The system will then take in the operations and form a valid subgraph for code generation.
Users can either provide a predefined fused kernel or directly call an external compiler like TorchInductor~\cite{inductor} to generate the kernel.
This approach is flexible enough to handle various fusion cases and provides a user-friendly interface for defining custom fusion rules.

For the BERT case in Fig.~\ref{fig:schedule}, users only need to write one line of code to find out all the \texttt{query}, \texttt{key}, \texttt{value} layers, and pass them to the schedule for fusion.
After fusion, the three \texttt{nn.Linear} layers will be replaced with a single \texttt{FusedQKV} layer.
The corresponding structure and dataflow representation both depict the changes.
In practice, we also fuse patterns like Bias$+$GeLU~\cite{shoeybi2019megatron}, which uses numerical approximation to calculate the compute-intensive kernel.
We use TorchScript~\cite{torchscript} for JIT compilation in this case to achieve better performance.
\end{comment}

\subsubsection{Pipeline Partitioning}
\label{sub:sch-compute-pipeline}

In \S~\ref{sec:sch-mod}, we demonstrated that {Slapo}\xspace can achieve tensor parallelism using parameter sharding at the module level. However, it is not possible to achieve pipeline parallelism with the same approach, as it requires rewriting the top module, including its \texttt{forward} method, to be a sequence of submodules.

SOTA dynamic graph-based DL frameworks support pipeline parallelism in two steps. First, unlike the native runtimes of DL frameworks that use a single process to execute the model graph, pipeline parallelism requires launching one process per pipeline stage. Therefore, DL frameworks with pipeline parallelism support must provide their own runtime. Second, the model must be rewritten to follow a specific API convention. The rewritten implementation consists of a sequence of submodules, with outputs connecting to inputs between two consecutive submodules. This allows the DL framework to assign each module to a stage for execution. However, this process can be tedious for developers to prepare a model for pipeline parallelism.

A recent work, PiPPy (Pipeline Parallelism for PyTorch)~\cite{pippy2022}, attempts to address this challenge by tracing the entire model to a static graph and partitioning the graph into a series of modules based on user annotations. However, this approach has limitations, as tracing the entire model can suffer from the limitations of the graph tracer, as discussed in \S~\ref{sec:intro}. If a part of the model cannot be transformed into a static graph, the entire model cannot be partitioned. In contrast, {Slapo}\xspace allows developers to configure the granularity and the form of the traced graph, meaning that developers can choose to only transform a few top-level modules into a static graph and use the \texttt{.pipeline\_split()} primitive to annotate the pipeline stage boundaries.

We use the example in Fig.~\ref{fig:pipeline} for illustration.
To evenly split a BERT model with 24 attention layers in its encoder into two partitions, we can use \texttt{.pipeline\_split()} primitive\footnote{Where to insert pipeline splits to achieve optimal throughput is out of scope in this paper, but developers can use auto-tuner in \S~\ref{sec:autotune} for exploration.} to annotate a stage boundary between layer 11 and 12 (0-based) in Fig.~\ref{fig:pipeline}(a).
In this case, only \texttt{encoder} module has to be traced, but not its submodules (e.g., \texttt{attention}) or siblings (e.g., \texttt{embeddings} and \texttt{pooler}).
We note that the untraceable, complex computation logic usually lies in core building block modules (e.g., \texttt{attention}), so limiting the tracing granularity makes our pipeline partitioning more applicable.

\begin{figure}[t]
    \centering
    \includegraphics[width=0.95\linewidth]{figs/pipeline.pdf}
    \caption{An example of partitioning a BERT model into two pipeline stages.}
    \label{fig:pipeline}
\end{figure}

However, since {Slapo}\xspace preserves the model structure hierarchy, it is non-trivial to partition a model into a sequence of modules with user-specified pipeline annotations.
Specifically, if we simply partition the model based on annotations as Fig.~\ref{fig:pipeline}(a), we fail to include \texttt{embeddings} and \texttt{pooler} modules in BERT architecture.
To generate the correct partitions, we propose an algorithm that propagates the pipeline annotations from the annotated submodule to the top module, so that all ancestor and descendant modules at all levels can be included. The algorithm is shown as follows.

\smallskip
\begin{lstlisting}[language=Python]
def partition(sch, common_parent_sch):
    seq_modules = partition_by_annotation(sch.module)
    parent_mod = sch.parent.module
    inline_and_annotate(seq_modules, parent_mod)
    if parent_mod != common_parent_sch.module:
        partition(parent_mod, common_parent_sch)

sub_schs = sch.get_all_sub_schedules_with_pipeline_split()
common_parent_sch = find_common_parent(sub_schs)
for sub_sch in schs:
    partition(sub_sch, common_parent_sch)
partition(common_parent_sch, None)
\end{lstlisting}

We first retrieve all the subschedules whose children have pipeline partition annotation (L8). For the example in Fig.~\ref{fig:pipeline}, only a single subschedule \texttt{sch["bert.encoder"]} is returned. We then find the common parent module (L9, \texttt{bert} for Fig.~\ref{fig:pipeline}), and partition each submodule with the annotations (L10-11). After partitioning the current module (L2), we inline the partitioned module sequence and propagate the pipeline split annotations to its parent module (L4) so that the parent module now also has the annotations.
We perform this process recursively until the common parent module. At this point, all submodules have been partitioned and inlined, and the common parent module is not partitioned yet but is annotated by its children. Finally, we partition from the common parent module (L12) to the top module to finish the process as depicted in Fig.~\ref{fig:pipeline}(b).

\section{{Slapo}\xspace Design}
\label{sec:name}

\begin{figure}[!htbp]
\centering
\includegraphics[width=0.9\linewidth]{figs/overview.pdf}
\caption{Overview of {Slapo}\xspace.}
\label{fig:overview}
\end{figure}

This section presents the design of {Slapo}\xspace, our schedule language to progressively optimize DL model training using proposed primitives.
{Slapo}\xspace decouples model definition from its training execution strategy for better portability.
{Slapo}\xspace also abstracts out the common optimization techniques using a set of primitives to apply (or un-apply) one by one, lowering the bar for model developers to try out different performance optimization ideas.
Furthermore, {Slapo}\xspace makes it possible to automate the performance tuning via hyperparameter search.

Fig.~\ref{fig:overview} illustrates the overview of {Slapo}\xspace.
{Slapo}\xspace accepts a deep learning model implementation in a DL framework with dynamic graphs (e.g., PyTorch~\cite{paszke2019pytorch}) and parses the original model execution as its default schedule.
Then, developers make use of the schedule primitives for optimizations on top of the default schedule.
We define the schedule language in \S~\ref{sec:sch}, and present the scheduling in various scenarios in \S~\ref{sec:sch-mod} and \S~\ref{sec:sch-compute}.
The scheduling strategy can be auto-tuned to search for a configuration that achieves the best performance (\S~\ref{sec:autotune}).
Meanwhile, {Slapo}\xspace adopts a verifier (\S~\ref{sec:verifier}) to ensure the functional correctness of all schedules.
After the schedule is determined and applied, the scheduled model can be trained on the runtime of the original DL framework (e.g., PyTorch), or if needed, on the dedicated runtime of existing distributed systems such as DeepSpeed~\cite{rasley2020deepspeed} pipeline.

\input{sections/3.1-sec-sch-design.tex}
\input{sections/3.2-sec-sch-mod.tex}
\input{sections/3.3-sec-sch-compute.tex}
\input{sections/3.4-sec-autotune.tex}
\input{sections/3.5-sec-verifier.tex}

\subsection{Auto-Tuning}
\label{sec:autotune}

Decoupling schedule from model definition enables auto-tuning.
The combinations of schedule primitives provided by {Slapo}\xspace can introduce a search space, which includes the decisions of the number of activation checkpoints and pipeline stages, whether to shard a parameter or replace a certain module/subgraph, and so on.
Consequently, {Slapo}\xspace provides an auto-tuner to explore the best combination given a particular training environment.

\begin{figure}[t]
    \centering
    \includegraphics[width=\linewidth]{figs/polygon.pdf}
    \caption{An example of a user-defined search space.}
    \label{fig:polygon}
\end{figure}

We provide symbolic variable constructions to help developers build a search space with arbitrary numbers of hyperparameters in a schedule.
Fig.~\ref{fig:polygon} depicts an example of a search space composed of batch size and the ratio of activation checkpointing.
In this example with two tunable parameters, we form the candidates of checkpoint ratio based on the batch size,
resulting in a polygon search space instead of a rectangle.
This allows developers to incorporate domain knowledge to prune invalid (the gray region) or inefficient configurations (the white region).
By avoiding the overlapping region of gray and orange regions plus the white region, the tuning efficiency can be improved.

With the search space constructed, the auto-tuner algorithm iteratively determines the values of all tunable parameters, schedules the model, and launches an evaluation script provided by developers to benchmark the performance and memory footprint.
Our auto-tuner leverages exhaustive search by default. Meanwhile, we also provide randomized coordinate descent for users to accelerate the tuning process.
We evaluate the efficiency of the search space and the tuning algorithm in \S~\ref{sub:ablation} to show the effectiveness of our auto-tuner.

\section{Conclusion}
\label{sec:conclusion}
In this paper, we propose a schedule language {Slapo}\xspace for progressive optimization of large deep learning model training.
{Slapo}\xspace decouples model execution from definition, and provides a set of schedule primitives for users to efficiently optimize the model execution.
Apart from scheduling modules and parameters, {Slapo}\xspace is capable of partially scheduling the computation of a model, making deep learning compilers more practical for training workloads.
Experimental results show {Slapo}\xspace can combine SOTA efforts to align or even outperform their performance with minimal programming effort.
\section{Limitations}
\label{sec:limitations}

Our tracer is inherited from the base tracer of \texttt{torch.fx}\xspace but adds support for control flow and hierarchical representation.
However, the control flow support is highly constrained by the random input data.
For one thing, we need advanced support for the constant propagation of the metadata.
We need to overload all the operators of the proxy to make sure the metadata can go through different modules.
For another, each time the tracer will only choose one path to execute and ignore others, which means if some features of the input data are dynamically changed (e.g., dynamic shapes), the traced model can only work for that specific kind of input.
We will need a quick transformation technique in order to trace and optimize the model just-in-time, so that it can suit different input data.

The functionality of the verifier is also limited.
Currently, we mainly depend on a large amount of randomly generated input data and test whether the output of the new module is the same as the original one.
However, it cannot work for a complex module with different branches.
Sampling-based verification cannot figure out whether the module is semantically correct.
In the future, we can incorporate symbolic shape propagation for the verifier instead of random inputs, so that we can further test if the shapes of the tensors match in the dataflow.

\section{Introduction}
\label{sec:intro}

The demand of large deep learning (DL) models is surging in recent years as they demonstrate dominating model accuracy on a range of tasks in natural language processing (NLP)~\cite{rishi2021foundation,devlin2018bert,brown2020language,chowdhery2022palm} and computer vision~\cite{zagoruyko2016wide,dosovitskiy2021vit,liu2022swinv2}.
These models are normally invented in user-friendly DL frameworks like PyTorch~\cite{paszke2019pytorch} with dynamic model graphs\footnote{Dynamic graph DL frameworks construct the model graph on the fly when executing its forward computation instead of constructing the graph ahead of time.}, which by design lacks sufficient optimization for high-performance execution.
This issue becomes more and more critical as the size of models grows exponentially and so does the time of training.

In order to reduce the model training time, developers propose various kinds of optimization.
The first type of optimization is implemented manually in different layers of model training, such as inserting high-performance kernels~\cite{dao2022flashattention,shoeybi2019megatron,apex,xFormers2022} for computationally intensive operators on specific devices (e.g., NVIDIA GPUs), employing data, tensor, and pipeline parallelism~\cite{rajbhandari2020zero,shoeybi2019megatron,narayanan2019pipedream}, as well as activation checkpointing~\cite{chen2016checkpoint,marisa2021dtr,jain2020checkmate}, to efficiently distribute the training across multiple devices.
However, manual optimization introduces the following two challenges.

\noindent\textbf{Challenge~1: Generality --}
Incorporating the above optimizations requires making direct changes to the model implementation, which means that the optimization is not easy to generalize to other models.
A new model, even with minimal change from the old one, may not be able to directly reuse the old optimization.
In addition, the optimized model becomes platform-specific, requiring developers to maintain multiple implementations to serve all requirements (e.g., training on different platforms and deploying on non-GPU devices).

\noindent\textbf{Challenge~2: Ease of Tuning --}
In practice, an optimization scheme has a number of configurations to tune (e.g., pipeline stages, number of activation checkpoints) to get a combination that results in the best performance.
Developers need to identify tunable configurations in the implementation and modify the model to expose them for effective tuning.
This process can be tedious and error-prone especially when the model definition is closely tied to optimizations.

In addition to manual optimization, the other set of optimization approaches converts the DL model into a number of \emph{static graphs} and leverages DL compilers to automatically apply optimizations.
For example, JAX~\cite{jax2018github} is a compiler-based DL framework powered by a DL compiler XLA~\cite{xla}.
JAX traces the entire model to obtain a whole graph statically, on top of which the compiler can perform aggressive optimizations such as operator fusion, expression simplification, and even 3D parallelism~\cite{zheng2022alpa}.
Similarly, the recent release PyTorch 2.0~\cite{pytorch20} provides a compiler interface to trace PyTorch dynamic graph executions and construct static graphs in \texttt{torch.fx}\xspace~\cite{reed2022fx} for optimizations.
While automatic optimization requires minimal engineering effort from model developers and addresses some of the challenges mentioned above, it also introduces two new challenges.

\noindent\textbf{Challenge~3: Programmability --}
Working on static model graphs is limited by the requirement that everything must be statically analyzable and deterministic.
Frameworks may impose constraints on the users to facilitate the conversion to static graphs.
For example, JAX programming model requires pure Python functions, no in-place updates, etc., so developers may need to rewrite the model to meet these constraints in order to make it runnable.
Moreover, it is usually non-trivial for developers to control or configure the optimizations in fine granularity, such as disabling certain rules, or excluding certain operators from a compiler pass.

\noindent\textbf{Challenge~4: Debuggability --} To make model implementation easy to understand and maintain, model developers usually implement layer modules (e.g., convolutional, fully connected, and attention layers) as building blocks, and use them to compose a model \emph{hierarchically}.
However, to expand the scope of optimization and improve performance, DL compilers operating on a static model graph often flatten the hierarchy to create a single-level dataflow graph, and rewrite certain operators (e.g., decomposing the \texttt{batch\_norm} op into a number of smaller ones).
This prevents developers from understanding and troubleshooting performance or convergence issues, as the optimized model may bear little resemblance to the original model implementation.

\smallskip
To address the challenges mentioned above, we propose \emph{{Slapo}\xspace}\footnote{Open source: \url{https://github.com/awslabs/slapo}.}, a \underline{\textbf{S}}chedule \underline{\textbf{LA}}nguage for \underline{\textbf{P}}rogressive \underline{\textbf{O}}ptimization, designed for DL frameworks with dynamic model graphs. {Slapo}\xspace has the following major features.

\noindent\textbf{Decouple model execution from its definition.} To address Challenge 1, we decouple model execution (named ``schedule'') from its definition.
As a result, the model implementation remains the same, and developers can optimize a model-specific and platform-specific schedule in a separate place.
This idea is inspired by well-known domain-specific compilers -- Halide~\cite{jonathan2013halide} and Apache TVM~\cite{chen2018tvm} -- which propose widely adopted schedule languages that decouple tensor operator scheduling from its arithmetic definition.

\noindent\textbf{Auto-tuning.} A separate schedule also enables massive auto-tuning opportunities. Similar to AutoTVM~\cite{chen2018learning}, {Slapo}\xspace provides a programming interface that allows developers to specify a set of tuneable knobs to form an efficient tuning space. The tuning space can then be explored by {Slapo}\xspace auto-tuner to realize the optimal configuration, which addresses Challenge 2. Along this direction, {Slapo}\xspace can also enable auto-scheduling as Ansor~\cite{zheng2020ansor}, and this is our planned future work.

\noindent\textbf{Progressive optimization.} {Slapo}\xspace incorporates a ``trace by need'' approach that only traces a desired module to be a static graph for compiler-based aggressive optimizations.
The traced part can be expanded or shrunk progressively as needed.
Developers explicitly call the scheduling primitives to realize this, addressing Challenge 3.

\noindent\textbf{Structure-preserving scheduling.} Model developers usually define building blocks (e.g., convolutional or attention layers), and then compose them together to form a complete model. Consequently, developers often leverage this structure to analyze and debug the model.
{Slapo}\xspace preserves this hierarchy when constructing the schedule (see \S~\ref{sec:name} for details), so that developers can easily locate the module and apply scheduling.
Also, as the model structure is preserved and optimization can be progressively applied, it facilitates the users to debug any performance and convergence issue, addressing Challenge 4.

In summary, we make the following contributions:
\begin{itemize}
    \setlength\itemsep{0.2em}

    \item We propose {Slapo}\xspace, a schedule language that decouples model execution from definition, and preserves model structure hierarchy to enable progressive optimization.

    \item We define a set of schedule primitives for {Slapo}\xspace to cover widely used optimizations in efficient training.
    \item We design and implement a lightweight auto-tuner for further reducing the efforts required to identify the optimal schedule configurations for training efficiency.

\end{itemize}

We evaluate {Slapo}\xspace by training popular deep learning models with billions of parameters and compare {Slapo}\xspace with SOTA distributed training frameworks such as DeepSpeed~\cite{rasley2020deepspeed} and Megatron-LM~\cite{shoeybi2019megatron}.
With minimal program effort, {Slapo}\xspace is capable of scheduling the existing hand-crafted optimizations to achieve up to 3.35$\times$ speedup on a single machine with 8 NVIDIA V100 GPUs, and up to 1.32$\times$ speedup on multiple machines with up to 64 NVIDIA V100 GPUs, when compared to the out-of-the-box best baseline.

\begin{comment}
For example, ZeRO~\cite{rajbhandari2020zero,rasley2020deepspeed} provide an advanced mechanism for data parallelism to shard the optimizer states, gradient, and parameters in the model.
Megatron-LM~\cite{shoeybi2019megatron} shards the parameters in Transformer models to explore tensor parallelism.

3D parallelism (data, model, and pipeline parallelism). We present more details in \S~\ref{sub:3dparallelism}.
Megatron-LM~\cite{shoeybi2019megatron} shards the parameters in Transformer models to explore tensor parallelism.
ZeRO-3~\cite{rajbhandari2020zero,rasley2020deepspeed} provide an advanced mechanism for data parallelism to shard the optimizer states, gradient, and weights in the model.
GPipe~\cite{huang2019gpipe} and PipeDream~\cite{narayanan2019pipedream} further scales up deep neural network training using pipeline parallelism.
These libraries are easy to try and enable quick deployment in production.
However, the speedup they provide is not without cost.
Users always need to rewrite the model in order to take advantage of the facilities offered by these libraries, which cannot be generalized to models other than Transformers and requires optimizing each model case by case.
Another optimization strategy is offloading efficient kernels.
Many libraries implement high-performance kernels for the basic components in the Transformer models~\cite{xFormers2022,dao2022flashattention}.
While replacing a single layer with these kernels is trivial, some advanced optimizations like operator fusion~\cite{tensorrt} cannot be easily implemented without refactoring the codebase, which can be a time-consuming and error-prone process that may even break the model.

On the other hand, there are emerging research on compiler for distributed workloads~\cite{zheng2022alpa,xla}.
Alpa~\cite{zheng2022alpa} is the first optimizing compiler that can automatically generate high-performance execution plan for 3D parallelism in distributed training.
However, its high barrier prevents developers from deploying it in actual production.
One of the main issues with these compilers is the lack of customizability.
Compiler passes are often based on heuristics with magic thresholds, and some of the optimization patterns are fixed.
For example, in XLA, there is a fixed fusion pattern for fusing convolution and batch normalization layers.
All the operators matching this pattern will be transformed, which means developers cannot conduct fine-grained optimizations for a specific operator or a specific pattern.
What makes the issue even worse, compilers always require operators to be fully offloaded.
If a user uses a custom operator that does not have a backend implementation, the compiler may fail to generate the executable.
Developing a new operator backend requires significant engineering effort, and deploying the model for training becomes a time-consuming process.
Furthermore, compiler optimization is not transparent to developers, making it difficult to understand the optimal strategies realized by compiler-based solutions, and hard to debug and locate the problem if the compiler crashes.
Conducting performance analysis to figure out the bottleneck of compiled models is even challenging, as the model hierarchy is broken by the compiler.

Instead of making a choice between fully manual optimization (e.g., PyTorch) and fully automatic optimization (e.g., Jax), we advocate for \emph{progressive optimization} that applies optimizations in a narrow scope step by step, to minimize the possible risk brought by external libraries or compilers.
A key observation is that those optimizations techniques specifying \emph{how} to execute the model is agnostic to the model itself that defines \emph{what} to execute.
Inspired by Halide~\cite{jonathan2013halide} and TVM~\cite{chen2018tvm}, we propose {Slapo}\xspace, a \underline{\textbf{S}}chedule \underline{\textbf{LA}}nguage for \underline{\textbf{P}}rogressive \underline{\textbf{O}}ptimization.
It \emph{decouples model execution from its definition} and defines a set of schedule primitives for describing the optimizations on model execution.
{Slapo}\xspace allows users to conduct optimizations without changing the original model, making it easier to maintain and enabling collaboration between model developer and system engineers.

We propose a hierarchical approach that uses both structural and dataflow representation to represent models, and build our scheduling language on top of that.
As a result, the schedule is still modular and preserves the hierarchical structure of the model.
{Slapo}\xspace enables users to compose the latest techniques on efficient kernels, whether they are hand-written or generated by compilers, into the model.
We can even offload part of the model to a compiler for optimization, while leaving the rest unchanged using the default backend.
The approach allows us to localize problems and focus on optimizing or debugging a small part of the model.

In summary, we make the following contributions:
\begin{itemize}
    \item We summarize the existing optimization techniques used in gigantic model training to achieve state-of-the-art (SOTA) performance and design a set of transformation primitives.
    \item We propose a schedule DSL that decouples model execution from definition, enabling users to describe the execution strategy without changing the model itself.
    \item We implement a system that enables end-to-end optimization and training, and provide a lightweight auto-tuner for further reducing the efforts required to find hyperparameters and improve performance.
    \item Experimental results show our DSL can optimize a vanilla PyTorch models from scratch and achieve the SOTA performance. We demonstrate a \fixme{?x} speedup compared to the SOTA solution for Megatron-LM on a single machine, and a \fixme{?x} speedup compared to DeepSpeed in a multi-server environment.
    \fixme{Describe more on experiment settings and results}
\end{itemize}
\end{comment}
\section{Background and Motivation}
\label{sec:background}
In this section, we first introduce common practices of improving a DL model training efficiency for dynamic graphs (e.g., PyTorch~\cite{paszke2019pytorch}), followed by an end-to-end motivating example to illustrate the challenges of this process.

\subsection{Efficient Model Training}
\label{sub:training}

\noindent\textbf{High-performance kernel libraries.} To achieve high efficiency in deep learning model training, it is straightforward to leverage efficient kernels specifically optimized for particular hardware platforms (e.g., NVIDIA GPUs).
For example, there are a number of libraries~\cite{dao2022flashattention,shoeybi2019megatron,apex,xFormers2022} that provide efficient CUDA kernel implementations.
These libraries encapsulate kernels as DL framework-compatible modules for developers to replace the native implementation in their models.
In the case where there are no existing CUDA implementations, developers may leverage compiler-based solutions, such as TorchScript~\cite{torchscript}, to generate a kernel.

\noindent\textbf{Activation checkpointing.} Apart from compute optimization techniques, memory footprint optimization is also essential for training large models.
A large portion of memory consumption in the training process is contributed by forward activation tensors that are stored for the later gradient calculation.
By checkpointing some activation tensors and re-computing the rest of them in backward propagation, we are able to significantly reduce memory footprint, and thus support a larger batch size and higher training throughput.
This approach is called activation checkpointing and is originally proposed by \cite{chen2016checkpoint}.
Furthermore, existing works~\cite{jain2020checkmate,marisa2021dtr,korthikanti2022reducing} also demonstrate that by carefully selecting which activations to checkpoint, we are capable of better utilizing device memory and achieving an even better throughput.

\noindent\textbf{Parallelism in distributed training.} When the model is too large to fit in a single device, training it in parallel in a distributed environment is inevitable.
The parallelism techniques are usually classified into three types: data parallelism, tensor parallelism, and pipeline parallelism. Both tensor and pipeline parallelism belong to a larger class called model parallelism.
\emph{Data parallelism} partitions training data, so each device trains the replicated model with different data~\cite{li2020ptddp,abadi2016tensorflow,rajbhandari2020zero}, and aggregates their partial gradients for parameter updating.
Since data parallelism replicates an entire model on each device, it is insufficient when the model size (i.e., total memory consumption of its parameters) is too large to fit on a single GPU.
In this case, \emph{tensor parallelism}, takes another approach by partitioning the model parameter tensor onto multiple devices~\cite{shoeybi2019megatron}.
However, it requires developers to explicitly use collective communication operators to scatter tensors and aggregate partial results.
For example, Megatron-LM~\cite{shoeybi2019megatron} is a widely used PyTorch-based framework that provides manual parallelized Transformer models~\cite{vaswani2017transformer} and is adopted to train extremely large models~\cite{zhang2022opt}.
Finally, \emph{pipeline parallelism}~\cite{huang2019gpipe,narayanan2019pipedream} partitions models by layers and groups them into a series of stages. By putting each stage on a different device, we can overlap the computation of multiple data batches.
To ensure correctness and performance, pipeline parallelism needs a specialized runtime to schedule and synchronize data.
These techniques are not mutually exclusive and can be combined. Combining all of them is known as \emph{3D parallelism}~\cite{narayanan2021megatronv2}.

\subsection{A Motivating Example}
\label{sub:example}
\begin{figure}[!htbp]
\centering
\includegraphics[width=\linewidth]{figs/attention.pdf}
\caption{An attention layer in BERT. The \texttt{Query}, \texttt{Key}, \texttt{Value}, and \texttt{Output} are \texttt{nn.Linear} modules containing learnable weight and bias. \Circled{1} - \Circled{4} indicates optimization points.}
\label{fig:attn}
\end{figure}

In this subsection, we use BERT~\cite{devlin2018bert} model implementation from HuggingFace Hub (v4.25.1)~\cite{wolf2019huggingface} to showcase how the above techniques are applied to a DL model for efficient training.
Fig.~\ref{fig:attn}(a) depicts the architecture of attention layer~\cite{vaswani2017transformer}, the core and most time-consuming module in BERT.
An attention layer is composed of two submodules -- \texttt{SelfAttention} and \texttt{Projection}.
We conduct a few steps to progressively improve the training efficiency of this attention layer and show the resulting implementation in Fig.~\ref{fig:attn}(b).

\noindent\textbf{\Circled{1} Fuse QKV.} By default, the \texttt{Query}, \texttt{Key}, \texttt{Value}
in \texttt{SelfAttention} are three standalone \texttt{nn.Linear} modules, which may incur extra kernel launch overheads.
We replace them with a single \texttt{nn.Linear} module with their parameters concatenated, as shown in the following code snippet.

\smallskip
\begin{lstlisting}[language=diff]
def __init__(self, ...):
-   self.query = nn.Linear(n_hidden, n_head)
-   self.key = nn.Linear(n_hidden, n_head)
-   self.value = nn.Linear(n_hidden, n_head)
+   self.qkv = nn.Linear(n_hidden, n_head * 3)
def forward(self, hidden_states, ...):
-   query = transpose(self.query(hidden_states))
-   key = transpose(self.key(hidden_states))
-   value = transpose(self.value(hidden_states))
+   qkv = transpose(self.qkv(hidden_states))
+   query, key, value = split(qkv, 1, dim=-1)
\end{lstlisting}

\noindent\textbf{\Circled{2} Use efficient kernels.} The pink blocks in Fig.~\ref{fig:attn} are the core attention computation, which is also the bottleneck of performance and memory footprint.
A recent work flash attention~\cite{dao2022flashattention} proposes to compute the attention in a block-by-block fashion, so only a block of the intermediate attention tensor is generated at a time, significantly reducing the peak memory consumption, and thus can improve the training efficiency by enlarging the batch size.
The following code snippet shows how we replace the existing attention with the flash attention implementation provided by xFormers~\cite{xFormers2022}. The transpose and reshape operations are simplified.

\smallskip
\begin{lstlisting}[language=diff]
def forward(self, hidden_states, ...):
-   query *= 1 / query.shape[-1] ** 0.5
-   attn = baddbmm(bias, query, key)
-   attn = dropout(softmax(attn), p)
-   attn = attn @ value
+   attn = xformers.ops.mem_eff_attention(...)
\end{lstlisting}

\noindent Note that this flash attention implementation only supports the latest NVIDIA GPUs with Volta, Ampere, and Hopper architectures, so once this kernel is used, the model is no longer compatible with other platforms.

Another optimization opportunity is the pattern in \texttt{Projection}. By default, the bias addition operation is contained in \texttt{Output} module.
A more aggressive and more efficient way is to use a DL compiler (e.g., TorchScript~\cite{torchscript}) to fuse the pattern \texttt{BiasAdd-Dropout-ResidualAdd-LayerNorm} into a single kernel, as suggested by \cite{vaidya2022nvfuserexample}.

\noindent\textbf{\Circled{3} Tensor parallelism.} We then partition the \texttt{FusedQKV} and \texttt{Output} parameters onto multiple devices.
Given the input of the attention module $X$, the weights of \texttt{FusedQKV} ($A$) and \texttt{Output} ($B$), we have self-attention function $f$:

\vspace{-1em}
\begin{equation*}
f(XA)B
=f\left(X\begin{bmatrix}A_1 & A_2\end{bmatrix}\right)
\begin{bmatrix}B_1 \\ B_2\end{bmatrix}
=f(XA_1)B_1 + f(XA_2)B_2\,.
\end{equation*}
\vspace{-1em}

Accordingly, we follow the convention of Megatron-LM \cite{shoeybi2019megatron} to shard the weights of \texttt{FusedQKV} in columns and shard the weights of \texttt{Output} in rows. We illustrate the latter in the following code snippet.

\smallskip
\begin{lstlisting}[language=diff]
def __init__(self, ...):
-   self.output = nn.Linear(n_hidden, n_hidden)
+   new_size = n_hidden // world_size
+   self.output = nn.Linear(new_size, n_hidden)
def forward(self, hidden_states):
    out = self.output(hidden_states)
+   dist.allreduce(out)
\end{lstlisting}

Since the output tensor only holds partial results after sharding, we need to conduct \texttt{all\_reduce} to aggregate outputs from different devices.

\noindent\textbf{\Circled{4} Pipeline parallelism.} To pipeline a BERT model and execute it on a SOTA pipeline runtime, we have to further manually partition the model to a sequence of sub-models (each of them includes a series of attention layers) by re-writing the top module\footnote{The code example is omitted due to page limit.}.

The above process poses a generality issue.
Although developers have spent efforts on identifying and optimizing the performance bottleneck of a model with semantics preserved, this effort is hard to be reused by another model.
Furthermore, the above improved model is no longer compatible with a single device, unless we add control logic to only enable model parallelism on multi-GPU environments or maintain a separate single-device version.
It also creates a barrier for the model deployment after training, because a model implementation with custom efficient kernels and parallelism may not be recognized by inference compilers.

In essence, the above pain points are the result of tightly coupling model definition and training/platform-specific optimizations.
This motivates us to propose a schedule language that decouples model execution (i.e., schedule) from definition and provides easy-to-use primitives for optimizing large model training.
In fact, the idea of decoupling optimization has been widely accepted in DL compilers~\cite{jonathan2013halide,chen2018tvm,baghdadi2019tiramisu}, and opens opportunities for auto-tuning~\cite{chen2018tvm} and auto-scheduling~\cite{zheng2020ansor}.

\begin{comment}
\subsection{Decoupling Schedule from Computation}

\yida{removing candidate for space saving. We can simply add one more sentence at the end of the previous paragraph}
The idea of decoupling optimizations (i.e., schedule) from the compute definition is already widely accepted in DL compilers~\cite{jonathan2013halide,chen2018tvm}. The following pseudo-code is an example of using Apache TVM to optimize matrix multiplication:

\smallskip
\begin{lstlisting}[language=Python]
# Compute definition.
compute([M, N], lambda i, j: \
    reduce(A[i, k] * B[k, j], axis=k))
# Basic schedule for GPUs.
oi, ii = split(i, factor=64)
oj, ij = split(j, factor=64)
reorder(oi, oj, ii, ij)
bind(oi, blockIdx.x)
bind(oj, blockIdx.y)
...
\end{lstlisting}

In the code snippet, L1-3 provides the arithmetic definition of matrix multiplication. Starting from L4 is the schedule to tile input tensors, reorder the computation, and bind certain dimensions of computations to GPU threads.
With this programming interface, developers can continuously provide efficient schedules for more and more hardware platforms without interference to each other.
In fact, the TVM community already contributes more than 15K lines of code for operator schedules.
Furthermore, since the decoupled schedule is clean and systematic, it opens opportunities for auto-tuning~\cite{chen2018tvm} and auto-scheduling~\cite{zheng2020ansor}.
\end{comment}

\subsection{Schedule Language}
\label{sec:sch}

\begin{figure}[t]
    \centering
    \includegraphics[width=\linewidth]{figs/example-1.pdf}
    \caption{An example of scheduling modules and parameters of a BERT model. \texttt{ffn} is the Feed-Forward Network. \texttt{eff\_attn} refers to the replaced attention module. The weight matrix has a shape of \texttt{(output\_dim, input\_dim)}. Thus, sharding the weight matrix at \texttt{axis=0} is equivalent to partitioning the output dimension.}
    \label{fig:schedule_mod}
\end{figure}

Motivated by \S~\ref{sub:example}, our goal is to let developers optimize models in a few lines of code without changing the model definition itself.
Fig.~\ref{fig:schedule_mod} presents the {Slapo}\xspace schedule language with the BERT model from HuggingFace Hub~\cite{wolf2019huggingface} as an example.
As shown in Fig.~\ref{fig:schedule_mod}(a), most DL model developers define models with a hierarchical structure for better readability and easy maintenance.
The \texttt{\_\_init\_\_} constructor defines the configurations, submodules, and learnable parameters, and the \texttt{forward} method defines the forward computation (the backward computation is generated by the framework with automatic differentiation~\cite{paszke2019pytorch}).
Developers can then pass the created model into {Slapo}\xspace and create a default \emph{schedule} that specifies \emph{how} to execute the model in the framework.
The schedule preserves the hierarchical structure, and \texttt{create\_schedule} is applied recursively to all the submodules, so that developers can easily apply schedule primitives at arbitrary levels.

\input{figs/primitives}

Table~\ref{tab:primitives} defines the schedule primitives.
It is worth noting that the primitive set can be easily extended in the future to incorporate other optimizations.
We categorize schedule primitives in two sets: whether or not it requires a static graph to be generated. On one hand, when scheduling at the module level, such as enabling activation checkpointing, replacing with an efficient alternative, and sharding learnable parameters, we do not change the computation specified in the \texttt{forward} method.
As a result, the schedule primitives in the left column of Table~\ref{tab:primitives} do not require a static graph.
We present the details of this scheduling in \S~\ref{sec:sch-mod}.
On the other hand, scheduling the computation, such as operator fusion and pipeline parallelism, has to manipulate the \texttt{forward} method.
Thus, the schedule primitives in the right column of Table~\ref{tab:primitives} require the computation to be in a static graph, so we have to use \texttt{.trace()} prior to applying these primitives, as presented in \S~\ref{sec:sch-compute}.

\section{Related Work}
\label{sec:related_work}

\noindent\textbf{Schedule Languages and Compilers.}
Many domain-specific languages (DSLs) leverage the idea of decoupling optimizations from algorithms~\cite{jonathan2013halide,chen2018tvm,baghdadi2019tiramisu,lai2019heterocl,xiang2022heteroflow}, allowing users to focus on customization and enabling compilers to perform complex optimizations.
TVM~\cite{chen2018tvm,feng2022tensorir} inherits the decoupling idea from Halide~\cite{jonathan2013halide} and builds an end-to-end compiler for deep learning inference.
Tiramisu~\cite{baghdadi2019tiramisu} further extends the scope of schedule to polyhedral compilation.
{Slapo}\xspace borrows a similar decoupling idea and applies it to the model execution level.

\noindent\textbf{Dynamic Graph Optimizations.}
Due to the dynamic nature and usability of PyTorch, many frameworks and libraries directly explore optimizations on top of it.
ZeRO~\cite{rajbhandari2020zero} is a three-stage data parallelism strategy that partitions optimizer states, gradients, and parameters to reduce memory usage, implemented first in DeepSpeed~\cite{rasley2020deepspeed} and then adopted by other frameworks~\cite{fsdp}.
MiCS~\cite{zhang2022mics} further improves ZeRO by minimizing the communication scale.
Megatron-LM~\cite{shoeybi2019megatron,narayanan2021megatronv2} takes a different approach to implement model parallelism for Transformer models, becoming one of the mainstream SOTA parallelism libraries.
{Slapo}\xspace applies those optimization techniques using primitives in a more user-friendly way.
We also use framework dialects described in \S~\ref{sec:impl} to train the scheduled model on these frameworks.

\noindent\textbf{Static Graph Optimizations.}
Some other DL frameworks adopt static graphs so that compiler optimizations can be easily involved.
JAX~\cite{jax2018github} is a popular framework that offers a programming model similar to NumPy, and is powered by XLA~\cite{xla} as the backend compiler.
Accordingly, it is able to achieve 3D parallelism with the corresponding sharding mechanism, GSPMD~\cite{xu2021gspmd} and GShard~\cite{lepikhin2021gshard}.
On top of that, Alpa~\cite{zheng2022alpa} is the first compiler based on JAX and XLA to achieve automatic 3D parallelism.
Besides, Unity~\cite{colin2022unity} is another compiler-based distributed training framework that automatically jointly optimizes model parallelism and operator fusion.
Their automation mechanisms are orthogonal to {Slapo}\xspace and could inspire {Slapo}\xspace's auto-scheduler in the future.

Moreover, PyTorch~\cite{paszke2019pytorch} eco-system also pays efforts on static graph optimization in recent years.
\texttt{torch.fx}\xspace~\cite{reed2022fx} is an IR to capture PyTorch dynamic graph
 so that it can be compiled by DL compilers such as PyTorch 2.0~\cite{pytorch20}.
Since {Slapo}\xspace adopts \texttt{torch.fx}\xspace as the IR for static graphs, we plan to integrate the PyTorch compiler in the future.

\begin{comment}
\textbf{Distributed Training.}
There are lots of existing works on accelerating distributed deep learning, where a large body of them explores 3D parallelism.
Data parallelism feeds different devices with different batches of data, and aggregates their gradients during the backward pass~\cite{jiang2020byteps,sergeev2018horovod}.
Popular deep learning frameworks like PyTorch~\cite{li2020ptddp} and TensorFlow~\cite{abadi2016tensorflow} both support Distributed Data Parallel (DDP) in their library.
However, it requires duplicating the model across the devices, causing redundant memory consumption.
ZeRO~\cite{rajbhandari2020zero} firstly proposes a three-stage sharding technique that shards optimizer state, gradients, and parameters to reduce memory usage.
DeepSpeed~\cite{rasley2020deepspeed} successfully implements the key idea of ZeRO to support different models.
FairScale~\cite{fairscale} proposes the Fully-Sharded Data Parallel (FSDP)~\cite{fsdp}, which is a variant of ZeRO-3 that shards the parameters across the devices and gathers them when needed.
MiCS~\cite{zhang2022mics} further improves ZeRO-3 by batching the parameter prefetching to reduce the payload and pre-allocating buffers for defragmentation.
ZeRO-Offload~\cite{ren2021zerooffload} and ZeRO-Infinity~\cite{samyam2021zeroinfinity} additionally leverage CPU and NVMe to relieve the burden of GPU memory.

Model parallelism does not replicate the model, but holds different layers on different devices.
ColocRL~\cite{azalia2017colocrl}, OptCNN~\cite{jia2018optcnn}, and FlexFlow~\cite{jia2019flexflow} all leverage heuristics or machine learning methods to automatically find the optimal placement.
Another model parallelism approach is also called tensor parallelism which shards the parameters onto different devices.
Megatron-LM~\cite{shoeybi2019megatron,narayanan2021megatronv2} first implements this idea on Transformer-based models and becomes the mainstream tensor parallelism library.
Pipeline parallelism can be viewed as an advanced model parallelism, which groups several consecutive layers together as a stage and places different stages onto different devices to form a pipeline.
GPipe~\cite{huang2019gpipe} leverages synchronized method, while PipeDream~\cite{narayanan2019pipedream,narayanan2021pipe2bw} and PipeMare~\cite{yang2021pipemare} leverage asynchronized forward and backward pass to further improve the throughput but at the cost of slow convergence.

Combining these 3D parallelism techniques together is also a hot topic~\cite{li2021colossal,miao2023galvatron}.
Megatron-LM~\cite{shoeybi2019megatron}, Colossal-AI~\cite{li2021colossal}, and OneFlow~\cite{yuan2021oneflow} are frameworks that support those parallelism techniques.
Alpa~\cite{zheng2022alpa} is the first compiler for generating optimal distributed plans.
GSPMD~\cite{xu2021gspmd}, GShard~\cite{lepikhin2021gshard}, and Unity~\cite{colin2022unity} also automates this process.
$P^2$~\cite{xie2022p2} uses program synthesis to achieve the same goal.
Our DSL is expressive enough to support all the 3D parallelism, providing a new interface for automation.
Moreover, our DSL not only supports 3D parallelism, but has other primitives to express operator fusion, gradient checkpointing, etc.

Efficient kernels are also important for distributed deep learning.
Apex~\cite{apex} provides mixed precision training and lots of useful fused kernels.
xformer~\cite{xFormers2022} uses Triton~\cite{tillet2019triton} to implement basic components in Transformer models, including the Flash Attention~\cite{dao2022flashattention} kernel.
Our DSL is orthogonal to the above solutions and provides easy-to-use API for users to integrate them into our system.

\noindent
\textbf{Scheduling Languages.}
There are many existing domain-specific languages (DSLs) that leverage the idea of decoupling optimizations from algorithms.
Halide~\cite{jonathan2013halide} first defines schedule primitives in image processing.
TVM~\cite{chen2018tvm} later extends it to support neural network compilation.
The idea of separation of concerns is then applied to different DSLs of deep learning~\cite{chen2018tvm,baghdadi2019tiramisu}, computer graphics~\cite{hu2019taichi}, graph processing~\cite{zhang2018graphit,chen2021krill}, and hardware customization~\cite{lai2019heterocl,xiang2022heteroflow,nitish2019t2s,yuka2022exo}.
Those decoupling makes the users better handle the program and focus more on the customization, and also enables the compiler to conduct complex optimizations.
Our DSL follows the same practice that decouples model execution from definition, and provides primitives for users to easily conduct optimizations for distributed training.

There are several major differences between our work and TVM:
(1) TVM can only boost model inference, but our DSL can speed up the training process.
(2) TVM only focuses on a single device, while our DSL can express parallelism across multiple devices and even multiple nodes.
(3) The schedule language in TVM is mainly designed for loop transformation, which is too low-level for large model training. Our DSL focuses on graph-level optimizations like operator fusion and 3D parallelism.
\end{comment}
\section*{Acknowledgments}

\subsection{Verification}
\label{sec:verifier}

To achieve high usability, {Slapo}\xspace primitives are flexible to schedule modules and computations. However, it is possible for developers to schedule a model incorrectly. For example, the replaced module may require a different data layout but developers do not provide the corresponding layout transformation logic, or developers insert the output aggregation to an invalid point when sharding a parameter tensor.

To guarantee the scheduled model is executable and maintains its numerical correctness, {Slapo}\xspace includes the following verification stages.
First, before applying the schedule, we validate the sequence of schedule primitives with a set of predefined rules in each primitive.
For example, a \texttt{.sync()} primitive must have a corresponding \texttt{.shard()} primitive beforehand; primitives for distributed training can only be specified in a distributed environment with multiple devices; primitives that require static graph must have a corresponding \texttt{.trace()} primitive in advance.
If any of the rules are violated, the verifier raises an error and stops the rest scheduling process.

Second, we need to assure the numerical correctness of the replaced or fused operators.
To achieve this, the verifier generates random inputs and compares the results of the replaced module against the original one. Developers can also configure the number of random inputs and provide a function to generate inputs when there are any constraints.

Lastly, we verify the end-to-end correctness also using random inputs. This verification includes two purposes: (1) whether the shape of the sharded parameters or the tensors in the replaced modules are correct, and (2) whether the scheduled model produces the same output as the vanilla one.