🚗 Тест-драйв PyTorch 2.0 и заглядываем под капот «двушки»

⚜️ Новейший метод для ускорения кода в PyTorch 2.0 — torch.compile(), который позволяет JIT‑компилировать код PyTorch в оптимизированные ядра, требуя минимальных изменений кода.

JIT‑компиляция (Just‑In‑Time compilation) — это процесс, при котором код на высокоуровневом языке преобразуется в код на низкоуровневом языке, который может быть быстрее и эффективнее исполнен процессором или графическим ускорителем. torch.compile() использует TorchDynamo и заданный бэкенд для JIT‑компиляции кода PyTorch.

Рассмотрим, насколько эффективен torch.compile() на практике, заглянем под капот «двушки» PyTorch 2.0, чтобы узнать, как работает этот метод и какие преимущества он дает.

Разработчики PyTorch заявляют, что torch.compile() может дать прирост производительности до 50% по сравнению с обычным кодом PyTorch. Для проверки этого заявления мы проведем ряд экспериментов на разных моделях и данных, выясним, есть ли реальная выгода от использования torch.compile().

📌 Для того чтобы понять, как работает torch.compile() рассмотрим основные его компоненты, которые отвечают за различные аспекты JIT‑компиляции кода PyTorch:

▪️TorchDynamo — это динамический компилятор, который анализирует код PyTorch и определяет, какие части кода могут быть скомпилированы в оптимизированные ядра. Отслеживает изменения в коде и перекомпилирует его при необходимости.

▪️AOT AutoGrad — это система автоматического дифференцирования, которая позволяет вычислять градиенты для скомпилированных ядер. Генерирует код для обратного распространения ошибки во время компиляции, а не во время исполнения, что ускоряет процесс обучения нейронных сетей.

▪️PrimTorch — это набор примитивных операций, которые используются для построения скомпилированных ядер. Включает в себя базовые математические и логические операции, а также операции над тензорами, такие как сложение, умножение, свертка и т. д.

▪️TorchInductor — это бэкенд для JIT‑компиляции кода PyTorch в оптимизированные ядра для разных устройств. Поддерживает разные бэкенды и адаптирует код PyTorch к специфике каждого устройства.

Далее подробнее рассмотрим, как работают компоненты TorchDynamo и TorchInductor и как они взаимодействуют друг с другом, чтобы обеспечить JIT‑компиляцию кода PyTorch.

📌Читать

@machinelearning_interview