Що таке TorchNeuron: AWS відкриває чипи Trainium для PyTorch

Більшість сучасних моделей ШІ написані на PyTorch, і більшість із них навчають на GPU від Nvidia. Amazon хоче змінити другу частину, не чіпаючи першої. Новий бекенд TorchNeuron дає змогу запускати код PyTorch нативно на AWS Trainium — власних ШІ-чипах Amazon — за допомогою, як каже компанія, «однієї зміни в налаштуваннях і без переписування коду», йдеться в її блозі. Відкритий код публікують 20 жовтня на PyTorch Conference North America.

Що таке TorchNeuron

У PyTorch є офіційний спосіб для виробників чипів підключати нове залізо — бекенд пристроїв PrivateUse1 і API torch.accelerator. TorchNeuron використовує цей механізм і реєструє Trainium як звичайний пристрій PyTorch під назвою neuron. За даними блогу PyTorch Foundation і документації AWS, він підтримує:

• eager-виконання для досліджень і налагодження, зокрема новий режим «adaptive eager»;
• torch.compile для швидкості в продакшені;
• стандартні API розподіленого навчання — FSDP, DTensor, DDP і тензорний паралелізм, а також фреймворк TorchTitan;
• Neuron Kernel Interface (NKI) для власних ядер;
• профілювання звичним torch.profiler у зв'язці з інструментом AWS Neuron Explorer.

Пакет називається torch-neuronx. Його опубліковано на GitHub під ліцензією Apache 2.0 у статусі публічної бети в межах релізу Neuron 2.33.

Що насправді змінюється для розробників

На практиці головна зміна — одне слово. Там, де код відправляв модель на GPU Nvidia через «cuda», тепер пишуть «neuron». За словами Amazon, її інженери запускали той самий скрипт навчання на звичайному чипі й на Trainium, змінивши лише це налаштування, і модель «навчалася тому самому з тією ж швидкістю». Також Amazon стверджує, що на кожній збірці перевіряє TorchNeuron на сотнях тисяч моделей. Hugging Face додала нативну підтримку Trainium у свої інструменти, тож і там специфічний для Amazon код не потрібен.

Серед перших користувачів Amazon називає компанію голосового ШІ Cartesia та Інститут розвитку мозку Лібера (Lieber Institute for Brain Development), який застосовує моделі ШІ в пошуку ліків.

Чого він не робить

• Це не для вашого ноутбука. Чипи Trainium є лише в дата-центрах AWS. TorchNeuron підтримує інстанси Trainium2 і Trainium3. Якщо хочете запускати моделі вдома, дивіться наш гід із вибору ноутбука для локального ШІ.
• Це бета. У бета-версіях бувають прогалини, а заяви про швидкість — це слова самої Amazon. Важливими будуть незалежні тести.
• Він сам не робить Trainium дешевшим. Ідея — знизити витрати на перехід, а не ціну.

Чому це важливо

Перевага Nvidia — не лише в чипах. Її програмний стек CUDA тримає команди на її залізі, бо перенесення коду забирає час інженерів. Amazon б'є саме по цій прив'язці й каже, що також надсилає зміни в основний код PyTorch, аби сам фреймворк менше залежав від конкретних чипів. І вона не одна: стартапи на кшталт Tenstorrent, що щойно залучив $1,5 млрд, полюють на ту саму можливість. Більша конкуренція серед ШІ-чипів може послабити тиск на обчислювальні потужності, які обмежує ще й електроенергія: це показав цього тижня вихід Microsoft із проєкту дата-центру у Британії.