ML Serving – Hamel’s Blog - Hamel Husain

Типы оптимизации

Эта таблица вдохновлена этим докладом, в котором Mark Saurofim описывает категории оптимизации, относящиеся к serving моделей.

Когда речь заходит о serving моделей, обсуждение обычно сосредоточено на inference-серверах. Однако одна недооценённая тема — это та роль, которую играют ML-компиляторы (обсуждается в следующих разделах).

Inference-серверы

Ниже перечислены inference-серверы, на которые стоит обратить внимание. Nvidia Triton, по мнению ~20+ профессионалов, с которыми я общался, выглядит самым популярным/надёжным.

Torch Serve TFServe KServe Nvidia Triton

Вот подробные заметки с примерами кода по inference-серверам, которые я изучал. Я рассматриваю два inference-сервера: TFServing и TorchServe, а также REST API сервер общего назначения для serving моделей с использованием fastapi.

ML-компиляторы

Почему вам стоит интересоваться компиляторами

Из этого поста:

Вычисления на edge-устройствах могут сэкономить кучу денег (и потенциально снизить сложность). Переписывание моделей под работу на edge-устройствах — дорогостоящая задача, потому что требует специалистов; хорошие компиляторы могут облегчить этот переход или вовсе сделать его ненужным, помогая транслировать код на высокоуровневом фреймворке в код, специфичный для конкретного железа. Иногда облачные вычисления — не вариант (приходится запускаться на устройстве). Новые типы архитектур могут очень плохо работать на железе без подходящих компиляторов, даже если эти модели существенно меньше и проще устоявшихся. Один такой пример описан исследователями из Google, которые пытались реализовать capsule networks. Перевод кода фреймворка на различное бэкенд-железо — нетривиальная задача, что иллюстрирует график ниже (растёт количество фреймворков + железа/устройств):

Терминология

В терминологии компиляторов легко запутаться. Вот некоторые важные для понимания термины:

Kernel Fusion: процесс объединения нескольких операций в одну. Например, свёртку, за которой следует ReLU, можно объединить в одну операцию. Intermediate Representation (IR): IR — это представление модели, не зависящее от фреймворка вроде PyTorch. Эти IR часто представляют собой JSON, YAML или строку, кодирующую структуру модели, веса и т. д. Эти IR в итоге транслируются в исполняемые программы и часто оптимизируются под конкретное железо. Graph Lowering: процесс перевода одного IR в другой. Например, модель, представленную в IR ONNX, можно понизить (lower) в модель в IR TensorRT. Смысл в том, чтобы перевести один IR в такой, который позволит конкретному компилятору выполнить оптимизации, специфичные для домена/железа. Разные компиляторы могут оптимизировать разные вещи и преследовать разные цели. Например, главная цель ONNX — быть переносимым IR, а главная цель TensorRT — оптимизация inference. «Graph lowering» назван так потому, что модель можно представить в виде графа, и вы последовательно «понижаете» представление этого графа в сторону того, из чего можно сделать машинный код. Компиляторы могут обрабатывать IR в разных фазах, так что IR упрощается/оптимизируется, и бэкенд использует его для генерации машинного кода. Front-End: часть Compiler Stack, которая транслирует язык программирования в IR. Back-End: часть Compiler Stack, которая транслирует IR в машинный код. Можно комбинировать разные фронтенды с разными бэкендами через «middle-end», которые транслируют IR из одного в другой (это также называют «lowering» или «graph lowering»). Graph Breaks: когда вы представляете модель в виде графа через IR, в коде могут быть участки, которые не вписываются в графовую форму, — например, конструкции if/else. В таких случаях возникает graph break, и программу приходится разбивать на несколько графов. Обычно компилятор может оптимизировать только те части, которые представимы в виде графа. В зависимости от используемого фреймворка либо вы, либо сам фреймворк должны будете «сшить» все графы так, чтобы они отражали исходную программу. Graph breaks обычно влекут потерю производительности, поэтому их важно минимизировать.

Частые источники путаницы

Compiler Stack

Часто бывает неоднозначно, что человек имеет в виду под словом «компилятор»: фронтенд, генерирующий IR, бэкенд, генерирующий/исполняющий код, или весь компиляторный стек целиком. Например, у ONNX есть и IR, и runtime. IR — это спецификация (строка), которая позволяет представить модель независимо от фреймворка. ONNX Runtime — это бэкенд, который позволяет исполнять модели, представленные в ONNX IR, на различном железе и из различных языков (Python, C++, Java, JS и т. д.).

