Amazon Web Services представила Arm-процессор нового поколения Graviton4 и ИИ-ускоритель Trainium2, предназначенный для обучения нейронных сетей. Всего к текущему моменту компания выпустила уже 2 млн Arm-процессоров Graviton, которыми пользуются более 50 тыс. клиентов.

Amazon Graviton4 (Изображение: AWS)

«Graviton4 представляет собой четвёртое поколение процессоров, которое мы выпустили всего за пять лет, и это самый мощный и энергоэффективный чип, который мы когда-либо создавали для широкого спектра рабочих нагрузок», — отметил Дэвид Браун (David Brown), вице-президент по вычислениям и сетям AWS. По сравнению с Graviton3 новый чип производительнее на 30 %, включает на 50 % больше ядер и имеет на 75 % выше пропускную способность памяти.

Изображение: AWS

Graviton4 будет иметь до 96 ядер Neoverse V2 Demeter (2 Мбайт L2-кеша на ядро) и 12 каналов DDR5-5600. Кроме того, новый чип получит поддержку шифрования трафика для всех своих аппаратных интерфейсов. Процессор изготавливается по 4-нм техпроцессу TSMC и, вероятно, имеет чиплетную компоновку. Возможно, это первый CPU компании, ориентированный на работу в двухсокетных платформах.

Изображение: AWS

Поначалу Graviton4 будет доступен в инстансах R8g (пока в статусе превью), оптимизированных для приложений, интенсивно использующих ресурсы памяти — высокопроизводительные базы данных, in-memory кеши и Big Data. Эти инстансы будут поддерживать более крупные конфигурации, иметь в три раза больше vCPU и в три раза больше памяти по сравнению с инстансами Rg7, которые имели до 64 vCPU и 512 Гбайт ОЗУ.

Amazon Trainium2 (Изображение: AWS)

В свою очередь, Trainium 2 предназначен для обучения больших языковых моделей (LLM) и базовых моделей. Сообщается, что ускоритель в сравнении с Trainium 1 вчетверо производительнее и при этом имеет в 3 раза больший объём памяти и в 2 раза более высокую энергоэффективность. Инстансы EC2 Trn2 получат 16 ИИ-ускорителей с возможностью масштабирования до 100 тыс. единиц в составе EC2 UltraCluster, которые суммарно дадут 65 Эфлопс, то есть по 650 Тфлопс на ускоритель. Как утверждает Amazon это позволит обучать LLM с 300 млрд параметров за недели вместо месяцев.

Со временем на Graviton4 заработает SAP HANA Cloud, портированием и оптимизацией этой платформы уже занимаются. Oracle также перенесла свою СУБД на Arm, а заодно перевела все свои облачные сервисы на чипы Ampere, в которую в своё время инвестировала. Microsoft же пошла по пути AWS и недавно анонсировала 128-ядерый Arm-процессор (Neoverse N2) Cobalt 100 и ИИ-ускоритель Maia 100 собственной разработки. Всё это может представлять отдалённую угрозу для AMD и Intel. С NVIDIA же все всё равно пока что продолжают дружбу — именно в инфраструктуре AWS, как ожидается, появится самый мощный в мире ИИ-суперкомпьютер на базе новых GH200.

Облако AWS объявило о доступности инстансов EC2 Trn1, которые впервые были анонсированы в конце прошлого года. Утверждается, что Trn1 предлагают лучшее соотношение цены и производительности в облаке для решения комплексных задач глубокого обучения.

Инстансы базируются на ускорителях Trainium собственной разработки Amazon: это решение специально оптимизировано для высокопроизводительного глубокого обучения. Отмечается, что по сравнению с другими GPU-инстансами Trn1 обеспечивает 50-% сокращение затрат на обучение моделей. По сравнению с EC2 P4d инстансы Trn1 обеспечивают 1,4-кратное увеличение производительности для BF16-расчётов, 2,5-кратное повышение быстродействия для TF32 и 5-кратное — для FP32. Пиковая производительность Trn1 достигает 3,4 Пфлопс (TF32/FP16/BF16).

Источник изображения: Amazon

В семейство EC2 Trn1 пока входят два типа Nitro-инстансов: trn1.2xlarge и trn1.32xlarge. В первом случае задействован один ускоритель Trainium, во втором — 16. Объём памяти HBM2e составляет соответственно 32 и 512 Гбайт. В составе trn1.2xlarge имеется локальный NVMe SSD ёмкостью 500 Гбайт, а trn1.32xlarge полагается сразу четыре NVMe-модуля на 2 Тбайт каждый. В старшем варианте ускорители связаны интерконнектом NeuronLink (768 Гбайт/с), а для самого инстанса предлагается EFA-подключение со скоростью до 800 Гбит/c. Младший вариант ограничен подключением 12,5 Гбит/с.

