# Pipe-BD: Pipelined Parallel Blockwise Distillation

Hongsun Jang2, Jaewon Jung4, Jaeyong Song3, Joonsang Yu5, Youngsok Kim3, and
Jinho Lee2∗ * Corresponding author 2Department of Electrical and Computer
Engineering, Seoul National University  4Department of Artificial
Intelligence, Yonsei University  3Department of Computer Science, Yonsei
University  5CLOVA ImageVision, CLOVA AI Lab, NAVER<EMAIL_ADDRESS>{jungjaewon<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Training large deep neural network models is highly challenging due to their
tremendous computational and memory requirements. Blockwise distillation
provides one promising method towards faster convergence by splitting a large
model into multiple smaller models. In state-of-the-art blockwise distillation
methods, training is performed block-by-block in a data-parallel manner using
multiple GPUs. To produce inputs for the student blocks, the teacher model is
executed from the beginning until the current block under training. However,
this results in a high overhead of redundant teacher execution, low GPU
utilization, and extra data loading. To address these problems, we propose
Pipe-BD, a novel parallelization method for blockwise distillation. Pipe-BD
aggressively utilizes pipeline parallelism for blockwise distillation,
eliminating redundant teacher block execution and increasing per-device batch
size for better resource utilization. We also extend to hybrid parallelism for
efficient workload balancing. As a result, Pipe-BD achieves significant
acceleration without modifying the mathematical formulation of blockwise
distillation. We implement Pipe-BD on PyTorch, and experiments reveal that
Pipe-BD is effective on multiple scenarios, models, and datasets.

###### Index Terms:

Distributed Training, Knowledge Distillation, Neural Architecture Search,
Model Compression

##  I Introduction

Modern deep neural network models are known to incur huge computational and
memory requirements, especially with large-scale datasets [1]. With the
continuing growth in model size, it takes tens, if not hundreds, of GPU days
to train them [2], and the model size often exceeds the GPU memory capacity.
Especially for methods that explore large solution spaces such as the neural
architecture search (NAS) [3, 4], the problem becomes even more significant.
This problem mandates the use of model parallelism [5, 6], which creates
substantial throughput loss with inevitable pipeline bubbles.

Blockwise distillation [7, 8, 9] is one promising approach to mitigate such
problems. As illustrated in Fig. 1, blockwise distillation splits the model
into multiple smaller blocks. As opposed to traditional knowledge distillation
methods that rely on input data and output labels from both ends, blockwise
distillation uses the intermediate activation values of pretrained blocks of a
‘teacher’ to train each ‘student’ block. As a result, each block converges
faster (i.e., fewer epochs) due to the smaller solution space.

Contrary to the earlier belief that teachers must be larger than students,
recent studies have revealed that smaller teachers can be used to train larger
students [10]. With such findings, blockwise distillation is used in various
fields such as model compression [7, 11] and NAS [9, 12]. Since training a
small teacher for a new task is quick and easy, blockwise distillation can be
applied in most cases where traditional training is used.

Figure 1: Conceptual diagram of blockwise distillation.

However, the existing state-of-the-art methods for blockwise distillation [9,
7] exhibit several inefficiencies. Relying on the traditional data-parallel
training scheme, they train each student block one by one independently. While
this fully exploits the independent nature of the blocks, it is not the best
choice for training throughput. First, to train a single intermediate student
block, the teacher blocks must be executed from the beginning to the
designated block. As a result, the teacher blocks exhibit substantial
redundant execution, especially with blocks closer to the output. Second, with
data parallelism, a batch of data is split among multiple GPUs, which leads to
a smaller batch size per GPU, often resulting in resource under-utilization.
Some approaches use a larger batch size to mitigate this [2], but it is known
to be difficult to ensure model convergence [13]. Last, the data must be
redundantly loaded for each student block. Unless the entire dataset fits into
the GPU memory, the data are loaded from the CPU memory or disks. As the
memory and disks are shared system-wide, the extra data loading becomes
another significant overhead in training.

To address the issues, we propose _Pipe-BD_ , a novel parallel training method
for blockwise distillation. We assign individual student blocks to different
devices and compute a teacher network in a relayed manner, which can reduce
teacher redundancy. Inspired by approaches with pipeline parallelism [5, 14,
15, 16], we restructure the training schedule of the student blocks such that
the training time is greatly improved.

Pipe-BD comprises three components: First, we propose _teacher relaying_.
Instead of relying on data parallelism, we spread the student model to
multiple training devices (i.e., GPUs) in a block granularity. Then, blocks of
the teacher model are executed by relaying the intermediate activation values
between the devices. This approach has the advantages of eliminating extra
data loading and increasing resource utilization from larger batch size per
device. Second, we propose _decoupled parameter update_ to remove the
scheduling bubbles and enhance the overall utilization. With teacher relaying,
devices have to wait for the intermediate activation values from previous
devices, creating scheduling bubbles. Decoupled parameter update performs
model parameter updates in a misaligned manner and starts the next step right
ahead, so those bubbles can be removed. Third, we suggest _automatic hybrid
distribution_. Achieving a balance between devices is difficult with blockwise
distillation because of the limited number of blocks available in typical
neural network structures. Automatic hybrid distribution enables fine-grained
balancing with further splitting blocks along the batch size dimensions.

Pipe-BD is implemented on PyTorch and can automatically make all scheduling
decisions to improve the throughput. Our extensive set of experiments shows
Pipe-BD achieves a significant speedup over the state-of-the-art methods on
multiple use cases and environments ranging from 2.37$\times$ to 7.38$\times$.

##  II Background and Related Work

###  II-A Blockwise Distillation

Blockwise distillation [9, 7, 8] is a promising direction for training a
neural network. In traditional knowledge distillation, a student model is
trained against a pre-trained teacher model. Because the solution space size
is identical to that of conventional supervised training, it faces convergence
and training time problems. Blockwise distillation splits the larger teacher
model into smaller ones and trains them blockwise as depicted in Fig. 1. Each
teacher block ($T_{i}$) and student block ($S_{i}$) pair obtains activation
values from the previous teacher block ($T_{i-1}$). This pair performs forward
pass using the activation as input and creates teacher output activation and
student output activation. Blockwise distillation minimizes a loss function
($L(\Delta output)$) which measures the difference between these two
activations, to distill knowledge from a teacher block to the dedicated
student block. This blockwise distillation process makes target problem spaces
smaller and is known to converge faster. Many applications such as NAS [9, 12]
and model compression [7, 11] use blockwise distillation because of these
characteristics.

###  II-B Parallelization Baseline of Blockwise Distillation

State-of-the-art methods of blockwise distillation [9] use the traditional
data-parallel scheme to further accelerate the training as illustrated in Fig.
3(a). This scheme trains a student block ($S_{i}$) with all devices in a data-
parallel manner for fixed $n$ epochs, then moves on to train the next student
block ($S_{i+1}$). It redundantly loads data multiple times because of this
iterative training. Each student block ($S_{i}$) requires the activation
values from the previous teacher ($T_{i-1}$), so it also entails redundant
teacher executions. Furthermore, it uses a smaller batch size per device which
leads to under-utilization. Due to these inefficiencies, the data-parallel
blockwise distillation suffers from poor scalability. An alternative scheme
[7] regards the training of each layer as a single task and adopts bin packing
algorithm to balance the workload. However, it still has redundant teacher
executions and suffers from workload imbalance when there are insufficient
layers in the model.

##  III Motivation

Figure 2: Motivational experiment. The breakdown demonstrates three major
inefficiencies of baseline; redundant teacher execution, extra data loading,
and low resource utilization.

In this section, we provide a motivational study highlighting the inefficiency
of the existing parallel blockwise distillation training scheme and the need
for a new approach. Fig. 2 depicts the breakdown of time spent in parallel
blockwise distillation with four RTX A6000 GPUs (NAS with Cifar-10; see
Section VI-B for the detailed setup). ‘Baseline’ refers to the state-of-the-
art parallel blockwise distillation method [9], where each block is trained
sequentially using four devices with data parallelism. As displayed in the
chart, the training time is spent on data loading, teacher execution (forward
pass), and student execution (forward/backward pass).

However, all three parts exhibit significant inefficiency, slowing down the
training. To demonstrate the inefficiencies, we plot the ‘ideal’ bar in Fig. 2
by measuring the training time of each part separately with a single GPU and
dividing each time by four. This represents an imaginary system with perfect
parallelization and infinite device memory.

The large gaps in teacher execution and data loading time occur because the
baseline has many redundant teacher executions and extra data loading. Because
each student block to train requires executing the teacher model from the
beginning, the earlier teacher blocks are redundantly executed multiple times
(see Fig. 3(a)). Similarly, block-by-block training forces loading data as
many as the number of blocks. In addition, data-parallelism leads to smaller
batch size per device, resulting in lower resource utilization. As
demonstrated in several empirical studies [17, 18], a sufficient per-device
batch size is critical for training throughput, which is the cause of the gap
on student execution time. Pipe-BD targets these inefficiencies. As presented
in Fig. 2, Pipe-BD reduces the training time close to the ideal case, with
only a small overhead (idle).

##  IV Pipe-BD Method

(a) Baseline

(b) w/ Teacher Relaying

(c) w/ Decoupled Parameter Update

(d) w/ Automatic Hybrid Distribution

Figure 3: Illustration of the techniques in Pipe-BD.

###  IV-A Teacher Relaying

Pipe-BD starts by restructuring the training pipeline of blockwise
distillation with teacher relaying. As opposed to the baseline (Fig. 3(a))
where a single block is fully trained in a data-parallel manner before moving
on to the next, teacher relaying exclusively distributes the teacher and
student blocks to all training devices. Then, each device relays the
intermediate teacher activation values to the next device as depicted in Fig.
3(b). The received activation is the input for both the teacher and the
student block. The teacher block is executed first, whose output activation is
sent to the next device such that the execution of the next block can start.
Overlapped with the transmission, the forward pass execution of the student
starts, taking the same input as the teacher block. After calculating the loss
by comparing the output activations of the teacher and the student, the
backward pass of the student follows. After all the backward passes are
finished, parameter updates are performed on each block, completing the
training step.

The teacher relaying scheme has two advantages over the existing approach.
First, each device executes the stages with larger batches and enjoys better
resource utilization. For example, in the baseline using four devices with an
effective batch size of 256, each device executes with a batch size of 64,
which is often too small to fully utilize the hardware resources. In contrast,
with teacher relaying, each device would run with a full batch size of 256,
increasing resource utilization. Second, the overhead of data loading is
reduced. When the dataset is large, the data must come from the main memory or
the disk, where both are system-wide shared resources. Because teacher
relaying does not go through multiple training passes, the number of data
loading decreases, leading to higher throughput.

One minor trade-off is communication overhead. In the baseline, gradient
sharing must occur after every backward pass. With teacher relaying, there is
some communication delay from relaying the intermediate activation values from
one device to another. However, the communication time is almost negligible in
our target settings of single-node multi-device training. Furthermore, in both
cases, most of the communications overlapped with computations.

###  IV-B Decoupled Parameter Update

Although teacher relaying removes the redundant teacher executions, the
removed redundancy is not directly translated to speedup. At the beginning of
each step, each device has to wait until the previous device delivers the
intermediate activation. Fig. 3(c) illustrates how decoupled parameter update
addresses this problem. As soon as the backward pass of each block is
complete, the parameter updates are performed without waiting for the other
devices. Then, the teacher execution of the next step can start earlier,
increasing the training throughput. This does not harm the training accuracy
by any means because the student blocks have no dependency on the weight
parameters of the other blocks, which is a special characteristic of blockwise
knowledge distillation training.

Decoupled parameter update successfully hides the teacher waiting time except
for the beginning of each epoch, where full synchronization is needed for
validating the whole model. Because there are usually tens to hundreds of
steps per epoch, such overhead is amortized to a negligible amount.

###  IV-C Automatic Hybrid Distribution

With the teacher relaying and decoupled parameter update, the system
throughput is determined by the throughput of the slowest device. Because of
this, load balancing between devices plays a critical role in performance. One
straightforward and intuitive load-balancing method is distributing the
workload in contiguous blocks. The distribution is simple because there are
only ${}_{B-1}C_{N-1}$ choices for $B$ blocks and $N$ devices. Unfortunately,
the naive distribution scheme often fails to provide a good balance. In
blockwise distillation, the number of blocks $B$ is determined by the neural
network architecture. Usually, $B$ is around ten [19, 3] and $N$ is four to
eight within a single server. Because there are not enough number of blocks to
distribute to the devices, the naive distribution is likely to end up in a
severe workload imbalance.

With automatic hybrid distribution, we provide another degree of freedom for
workload distribution as presented in Fig. 3(d). Instead of relying on the
block granularity, we allow further splitting of each block along the batch
dimension. Thus, when a block is too long, it can be split into two or more
smaller effective blocks. Because a batch is split, the total workload can
become larger because of GPU under-utilization. However, sometimes a slight
increase in the total workload is dwarfed by the gain from workload balancing.

Automatic hybrid distribution introduces a larger design space to workload
distribution, which is difficult to tune manually. To estimate throughputs of
possible schedules, we measure consumed time of a few test execution for each
block under feasible batch sizes. Then, considering the practical problem size
of both $B$ and $N$ at around ten, the optimal solution can be found using an
exhaustive search. Because the decision is made only once at the beginning,
its overhead is amortized over the entire training and is negligible in our
experiments.

##  V Pipe-BD Framework

###  V-A Overall Procedure

Algorithm 1 Pipe-BD procedure

1:

2: $G$: # of devices, $D_{i}$: $i$-th device

3: $T_{i}$: Teacher blocks assigned to $D_{i}$

4: $S_{i}$: Student blocks assigned to $D_{i}$

5:Initialization: Decide $T_{i}$ and $S_{i}$ of each device // AHD

6:for each epoch do

7: for parallel $i=0,1,$ … $,G-1$ do

8: for each step do

9: if $D_{i}$.prev == $\emptyset$ then $act_{in_{i}}$ = load_data()

10: else $act_{in_{i}}$ = receive(from=$D_{i}$.prev) // TR

11: $act_{t\\_out_{i}}$ = $T_{i}$.forward($act_{in_{i}}$)

12: if $D_{i}$.next != $\emptyset$ then send($act_{t\\_out_{i}}$,
to=$D_{i}$.next) // TR

13: $act_{s\\_out_{i}}$ = $S_{i}$.forward($act_{in_{i}}$)

14: $S_{i}$.backward($\mathcal{L}(act_{s\\_out_{i}},act_{t\\_out_{i}}$))

15: if AHD_enabled then $S_{i}$.share_gradient() // AHD

16: if $\sim$DPU_enabled then wait_all_devices() // DPU

17: $S_{i}$.update_weight()

18: end for

19: end for

20:end for

Algorithm 1 displays the overall procedure for Pipe-BD. At initialization, the
optimal schedule is decided from the profiled results, and the blocks are
assigned to the devices (line 4). At the beginning of each step, each device
receives the intermediate activation from the previous device (line 9). If
$T_{i}$ and $S_{i}$ contain the first block, the device instead starts with
loading the data (line 8). It mostly overlaps with the computation except for
the first step in each epoch. After the teacher forward pass is completed
(line 10), the result is sent out such that the next device can execute
$T_{i+1}$. (line 11). Then $S_{i}$ is executed (lines 12-13). If automatic
hybrid distribution made a decision to share the block with other devices,
gradient sharing is performed (line 14). Finally, decoupled parameter update
(line 15) removes barrier operation, enabling each device to update its
student weight without waiting for the other devices.

###  V-B Implementation

We used a native PyTorch distributed package for point-to-point
communications. All communications are implemented to overlap with
computations as much as possible. We used Pytorch DistributedDataParallel
class for data-parallel communications. For automatic hybrid distribution, the
profiling function is called before training, which runs 100 steps of each
block with feasible batch sizes to obtain execution times under the current
environment. Based on these profiled execution times, Pipe-BD determines the
best scheduling and starts training. The implementation of Pipe-BD is
available at https://github.com/hongsunjang/Pipe-BD.

##  VI Experimental Setup

###  VI-A Workload

To demonstrate the advantage of Pipe-BD, we applied it to two popular
blockwise distillation applications.

Neural Architecture Search. NAS is the current de facto standard for building
a new neural network architecture. To search for a final architecture,
multiple candidate operations in each layer are associated with a trainable
_architecture parameter_ , representing the probability of selecting the
operation every step. After the entire network is trained, the operation with
the highest probability within each layer is selected as the final
architecture. For an efficient search, blockwise distillation is a popular
method [9, 12] for a smaller solution space. One notable aspect of NAS is that
each step periodically requires two rounds of forward/backward passes for
students: one for the architecture parameters and another for the weight
parameters. However, this does not cause any difference to blockwise
distillation or Pipe-BD because each round can be regarded as a single
training step. We used ProxylessNAS [3] as the search backbone. For the
teacher model, we used pre-trained MobileNetV2 [20]. For other settings, we
followed the values suggested from the official implementations of DNA [9].

Model Compression. Model compression is also another popular application of
blockwise knowledge distillation [7, 11]. A small student neural network model
is trained from a larger pretrained teacher model. We follow the tradition and
use layers of VGG-16 [21] as the teacher model and depth-wise separable
convolution (DS-Conv) [22] layers as replacements. We follow the settings from
[7] for the training.

###  VI-B Experimental Environment

For the experiments, we use two types of environments. By default, four RTX
A6000 GPUs (Ampere) are attached to an AMD EPYC 7302 CPU. For additional
experiments on a slightly low-cost configuration, four RTX 2080Ti GPUs
(Turing) are attached to two Intel Xeon Silver 4214 CPUs. We used two
datasets, CIFAR-10 [23] and ImageNet [1]. For the model compression, we used
stochastic gradient descent (SGD) optimizer with a learning rate of 0.1 for
compressing and 0.0001 for finetuning. For the NAS, we used SGD optimizer with
a learning rate of 0.005 for neural network architecture searching and 0.1 or
0.05 for retraining the final architecture.

TABLE I: Experimental Environment HW |  Default (w/ A6000) | GPU |  4$\times$ NVIDIA RTX A6000
---|---|---|---
CPU |  1$\times$ EPYC 7302, 16 cores
Memory | 256 GB DDR4 ECC
Interconnect | PCIe 4.0
Alternative (w/ 2080Ti) | GPU |  4$\times$ NVIDIA RTX 2080Ti
CPU |  2$\times$ Xeon 4214 Silver, 12 cores
Memory | 256 GB DDR4 ECC
Interconnect | PCIe 3.0
SW | Common | Python | 3.10
CUDA | 11.6
PyTorch | 1.13
NAS | Teacher Model | MobileNetV2
Kernel Size | 3,5,7
Expansion Ratio | 3,6
Model Compression | Teacher | VGG-16
Replacement | DS-Conv

###  VI-C Baselines

Based on the prior work mentioned in Section II-B, we used two baselines for
our experiments. The first baseline (DP) is the traditional _data-parallel_
blockwise distillation used in [9] official implementation. The second
baseline (LS) is the _layerwise scheduling_ introduced in [7]. Each baseline
targets either one of neural architecture search or model compression, so we
implemented these baselines to both of our target workloads in PyTorch.

##  VII Experimental Results

###  VII-A Speedup and Ablation

Figure 4: Speedup and ablation of baselines and Pipe-BD.

Fig. 4 shows the speedup of Pipe-BD over the baselines with an ablation study
of the proposed techniques using four RTX A6000 GPUs. Each colored bar shows
the speedup of Pipe-BD where 1) only teacher relaying is applied (TR), 2)
decoupled parameter update is further applied (TR+DPU), and 3) all three
schemes are applied, including automatic hybrid distribution (TR+DPU+AHD). In
addition, we tested an alternative method named _Internal Relaying_ (TR+IR).
With internal relaying, each device trains all existing blocks in every step,
and parallelization is obtained via data parallelism. Instead of re-executing
the teacher blocks or relaying activations between devices, the teacher
activations are internally stored in memory and are retrieved for the next
block. This approach allows for removing the redundancies of teacher and data
loading as well as the load imbalance. However, it has the disadvantage of
using a small batch size per device. In fact, internal relaying is a special
case of Pipe-BD with TR+DPU+AHD when all blocks are only split along the batch
dimension.

