Пленарный доклад на конференции Google I/O в нынешнем году был посвящен машинному обучению и интеллектуальным системам помощи. Google уже несколько лет разрабатывает собственные аппаратные решения для систем глубокого обучения - Tensor Processing Units или TPUs.
TPU представляет собой интегральную схему специального назначения ASIC (Application Specific Integrated Circuit), которая оптимизирована под операции машинного обучения TensorFlow. Второе поколение TPU опиралась на чипы с двумя ядрами с 128 × 128 блоками Mixed Multiply Unit (MXU) каждое. Подключенные 16 Гбайт памяти HBM к каждому чипу обеспечивают пропускную способность 600 Гбайт/с. Google указывает вычислительную производительность TPU на уровне 180 TFLOPS.
Теперь Google представила третье поколение TPU. Однако со стороны ASIC и аппаратных ресурсов изменений нет. Мы по-прежнему получаем четыре 2-ядерных чипа в одном TPU. Производительность TPU составляет 180 TFLOPs. Внутри чипа работают скалярные, векторные и матричные (MXU) блоки.
Для тренировки сетей глубокого обучения важны вычисления FP32, то есть расчеты с плавающей запятой с одинарной точностью. Чтобы проиллюстрировать сложность матричных вычислений, Google привела следующую диаграмму.
Как можно видеть, сложность расчетов FP32 существенно повышается, причем системе приходится вычислять миллиарды чисел в секунду. Поэтому Google перешла на новый формат чисел bfloat16. Он должен обеспечивать диапазон, идентичный FP32, с пусть и с меньшей точностью.
Без новых чипов - но с более плотной упаковкой и водяным охлаждением
Непосредственно у самих чипов изменений немного. По всей видимости, Google теперь более плотно располагает чипы и использует новый интерконнект. К сожалению, технических подробностей о пропускной способности и числе каналов не приводится, но теперь Google может связывать больше TPU. То есть в одной ферме теперь могут работать десятки или даже тысячи TPU. До сих пор каждая ферма содержала до 64 TPU. В результате новые фермы могут быть крупнее, вычислительная производительность может достигать 100 POPS для вычислений INT8, важных для машинного обучения.
В случае третьего поколения TPU Google отказалась от воздушного охлаждения. Причем причина кроется не в увеличении производительности TPU, а в более плотном их расположении в фермах, когда выделяющееся тепло намного эффективнее отводить с помощью СВО.
Google использует TPU не только в собственных задачах, но и предлагает доступ к ним в своем облаке. Минимальная облачная конфигурация TPU состоит из четырех 2-ядерных чипов и 64 Гбайт памяти HBM2. Цена составляет $6,50 за TPU в час.