Ускорители AWS Trainium включают специализированные скалярные, векторные и тензорные движки, оптимизированные под задачи глубокого обучения. Поддерживается работа с различными типами данных, включая FP32, TF32, BF16, FP16 и UINT8. Кроме того, обеспечена совместимость с новым форматом cFP8, который хорошо подходит для крупных моделей. Аппаратно поддерживается функция стохастического округления. Наконец, AWS Trainium поддерживают и кастомные режимы работы с данными.

Возможно развёртывание Trn1 в кластерах EC2 UltraCluster с поддержкой до 30 000 ускорителей AWS Trainium для решения наиболее сложных задач. Архитектура кластеров предусматривает использование петабитной неблокированной сети и масштабируемого хранилища с низкой задержкой Amazon FSx for Lustre. Ускорители используют тот же набор AWS Neuron SDK, что и AWS Inferentia. Кроме того, новинки будут доступны для Amazon SageMaker, EKS, ECS и AWS Batch.

GPU давно применяются для ускорений вычислений и в последние годы обросли поддержкой специфических форматов данных, характерных для алгоритмов машинного обучения, попутно практически лишившись собственно графических блоков. Но в ближайшем будущем их по многим параметрам могут превзойти специализированные ИИ-процессоры, к числу которых относится и новая разработка AWS, чип Trainium.

На мероприятии AWS Re:Invent компания рассказала о прогрессе в области машинного обучения на примере своих инстансов P3dn (Nvidia V100) и P4 (Nvidia A100). Первый вариант дебютировал в 2018 году, когда модель BERT-Large была примером сложности, и благодаря 256 Гбайт памяти и сети класса 100GbE он продемонстрировал впечатляющие результаты. Однако каждый год сложность моделей машинного обучения растёт почти на порядок, а рост возможностей ИИ-ускорителей от этих темпов явно отстаёт.

Сложность моделей машинного обучения будет расти всё быстрее

Когда в прошлом году был представлен вариант P4d, его вычислительная мощность выросла в четыре раза, а объём памяти и вовсе на четверть, в то время как знаменитая модель GPT-3 превзошла по сложности BERT-Large в 500 раз. А теперь и 175 млрд параметров последней — уже ничто по сравнению с 10 трлн в новых моделях. Приходится наращивать и объём локальной памяти (у Trainium имеется 512 Гбайт HBM с суммарной пропускной способностью 13,1 Тбайт/с), и активнее использовать распределённое обучение.

Для последнего подхода узким местом стала сетевая подсистема, и при разработке стека Elastic Fabric Adapter (EFA) компания это учла, наделив новые инстансы Trn1 подключением со скоростью 800 Гбит/с (вдвое больше, чем у P4d) и с ультранизкими задержками, причём доступен и более оптимизированный вариант Trn1n, у которого пропускная способность вдвое выше и достигает 1,6 Тбит/с. Для связи между самими чипами внутри инстанса используется интерконнект NeuroLink со скоростью 768 Гбайт/с.

Прогресс подсистем сети и памяти в ИИ-инстансах AWS

Но дело не только в возможности обучить GPT-3 менее чем за две недели: важно и количество используемых для этого ресурсов. В случае P3d это потребовало бы 600 инстансов, работающих одновременно, и даже переход к архитектуре Ampere снизил бы это количество до 200. А вот обучение на базе чипов Trainium требует всего 130 инстансов Trn1. Благодаря оптимизациям, затраты на «общение» у новых инстансов составляют всего 7% против 14% у Ampere и целых 49% у Volta.

Меньше инстансов, выше эффективность при равном времени обучения — вот что даст Trainium

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

В Trainium, по словам AWS, гибкость сохранена — ускоритель имеет 16 полностью программируемых (на С/С++) обработчиков. Есть и у него и другие оптимизации. Например, аппаратное ускорение стохастического округления, которое на сверхбольших моделях становится слишком «дорогим» из-за накладных расходов, хотя и позволяет повысить эффективность обучения со смешанной точностью. Всё это позволяет получить до 3,4 Пфлопс на вычислениях малой точности и до 840 Тфлопс в FP32-расчётах.

AWS постаралась сделать переход к Trainium максимально безболезненным для разработчиков, поскольку SDK AWS Neuron поддерживает популярные фреймворки машинного обучения. Впрочем, насильно загонять заказчиков на инстансы Trn1 компания не собирается и будет и далее предоставлять на выбор другие ускорители поскольку переход, например, с экосистемы CUDA может быть затруднён. Однако в вопросах машинного обучения для собственных нужд Amazon теперь полностью независима — у неё есть и современный CPU Graviton3, и инфереренс-ускоритель Inferentia.