torch.compile для инференса

torch.compile компилирует код PyTorch в оптимизированные ядра, что значительно ускоряет инференс. Эта функция использует TorchDynamo для компиляции кода в графы и TorchInductor для дальнейшей компиляции графов в оптимизированные ядра. Это мощный инструмент оптимизации, который во многих случаях требует лишь добавления одной строки кода.

Оберните модель с помощью torch.compile, чтобы скомпилировать её и вернуть оптимизированную модель.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model)

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

Существует несколько параметров для настройки процесса компиляции. Ниже приведены два из наиболее важных. Полный список параметров см. в документации torch.compile.

Режимы

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

• default — это сбалансированный вариант между скоростью и потреблением памяти.
• reduce-overhead снижает накладные расходы Python за счёт небольшого увеличения потребления памяти, но может работать быстрее.
• max-autotune обеспечивает максимальную скорость, но компиляция занимает больше времени.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead")

Fullgraph

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

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead", fullgraph=True)

Бенчмарки

См. таблицу ниже с бенчмарками производительности, сравнивающими среднее время инференса в миллисекундах при включённом и отключённом torch.compile на различных GPU и размерах батчей для одного и того же изображения при выполнении разных задач компьютерного зрения.

Выберите Subset в таблице ниже, чтобы переключаться между различными GPU, а также между бенчмарками на PyTorch nightly 2.1.0dev и torch.compile с включенным режимом reduce-overhead.

Дальнейшие шаги

• См. руководство torch.compile для обучения, чтобы узнать, как использовать torch.compile для обучения.