AWS открыла исходный код PyTorch-бэкенда TorchNeuron для ускорителей Trainium

TorchNeuron позволяет запускать приложения на PyTorch с ускорителями AWS Trainium2 и Trainium3 через стандартные интерфейсы фреймворка и с минимальными изменениями кода. Публичная бета поддерживает два режима выполнения, распределенное обучение и разработку собственных ядер.

Amazon Web Services (AWS), облачное подразделение Amazon, выпустила публичную бета-версию TorchNeuron — нативного бэкенда PyTorch для собственных ИИ-ускорителей Trainium. Исходный код опубликован под лицензией Apache 2.0.

TorchNeuron подключает Trainium через предусмотренные в PyTorch механизмы для стороннего оборудования. В типичном случае разработчику достаточно выбрать устройство neuron, например с помощью .to('neuron'); код, использующий аппаратно-независимый интерфейс torch.accelerator, может работать без изменений.

Бэкенд поддерживает Trainium2 и Trainium3, непосредственное выполнение операций для исследований и отладки, компиляцию через torch.compile, стандартные распределенные API PyTorch, профилирование и собственные ядра на базе Neuron Kernel Interface.

AWS также интегрирует Trainium с TorchTitan, библиотеками Hugging Face и проектом vLLM. Компания намерена передавать изменения в основной проект PyTorch, в том числе устранять жесткие привязки фреймворка к отдельным типам ускорителей.