Amazon открыла TorchNeuron для запуска PyTorch на ускорителях Trainium

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

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.