Что такое 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 из проекта дата-центра в Великобритании.