На этой неделе Amazon Web Services представила следующее поколение ускорителя Trainium3 для обучения и инференса ИИ. Как заявляет AWS, новый процессор вдвое быстрее своего предшественника и в четыре раза эффективнее. Это делает его одним из лучших решений для обучения и инференса ИИ с точки зрения стоимости. В абсолютных цифрах Trainium3 обеспечивает до 2517 MXFP8 TFLOPS, что почти вдвое ниже, чем у Nvidia’s Blackwell Ultra. Однако, Trn3 UltraServer от AWS оснащен 144 чипами Trainium3 и обеспечивает 0,36 ExaFLOPS производительности FP8, тем самым сопоставив производительность с NVL72 GB300 от Nvidia. Это очень важное событие, поскольку лишь немногие компании могут бросить вызов системам ИИ с масштабированием до уровня стойки от Nvidia.
AWS Trainium3 – это двухчиплетный ускоритель ИИ, оснащенный 144 ГБ памяти HBM3E с использованием четырех стеков, что обеспечивает пиковую пропускную способность памяти до 4,9 ТБ/с. Каждый вычислительный чиплет, предположительно изготовленный TSMC с использованием технологии 3-нм класса, содержит четыре ядра NeuronCore-v4 (которые оснащены расширенной ISA по сравнению с предшественниками) и соединяет два стека памяти HBM3E. Два чиплета соединены с помощью фирменного высокоскоростного интерфейса и разделяют 128 независимых аппаратных движков перемещения данных (которые являются ключевым элементом архитектуры Trainium), ядра коллективной связи, которые координируют трафик между чипами, и четыре интерфейса NeuronLink-v4 для масштабируемости.
NeuronCore-v4 интегрирует четыре блока выполнения: тензорный движок, векторный движок, скалярный движок, блок GPSIMD и 32 МБ локальной SRAM, которая управляется компилятором, а не кэшируется. С точки зрения разработки программного обеспечения, ядро построено вокруг модели потока данных, определяемой программным обеспечением, в которой данные поступают в SRAM с помощью DMA-движков, обрабатываются единицами выполнения и затем записываются обратно в виде накопления памяти, близкой к памяти, что позволяет DMA выполнять операции чтения-добавления-записи в одной транзакции. SRAM не когерентна между ядрами и используется для разбивки, этапов и накопления, а не для общего кэширования.
Пожалуй, наиболее интересным компонентом NeuronCore-v4 является блок GPSIMD, который интегрирует восемь полностью программируемых 512-битных векторных процессоров, которые могут выполнять код общего назначения, написанный на C/C++, при доступе к локальной SRAM. GPSIMD интегрирован в NeuronCore, потому что не все в реальных моделях ИИ четко подходят к тензорному движку. Современные рабочие нагрузки ИИ содержат много кода для необычных расположений данных, логики постобработки, индексирования и математики, специфичной для модели. Эти операции трудно или неэффективно выразить как матричные операции, и запуск их на хост-CPU приведет к задержкам и дорогостоящей передаче данных. GPSIMD решает эту проблему, предоставляя реальные универсальные программируемые векторные блоки внутри ядра, чтобы такая логика работала непосредственно рядом с тензорами на полной скорости и с использованием той же локальной SRAM.
Вкратце, NeuronCore-v4 работает как плотно связанный движок потока данных, в котором тензорная математика, векторные преобразования, скалярное управление и пользовательский код совместно используют 32 МБ локальной области подкачки и оркеструются компилятором Neuron, а не планировщиком warp, используемым на оборудовании Nvidia.
С точки зрения производительности, Trainium3 превосходит своего непосредственного предшественника в вычислениях FP8 (в частности, MXFP8) почти вдвое и достигает 2,517 PFLOPS на пакет (очевидно опережая Nvidia’s H100/H200, но отставая от Blackwell B200/B300) и добавляет поддержку MXFP4. Однако производительность Trainium3 в BF16, TF32 и FP32 остается на уровне Trainium2, что ясно показывает, что AWS делает ставку на MXFP8 для обучения и инференса в будущем. В связи с этим, он не разрабатывает свои возможности BF16 (которые широко используются для обучения в настоящее время) и FP32, поскольку чувствует себя комфортно с производительностью, которую он имеет, учитывая, что эти форматы теперь используются в основном для накопления градиентов, мастерских весов, состояний оптимизатора, масштабирования потерь и некоторых операций, чувствительных к точности.
Одна интересная возможность, которой обладает Trainium3, и которую стоит упомянуть, – это функция Logical NeuronCore Configuration (LNC), которая позволяет компилятору Neuron объединять четыре физических ядра в более широкое логическое ядро, которое автоматически синхронизируется с объединенными вычислительными, SRAM и HBM, что может быть полезно для очень широких слоев или больших длин последовательностей, которые обычно используются в очень больших моделях ИИ.
В последние кварталы значительная часть успеха Nvidia была обусловлена ее решениями NVL72 на базе Blackwell, включающими 72 графических процессора Blackwell. Поддержка масштабного наращивания world size и топологии all-to-all, что особенно важно для Mixture-of-Experts (MoE) и инференса авторегрессии. Это дает Nvidia огромное преимущество перед AMD и разработчиками пользовательских ускорителей, такими как AWS. Чтобы обеспечить эту возможность, Nvidia пришлось разработать NVLink switches, сложные сетевые карты и DPU, что является огромным кремниевым усилием. Однако, похоже, что Trn3 UltraServer от AWS даст Nvidia’s GB300 NVL72 достойный ответ.
Trn3 UltraServer, работающие на ускорителях Trainium3, будут предлагаться в двух размерах: одна конфигурация оснащена 64 ускорителями и предположительно процессором Intel Xeon, в то время как большая версия объединяет 144 ускорителя и процессор на базе Arm Graviton в одном решении для масштабирования до уровня стойки. В большей системе 144 ускорителя Trainium3 распределены по 36 физическим серверам, на каждом из которых установлен один процессор Graviton и четыре чипа Trainium3. Во многих отношениях такая схема напоминает подход Nvidia NVL72, который использует CPU, GPU и кремниевую связность Nvidia, подчеркивая направление AWS по созданию интегрированных вертикально платформ ИИ.
Внутри сервера ускорители Trainium3 соединены через первый слой NeuronSwitch-v1 с использованием NeuronLink-v4 (со скоростью 2 Гбит/с на устройство, хотя неясно, говорим ли мы об одной в направлении полосе пропускания или агрегированной двунаправленной полосе пропускания), и связь между разными серверами осуществляется через два дополнительных слоя NeuronSwitch-v1, также выполняемых через NeuronLink-v4. К сожалению, AWS не публикует агрегированную пропускную способность NeuronSwitch-v1 на уровне домена.
С точки зрения производительности, большая конфигурация с 144 Trainium3 обеспечивает производительность 362,5 MXFP8/MXFP4 PetaFLOPS (плотная), которая (на уровне с GB300 NVL72), 96,624 PFLOPS пропускной способности BF16/FP16/TF32 и 26,352 PFLOPS в FP32. Система также оснащена 21 ТБ памяти HBM3E, обеспечивающей общую пропускную способность памяти 705,6 ТБ/с, что отстает от Nvidia’s GB300 NVL72 в этом показателе.
В целом, Trn3 Gen2 UltraServer выглядит очень конкурентоспособным по сравнению с Nvidia’s GB300 NVL72 с точки зрения производительности FP8. FP8 вот-вот станет более популярным для обучения, поэтому ставка на этот формат имеет большой смысл. Конечно, у Nvidia есть козырь в рукаве в виде NVFP4, который позиционируется как для инференса, так и для обучения, и с помощью этого формата машины на базе Blackwell компании являются непобедимыми. То же самое относится и к BF16, который стал быстрее по сравнению с Trainium2, но не настолько, чтобы превзойти Blackwell от Nvidia.
В целом, хотя Trn3 Gen2 UltraServer с 144 ускорителями Trainium3 выглядит довольно конкурентоспособным по сравнению с машинами NVL72 на базе Blackwell от Nvidia, решение Nvidia является более универсальным в целом.
В дополнение к выпуску нового оборудования ИИ, AWS объявила о широком расширении программного стека AWS Neuron на своей ежегодной конференции re:Invent на этой неделе. AWS позиционирует этот релиз как переход к открытости и доступности для разработчиков, поэтому обновление обещает сделать платформы на базе Trainium более простыми в использовании, позволить стандартным фреймворкам машинного обучения работать непосредственно на оборудовании Trainium, предоставить пользователям больший контроль над производительностью и даже открыть доступ к путям оптимизации на низком уровне для экспертов.
Основным дополнением является интеграция PyTorch нативно через открытый исходный код backend под названием TorchNeuron. Используя механизм PrivateUse1 PyTorch, Trainium теперь отображается как тип устройства по умолчанию, что позволяет существующему коду PyTorch выполняться без изменений. TorchNeuron также поддерживает интерактивное немедленное выполнение, torch.compile и распределенные функции, такие как FSDP и DTensor, и работает с популярными экосистемами, включая TorchTitan и Hugging Face Transformers. Доступ к этой функции в настоящее время ограничен избранными пользователями в рамках частной программы предварительного просмотра.
AWS также представила обновленный интерфейс ядра Neuron (NKI), который дает разработчикам прямой контроль над поведением оборудования, включая программирование на уровне инструкций, явное управление памятью и гранулярное планирование, раскрывая набор инструкций Trainium для разработчиков ядра. Кроме того, компания выпустила NKI Compiler под лицензией Apache 2.0 с открытым исходным кодом. Интерфейс программирования доступен публично, в то время как компилятор остается в ограниченном предварительном просмотре.
AWS также выпустила свой Neuron Explorer, набор инструментов отладки и настройки, который позволяет разработчикам программного обеспечения и инженерам по производительности улучшать работу своих моделей на Trainium. Это достигается путем отслеживания выполнения от вызовов высокоуровневых фреймворков до отдельных инструкций ускорителя, предлагая послойный профилирование, видимость на уровне исходного кода, интеграцию со средами разработки и управляемые ИИ предложения для настройки производительности.
Наконец, AWS представила свой Neuron Dynamic Resource Allocation (DRA) для интеграции Trainium непосредственно в Kubernetes без необходимости в пользовательских планировщиках. Neuron DRA опирается на встроенный планировщик Kubernetes и добавляет осведомленность об аппаратной топологии, чтобы полностью интегрировать UltraServer как единый ресурс и затем гибко назначать оборудование для каждой рабочей нагрузки. Neuron DRA поддерживает Amazon EKS, SageMaker HyperPod и развертывания UltraServer и предоставляется как программное обеспечение с открытым исходным кодом с образами контейнеров, опубликованными в общедоступном реестре AWS ECR.
И Neuron Explorer, и Neuron DRA предназначены для упрощения управления кластером и предоставления пользователям детального контроля над тем, как выделяются и используются ресурсы Trainium. В целом, AWS стремится сделать свои платформы на базе Trainium гораздо более распространенными, чем сегодня, в попытке сделать их более конкурентоспособными по сравнению с предложениями на базе CUDA от Nvidia.
На этой неделе Amazon Web Services выпустила 3-го поколения ускоритель Trainium для обучения и инференса ИИ, а также сопутствующие решения для масштабирования до уровня стойки Trn3 UltraServers. Впервые Trn3 Gen2 UltraServers, машины для масштабирования до уровня стойки, будут полагаться исключительно на внутреннее оборудование AWS, включая процессор, ускорители ИИ, коммутационную аппаратуру и ткани для подключения, сигнализируя о том, что компания приняла аппаратную стратегию вертикальной интеграции Nvidia.
AWS утверждает, что процессор Trainium3 предлагает примерно в 2 раза более высокую производительность и в 4 раза лучшую энергоэффективность по сравнению с Trainium2, поскольку каждый ускоритель обеспечивает до 2,517 PFLOPS (MXFP8) — опережая Nvidia’s H100, но отставая от B200 — и сопровождается 144 ГБ HBM3E с пропускной способностью 4,9 ТБ/с. Тем временем Trn3 Gen2 UltraServers масштабируются до 144 ускорителей для производительности FP8 примерно в 0,36 ExaFLOPS, что сопоставимо с решением для масштабирования до уровня стойки Nvidia’s GB300 NVL72. Тем не менее, оборудование Nvidia по-прежнему выглядит более универсальным, чем AWS.
Чтобы догнать Nvidia, Amazon также объявила о крупных обновлениях своего программного стека Neuron, чтобы сделать платформы на базе Trainium более простыми в использовании, разрешить стандартным фреймворкам машинного обучения работать нативно на оборудовании, предоставить разработчикам больший контроль над производительностью и открыть доступ к настройке на низком уровне для экспертов.
Автор – Anton Shilov