Among the baselines, LS performs better than DP on Cifar-10 but worse on
ImageNet. Because the composition of the neural networks for ImageNet
typically has a few heavy blocks, LS suffers from severe load imbalance.
Nevertheless, they both perform inferior to Pipe-BD. TR provides speedup for
all cases due to eliminating extra data loading, redundant teacher execution,
and enhancing resource utilization. Further, DPU provides additional speedup
by removing synchronization barriers, which improves the overlapping of the
teacher waiting time with student executions. Additionally, AHD removes the
pipeline bubbles by balancing workloads, which drives an additional speedup
over TR+DPU.

With the ImageNet dataset, its larger spatial dimension of the images
(224$\times$224 vs 32$\times$32) leads to heavy workloads in the first block.
As a result, with TR only, the execution time of block 0 dominates all the
others. Because of this, DPU has little room for improvement, whereas
splitting the workload of the first block with AHD has a large impact on
reducing the bubbles. In contrast, in the Cifar-10 case, the workload is
already well-balanced only with TR+DPU version, and the gain from more
balancing is offset by the loss from lower resource utilization caused by AHD.

TABLE II: Parallel Blockwise Distillation Training Results

Task |  Dataset | Teacher | Student | Elapsed Time (1 epoch)
---|---|---|---|---
Model | #Params | FLOPs | Acc. (%) | Backbone | #Params | FLOPs | Acc. (%) | DP | LS | Pipe-BD
NAS | Cifar-10 | MobileNetV2 | 2.24 M | 87.98 M | 95.42% |  ProxylessNAS [3] | 1.40 M | 76.10 M | 95.48% | 31.52s. | 16.33s. | 10.23s.
ImageNet | MobileNetV2 | 3.50 M | 300.77 M | 72.00% | ProxylessNAS | 4.22 M | 420.20 M | 74.54% | 62m 21s. | 125m 26s. | 14m 15s.
Compression | Cifar-10 | VGG-16 | 14.72 M | 0.63 B | 91.85% |  DS-Conv [22] | 7.25 M | 0.39 B | 91.51% | 13m 18s. | 6m 37s. | 1m 49s.
ImageNet | VGG-16 | 138.36 M | 30.98 B | 71.59% | DS-Conv | 138.09 M | 26.15 B | 71.32% | 229m 23s. | 566m 49s. | 60m 39s.

