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.
翻译:训练大型深度神经网络模型因其巨大的计算和内存需求而极具挑战性。逐块蒸馏通过将大型模型拆分为多个较小模型,为加速收敛提供了一种有前景的方法。在最先进的逐块蒸馏方法中,训练以数据并行的方式在多GPU上逐块进行。为了生成学生块的输入,教师模型需从起始执行到当前训练块。然而,这会导致教师模型冗余执行的高开销、GPU利用率低以及额外数据加载问题。为解决这些问题,我们提出Pipe-BD,一种针对逐块蒸馏的新型并行化方法。Pipe-BD主动利用流水线并行进行逐块蒸馏,消除了教师块的冗余执行,并增加了每设备批次大小以提升资源利用率。我们还扩展了混合并行以实现高效的工作负载平衡。因此,Pipe-BD在不修改逐块蒸馏数学公式的前提下实现了显著加速。我们在PyTorch上实现了Pipe-BD,实验表明Pipe-BD在多种场景、模型和数据集上均有效。