Amazon анонсировала TorchNeuron, собственный бэкэнд PyTorch с открытым исходным кодом для своего семейства микросхем искусственного интеллекта Trainium, что устраняет одно из самых больших практических препятствий на пути внедрения специального кремния для рабочих нагрузок машинного обучения. Согласно официальному сообщению в блоге PyTorch, опубликованному в преддверии конференции PyTorch в Северной Америке 20 октября, разработчики теперь могут запускать существующий код модели PyTorch на Trainium, внося однострочное изменение, вместо того, чтобы переписывать свои рабочие процессы с использованием проприетарных инструментов.

Это объявление сразу же привлекло внимание, поскольку оно направлено на то, что аналитики уже давно называют самым устойчивым преимуществом Nvidia: не производительность чипов, а программную экосистему CUDA, которая привязывает разработчиков к ее графическим процессорам. Reuters представило этот релиз как заявление Amazon о том, что ее чипы искусственного интеллекта будут работать с самой популярной средой разработки в области искусственного интеллекта, а TechCrunch назвал TorchNeuron самым большим шансом компании на Nvidia. Чтобы получить более широкое представление о том, как меняется рынок ускорителей, мы отслеживаем эти события по мере их развития.

Что на самом деле меняет TorchNeuron

Основное обещание TorchNeuron — простота. В примере, опубликованном в блоге PyTorch, существующий цикл обучения запускается в Trainium после всего лишь перемещения модели на новое устройство с помощью `.to('neuron')`. Согласно сообщению, код, написанный на основе `torch.accelerator`, рекомендованного PyTorch независимого от устройства API, работает без изменений. Для скомпилированных рабочих нагрузок разработчики один раз устанавливают серверную часть по умолчанию с помощью `torch.compiler.set_default_backend("neuron")`, а затем вызывают `model.compile()` без дальнейших изменений.

Под капотом TorchNeuron регистрирует Trainium через PrivateUse1, официальную точку расширения PyTorch для пользовательских аппаратных серверных частей, что означает, что чип ведет себя как стандартное устройство PyTorch, а не как придаток, прикрепленный к платформе. AWS описывает интеграцию как охватывающую несколько режимов выполнения:

  • Ускоренное выполнение для исследований и отладки с экспериментальным «адаптивным нетерпеливым режимом», который снижает затраты на запуск.
  • torch.compile для повышения производительности: преобразование моделей в графики, оптимизируемые компилятором Neuron, со встроенной поддержкой FlexAttention.
  • Интерфейс ядра нейронов (NKI) для исследователей, которым нужны специальные низкоуровневые ядра.

Согласно объявлению, бэкэнд работает как на Trainium2, так и на более новых чипах Trainium3.

Прямой удар по программному рву Nvidia

Доминирование Nvidia в сфере ИИ-вычислений никогда не опиралось исключительно на более быстрое оборудование. CUDA, программная платформа компании, созданная уже два десятилетия, — это место, где пишется, настраивается и развертывается большая часть исследовательского кода искусственного интеллекта, и каждому чипу-претенденту — от линейки AMD Instinct до TPU Google и Trainium Amazon — приходится сталкиваться с затратами, требующими от разработчиков отказаться от нее. Освещение запуска TorchNeuron в GeekWire, Benzinga и GuruFocus однозначно представило это объявление как нападение на этот ров.

Amazon не пытается вытащить разработчиков из PyTorch, чтобы привлечь их на свою сторону. Вместо этого компания внедряет Trainium в PyTorch, чтобы огромная база пользователей платформы могла внедрить этот чип, не меняя своих привычек. Сообщение в блоге PyTorch содержало одобрение со стороны Meta, которая поддерживает эту структуру: «Встроенная поддержка AWS Trainium расширяет этот выбор, позволяя исследователям и инженерам экспериментировать и выполнять итерации, используя уже знакомые им быстрое выполнение и распределенные API», — сказала Яна ван Греунен, директор по разработке PyTorch в Meta, в своем посте.

