Amazon представила TorchNeuron — публичный программный модуль для переноса моделей машинного обучения на фирменные ускорители Trainium 2 и Trainium 3. Разработчикам при этом доступны привычные инструменты для обучения, запуска и отладки нейросетей.

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

Amazon отмечает, что с помощью TorchNeuron модели можно переносить без правок, но пока с оговорками. Например, torch.compile требует фиксированных размеров входных данных, а при новой форме тензоров заново начинает компиляцию. Также некоторые функции PyTorch пока не реализовали, а часть операций платформа выполняет на CPU.

TorchNeuron уже протестировали разработчики компании Deepgram и нейросети для генерации голоса Cartesia. Представитель Deepgram отметили, что в тестах чипы семейства Trainium обрабатывают на 60% больше секунд входящего аудио на каждый потраченный доллар, чем ускорители, которые использовали в компании ранее. Впрочем, точную схема ускорителей не называют.

Код TorchNeuron открыт и опубликован на GitHub. Hugging Face также добавила поддержку оборудования Amazon в свои библиотеки, в том числе Transformers и Diffusers.