Мощный чип, специализирующийся на машинном обучении

TL; DR - это ускоритель, созданный на основе 16-битного блока систолической матрицы с умножением и накоплением 128x128 («MXU»). Если вам этого достаточно - отлично! Если нет, читайте дальше ...

Введение

Возможно, вы слышали, что у Google есть специальный чип для машинного обучения. Он называется TPU («Tensor Processing Unit») и представляет собой все возможное, чтобы Google вложил столько возможностей машинного обучения, сколько они могут уместить в одном чипе. Google Cloud дает разработчикам возможность использовать эту мощь, но сам чип - это по большей части черный ящик… или нет? Давайте отодвинем несколько слоев, чтобы увидеть, что внутри волшебства.

В этой статье я попытаюсь рассмотреть системную архитектуру TPU, сохраняя при этом достаточно простые вещи, чтобы разработчик программного обеспечения с минимальным опытом работы с оборудованием мог следовать. Однако я не скуплюсь на технические детали. Я собираюсь дать как можно более полный обзор высокого уровня.

— — —

Я не связан с Google, я всего лишь временный инженер-компьютерщик, пытающийся обобщить и объяснить, какая документация существует.

Вывод с высокой мощностью

Обучение и запуск нейронных сетей требует больших вычислительных мощностей. Еще в 2013 году Google¹ провел несколько простых расчетов, чтобы узнать, что им нужно для голосового поиска, и результаты оказались неожиданными:

Если бы мы рассмотрели сценарий, в котором люди используют голосовой поиск Google всего три минуты в день, и запустили глубокие нейронные сети для нашей системы распознавания речи на процессорах, которые мы использовали, нам пришлось бы удвоить количество центров обработки данных Google!

- Норм Джуппи

ASIC

Нейронные сети - мощные инструменты, но они будут стоить слишком дорого (даже для Google), чтобы их можно было везде запускать на стандартных компьютерах.

К счастью, для выполнения тяжелой работы вам не нужен стандартный компьютер. В какой-то момент становится рентабельным разработать специальный чип, способный выдержать такой вес. Этот заказной чип представляет собой специализированную интегральную схему (ASIC).

Обычно ASIC доставляют больше проблем, чем они того стоят. На их разработку уходит много времени: Google потратил 15 месяцев на TPUv1, и это было поразительно быстро. Они изначально дороги, требуют специализированных инженеров и затрат на производство, которые начинаются от миллиона долларов. И они негибкие: нет возможности заменить чип, когда он будет готов.

Но если вы знаете, что будете выполнять одну конкретную работу в достаточном объеме, повторяющиеся преимущества могут компенсировать первоначальные недостатки. ASIC обычно являются самым быстрым и наиболее энергоэффективным способом выполнения задачи. Google хотел, чтобы эта производительность работала в нейронных сетях, и результатом стал TPU.

Скаляр, Вектор, Матрица

Нейронная сеть требует много математики, но большая часть математики довольно проста: умножьте кучу чисел и сложите результаты. Вы можете соединить эти два значения вместе с помощью одной операции, называемой умножение-накопление (MAC). И если нам больше ничего не нужно делать, мы можем очень, очень быстро умножать-накапливать.

Без нового чипа мы бы сделали это либо с центральным, либо с графическим процессором. ЦП - это скалярная машина, что означает, что он обрабатывает инструкции пошагово. Это отлично подходит для приложений общего назначения, таких как ваш ноутбук или сервер, но мы можем добиться большей производительности, специализируясь.

Графический процессор - это векторная машина. Вы можете предоставить ему длинный список данных - одномерный вектор - и одновременно выполнять вычисления для всего списка. Таким образом, мы можем выполнять больше вычислений в секунду, но мы должны выполнять те же вычисления с вектором данных параллельно. Этот вид вычислений хорошо подходит для графики или крипто-майнинга, когда одно задание должно выполняться много раз.

