In this paper, we present a novel technique to search for hardware architectures of accelerators optimized for end-to-end training of deep neural networks (DNNs). Our approach addresses both single-device and distributed pipeline and tensor model parallel scenarios, latter being addressed for the first time. The search optimized accelerators for training relevant metrics such as throughput/TDP under a fixed area and power constraints. However, with the proliferation of specialized architectures and complex distributed training mechanisms, the design space exploration of hardware accelerators is very large. Prior work in this space has tried to tackle this by reducing the search space to either a single accelerator execution that too only for inference, or tuning the architecture for specific layers (e.g., convolution). Instead, we take a unique heuristic-based critical path-based approach to determine the best use of available resources (power and area) either for a set of DNN workloads or each workload individually. First, we perform local search to determine the architecture for each pipeline and tensor model stage. Specifically, the system iteratively generates architectural configurations and tunes the design using a novel heuristic-based approach that prioritizes accelerator resources and scheduling to critical operators in a machine learning workload. Second, to address the complexities of distributed training, the local search selects multiple (k) designs per stage. A global search then identifies an accelerator from the top-k sets to optimize training throughput across the stages. We evaluate this work on 11 different DNN models. Compared to a recent inference-only work Spotlight, our method converges to a design in, on average, 31x less time and offers 12x higher throughput. Moreover, designs generated using our method achieve 12% throughput improvement over TPU architecture.

翻译：本文提出了一种新颖的技术，用于搜索针对深度神经网络（DNN）端到端训练优化的加速器硬件架构。我们的方法同时覆盖单设备场景以及分布式流水线和张量模型并行场景，后者是首次被研究。该搜索在固定面积和功耗约束下，优化了与训练相关的指标（如吞吐量/热设计功耗）。然而，随着专用架构和复杂分布式训练机制的普及，硬件加速器的设计空间探索变得极其庞大。此前的研究通过将搜索空间缩小到仅针对推理的单加速器执行，或针对特定层（如卷积层）调整架构来应对这一挑战。相反，我们采用了一种独特的基于启发式关键路径的方法，以确定可用资源（功耗和面积）的最佳利用方式——无论是针对一组DNN工作负载，还是针对每个工作负载单独优化。首先，我们执行局部搜索，为每个流水线和张量模型阶段确定架构。具体而言，系统迭代生成架构配置，并使用一种新颖的基于启发式的方法进行设计调整，该方法优先为机器学习工作负载中的关键算子分配加速器资源和调度。其次，为应对分布式训练的复杂性，局部搜索为每个阶段选择多个（k个）设计。随后，全局搜索从top-k集合中识别出加速器，以优化跨阶段的训练吞吐量。我们在11个不同的DNN模型上评估了该方法。与近期仅专注于推理的Spotlight工作相比，我们的方法平均收敛速度快31倍，吞吐量提升12倍。此外，使用我们的方法生成的设计相比TPU架构实现了12%的吞吐量提升。