Термин «компилятор» часто перегружен. Понимание контекста, в котором используется этот термин, помогает разобраться в документации. Например, Torch Dynamo — это фронтенд (его часто называют «graph acquisition tool»), который производит IR. Пользователь затем может понизить (lower) этот IR в бэкенд другого компиляторного стека, например в C++ (для CPU) или OpenAI Triton (для GPU), который в итоге и будет исполнен.

Nvidia против OpenAI Triton

Triton от Nvidia — это inference-сервер. Triton от OpenAI — это высокоуровневый язык программирования для CUDA и компиляторный стек. Эти два проекта не связаны.

Обучение: TorchDynamo + JIT + Triton

Компиляторы часто воспринимают как оптимизацию для inference. Однако есть компиляторы, которые помогают и в обучении. Самый заметный сегодня — TorchDynamo + JIT Compiler. TorchDynamo — это фронтенд, который позволяет захватить модель PyTorch в виде IR1. Поскольку этот фронтенд поддерживается командой PyTorch, у него будет лучшая поддержка моделей PyTorch. Лучшее в стеке TorchDynamo+JIT — то, что он бесшовно обрабатывает graph breaks за вас (сам сшивает разные подграфы и т. п.). JIT-компилятор заменяет интерпретатор CPython, который обычно «эагерно» (eagerly) исполняет ваш PyTorch-код, ускоряя исполнение2. Он работает, динамически модифицируя байткод Python прямо перед его исполнением. Всё, что вам нужно сделать, — вызвать одну строчку кода: torch.compile(...), чтобы получить преимущества. Andrej Karpathy использует это в своих туториалах по NanoGPT:

Это так недооценено! Смешной прирост скорости от вызова torch.compile(...) Код: https://t.co/9mMGGyEcfzСвязанные посты:1. https://t.co/JOLZ8BQdet2. https://t.co/ci8KAzLW1p (этот блог от @marksaroufim — настоящая жемчужина, кстати) https://t.co/E0oRVC4Ow8 — Hamel Husain (@HamelHusain) February 7, 2023

JIT-компилятор теоретически может использовать разные бэкенды для исполнения, но на данный момент проторённый путь — это TorchInductor, который понижает IR до бэкенда OpenAI Triton. Поэтому весь «compiler stack» выглядит так:

TorchDynamo захватывает модель PyTorch в виде IR (FX-графы, несколько графов, если есть graph breaks) JIT-компилятор понижает FX-граф (через другие промежуточные компиляторные стеки) до TorchInductor TorchInductor понижает IR до OpenAI-Triton для GPU или C++ для CPU OpenAI Triton компилирует IR и исполняет его (возможно, сначала передав его в какой-то другой бэкенд, например CUDA)

Это много шагов! Хотя конечному пользователю необязательно знать обо всех этих шагах, документация по Torch Dynamo может быть очень запутанной, если не понимать эти разные сущности.

Заметки по конкретным компиляторным стекам

TensorRT

Это компилятор/среда исполнения. В отличие от JIT (just-in-time), это AOT (ahead-of-time) компилятор. Он совместим с PyTorch через torch-TensorRT. Из документации: основное средство импорта обученной модели из фреймворка в TensorRT — это формат обмена ONNX, поэтому перед использованием TensorRT нужно сконвертировать модель в ONNX (это происходит автоматически при использовании torch-TensorRT).

ONNX

ONNX в основном сосредоточен на переносимости. ONNX предоставляет собственный набор операций, которые кросс-платформенны и могут быть запущены на ONNX Runtime. ONNX runtime предоставляет клиенты на множестве языков — Python, C, C++, Java, JS и т. д., позволяя загрузить модель как для inference, так и для обучения. Есть железо-специфичные клиенты, оптимизированные под ОС (Linux, Windows, Mac и т. д.), аппаратное ускорение (CUDA, CoreML) и т. п. Клиента можно выбрать здесь: https://onnxruntime.ai/.

Вычислительные графы можно строить из встроенных операций ONNX; однако это непрактичный способ построить модель. На практике вы хотите использовать любимый ML-фреймворк вроде PyTorch или Tensorflow и применить конвертер для конвертации модели в ONNX, например torch.onnx

Особенности PyTorch

Однако бесплатных обедов не бывает. Для PyTorch есть длинный список ограничений и подводных камней, где что-то может пойти не так. Нужно быть аккуратным при экспорте модели и проверять экспортированную ONNX-модель на соответствие исходной. Экспортированная модель выглядит как JSON-файл, который описывает граф вместе с весами.