###  VII-B Sensitivity and Scheduling

(a) Speedup.

(b) 2080Ti schedule.

(c) A6000 schedule.

Figure 5: GPU type sensitivity of Pipe-BD on NAS.

Fig. 5b and Fig. 5c show how Pipe-BD automatically determines the appropriate
schedule according to two different environmental settings with the same NAS
on ImageNet workload. While the speedup trends are similar in Fig. 5a, they
are from different schedules. The execution time of block 0 is the longest
among the six blocks in both settings. However, the gap is more extensive on
A6000 than on 2080Ti. To mitigate the imbalance, Pipe-BD settles at a schedule
where the first three blocks (0-2) are shared on three devices (0-2) for
A6000, while block 0 on 2080Ti is shared among two devices (0-1) and two
blocks (1-2) are assigned to device 2.

In Fig. 6, we demonstrate the sensitivity to the batch size on the NAS
workload, normalized against DP of each batch size. In general, the advantage
of Pipe-BD is not very sensitive to the batch size. One common trend is that
the speedup is better in smaller batch sizes because the resource utilization
difference becomes more significant with smaller batch sizes. One exception is
AHD for ImageNet, where the speedup is better on larger batch sizes. The
reason is found in the schedule depicted in Fig. 5(c) which uses three-way
data parallelism to balance workloads. Because the training time for the
student is shorter in both the baseline and Pipe-BD with larger batch sizes,
reduction in the teacher redundancy and extra data loading account more for
the overall speedup.

