Amazon анонсувала TorchNeuron, власний бекенд PyTorch з відкритим вихідним кодом для сімейства чіпів ШІ Trainium, усуваючи одну з найбільших практичних перешкод на шляху використання спеціального кремнію для робочих навантажень машинного навчання. Згідно з офіційною публікацією в блозі PyTorch, опублікованій напередодні конференції PyTorch у Північній Америці 20 жовтня, розробники тепер можуть запускати існуючий код моделі PyTorch на Trainium із однорядковою зміною, а не переписувати свої робочі процеси навколо власних інструментів.

Оголошення одразу привернуло увагу, оскільки воно спрямоване на те, що аналітики давно назвали найбільш довгостроковою перевагою Nvidia: не сиру продуктивність чіпа, а екосистему програмного забезпечення CUDA, яка прив’язує розробників до графічних процесорів. Агентство Reuters озвучило реліз як заяву Amazon про те, що її чіпи AI працюватимуть із найпопулярнішою структурою розробки в AI, а TechCrunch назвав TorchNeuron найбільшим успіхом компанії в Nvidia. Щоб отримати ширше уявлення про те, як змінюється ринок прискорювачів, ми постійно висвітлюємо новини про апаратне забезпечення штучного інтелекту відстежуємо ці події по мірі їх розвитку.

Що насправді змінює TorchNeuron

Основна обіцянка TorchNeuron — простота. У прикладі, опублікованому в блозі PyTorch, існуючий навчальний цикл виконується на Trainium лише після переміщення моделі на новий пристрій за допомогою `.to('neuron')`. Згідно з дописом, код, написаний на основі `torch.accelerator`, рекомендованого PyTorch агностичного API для пристроїв, працює без змін. Для скомпільованих робочих навантажень розробники встановлюють бекенд за умовчанням один раз за допомогою `torch.compiler.set_default_backend("neuron")`, а потім викликають `model.compile()` без подальших змін.

Під капотом TorchNeuron реєструє Trainium через PrivateUse1, офіційну точку розширення PyTorch для призначених для користувача апаратних серверів, що означає, що чіп поводиться як стандартний пристрій PyTorch, а не як додаток, прикріплений до каркасу. AWS описує інтеграцію як охоплюючу кілька режимів виконання:

  • Енергійне виконання для досліджень і налагодження з експериментальним «адаптивним режимом нетерпіння», який зменшує накладні витрати на запуск
  • torch.compile для продуктивності виробництва, записуючи моделі в графіки, які оптимізує компілятор Neuron, із вбудованою підтримкою FlexAttention
  • Інтерфейс ядра Neuron (NKI) для дослідників, яким потрібні спеціальні ядра низького рівня

Згідно з оголошенням, серверна частина працює як на Trainium2, так і на нових чіпах Trainium3.

Прямий удар по програмному рову Nvidia

Контроль Nvidia над обчисленнями зі штучним інтелектом ніколи не спирався лише на швидше обладнання. CUDA, програмна платформа компанії, якій вже два десятиліття, — це місце, де пишеться, налаштовується та розгортається більшість дослідницького коду штучного інтелекту, і кожен чіп-претендент — від лінійки Instinct від AMD до TPU від Google і Trainium від Amazon — змушений був зіштовхнутися з витратами, пов’язаними з проханням розробників залишити його. Висвітлення запуску TorchNeuron у GeekWire, Benzinga та GuruFocus однозначно оформило це оголошення як напад на цей рів.

Amazon не намагається відтягнути розробників від PyTorch, щоб завоювати їх. Натомість він додає Trainium до PyTorch, тож величезна база користувачів фреймворку може прийняти чіп, не змінюючи звичок. Повідомлення в блозі PyTorch містило схвалення від Meta, яке підтримує структуру: «Рідна підтримка AWS Trainium розширює цей вибір, дозволяючи дослідникам і інженерам експериментувати та повторювати, використовуючи готове виконання та розподілені API, які вони вже знають», — сказала Яна ван Греунен, директор PyTorch Engineering у Meta, у дописі.