Такая структура имеет значение для облачного бизнеса AWS. Экземпляры Trainium продаются как более дешевая альтернатива графическим процессорам Nvidia для обучения и вывода, и эти чипы уже привлекли крупных клиентов в области ИИ. Создание ощущения, что чипы являются родными для PyTorch, напрямую устраняет возражение, которое замедляет более широкое внедрение: усилия по портированию.

Технические возможности: от ноутбуков до обучения в масштабе кластера

Возможности, перечисленные в блоге PyTorch, позволяют предположить, что Amazon стремится выйти далеко за рамки экспериментов для любителей. TorchNeuron охватывает наборы операторов, наблюдаемые в текстовых, мультимодальных и визуально-языковых моделях, а также при генерации изображений и видео, моделях пространства состояний и моделях мира с точностью BF16, FP32 и FP16; Поддержка TorchAO и FP8 находится в стадии закрытого бета-тестирования, а в следующем выпуске они станут общедоступными.

Распределенное обучение — дисциплина, которая превращает код с одним графическим процессором в прогоны в масштабе кластера — поддерживается через стандартные API PyTorch — DDP, FSDP2, SimpleFSDP, тензорный и конвейерный параллелизм, а также сегментирование в стиле ZeRO, а также DTensor и DeviceMesh для многомерных групп процессов. Потоки и события сопоставляют семантику PyTorch с аппаратными очередями Trainium, говорится в сообщении.

Можно сделать вывод, что стеки Amazon Neuron построены на одной и той же основе: vLLM Neuron обслуживает большие языковые модели, а vLLM Omni Neuron обслуживает диффузионные модели для генерации изображений и видео. Интерфейс ядра Neuron позволяет исследователям писать собственные операторы как для быстрого, так и для скомпилированного выполнения, а AWS представляет собой NKI-lib с открытым исходным кодом, набор оптимизированных для Trainium ядер, охватывающих блочные вычисления FlashAttention и MoE для основных сценариев LLM и мультимодальных сценариев.

Инструменты повышения производительности также включены. Neuron Explorer расширяет стандартный torch.profiler для отслеживания выполнения от платформы до аппаратного обеспечения, связывая каждое событие обратно с исходной строкой, и может экспортировать метрики в формате Parquet, которые агенты могут использовать в почтовых заметках для автоматического анализа узких мест.

Доступность и что будет дальше

Согласно сообщению в блоге, запуск TorchNeuron с открытым исходным кодом запланирован на PyTorch Conference North America 20 октября, при этом совместимость распространяется на Hugging Face Transformers v5. AWS описывает Trainium как портфель микросхем, разработанных с нуля для обучения и вывода ИИ в облачном масштабе.

Стратегическим контекстом является рынок, на котором каждый крупный поставщик облачных услуг одновременно является крупнейшим покупателем ускорителей Nvidia и потенциальным конкурентом. Microsoft, Google и Meta разрабатывают собственные микросхемы, а китайские фирмы, включая Huawei, создают свои собственные программные стеки, чтобы избежать гравитации CUDA. Встречаясь с разработчиками внутри PyTorch, а не предлагая им мигрировать, Amazon выбрала наименее разрушительный путь в этом соревновании — и наиболее вероятный, если опыт разработчиков сохранится, чтобы перенести реальные рабочие нагрузки обучения с оборудования Nvidia.

Действительно ли TorchNeuron подрывает преимущество CUDA, будет зависеть от производительности в масштабе и длинного хвоста реальных моделей — вопросы, на которые материалы запуска не дают ответа. Но барьер, который удерживал большинство разработчиков PyTorch даже от попытки переписать Trainium, теперь снижен до одной строки кода.

---

Будьте впереди ИИ

Получайте последние новости, анализ и открытия в области искусственного интеллекта — все в одном месте.

Подробнее новости AI →