Ускорители нейронных сетей

CPU играет первостепенную роль в обучении и инференсе нейронных сетей. Остальные ускорители выступают как сопроцессоры: часть из них массивна и самостоятельна (как, например, GPU), часть может выполнять только особые операции, поэтому для их полноценного использования нейросеть будет должна пройти сложные этапы компиляции и перестройки графа. Характерный пример таких ускорителей — Google TPU. И как раз про него мы и расскажем в этой статье.

Графические ускорители чрезвычайно эффективны, но у них есть узкие места, которые можно оптимизировать:

1. Хоть GPU-ядра значительно проще, чем у CPU, они все же не упрощены до нужного нам предела: каждый потоковый мультипроцессор (Streaming Multiprocessor, SM), а их несколько десятков на GPU, содержит несколько скалярных ядер (обычно кратно восьми, они также называются CUDA-ядрами), модули специальных функций, блок управления командами и собственную память. Все это несильно нужно для выполнения матричных умножений.
2. Устройство памяти также избыточно. И хоть некоторые ее части (например, текстурная память) отсутствуют в серверных GPU, ее можно упростить еще сильней.
3. Из-за множества модулей усложняется топология устройства: необходимо выделять больше места для связей между модулями (так, например, SM связан с памятью и SM связаны между собой)
4. Из-за сложных ядер и топологии у GPU высокое энергопотребление и для размещения всех модулей необходимы большие кристаллы. Это означает, что у GPU дорогие эксплуатация и производство.

На изображении ниже легко можно увидеть, сколько места занимают память и шины, а сколько — SM (который тоже избыточен по своему функционалу).

Таким образом, для эффективного исполнения нейронных сетей необходимо только несколько модулей:

1. Некоторый аналог скалярных ядер для выполнения операций умножения и сложения;
2. Коммуникационная система;
3. Память устройства с высокой скоростью доступа.

Именно по этим принципам построен Google TPU v1.

Google TPU (Tensor Process Unit)

Давайте рассмотрим его с точки зрения аппаратной архитектуры:

Matrix Multiply Unit (MXU)

Основное и самое главное отличие TPU от GPU — Matrix Multiply Unit, представляющий собой двумерный систолический массив из \(256х256\) вычислительных единиц (systolic computing unit). Каждая вычислительная единица выполняет операцию умножения-сложения за один такт. Таким образом, при полной нагрузке каждый столбец систолического массива MXU делает 256 умножений-сложений за такт — то есть, формирует один элемент выходного вектора.

Как видно на изображении 2, MXU занимает 24 процента схемы, в то же время CUDA-ядра, если судить по изображению 1, занимают около 5 процентов.

Проиллюстрируем работу систолического массива примером умножения двух матриц \( А \) и \( В \) размера \( 3х3 \):

1. На нулевом такте мы загружаем матрицу \( B \) в память массива.
2. На первом такте происходит одно умножение-сложение, накопление частичной суммы в аккумуляторе, распределенном по систолическому массиву, и передача данных дальше по линиям массива.

3. На втором такте происходят уже три умножения-сложения. Легко заметить, что на третьем такте мы получим результат для первого элемента матрицы \( АB \), а на пятом — систолический массив будет полностью нагружен.

MXU может выполнять только матричные умножения и свертки im2col. Остальные операции (функция активации, нормализация и, если нужно, пулинг) производятся в специальных модулях. Все промежуточные результаты сохраняются в буффер (Local Unified Buffer), если же буффера не хватает — всегда есть 8 GB DRAM. Google TPU v1 поддерживает только int8 арифметику — то есть, предполагается, что все веса и активации сетей квантованы в этот тип данных (в более поздних версиях добавили поддержку других типов).

С точки зрения программной части TPU тоже устроен просто. У него всего пять CISC инструкций, функционал которых понятен из названия:

• Read_Host_Memory
• Read_Weights
• MatrixMultiply/Convolve
• Activate
• Write_Host_Memory

Эти инструкции передаются из хоста на TPU по PCIe шине.

Однако хоть TPU и спроектирован для запуска нейронных сетей, возникает несколько вопросов к архитектуре:

1. Что если размерности матрицы В не равны 256х256?
2. У современных сетей множество специфичных слоев: начиная от билинейной интерполяции и заканчивая свертками Фурье. Как вычислять их?

Именно поэтому Google разработала компилятор нейронных сетей — XLA (Accelerated Linear Algebra). Но о нем и о различных усовершенствованиях архитектуры TPU вы узнаете уже в следующей статье.

Статьи и книги:

(1) E. Lindholm et al. Nvidia Tesla: A Unified Grafics and Computer Architecture

(2) In-Datacenter Performance Analysis of a Tensor Processing Unit

(3) Ten Lessons From Three Generations Shaped Google’s TPUv4i

(4) X. Liang Ascend AI Processor Architecture and Programming