Но мы все еще можем добиться большего. Данные нейронной сети организованы в матрицу, двумерный вектор. Итак, построим матричную машину. И мы действительно заботимся только о умножении-накоплении, поэтому мы будем отдавать ему приоритет перед другими инструкциями, которые обычно поддерживает процессор. Мы посвятим большую часть нашего чипа MAC, которые выполняют умножение матриц, и в большинстве случаев игнорируем другие операции².

Введите систолический массив

Способ достижения такой производительности матрицы заключается в архитектуре, называемой систолическим массивом. Это интересный момент, и именно поэтому TPU так эффективен. Систолический массив - это своего рода аппаратный алгоритм, который описывает структуру ячеек на микросхеме, которая вычисляет матричное умножение. «Систолический» описывает, как данные перемещаются по чипу волнами, как биение человеческого сердца³.

Внимание: впереди подробности. Если вас не волнует, как работает систолический массив, прокрутите вниз, пока не увидите цветок.

Есть несколько вариантов дизайна систолического массива; здесь я говорю о версии, реализованной в ТПУ.

Рассмотрим операцию умножения матриц:

Для входов 2x2 каждый член на выходе представляет собой сумму двух произведений. Ни один продукт не используется повторно, но используются индивидуальные условия.

Мы реализуем это, построив сетку 2x2. (Это на самом деле сетка, а не просто абстракция - оборудование - это весело). Обратите внимание, что 2x2 - это игрушечный пример, а полноразмерный MXU - чудовищный 128x128.

Допустим, AB / CD представляют наши активации, а EF / GH - наши веса. Для нашего массива мы сначала загружаем веса следующим образом:

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

Каждый цикл часов каждая ячейка будет выполнять следующие шаги, все параллельно:

  1. Умножьте наш вес на активацию слева. Если слева нет ячейки, возьмите из очереди ввода.
  2. Добавьте этот продукт к частичной сумме, поступающей сверху. Если ячейки сверху нет, то частичная сумма сверху равна нулю.
  3. Пройдите активацию в ячейку справа. Если справа нет ячейки, выбросьте активацию.
  4. Передайте частичную сумму в ячейку внизу. Если внизу нет ячейки, выведите частичную сумму.

(Мокап этих правил на языке Python можно найти здесь.)

По этим правилам вы можете видеть, что активация начинается слева и перемещается на одну ячейку вправо за цикл, а частичные суммы начинаются сверху и перемещаются на одну ячейку вниз за цикл.

Поток данных будет выглядеть так:

И это наш массив! Давайте рассмотрим выполнение первого вывода:

Цикл 1

  1. Вверху слева считывается A из входной очереди, умножается на вес E, чтобы получить продукт AE.
  2. AE, добавленная к частичной сумме 0 сверху, дает частичную сумму AE.
  3. Активация A передана в ячейку справа вверху.
  4. Частичная сумма AE передана в ячейку слева внизу.

Цикл 2

  1. Слева внизу считывает B из входной очереди, умножается на вес G для получения продукта BG.
  2. BG, добавленный к частичной сумме AE сверху, дает частичную сумму AE + BG
  3. Активация B передана в ячейку справа внизу
  4. Частичная сумма AE + BG является выходом.

И вы можете видеть, что мы правильно вычислили первый член нашей выходной матрицы. Между тем, в цикле 2 мы также вычисляем CE в верхнем левом углу и AF в верхнем правом углу. Они будут распространяться по ячейкам, и к 4-му циклу мы произведем весь результат 2x2.

Циклы 3 и 4 оставлены в качестве упражнения для читателя.

Вот несколько диаграмм от Google, которые дают немного больше наглядной картины.

Вы увидите, что входные активации смещены нулями, чтобы гарантировать, что они попадают в массив в нужный момент, и что выходы, которые покидают массив, смещены таким же образом. Для полного вычисления матрицы результатов требуется 3n-2 цикла, тогда как стандартное последовательное решение - n³. Неплохо!

Мы можем это сделать, потому что мы выполняем операции MAC 128x128 параллельно. Множители обычно большие и дороги в аппаратной реализации, но высокая плотность систолических массивов позволяет Google упаковать 16 384 из них в MXU. Это напрямую связано с тренировкой скорости и запуском вашей сети.

