Amazon и Hugging Face упростили запуск моделей на чипах Trainium

Модели из библиотек Hugging Face запускаются и дообучаются на чипах Amazon Trainium тем же кодом, что на видеокартах, без отдельной библиотеки Optimum Neuron.

С 20 октября TorchNeuron, дополнение к PyTorch¹ для чипов Amazon Trainium², выходит в открытую бета-версию. Главные библиотеки Hugging Face работают с Trainium как с любым другим ускорителем.

Почему это важно: раньше модели Hugging Face попадали на Trainium через отдельную библиотеку Optimum Neuron с собственными версиями моделей и кода для обучения. Теперь на чипах Amazon запускается тот же код, что и на видеокартах, а обученная модель не привязана к железу.

Предыстория: AWS представила TorchNeuron в декабре 2025 года на конференции re:Invent. По документации AWS, до сих пор он был доступен только участникам закрытой бета-программы.

Как это работает: Optimum Neuron опиралась на PyTorch XLA, и код приходилось писать с оглядкой на компилятор. Вычисления не выполнялись сразу: их сначала собирали в общую схему и только потом запускали. Размеры данных были жёстко заданы, и любой участок кода, который не укладывался в схему, запускал новую компиляцию. TorchNeuron делает Trainium обычным устройством PyTorch: команды выполняются сразу, а для скорости есть штатная компиляция torch.compile и встроенные средства PyTorch, которые распределяют модель по чипам.

Детали: изменения коснулись текстовых моделей, генерации картинок и обучения:

В Transformers, библиотеке для языковых моделей, 89% из 100 самых популярных архитектур для генерации текста работают без доработок, а с Optimum Neuron v0.4.5 таких было 46%.
Для дообучения хватает обычных библиотек Hugging Face — TRL, PEFT и Accelerate, а особый модуль обучения из Optimum Neuron больше не нужен. Пример дообучения авторы запустили на восьми чипах Trainium2 — половине сервера trn2.48xlarge.
В Diffusers, библиотеке для генерации картинок, модели FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 и H3 можно разделить между несколькими чипами.
Одно из ядер³ Neuron на площадке Hugging Face ускоряет визуальные трансформеры до 2,4 раза.

Но: на Trainium по-прежнему лучше держать размеры данных постоянными. Поэтому TRL начиная с версии 1.16 дополняет пакеты данных до длины, кратной 128 токенам, — так модель реже приходится компилировать заново. Асинхронная загрузка весов в Transformers на этих чипах не работает, её нужно отключать.

Что дальше: Hugging Face советует пользователям чипов Amazon перейти с Optimum Neuron на основные библиотеки. Команда обещает, что всё больше моделей Diffusers можно будет делить между чипами, а на площадке появятся новые ядра Neuron.

Если шире: 12 октября WSJ писала, что Amazon обсуждает покупку стартапа Decart за 7 млрд долларов. Decart помогает переносить ИИ-модели между чипами.