Це формування має значення для хмарного бізнесу AWS. Екземпляри Trainium продаються як дешевша альтернатива графічним процесорам Nvidia для навчання та висновків, і ці чіпи вже привернули увагу клієнтів ШІ. Зробити чіпи рідними для PyTorch безпосередньо вирішує заперечення, яке сповільнило ширше впровадження: спроби портування.

Технічний обсяг, від ноутбуків до навчання в кластерному масштабі

Можливості, перелічені в блозі PyTorch, свідчать про те, що Amazon націлений далеко за рамки експериментів для любителів. TorchNeuron охоплює набори операторів, які можна побачити в текстових, мультимодальних і мовних моделях, а також генерацію зображень і відео, моделі простору станів і моделі світу з точністю BF16, FP32 і FP16; Підтримка TorchAO і FP8 знаходиться в приватній бета-версії перед загальною доступністю в наступному випуску.

Розподілене навчання, дисципліна, яка перетворює код на одному графічному процесорі в виконання кластерного масштабу, підтримується через стандартні API PyTorch — DDP, FSDP2, SimpleFSDP, паралелізм тензорів і конвеєрів, а також шардинг у стилі ZeRO, а також DTensor і DeviceMesh для багатовимірних груп процесів. Потоки та події відображають семантику PyTorch на апаратних чергах Trainium, йдеться у дописі.

Для висновку, стеки Amazon Neuron побудовані на тій самій основі: vLLM Neuron обслуговує великі мовні моделі, тоді як vLLM Omni Neuron обслуговує дифузійні моделі для створення зображень і відео. Інтерфейс ядра Neuron дозволяє дослідникам писати спеціальні оператори як для швидкого, так і для скомпільованого виконання, а AWS є бібліотекою NKI-lib з відкритим кодом, набором оптимізованих для Trainium ядер, що охоплює поблокові обчислення FlashAttention і MoE для стандартних LLM і мультимодальних сценаріїв.

Також включено інструменти підвищення продуктивності. Neuron Explorer розширює стандартний `torch.profiler`, щоб відстежувати виконання від фреймворка до апаратного забезпечення, пов’язуючи кожну подію з вихідним рядком, і може експортувати метрики у форматі Parquet, які можуть використовуватися агентами для автоматичного аналізу вузьких місць.

Доступність і що буде далі

TorchNeuron планується запустити з відкритим кодом на конференції PyTorch у Північній Америці 20 жовтня, згідно з повідомленням у блозі, із сумісністю з Hugging Face Transformers v5. AWS описує Trainium як портфоліо мікросхем, розроблених з нуля для навчання штучного інтелекту в хмарному масштабі та висновків.

Стратегічний контекст — це ринок, на якому кожен великий хмарний провайдер є одночасно найбільшим покупцем прискорювачів Nvidia і потенційним конкурентом. Microsoft, Google і Meta розробляють нестандартний кремній, а китайські фірми, включаючи Huawei, створюють власні стеки програмного забезпечення, щоб уникнути гравітації CUDA. Зустрівшись із розробниками в PyTorch, замість того, щоб попросити їх перейти, Amazon вибрала найменш руйнівний шлях у цьому конкурсі — і, найімовірніше, той, якщо досвід розробника витримає, щоб перенести реальне навчальне навантаження з апаратного забезпечення Nvidia.

Чи справді TorchNeuron послабить перевагу CUDA, залежатиме від продуктивності в масштабі та від довгого хвоста моделей реального світу, на це питання не дають відповіді матеріали для запуску. Але бар’єр, який утримував більшість розробників PyTorch навіть від спроби Trainium — переписання — тепер опустився до одного рядка коду.

---

Будьте попереду ШІ

Отримуйте останні новини штучного інтелекту, аналіз і прориви — усе в одному місці.

Читати більше новин AI →