Веса загружаются почти так же, как и активации - через входную очередь. Мы просто отправляем специальный управляющий сигнал (красный на приведенной выше диаграмме), чтобы сообщить массиву о необходимости сохранять веса по мере их прохождения вместо выполнения операций MAC. Веса остаются в тех же элементах обработки, поэтому мы можем отправить весь пакет перед загрузкой нового набора, что сокращает накладные расходы.

Вот и все! Остальная часть микросхемы важна и заслуживает рассмотрения, но основным преимуществом TPU является его MXU - блок умножения систолической матрицы.

Сделай передышку

Давайте на минутку подумаем о чем-нибудь приятном, например, о тосте с авокадо или солнечной улыбке Эндрю Нг.

Остальная часть этой статьи будет сосредоточена на реализации TPU, и я не буду вдаваться в такие подробности, как ситолический массив. Если вы готовы, мы продолжим.

Остальная сова

У нас есть замечательный систолический массив, но предстоит еще много работы, чтобы создать поддержку и инфраструктуру, чтобы заставить его работать. Во-первых, нам нужен способ передачи данных в сам чип и из него. Затем нам нужно в нужное время вводить и выводить его из массива. И, наконец, нам нужен способ обрабатывать в нейронной сети все, что не умножения матриц. Давайте посмотрим, как все это происходит в аппаратном обеспечении.

Полная система

Ниже вы найдете схему системы и макет для более старого TPUv1. Это не обязательно применимо к новым версиям TPUv2 и TPUv3 - Google внесла некоторые архитектурные изменения с момента выхода этой документации. Однако мне не известны новые справочные материалы или документация по новым чипам, поэтому мы будем придерживаться TPUv1. Я рассмотрю наиболее заметные различия между версиями позже.

На высшем уровне ТПУ спроектирован как ускоритель. Это означает, что он будет подключаться к хост-системе, и хост будет выгружать данные и инструкции, которые будут вычисляться на ускорителе. Результаты возвращаются на хост через тот же интерфейс. По этой модели ускоритель может ускорить длительные и дорогостоящие операции с матрицей, в то время как хост может справиться со всем остальным.

Давайте проверим, что внутри ускорителя, по блок-схемам⁴. Мы рассмотрим их шаг за шагом.

Хост-интерфейс будет подключаться к нашей управляющей машине через PCIe. Мы видим, что через этот интерфейс поступают 3 формы данных: веса (для DDR3), активации (для унифицированного буфера) и инструкции по управлению (для красного Контрольный путь).

Веса одинаковы для всех входных данных в пакете, поэтому они загружаются один раз для каждого пакета. Это нечасто по сравнению с активациями, поэтому веса можно сохранять в медленной DDR3 DRAM вне кристалла. Мы можем подключить хост-интерфейс для записи в эту DDR3, и мы поместим туда все наши веса при загрузке модели. Перед вычислением веса считываются из DDR3 в FIFO весов, что означает, что мы можем предварительно выбрать следующий набор весов во время вычисления текущего пакета.

Унифицированный буфер хранит наши активации. Во время работы хосту необходимо быстро получить доступ к этому буферу, чтобы прочитать результаты и записать новые входные данные. Унифицированный буфер напрямую подключен к MXU, и эти два компонента занимают львиную долю площади чипа (53%). Буфер вмещает до 384 матриц активации с размерами 256x256, что является самым большим пакетом, поддерживаемым чипом. Эти активации считываются и записываются в очень часто, поэтому мы выделяем 30% нашего макета на внутреннюю память и шину 167 ГиБ / с вокруг MXU и пути унифицированного буфера.

MXU записывает обратно в объединенный буфер через аккумуляторы, а затем через конвейер активации. Сначала аккумуляторы собирают данные из MXU. Затем конвейер активации применяет стандартные функции нейронной сети (такие как ReLU и Maxpool⁵), которые не так затратны в вычислительном отношении, как умножение матриц. Когда это будет сделано, мы можем поместить вывод обратно в унифицированный буфер, готовый к следующему этапу обработки.

