Компиляция

Обзор

Pytorch 2.0 представил torch.compile — мощную возможность, которая ускоряет выполнение кода PyTorch за счёт JIT-компиляции в оптимизированные ядра. Ключевые возможности torch.compile включают:

• Улучшение производительности: значительное ускорение выполнения модели за счёт оптимизации вычислительного графа.
• Простота использования: Для внедрения требуется минимальное количество изменений в коде, что делает его высокодоступным.
• Совместимость: Бесшовно работает с существующим кодом и моделями PyTorch.

При использовании с Accelerate torch.compile легко встраивается в рабочие процессы распределённого обучения, позволяя одновременно извлекать выгоду из распределённого выполнения и оптимизаций компиляции.

Первое выполнение скомпилированного кода обычно занимает больше времени, так как включает время компиляции, но последующие запуски значительно быстрее. Для оптимальной производительности в различных сценариях torch.compile предлагает различные режимы, такие как "default", "reduce-overhead" (использует CUDA graphs для дальнейшего снижения накладных расходов) и "max-autotune" (выполняет обширную автонастройку для поиска лучших ядер для вашей модели).

Использование torch.compile с Accelerate

Accelerate предоставляет TorchDynamoPlugin для простой и бесшовной интеграции torch.compile в ваши скрипты обучения.

from accelerate import Accelerator
from accelerate.utils import TorchDynamoPlugin

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    backend="inductor",  # Options: "inductor", "aot_eager", "aot_nvfuser", etc.
    mode="default",      # Options: "default", "reduce-overhead", "max-autotune"
    fullgraph=True,
    dynamic=False
)

# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply torch.compile to your model
model = accelerator.prepare(model)

Он совместим со всеми остальными функциями и плагинами Accelerate, включая смешанную точность, распределённое обучение (DDP, FSDP, Deepspeed) и т. д.

Региональная компиляция

Вместо попытки скомпилировать всю модель, что обычно приводит к большой области для оптимизации, региональная компиляция нацелена на повторяющиеся блоки одного класса и компилирует их последовательно, чтобы использовать кэш компилятора. Например, в GPT2LMHeadModel повторяющимся блоком/классом является GPT2Block, к которому можно обратиться как model.transformer.h[0]. Остальная часть модели (например, model.lm_head) компилируется отдельно.

Это позволяет ускорить время компиляции и холодный запуск моделей, таких как LLM и Transformers в целом. Подробнее см. https://pytorch.org/tutorials/recipes/regional_compilation.html.

Как использовать региональную компиляцию

Его можно включить, установив use_regional_compilation=True в конфигурации TorchDynamoPlugin:

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    use_regional_compilation=True,
    ... # other parameters
)
# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply compile_regions to your model
model = accelerator.prepare(model)

Вы также можете использовать утилиту accelerate.utils.compile_regions напрямую так же, как вы используете torch.compile.

Преимущества региональной компиляции

Мы провели обширное тестирование, сравнив полную компиляцию и региональную компиляцию с использованием функции torch.compile в PyTorch. Полные результаты доступны в репозитории accelerate. Ключевые выводы из нашего тестирования:

1. Сопоставимая производительность: Региональная компиляция обеспечивает ускорение производительности, сопоставимое с полным компилированием, особенно для крупных моделей.
2. Более быстрая компиляция: Региональная компиляция значительно сокращает время, необходимое для компиляции моделей, что делает её более эффективным выбором для развёртывания.
3. Влияние размера батча: Разница в производительности между стратегиями компиляции уменьшается при увеличении размера батча, что указывает на меньшее влияние накладных расходов на компиляцию в таких сценариях.
4. Учёт размера модели: Преимущества региональной компиляции более выражены в больших моделях, где экономия времени на компиляцию может быть значительной.
5. Практическое применение: Для реальных задач региональная компиляция является практичным выбором для оптимизации времени холодного запуска обучения, особенно при работе с большими моделями.

Заключение

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