Amazon Web Services (AWS) с 19 октября расширяет доступ к инструменту TorchNeuron — из частного тестирования он переходит в публичное. С его помощью разработчики смогут запускать модели PyTorch на чипах Trainium, внеся лишь незначительные изменения в настройки.
TorchNeuron предлагает нативный аппаратный бэкенд для PyTorch. Инструмент поддерживает режим немедленного выполнения (eager mode), компиляцию torch.compile, функцию распределённого обучения FSDP и другие возможности. Чтобы перенаправить код с NVIDIA GPU на Trainium, разработчикам достаточно заменить .to(‘cuda’) на .to(‘neuron’). Платформа Hugging Face уже добавила нативную поддержку TorchNeuron в такие инструменты, как Transformers и Diffusers.
Ранее AWS уже пыталась обеспечить совместимость Trainium с PyTorch — с помощью Neuron SDK и PyTorch/XLA. Однако эти решения требовали от разработчиков освоения иных способов компиляции и выполнения кода. Например, XLA сначала преобразовывал вычисления в вычислительный граф — это усложняло отладку и снижало совместимость с некоторыми инструментами PyTorch.
Цель TorchNeuron — упростить для разработчиков работу с чипами AWS, не заставляя их переучиваться. AWS также вносит код в основные проекты PyTorch и сотрудничает с сообществом открытого ПО, чтобы уменьшить привязку фреймворка к конкретным чипам. Это может ослабить доминирование экосистемы CUDA от NVIDIA — в неё входят не только интерфейсы для вычислений на GPU, но и библиотеки вроде cuDNN и NCCL, а также множество оптимизированных под NVIDIA GPU программ.
По словам представителя AWS Колина Брейса (Colin Brace), уже более 30 компаний начали использовать TorchNeuron — среди них игроки из сфер робототехники, языковых моделей и голосового ИИ. Например, компания Deepgram совместно с AWS развернула модель распознавания речи Nova-3 на Trainium2: в определённых задачах обработка аудио на доллар выросла на 11–70 %. Но для достижения нужной производительности инженерам обеих компаний потребовалось несколько месяцев на перенос модели и настройку среды выполнения.
AWS сотрудничает с крупными игроками вроде Anthropic и OpenAI, однако остаётся открытым вопрос, готовы ли средние компании, не имеющие тесных связей с AWS, перейти на Trainium благодаря TorchNeuron.
Подписаться в Макс