И еще есть поток управления: красный путь, который принимает инструкции от хоста. Этот путь будет делать такие вещи, как контроль, когда MXU будет умножаться, или какие веса будет загружать FIFO, или какие операции будет выполнять конвейер активации. Думайте об этом как о капитане, который говорит остальной части чипа, что делать.

Вот полная картина казни:

  1. Чип запускается, буферы и DDR3 пустые
  2. Пользователь загружает скомпилированную TPU модель, помещая веса в память DDR3.
  3. Хост заполняет буфер активации входными значениями.
  4. Управляющий сигнал отправляется для загрузки слоя весов в MXU (через FIFO весов)
  5. Хост запускает выполнение, и активации передаются через MXU в аккумуляторы.
  6. Выходные данные проходят через конвейер активации, и новые слои заменяют старые в буфере.
  7. Повторяем с 4 по 6, пока не дойдем до последнего слоя.
  8. Активации последнего уровня отправляются обратно на хост-машину.

Вот оно! Полная картина логического вывода на TPUv1. Нечто подобное происходит в новом TPUv2…

… Но мы не можем подробно рассказать об этом. (Если вы тренируетесь на облачном TPU, вы работаете на TPUv2.) Мы можем сказать одно: TPU нового поколения допускают обучение (т.е. обновление весов), поэтому должен быть путь данных от MXU к хранение веса. В блок-схеме TPUv1 дело обстоит иначе.

Однако, просто зная, на что способен TPUv2, мы можем угадать несколько отличий:

  1. MXU в TPUv1 представлял собой 8-битный целочисленный массив 256x256, больший и менее точный, чем 16-битный bfloat16 128x128 MXU в TPUv2.
  2. Конвейер активации в TPUv1 был заменен полными векторными и скалярными модулями в TPUv2. Это дает ему гораздо более широкий диапазон доступных функций, чем ограниченные Активация и Нормализация / объединение из TPUv1.
  3. Единый буфер, похоже, заменен памятью с высокой пропускной способностью. Это освободит место для остальной части чипа за счет задержки.

Эти различия в первую очередь связаны с тем, что TPUv1 был разработан для вывода, а не обучения, поэтому допустимы арифметические операции с низкой точностью и меньшее количество операций. Обновления означают, что TPU нового поколения намного более гибкие - достаточно, чтобы Google было удобно размещать их в своем облаке.

Почти готово

Еще тост с авокадо. Подробнее Эндрю Нг.

Остальная часть этой статьи охватывает еще несколько концепций, которые важны для работы TPU, но не критичны для архитектуры. Осталось немногое!

Шансы и концы

bfloat16

Большинство вычислений машинного обучения CPU / GPU выполняется с 32-битными числами с плавающей запятой. Когда мы опускаемся до 16 бит, инженеров машинного обучения больше беспокоит диапазон их чисел, чем точность. Можно округлить несколько долей десятичной запятой, но выходить за рамки максимума или минимума числового представления - головная боль. Традиционные float16 обладают высокой точностью и недостаточным диапазоном. Ответ Google на это - формат bfloat16, в котором больше битов посвящено экспоненте, а меньше - значимости.

Для сравнения биты в числе с плавающей запятой:

  • float32: 1-битный знак, 8-битная экспонента, 23-битное значение
  • float16: 1-битный знак, 5-битная экспонента, 10-битная величина
  • bfloat16: 1-битный знак, 8-битная экспонента, 7-битная величина

Поскольку у него одинаковое количество битов экспоненты, bfloat16 - это всего лишь два старших байта float32. Это означает, что он имеет примерно тот же диапазон, что и float32, но с меньшей точностью. На практике эта стратегия работает хорошо. В TPU большая часть данных по-прежнему хранится в формате float32. Однако в MXU есть множители bfloat16 и аккумуляторы float32. Меньшие множители сокращают требования к мощности и площади чипа, оставляя место для большего количества множителей, работающих с более высокими тактовыми частотами. Это улучшает производительность чипа, не теряя при этом точности.

