Trainium

Trainium — это семейство специализированных микропроцессоров (ASIC), разработанных американской компанией Amazon Web Services (AWS) для ускорения задач машинного обучения, в частности, для обучения и инференса (вывода) нейронных сетей. Чипы Trainium являются частью инфраструктуры облачных вычислений AWS и предназначены для обеспечения высокой производительности при низком энергопотреблении по сравнению с универсальными графическими процессорами (GPU), традиционно используемыми в этой области.

История

Разработка Trainium была начата Amazon в рамках стратегии по снижению зависимости от сторонних поставщиков чипов, таких как NVIDIA, и оптимизации затрат на собственные вычислительные ресурсы. Первое поколение чипов, AWS Trainium (Trainium1), было анонсировано в декабре 2020 года. Коммерческая доступность инстансов (виртуальных серверов) на базе Trainium1, получивших название Amazon EC2 Trn1, началась в 2022 году.

Второе поколение, AWS Trainium2, было представлено в 2023 году и поступило в широкий доступ в 2024 году. Оно обеспечило значительный прирост производительности и эффективности. В 2024 году также было анонсировано третье поколение — Trainium3, выход которого ожидается в 2025 году.

Архитектура и устройство

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

• Вычислительные ядра (NeuronCores): Каждый чип содержит несколько (обычно два) вычислительных ядра, оптимизированных для операций с плавающей запятой (FP32, BF16, FP16) и целочисленных операций (INT8). Ядра поддерживают разреженные вычисления (sparsity), что позволяет ускорить обработку моделей с большим количеством нулевых весов.
• Высокоскоростная память (HBM): Чипы оснащены памятью типа HBM (High Bandwidth Memory) с высокой пропускной способностью, необходимой для быстрой загрузки весов и промежуточных данных.
• Внутреннее соединение (NeuronLink): Для масштабирования на множество чипов используется собственный высокоскоростной интерконнект NeuronLink, позволяющий объединять тысячи чипов в единую вычислительную систему (кластер). Пропускная способность NeuronLink в Trainium2 достигает 448 ГБ/с на устройство.
• Сеть (Elastic Fabric Adapter — EFA): Для связи между серверами в кластере используется технология EFA, обеспечивающая низкую задержку и высокую пропускную способность.

Сравнение поколений

Классификация

Trainium относится к классу ускорителей нейронных сетей (Neural Processing Units, NPU), или тензорных процессоров. В отличие от универсальных CPU и GPU, которые могут выполнять широкий спектр задач, Trainium жёстко оптимизирован для линейной алгебры и операций, лежащих в основе глубокого обучения. По модели использования это облачный чип — он не поставляется отдельно, а доступен исключительно через сервисы AWS (Amazon EC2).

Применение

Основное применение Trainium — обучение и инференс крупномасштабных моделей машинного обучения, включая:

• Большие языковые модели (LLM): Обучение моделей типа GPT, LLaMA, BERT и их вариаций. AWS сообщает о возможности обучения моделей с сотнями миллиардов параметров на кластерах Trainium2.
• Модели компьютерного зрения: Обучение свёрточных нейронных сетей (CNN) и трансформеров для задач классификации, детекции и сегментации изображений.
• Рекомендательные системы: Обучение моделей, используемых в электронной коммерции, стриминговых сервисах и социальных сетях.
• Генеративные модели: Обучение диффузионных моделей (например, Stable Diffusion) и других генеративных архитектур.

В экосистеме AWS чипы Trainium используются в рамках сервиса Amazon SageMaker для автоматизированного обучения и развёртывания моделей, а также через специализированные инстансы Amazon EC2 Trn1 (Trainium1) и Trn2 (Trainium2).

Программное обеспечение

Для работы с Trainium AWS разработала набор инструментов и библиотек под общим названием AWS Neuron SDK. Он включает:

• Neuron Compiler: Компилятор, который преобразует модели из популярных фреймворков (PyTorch, TensorFlow, JAX) в оптимизированный код для Trainium.
• Neuron Runtime: Среда выполнения, управляющая работой чипов.
• Neuron Framework: Интеграции с популярными фреймворками, позволяющие использовать Trainium с минимальными изменениями кода.
• Neuron Monitor: Инструменты для мониторинга производительности и загрузки чипов.

Критика и ограничения

Несмотря на заявленную высокую производительность и энергоэффективность, Trainium имеет ряд ограничений и критических замечаний:

• Привязка к экосистеме AWS: Чипы доступны исключительно в облаке Amazon, что создаёт зависимость от одного поставщика. Миграция на другие платформы (Google Cloud, Microsoft Azure, локальные серверы) невозможна без полной замены вычислительной инфраструктуры.
• Ограниченная поддержка фреймворков: На момент выхода Trainium1 поддержка фреймворков была неполной (например, отсутствовала поддержка JAX). Со временем ситуация улучшилась, но полная совместимость с новыми версиями и специфическими операциями может задерживаться.
• Сложность оптимизации: Для достижения максимальной производительности часто требуется ручная настройка компилятора и кода, что повышает порог входа для разработчиков.
• Конкуренция: Основные конкуренты — NVIDIA (GPU серий A100, H100, B200) и Google (TPU v5p). Trainium часто сравнивают с ними по показателям производительности на доллар и на ватт. В некоторых тестах (например, обучение больших LLM) Trainium2 демонстрирует сопоставимые или лучшие результаты, но в других (например, инференс с малыми задержками) уступает GPU.

Перспективы

Третье поколение, Trainium3, анонсированное в 2024 году, должно выйти в 2025 году. Ожидается, что оно будет производиться по 3-нм техпроцессу и обеспечит четырёхкратный прирост производительности по сравнению с Trainium2. Также анонсирован проект UltraCluster, который позволит объединять десятки тысяч чипов Trainium3 для обучения сверхбольших моделей.

Источники

• Официальная документация AWS по Trainium и Neuron SDK (aws.amazon.com/neuron).
• Презентации AWS re:Invent 2020, 2023, 2024 годов.
• Технические обзоры инстансов Amazon EC2 Trn1 и Trn2.
• Статьи в специализированных изданиях (The Next Platform, AnandTech, SemiAnalysis) за 2022–2024 годы.