###  VII-C Memory Overhead

Fig. 7 depicts the memory overhead of Pipe-BD on the NAS task for each rank
(GPU). Due to the characteristics of CNN-based models, lower-indexed teacher
blocks generally have larger feature map sizes. TR and DPU consume more memory
than DP because of this characteristic, especially on rank 0. This outcome is
also demonstrated in Fig. 7b because models for ImageNet contain even larger
feature map sizes in lower-indexed blocks. However, AHD successfully addresses
this issue using data parallelism in a hybrid manner, which lessens the memory
overhead of earlier ranks as depicted in Fig. 5(c). As a result, Pipe-BD
provides superior multi-fold speedups with a minor 8.7% and 21.3% additional
memory overheads over DP on average for Cifar-10 and ImageNet, respectively.

Figure 6: Batch size sensitivity of Pipe-BD on NAS.

###  VII-D Training Quality

Pipe-BD has no component that can hurt the accuracy because it only alters the
scheduling strategy. Nonetheless, we report the accuracy in Table II to
demonstrate that the Pipe-BD framework faithfully reproduces the end-to-end
training results in the prior art, with much shorter training time. For all
use cases under evaluation, Pipe-BD achieves significant speedup with the same
accuracy.

Figure 7: Memory overhead of Pipe-BD on NAS.

