  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.
