Новые TPU восьмого поколения от Google. Обучение и инференс теперь отдельно?

На конференции Google Cloud Next были представлены ^[1] новые Tensor Processor Unit (TPU) восьмого поколения. В этом поколении чипы анонсированы в двух различных архитектурах TPU — 8t и 8i:

• 8t (Sunfish) — для обучения ^[2],

• 8i (Zebrafish) — для инференса.

По сути, это дополнительное подтверждение, что подход «один GPU под все задачи» — устаревает. В этом поколении Google глубоко разносит обучение и инференс по разным архитектурам — вплоть до разной топологии сети. 

Как следует из поста Амина Вахдата (Amin Vahdat), старшего вице-президента и главного технолога по искусственному интеллекту ^[3] и инфраструктуре: TPU 8t и TPU 8i были разработаны в сотрудничестве с Google DeepMind для выполнения самых ресурсоемких задач в области искусственного интеллекта и адаптации к масштабируемым архитектурам моделей. TPU восьмого поколения — это результат более чем десятилетней работы (первое поколение было создано в 2015 году).

Как сообщает Вахдат: «Несколько лет назад мы предвидели растущий спрос на логические выводы со стороны клиентов по мере внедрения передовых моделей искусственного интеллекта в производственные процессы. А с появлением AI-агентов мы пришли к выводу, что сообществу будут полезны чипы, индивидуально адаптированные под нужды обучения и обслуживания». 

8t для обучения

TPU 8t ориентирован на ресурсоемкие задачи обучения. По заверениям автора: «Благодаря сочетанию максимально возможной вычислительной мощности, общей памяти ^[4] и пропускной способности межчиповых соединений с максимально возможной энергоэффективностью и продуктивным временем вычислений мы создали систему, которая обеспечивает почти в три раза более высокую вычислительную производительность на модуль по сравнению с предыдущим поколением. Это позволяет быстрее внедрять инновации и гарантирует, что наши клиенты будут задавать темп развития отрасли».

В цифрах это выглядит так: векторные, матричные и SparseCore-ядра, дополненные 128 МБ SRAM и 216 ГБ HBM3e. Для вертикального масштабирования используется межчиповый интерконнект ICI со скоростью 19,2 Тбит/с, для горизонтального — 400 Гбит/с. Это позволяет самым большим моделям использовать единый огромный пул памяти.

Кроме того, TPU 8t получили расширенные возможности RAS (Remote Access Service). Телеметрия в реальном времени, автоматическое обнаружение неисправных каналов ICI и перенаправление трафика без прерывания задания, а также оптическая коммутация каналов (OCS), которая перенастраивает оборудование в случае сбоев без участия человека… Все это позволяет довести уровень утилизации чипа до 97%.

Superpod состоит из 9 600 чипов с 2 ПБ общей памяти: 121 экзафлопс вычислений. А новая сеть Virgo ^[5], которая использует плоскую двухуровневую неблокирующую топологию, через JAX и Pathways позволяет объединить до миллиона TPU в рамках нескольких ЦОД.

Дополнительно были представлены ^[6] TPUDirect RDMA и TPUDirect Storage.

• TPUDirect RDMA — это прямая передача данных между памятью TPU (HBM) и сетевыми картами (NIC).

• TPUDirect Storage — прямой доступ к памяти между TPU и высокоскоростным управляемым хранилищем, таким как 10T Lustre.

Арендуйте GPU за 1 рубль!

Выберите нужную конфигурацию в панели управления Selectel. *

Подробнее → ^[7]

8i для инференса

TPU 8i, в свою очередь, ориентирован на инференс. Чип разработан для обработки сложной совместной итеративной работы множества специализированных агентов, часто объединяющихся в сложные потоки для предоставления решений и аналитических данных по самым сложным задачам, — сообщает Вахдат.

TPU 8i получил 288 ГБ памяти HBM3e в паре с 384 МБ SRAM — такой объем SRAM помогает удерживать активные веса модели на самом чипе. А новый механизм ускорения коллективных операций (Collectives Acceleration Engine, CAE), по информации Google, разгружает глобальные операции, сокращая внутреннюю задержку в пять раз

Как и у 8t, пропускная способность межчиповых взаимодействий удвоена до 19,2 Тбит/с. Однако сетевая топология отличается — в данном случае это Boardfly. 

В основе лежит Building Block (BB), в котором четыре чипа объединены в кольцо с полной связностью. BB масштабируются в группу из восьми штук. 36 групп замыкаются через оптические коммутаторы в под. Сетевой диаметр сокращен до семи хопов, поскольку для MoE-инференса важно меньшее количество сетевых переходов между чипами.

TPU 8i масштабируется в Superpod до 1 152 чипов в одном кластере с производительностью 11,6 экзафлопс. По заявлению Google, TPU 8i на 80% выгоднее по соотношению производительности и цены, чем Ironwood — за те же деньги можно обслужить почти вдвое больше пользователей.

Итоги

Оба чипа впервые работают на собственных ARM-процессорах Axion в качестве хост-CPU и поддерживают СЖО

Удобная таблица для сравнения спецификаций:

Как сообщает Google: «Оба чипа поступят в продажу в конце этого года и могут быть использованы в составе гиперкомпьютера Google с искусственным интеллектом, который объединяет специализированное оборудование (вычислительные ресурсы, системы хранения данных, сетевые ресурсы), открытое программное обеспечение (фреймворки, механизмы логического вывода) и гибкие модели потребления (оркестрация, управление кластерами и модели доставки) в единый стек».

Коммерческий контекст этого анонса очевиден. Google несколько раз в пресс-релизе указывает на возросшую продуктивность и очевидную экономию, связанную с этим. Кроме того, теперь это не просто альтернатива GPU, а усиленный сдвиг на отдельный ускоритель для разных AI-задач. Еще одна из увлекательных новостей: TPU 8t масштабируется до миллиона чипов в одном кластере — это больше, чем публично известные размеры кластеров других компании.

Что думаете насчет новых TPU? Делитесь мнением в комментариях.