(TPUv1 использовал 8-битные целочисленные операции, что непросто перенести на предварительно обученную модель.)

XLA

XLA - это экспериментальный JIT-компилятор для бэкэнда Tensorflow. Он превращает ваш граф TF в линейную алгебру, и у него есть собственные серверные части для работы на процессорах, графических процессорах или TPU. Чтобы запустить модель в облачном TPU, вы настраиваете ее с помощью tf.contrib.tpu.TPUEstimator, а затем… позвольте Google перенести ее в волшебную страну TPU для вас.

Похоже, они пока хранят набор инструкций, драйверы и бэкэнд XLA TPU только для внутреннего использования. Это понятно, но это означает, что мы не можем смотреть по сторонам, чтобы увидеть, как выглядит интерфейс CPU / TPU. Надеюсь, эта секретность не приведет к проблемам совместимости, поскольку другие компании разрабатывают собственные ускорители машинного обучения.

Стручки

TPU в производстве находятся в модулях, которые представляют собой большие стойки с огромной вычислительной мощностью. Каждый модуль вмещает 64 платы TPUv2 на 11,5 петафлопс.

TPU организованы в блоки для поддержки обучения. Одного TPU часто недостаточно для обучения большой модели с желаемой скоростью, но обучение включает в себя частые обновления веса, которые необходимо распределить между всеми задействованными чипами. Модули обеспечивают тесную связь между TPU и модулями, поэтому обучение, распределяемое между всеми чипами модуля, масштабируется почти линейно.

Это линейное масштабирование - важный принцип. Если специалисту по обработке данных требуется больше мощности, он добавляет больше микросхем, не задавая вопросов. Оборудование прозрачно для пользователя, поэтому он может сосредоточиться на выполняемой работе.

В этой статье больше говорится о самих стойках и о роли чипа в бизнесе.

Выводы

Это была вся информация, которую мне удалось найти о работе чипа. Это, конечно, неполно, но идею вы поняли.

TPU - действительно хорошее оборудование, но он мог существовать за годы до выхода версии 1. В основном это связано с тем, что Google так много занимается линейной алгеброй (управляемой машинным обучением), что этот вид микросхем стал коммерчески жизнеспособным. Другие компании подхватили эту тенденцию: Groq является побочным продуктом проекта TPU, Intel надеется вскоре начать продавать чипы Nervana, а там есть многие еще надеются войти. гонка. Специальное оборудование обещает снизить затраты на обучение и запуск моделей; Надеюсь, это ускорит темпы внедрения инноваций.

Больше информации:

Сноски:

  1. В частности, возможно, Джеффа Дина.
  2. тензор TensorFlow - это следующий шаг измерения размерности от матрицы. Мы обрабатываем тензоры, повторяя матричные операции партиями.
  3. Если вы посмотрите систолические массивы, вы, вероятно, получите диаграммы массива, реализующего алгоритм Кэннона, но схема здесь немного другая. В стандартной реализации частичные суммы остаются статичными, а веса перемещаются вниз; в TPU веса остаются неизменными, а частичные суммы движутся вниз. Обычно мы хотим вычислять результаты партиями, и если мы сохраняем веса там, где они есть, нам нужно загружать их только один раз за партию.
  4. Многие объяснения, которые я даю этой системе, являются умозрительными. Google прямо не сообщает, что происходит, поэтому я ни в чем не могу быть уверен, но это довольно простая система, и я уверен, что довольно точен в своих оценках.
  5. Мне неизвестен полный список функций, поддерживаемых TPUv1. Google по-прежнему использует их для многих своих выводов, поэтому я предполагаю, что список не слишком ограничивает.

Антон - недавно начавший работу инженер-электрик, которому нравится машинное обучение и компьютерная архитектура. Если вам нравится моя работа, я сейчас ищу свой следующий проект.