Integrating deep learning accelerators with TensorFlow

Description

Deep learning computation uses mixed-precision linear algebra and can benefit from specialized accelerators, such as NVIDIA’s GPU, Google’s tensor processing unit (TPU), and AWS’s Inferentia, to optimize performance, cost, and power. These accelerators need to be integrated with deep learning frameworks and exposed via existing programming interfaces within the frameworks so that developers and data scientists can leverage hardware acceleration with minimal changes to their existing model pipeline code.

Sudipta Sengupta dives into his experience with integrating Amazon Elastic Inference and AWS Inferentia with TensorFlow in the AWS cloud. Accelerators often come with runtimes that optimize model execution by performing whole graph optimization during compilation—ahead of time (AOT) or just in time (JIT). These runtimes support only a subset of the operators available in the framework while expanding operator coverage over time.

Sudipta builds upon and extends TensorFlow mechanisms to use a sub-graph-level interface to the respective accelerator runtimes in the AWS ecosystem, while respecting restricted operator coverage of the accelerator runtime and resource constraints of individual accelerators. The optimizations built into TensorFlow determine efficient ways to run the arithmetic operations in the model and distribute graph execution intelligently across the accelerator(s) and CPU on the instance. This is achieved by analyzing the computation needs of each model along with the accelerator and CPU resources that are available and optimizing the placement of operators across host and accelerator(s).

Going forward, there needs to be standard and generic mechanisms in TensorFlow for plugging in accelerator runtimes, with flexibility to customize for operator coverage, graph partitioning, support multiple devices and multiple compute elements within same device, and compilation mode (AOT versus JIT).