torch.onnx полагается на TorchScript для экспорта модели в IR. Приходится использовать либо Tracing, либо Scripting в зависимости от потока управления (control flow) в модели.

Tracing: если в модели нет управляющих конструкций (if-выражений, циклов и т. д.), используйте Tracing. Tracing работает за счёт записи модели во время её исполнения. Если в графе есть динамические элементы, они будут записаны как константы. Scripting: если есть управляющие конструкции, используйте Scripting. Mixed: если в коде есть модули без управляющих конструкций и модули с ними, можно скомпилировать эти модули по отдельности с помощью Tracing или Scripting и затем объединить их в одну модель.

Прочитайте туториал по TorchScript и Tracing vs. Scripting, чтобы узнать больше.

Используйте Torch-TensorRT вместо ONNX

Я поговорил с надёжным источником из команды PyTorch, и мне сказали, что TorchScript активно не поддерживается и что мне стоит посмотреть в сторону pytorch - TensorRT.

OpenAI Triton

Только для GPU, из этого доклада:

ApacheTVM

ApacheTVM — это компилятор, у которого, по всей видимости, один из самых широких диапазонов поддерживаемого железа.

Адаптивные компиляторы (TVM)

autoTVM, входящий в состав фреймворка ApacheTVM, относится к классу адаптивных компиляторов, которые запускают тесты на железе и автоматически пытаются оптимизировать код под него, перебирая множество различных конфигураций графа. Из введения Chip в компиляторы:

autoTVM измеряет реальное время выполнения каждого пути, по которому он идёт, что даёт ему ground truth данные для обучения cost model, которая предсказывает, сколько займёт будущий путь. Плюс этого подхода в том, что модель обучается на данных, сгенерированных во время выполнения, и может адаптироваться к любому типу железа, на котором запущена. Минус — cost model требуется больше времени, чтобы начать улучшаться.

Один пример — выход чипов M1 от Apple в ноябре 2020 года. M1 — это ARM-based system on a chip, а ARM-архитектуры более или менее хорошо изучены. Однако у M1 много новых компонентов в его реализации ARM, и требуется значительная оптимизация, чтобы различные ML-модели работали на нём быстро. Через месяц после релиза ребята из OctoML показали, что оптимизация autoTVM почти на 30% быстрее, чем ручная оптимизация команды Apple Core ML.

Цитата эксперта по компиляторам:

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

ML также можно использовать для оптимизации вычислительного графа — то есть «ML оптимизирует ML». Ниже — иллюстрация из работы Chen et al.; TVM: An Automated End-to-End Optimizing Compiler for Deep Learning, где авторы отмечают:

Возможные оптимизации образуют большое пространство, поэтому мы используем cost model на основе ML для поиска оптимизированных операторов.

Адаптивные компиляторы интересны тем, что не зависят от вручную закодированных правил, как традиционные компиляторы, и кажутся более user-friendly и обобщаемыми на новые виды железа. Ещё один адаптивный компилятор — cuDNN autotune, который можно включить через torch.backends.cudnn.benchmark=True, но, видимо, он работает только со свёртками. Недостаток адаптивных компиляторов — то, что им требуется много времени, чтобы «выучить» или «найти» хорошие оптимизации. Однако стоимость оптимизации модели можно амортизировать по многим устройствам с этим железом, а полученный чекпойнт можно даже использовать как стартовую точку для будущих сессий настройки.

Я нахожу адаптивную компиляцию более интересным и многообещающим направлением будущего компиляторов.

WASM

Таргетирование WASM теоретически позволило бы вам работать на любом устройстве, способном запустить браузер, — а браузеры могут запустить очень многие устройства/железо! Однако один из недостатков — то, что это всё ещё медленно. См. этот пост для более подробного обсуждения.

За кем стоит следить

Если хотите узнать больше о компиляторах, рекомендую следить за этими людьми:

Chris Lattner: смотрите, что он строит в Modular Horace He Mark Saroufim

Материалы

Я многое узнал, читая эту отличную статью Chip Huyen. Этот post-mortem попытки запустить capsule networks на современном железе хорошо иллюстрирует важность компиляторов.

Сноски

IR называется FX, а Inductor (другой компилятор) транслирует FX-графы для исполнения в одной из двух сред: OpenAI Triton для GPU и C++/OpenMP для CPU. Документация называет inductor бэкендом, но на самом деле это middle-layer, который понижает IR до другого компиляторного стека.↩︎ PyTorch использует интерпретатор Python, потому что PyTorch нативно работает в «eager mode» и допускает динамизм и поток управления Python (за что его и любят — он очень хакабельный и удобный для отладки).↩︎