##  VIII Conclusion

We propose Pipe-BD, a novel parallelization method for blockwise distillation.
By restructuring the existing parallelization scheme, we achieve a multi-fold
speedup on various use cases. In this study, we focused on a single-node,
multi-GPU setting since it is the most common setup. However, if the method
has to be scaled for a multi-node setting, the communication overhead needs to
be addressed. Along with the heterogeneous GPU/servers, this will be our
future direction.

## Acknowledgement

This work was partly supported by the National Research Foundation of Korea
(NRF) grants (2022R1C1C1011307, 2022R1C1C1008131) and Samsung Electronics Co.,
Ltd (IO221213-04119-01) and Institute of Information & communications
Technology Planning & Evaluation (IITP) grants (2020-0-01361) funded by the
Korean government (MSIT).

## References

  * [1] J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei, “ImageNet: A large-scale hierarchical image database,” in _CVPR_ , 2009.
  * [2] P. Goyal, P. Dollár, R. Girshick, P. Noordhuis, L. Wesolowski, A. Kyrola, A. Tulloch, Y. Jia, and K. He, “Accurate, Large Minibatch SGD: Training Imagenet in 1 Hour,” _arXiv preprint arXiv:1706.02677_ , 2017.
  * [3] H. Cai, L. Zhu, and S. Han, “ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware,” in _ICLR_ , 2019.
  * [4] H. Liu, K. Simonyan, and Y. Yang, “DARTS: Differentiable Architecture Search,” in _ICLR_ , 2019.
  * [5] Y. Huang, Y. Cheng, A. Bapna, O. Firat, D. Chen, M. Chen, H. Lee, J. Ngiam, Q. V. Le, Y. Wu, and z. Chen, “GPipe: Efficient Training of Giant Neural Networks Using Pipeline Parallelism,” in _NeurIPS_ , 2019.
  * [6] S. Zhao, F. Li, X. Chen, T. Shen, L. Chen, S. Wang, N. Zhang, C. Li, and H. Cui, “Naspipe: high performance and reproducible pipeline parallel supernet training via causal synchronous parallelism,” in _ASPLOS_ , 2022\.
  * [7] C. Blakeney, X. Li, Y. Yan, and Z. Zong, “Parallel Blockwise Knowledge Distillation for Deep Neural Network Compression,” _IEEE TPDS_ , vol. 32, no. 07, pp. 1765–1776, 2021.
  * [8] H. Wang, H. Zhao, X. Li, and X. Tan, “Progressive Blockwise Knowledge Distillation for Neural Network Acceleration,” in _IJCAI_ , 2018.
  * [9] C. Li, J. Peng, L. Yuan, G. Wang, X. Liang, L. Lin, and X. Chang, “Block-Wisely Supervised Neural Architecture Search With Knowledge Distillation,” in _CVPR_ , 2020.
  * [10] L. Yuan, F. E. Tay, G. Li, T. Wang, and J. Feng, “Revisiting knowledge distillation via label smoothing regularization,” in _CVPR_ , 2020.
  * [11] J. Yu, S. Kang, and K. Choi, “Network recasting: a universal method for network architecture transformation,” in _AAAI_ , 2019.
  * [12] B. Moons, P. Noorzad, A. Skliar, G. Mariani, D. Mehta, C. Lott, and T. Blankevoort, “Distilling Optimal Neural Networks: Rapid Search in Diverse Spaces,” in _ICCV_ , 2021.
  * [13] N. S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang, “On large-batch training for deep learning: Generalization gap and sharp minima,” in _ICLR_ , 2017.
  * [14] D. Narayanan, A. Harlap, A. Phanishayee, V. Seshadri, N. R. Devanur, G. R. Ganger, P. B. Gibbons, and M. Zaharia, “PipeDream: Generalized Pipeline Parallelism for DNN Training,” in _SOSP_ , 2019.
  * [15] M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro, “Megatron-LM: Training Multi-billion Parameter Language Models Using Model Parallelism,” _arXiv preprint arXiv:1909.08053_ , 2019.
  * [16] J. Song, J. Yim, J. Jung, H. Jang, H.-J. Kim, Y. Kim, and J. Lee, “Optimus-CC: Efficient Large NLP Model Training with 3D Parallelism Aware Communication Compression,” in _ASPLOS_ , 2023.
  * [17] D. Kinghorn, “Gpu memory size and deep learning performance (batch size),” 2018\. [Online]. Available: https://www.pugetsystems.com/labs/hpc/GPU-Memory-Size-and-Deep-Learning-Performance-batch-size-12GB-vs-32GB----1080Ti-vs-Titan-V-vs-GV100-1146/
  * [18] “Deep learning frameworks speed comparison,” 2017. [Online]. Available: https://wrosinski.github.io/deep-learning-frameworks/
  * [19] K. He, X. Zhang, S. Ren, and J. Sun, “Deep Residual Learning for Image Recognition,” in _CVPR_ , 2016.
  * [20] M. Sandler, A. Howard, M. Zhu, A. Zhmoginov, and L.-C. Chen, “MobileNetV2: Inverted Residuals and Linear Bottlenecks,” in _CVPR_ , 2018.
  * [21] K. Simonyan and A. Zisserman, “Very Deep Convolutional Networks for Large-Scale Image Recognition,” in _ICLR_ , 2015.
  * [22] A. G. Howard, M. Zhu, B. Chen, D. Kalenichenko, W. Wang, T. Weyand, M. Andreetto, and H. Adam, “MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications,” _arXiv preprint arXiv:1704.04861_ , 2017.
  * [23] A. Krizhevsky _et al._ , “Learning multiple layers of features from tiny images,” 2009.