TorchNeuron от AWS приносит нативный бэкенд PyTorch на Trainium

В примечаниях к выпуску Neuron 2.33 от AWS публичная бета-версия TorchNeuron датирована 19 октября. Анонс PyTorch, также датированный 19 октября, сообщает, что публичная бета-версия и открытый исходный код станут доступны 20 октября. Пакет torch-neuronx регистрирует в PyTorch устройство neuron для Trainium2 и Trainium3, позволяя поддерживаемым операциям моделей и тензоров использовать привычные API PyTorch. Для инженеров, рассматривающих эти чипы, полезное изменение — более короткий путь от существующей программы PyTorch к пробному запуску на устройстве. Будет ли программа работать правильно и эффективно, по-прежнему зависит от её операций, форм и программного окружения.

Прежний путь PyTorch для Trainium использовал библиотеку torch-neuronx на основе XLA и код устройства torch_xla. По словам AWS, нативный бэкенд использует расширение PyTorch PrivateUse1 и заменяет xm.xla_device() размещением на устройстве neuron. Для проектов, уже использующих Neuron, это существенное изменение интерфейса, хотя название пакета знакомо. Выпуск относится к обучению и инференсу на инстансах Trn2 и Trn3; в матрице поддержки Trn1 не указан.

Что меняется в программе PyTorch

В примечаниях к выпуску AWS показан перенос модели с помощью .to("neuron") и создание входного тензора на device="neuron". Поддерживаемые операции torch.nn, autograd и оптимизаторов затем выполняются через бэкенд. Код для не привязанного к конкретному устройству torch.accelerator API PyTorch может обнаруживать Neuron без переписывания под определённое устройство, согласно обзору TorchNeuron.

Есть два пути выполнения. Режим непосредственного выполнения, eager, запускает операции PyTorch по мере вызова из Python, сохраняя обычный поток управления, вывод и точки останова. Каждая операция компилируется при первом использовании, а последующие совпадающие вызовы могут использовать кэш. Инженеры также могут скомпилировать модель или выбранный участок с помощью torch.compile(model, backend="neuron", dynamic=False). Нативный бэкенд компилирует поддерживаемые графы прямого и обратного проходов; код вне захваченного участка продолжает работать в eager-режиме. Анонс PyTorch также описывает однократную установку torch.compiler.set_default_backend("neuron") перед компиляцией, чтобы в последующих вызовах компиляции можно было опускать аргумент бэкенда.

Выпуск включает бэкенд neuron для torch.distributed с документированными коллективными операциями и операциями «точка-точка». AWS перечисляет DistributedDataParallel, FSDP2, DTensor и DeviceMesh, а также конвейерный параллелизм среди поддерживаемых API. Это полезные направления для изучения в проекте распределённого обучения, но примечания к выпуску не доказывают, что существующая многоузловая задача будет работать без изменений или достигнет конкретной целевой пропускной способности.

Оценка начинается с подходящего окружения

В кратком руководстве по установке AWS требуется инстанс Trn2 или Trn3 и либо Neuron Base DLAMI, либо поддерживаемый Ubuntu AMI. Матрица выпуска указывает Linux x86_64 на Ubuntu 24.04 или 26.04. Поддержка Ubuntu 22.04 для этого выпуска завершена. В списке указаны Python 3.12, 3.13 и 3.14, а также torch-neuronx 2.13.3.0 для PyTorch 2.13.0 или torch-neuronx 2.14.3.0 для PyTorch 2.14.0. Семейство пакета должно соответствовать версии PyTorch.

Краткое руководство устанавливает torch-neuronx==2.14.* вместе с neuronx-cc==2.* и nki из индексов пакетов Neuron и PyTorch. Затем оно проверяет доступность устройства и сравнивает небольшую матричную операцию на Neuron с результатом CPU. Это проверка пути к установленному устройству в примере AWS. Она не подтверждает работоспособность модели читателя. Разумная следующая оценка — запустить прямой и обратный проходы реальной модели, изучить зарегистрированные переходы на CPU с помощью torch_neuronx.get_fallback_ops(), сравнить выходные данные с текущим окружением и отдельно измерить компиляцию и выполнение в установившемся режиме. Эти последние шаги — предложенные BIG CHANGE вопросы для оценки, а не проведённые нами испытания.

Где бета-версия ещё может изменить работу

В eager-режиме неподдерживаемая операция может перейти на CPU и выдать предупреждение. При использовании torch.compile операция в захваченном графе, которую бэкенд Neuron не может преобразовать в низкоуровневое представление, по умолчанию приводит к сбою компиляции. Для проверки поддержки каждой операции AWS направляет разработчиков к своей таблице совместимости операторов. Поэтому завершившийся eager-запуск сам по себе не доказывает ни полное выполнение на устройстве, ни совместимость с компиляцией.

Выпуск также указывает, что torch.compile требует dynamic=False; новая форма входных данных может потребовать ещё одной компиляции. AWS предлагает дополнение или группировку по формам, чтобы уменьшить их вариативность. Первый запуск включает компиляцию, если подходящего артефакта нет в кэше. Эти детали важны при измерении времени рабочей нагрузки: задержка компиляции и характер форм входных данных могут изменить результат оценки.

Другие явно обозначенные ограничения сужают утверждение о совместимости. Бета-версия не поддерживает torch.sparse, torch.func, torch.nested или torch.distributions. Только некоторые преобразования torch.fft выполняются на Neuron, а float64 выполняется как float32 с предупреждением. В этом выпуске нет интеграции TorchAO и поддержки составного типа данных MXFP8. Команде, использующей эти интерфейсы, нужно проверить собственный путь кода, прежде чем считать размещение на устройстве планом миграции.

Главное изменение

TorchNeuron в публичной бета-версии подключает Trainium через нативные интерфейсы устройств и компиляции PyTorch. Это делает первую пробу поддерживаемых программ PyTorch более прямой и даёт существующим пользователям Neuron иной путь по сравнению с кодом на основе XLA. Решающее свидетельство в пользу внедрения должна дать каждая рабочая нагрузка: остаются ли операторы на нужном устройстве, приемлемы ли численные результаты и оправдывают ли измеренные производительность и стоимость инстанс и работу по миграции. Для этого анализа BIG CHANGE изучил документацию и анонсы; мы не запускали TorchNeuron и не проверяли независимо заявления AWS о производительности.

Источники и дополнительное чтение

• Примечания к выпуску AWS Neuron 2.33 TorchNeuron: дата выпуска, статус бета-версии, матрица версий и операционных систем, API и известные ограничения.
• Обзор TorchNeuron от AWS: интеграция устройств, непосредственное и скомпилированное выполнение, поведение при переходе на CPU и программный путь через Neuron.
• Краткое руководство AWS по установке и проверке: требования к инстансу и образу, установка пакетов и документированная проверка устройства.
• Блог PyTorch: представляем TorchNeuron: объяснение смены устройств, компиляции и распределённых API со стороны PyTorch. Статья — анонс участников AWS в блоге PyTorch, а не независимый бенчмарк.
• Анонс Amazon: изложение AWS своих целей и первых пользователей. Заявления о миграции, производительности и экономических результатах требуют проверки для конкретной рабочей нагрузки.