ML Serving – Hamel’s Blog - Hamel Husain

Types of Optimization

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

This table is inspired by this talk where Mark Saurofim outlines categories of optimization relevant to model serving.

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

When considering model serving, the discussion tends to be focused on inference servers. However, one underrated topic is the importance ML compilers play (which is discussed in later sections).

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

Inference Servers

Inference-серверы

Below are the inference servers I would pay attention to. Nvidia Triton seems to be the most popular/robust according to ~20+ professionals I’ve spoken with.

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

Torch Serve TFServe KServe Nvidia Triton

Here are detailed notes with code examples on inference servers I’ve explored. I explore two inference servers: TFServing and TorchServe, as well as a general-purpose REST API server for model serving with fastapi.

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

ML Compilers

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

Why You Should Care About Compilers

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

From this blog post:

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

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

Terminology

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

It’s easy to get lost in the terminology of compilers. Here are some terms that are important to understand:

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

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 обычно влекут потерю производительности, поэтому их важно минимизировать.

Common Sources of Confusion

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

The Compiler Stack

Compiler Stack

It can often be ambiguous when someone refers to a compiler if they are referring to a front-end that generates the IR, the back-end that generates/executes the code, or the entire compiler stack as a whole. For example, ONNX has both an intermediate representation IR and a runtime. The IR is a specification (a string) that allows you to represent a model in a way that’s independent of a framework. The ONNX Runtime is a backend that allows you to execute models represented by the ONNX IR on a variety of hardware and from various languages (Python, C++, Java, JS, etc.).

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

The term “compiler” is often overloaded. Understanding the context in which the term is used can be helpful for understanding documentation. For example, Torch Dynamo is a front-end (often referred to as a “graph acquisition tool”) that produces an IR. The user can then lower this IR to a back-end to another compiler stack like C++ (for CPUs) or OpenAI Triton (for GPUs) that eventually gets executed.

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

Nvidia vs OpenAI Triton

Nvidia против OpenAI Triton

Triton by Nvidia is an inference server. Triton by OpenAI is a high-level CUDA programming language and compiler stack. The two are not related.

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

Training: TorchDyamo + JIT + Triton

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

Compilers are often thought of as optimizations for inference. However, there are compilers that help with training too. The most notable these days is TorchDynamo + JIT Compiler. TorchDynamo is a front-end that allows you to capture a PyTorch model as an IR¹. Since this front-end is maintained by the Pytorch team, it is going to have the greatest support for Pytorch models. The best thing about the TorchDynamo+JIT stack is that it seamlessly handles graph breaks for you (stitches various subgraphs together for you, etc). The JIT compiler supplants the CPython interpreter that normally “eagerly” runs your PyTorch code for faster execution². It works by dynamically modifying Python bytecode right before it is executed. All you have to do is to call one line of code: torch.compile(...) to see the benefits. Andrej Karpathy is using this in his NanoGPT tutorials:

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

This is so underrated! Ridiculous speedup by calling torch.compile(...)

Code: https://t.co/9mMGGyEcfz

Related blog posts:
1. https://t.co/JOLZ8BQdet
2. https://t.co/ci8KAzLW1p (this blog by @marksaroufim is a real gem, BTW) https://t.co/E0oRVC4Ow8

Это так недооценено! Смешной прирост скорости от вызова 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

The JIT Compiler can theoretically leverage different backends for execution, but at this time, the paved path is to use the TorchInductor which lowers the IR to the OpenAI Triton backend. So the entire “compiler stack” looks like this:

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)

That is a lot of steps! While the end-user doesn’t necessarily need to be aware of all these steps, the documentation on Torch Dynamo can be really confusing unless you are aware of these different things.

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

Notes on Specific Compiler Stacks

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

TensorRT

TensorRT

This is a compiler/execution environment. Unlike the JIT (just-in-time), it is an AOT (ahead-of-time) compiler. It is compatible with PyTorch via torch-TensorRT. From the docs: TensorRT’s primary means of importing a trained model from a framework is through the ONNX interchange format, so you need to convert your model to ONNX before getting to TensorRT (this happens automatically when you use torch-TensorRT).

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

