Автор: Александр Наздрюхин Редактура: Ксения Рябинова

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

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

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

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

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

Изображение 1. GeForce 8800 Ultra (в целом, устройство видеокарт не сильно изменилось с 2006 года, разве что экстенсивно)

Изображение 1. GeForce 8800 Ultra (в целом, устройство видеокарт не сильно изменилось с 2006 года, разве что экстенсивно)

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

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

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

Google TPU (Tensor Process Unit)

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