Workload-Aware Hardware Accelerator Mining for Distributed Deep Learning Training

TL;DR

The paper tackles the challenge of designing hardware accelerators optimized for end-to-end distributed DNN training, where both forward and backward passes and weight updates must be considered. It introduces WHAM, a workload-aware search framework that uses a critical-path-based heuristic, ILP, and a top-k global strategy to explore a tunable architectural template combining Tensor Core and Vector Core resources under area and power constraints. WHAM demonstrates substantial gains in throughput and energy efficiency compared to inference-focused baselines and hand-tuned TPU/NVDLA-style designs, and it scales to distributed training with pipeline and tensor-model parallelism. This approach enables faster convergence in architecture search (up to 174× faster than some baselines) and improved end-to-end performance (up to 12–23% higher throughput and up to 8.1× Perf/TDP improvements) across a diverse set of DNN workloads, offering a practical path to optimized training accelerators. The framework’s ability to produce both per-model specialized designs and common designs for multiple workloads highlights its potential to guide accelerator development for large-scale distributed training deployments.

Abstract

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.

Figures (14)

• Figure 1: Design space exploration with WHAM.
• Figure 2: The per-layer utilization of tensor cores and vector cores in Inception_v3 model using a single tensor (256 $\times$ 256) and vector core (256 wide). The y-axis is capped at 50%. Layers with fewer channels have lower utilization.
• Figure 3: Architectural template. WHAM explores TC dimensions, VC width, and number of TCs and VCs.
• Figure 4: WHAM's accelerator search takes an algorithmic approach to prune the large search space for training.
• Figure 5: ASAP and ALAP schedules. For simplicity, each operator executes with a unit cycle.