ONNX

ONNX

ONNX is mainly focused on portability. ONNX provides its own set of ops that are cross-platform and can be run on an ONNX Runtime. The ONNX runtime provides clients in many languages, like Python, C, C++, Java, JS, etc., allowing you to load a model for either inferencing or training. There are hardware-specific clients that are optimized for OS (i.e., Linux, Windows, Mac, etc.), Hardware acceleration (CUDA, CoreML) and so forth. You can select a client from here: https://onnxruntime.ai/.

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

You can construct your computation graphs with ONNX’s built-in ops; However, this is an impractical way to build a model. In practice, you want to use your favorite ML framework like PyTorch or Tensorflow and use a converter to convert that model to ONNX like torch.onnx

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

Pytorch Considerations

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

However, there is no free lunch. For PyTorch, there is a long list of limitations and gotchas where things can go wrong. You must be careful when exporting a model and should test the exported ONNX model against the native model for consistency. The exported model looks like a JSON file that describes the graph along with weights.

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

torch.onnx relies on TorchScript to export the model into an IR. You have to either use Tracing or Scripting depending upon the control flow in your model.

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

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

Read the TorchScript Tutorial and Tracing vs. Scripting to learn more.

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

Use Torch-TensorRT instead of ONNX

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

I talked to a trusted source on the Pytorch team, and they said that TorchScript is not actively maintained and that I should look at pytorch - TensorRT instead.

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

OpenAI Triton

OpenAI Triton

Only for GPUs, from this talk:

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

ApacheTVM

ApacheTVM

ApacheTVM is a compiler that seems to have one of the widest ranges of hardware support.

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

Adaptive Compilers (TVM)

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

autoTVM, part of ApacheTVM compiler framework are part of a class of compilers that are adaptive that run tests on hardware and automatically try to optimize it for the hardware by trying many different graph configurations. From Chip’s intro to compilers:

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

autoTVM measures the actual time it takes to run each path it goes down, which gives it ground truth data to train a cost model to predict how long a future path will take. The pro of this approach is that because the model is trained using the data generated during runtime, it can adapt to any type of hardware it runs on. The con is that it takes more time for the cost model to start improving.

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

One example is when Apple released their M1 chips in Nov 2020. M1 is an ARM-based system on a chip, and ARM architectures are more or less well-understood. However, M1 still has a lot of novel components of its ARM implementation and requires significant optimization to make various ML models run fast on it. A month after the release, folks at OctoML showed that the optimization made by autoTVM is almost 30% faster than hand-designed optimization by Apple’s Core ML team.

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

Quote from a compiler expert:

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

Pros of dynamic compilers is that we maximize performance across a broader set of models/ops. Downside of dynamic compilers is search time. AutoTVM generally takes hours, even more than a day to compile because it’s trying every possible combination.

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

ML can also be used to optimize the computation graph, or “ML is optimizing ML”. Below is an illustration from Chen, et.al; TVM: An Automated End-to-End Optimizing Compiler for Deep Learning, where they comment that:

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

Possible optimizations form a large space, so we use an ML-based cost model to find optimized operators.

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

Adaptive compilers are interesting because they do not depend on hand-coded rules like traditional compilers and seem more user-friendly and generalizable to new kinds of hardware. Another adaptive compiler is cuDNN autotune, which you can enable with torch.backends.cudnn.benchmark=True, but apparently this only works with convolutions. The drawback of adaptive compilers is that they can take a long time to “learn” or “search” for good optimizations. However, you can amortize the cost of optimizing the model over many devices on that hardware, and can even take that checkpoint as a starting point for future tuning sessions.

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

I find adaptive compilation more exciting and promising as a future direction in compilers.

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

WASM

WASM

Targeting WASM would theoretically enable you to work on any device that can run a browser - and many devices/hardware can run browsers! However, one drawback is that it’s still slow. See this blog post for a more detailed discussion.

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

People to follow

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

If you want to learn more about compilers, I recommend following these people:

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

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

Resources

Материалы

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

Footnotes

Сноски

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