NVIDIA представила AITune: инструмент для оптимизации инференса моделей PyTorch

NVIDIA представила новый инструмент AITune, предназначенный для оптимизации инференса и развертывания моделей глубокого обучения. Этот открытый инструмент упрощает процесс настройки моделей, позволяя пользователям сосредоточиться на разработке, а не на сложной интеграции различных технологий. AITune доступен под лицензией Apache 2.0 и может быть установлен через PyPI, что делает его доступным для команд, стремящихся автоматизировать оптимизацию инференса без необходимости переписывать существующие пайплайны PyTorch.

Основная функция AITune заключается в автоматизации выбора наилучшего бэкенда для каждой модели. Он работает на уровне nn.Module и позволяет значительно улучшить скорость и эффективность инференса в различных областях, таких как компьютерное зрение, обработка естественного языка, распознавание речи и генеративный ИИ. Инструмент автоматически оценивает доступные бэкенды, включая TensorRT, Torch-TensorRT и TorchAO, и выбирает наиболее эффективный, что исключает необходимость ручной настройки.

AITune поддерживает два режима настройки: предварительная настройка (AOT) и настройка в реальном времени (JIT). В режиме AOT пользователи могут предоставить модель и набор данных, а AITune автоматически определит модули, которые можно оптимизировать. В режиме JIT инструмент позволяет оптимизировать модули на лету, что удобно для быстрой проверки перед окончательным развертыванием.

Кроме того, AITune поддерживает кэширование, что позволяет избежать повторной сборки ранее настроенных артефактов. Это значительно ускоряет процесс развертывания, так как пользователи могут загружать уже подготовленные модели без дополнительных затрат времени.

Стратегии выбора бэкенда в AITune также интересны: инструмент предлагает три стратегии, включая FirstWinsStrategy, которая ищет первый успешный бэкенд, и HighestThroughputStrategy, которая выбирает самый быстрый бэкенд, но требует больше времени на первоначальную настройку. Это делает AITune гибким инструментом для различных сценариев использования и позволяет пользователям находить оптимальные решения для своих задач.