Amazon представила TorchNeuron для запуска PyTorch на ускорителях Trainium
Amazon представила TorchNeuron — открытый программный модуль для переноса моделей машинного обучения на фирменные ускорители Trainium 2 и Trainium 3. Инструмент предназначен для разработчиков, которые используют PyTorch и хотят обучать, запускать и отлаживать нейросети на оборудовании Amazon.
TorchNeuron подключает Trainium к PyTorch через стандартный механизм расширений фреймворка. Благодаря этому разработчикам не требуется переходить на полностью отдельный программный стек: основные инструменты и привычные подходы к работе с моделями сохраняются.
Два режима работы и распределённое обучение
Модуль поддерживает два режима выполнения. В режиме eager операции запускаются по отдельности. Это упрощает поиск ошибок, проверку отдельных участков кода и проведение экспериментов с архитектурой модели.
Второй вариант основан на torch.compile. В этом режиме вычисления объединяются в граф, после чего оптимизируются компилятором Neuron. TorchNeuron также поддерживает распределённое обучение на нескольких ускорителях, что позволяет использовать его в более крупных вычислительных конфигурациях.
Amazon заявляет, что модели можно переносить без изменений, однако на практике у решения пока есть ограничения. В частности, torch.compile требует фиксированных размеров входных данных. Если форма тензоров меняется, системе приходится запускать компиляцию заново. Кроме того, часть функций PyTorch ещё не реализована, а некоторые операции выполняются на CPU, а не на ускорителе.
Результаты первых тестов
TorchNeuron уже опробовали разработчики компании Deepgram и голосовой нейросети Cartesia. По словам представителя Deepgram, в их тестах ускорители семейства Trainium обрабатывали на 60% больше секунд входящего аудио на каждый потраченный доллар по сравнению с оборудованием, которое компания использовала ранее. Какая именно модель прежних ускорителей применялась для сравнения, не уточняется.
Исходный код TorchNeuron опубликован в открытом доступе на GitHub. Параллельно Hugging Face добавила поддержку оборудования Amazon в свои библиотеки, включая Transformers и Diffusers. Это может упростить запуск распространённых моделей на Trainium, хотя совместимость отдельных операций и требования к формам входных данных по-прежнему нужно проверять при переносе.
