Ускорьте инференс с помощью кэширования torch.compile

Сейчас мы кэшируем артефакты torch.compile, чтобы сократить время загрузки моделей, использующих PyTorch. Модели, такие как black-forest-labs/flux-kontext-dev, prunaai/flux-schnell и prunaai/flux.1-dev-lora, теперь запускаются в 2-3 раза быстрее. Мы опубликовали руководство по улучшению производительности моделей с помощью torch.compile, которое охватывает больше деталей.

Что такое torch.compile? Многие модели, особенно из семейства FLUX, применяют различные техники и трюки torch.compile для повышения скорости инференса. Первый вызов скомпилированной функции трассирует и компилирует код, что добавляет накладные расходы. Последующие вызовы выполняют оптимизированный код и значительно быстрее.

В наших тестах скорости инференса с black-forest-labs/flux-kontext-dev скомпилированная версия работает более чем на 30% быстрее, чем нескомпилированная. Улучшения производительности достигаются за счет кэширования скомпилированных артефактов на протяжении жизненного цикла контейнеров моделей, что приводит к резкому сокращению времени холодной загрузки: black-forest-labs/flux-kontext-dev: ~120с → ~60с (на 50% быстрее), prunaai/flux-schnell: ~150с → ~70с (на 53% быстрее), prunaai/flux.1-dev-lora: ~400с → ~150с (на 62% быстрее).

Кэш также улучшает время от запуска контейнера до успешного первого предсказания для всех моделей, использующих torch.compile. Как это работает? Система кэширования работает как многие системы кэширования CI/CD: когда контейнер модели запускается, он ищет кэшированные скомпилированные артефакты. Если они найдены, Torch повторно использует их вместо повторной компиляции с нуля. Когда контейнеры корректно останавливаются, они обновляют кэш при необходимости. Файлы кэша индексируются по версии модели и хранятся близко к узлам GPU.

Чтобы узнать больше о том, как использовать torch.compile, ознакомьтесь с нашей документацией и официальным учебником по torch.compile от PyTorch.