Что на самом деле управляет нейронной сетью: компиляторы ИИ под Python

Любой человек с опытом работы с компиляторами и базовыми знаниями в области машинного обучения, возможно, никогда не связывал эти две области. Лексеры, парсеры, деревья синтаксического анализа и IR-язык LLVM относятся к одному миру, тогда как нейронные сети — к другому. Они встречаются в компиляторах для ИИ/машинного обучения, и первым сюрпризом является следующее: модель глубокого обучения на самом деле не работает на Python.

Это звучит странно, ведь строка вроде следующей кажется доказательством того, что Python выполняет всю работу:

output = model(input)

Однако Python сам по себе не выполняет миллиарды умножений матриц. Под капотом происходит гораздо более сложный процесс.

От формулы к инструкциям аппаратного обеспечения

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

Input
  ↓
Matrix Multiplication
  ↓
ReLU
  ↓
Matrix Multiplication
  ↓
Output

В Python его основная часть может выглядеть так:

y = relu(x @ W + b)

Это кажется простым, но подумайте о том, что нужно устройству. У процессора или графического процессора нет понятия этой функции:

relu()

или этого оператора:

x @ W

как они представлены в коде. Аппаратное обеспечение выполняет инструкции. Где-то высокоуровневое описание модели должно быть преобразовано во что-то, что чип может эффективно обрабатывать, и именно эту задачу выполняют компиляторы ИИ/машинного обучения.

Упрощённая схема цепочки от модели до процессора:

AI Model
   ↓
PyTorch / TensorFlow
   ↓
Computational Graph
   ↓
Intermediate Representation (IR)
   ↓
Optimisations
   ↓
Lowering
   ↓
Hardware-specific Code
   ↓
Machine Instructions
   ↓
CPU / GPU / NPU

Это много этапов, поэтому стоит пройти их по одному.

Что делает фреймворк

В PyTorch та же идея реализуется следующим образом:

y = torch.relu(x @ W + b)

Разработчик не хочет заниматься регистрами GPU, структурой памяти, планированием инструкций или тысячами аппаратных потоков. Он просто описывает, какие вычисления ему нужны. Python и PyTorch предоставляют удобный способ сформулировать это описание, в то время как другие инструменты должны преобразовать его в формат, понятный аппаратному обеспечению.

Граф вычислений

Нейронную сеть можно рассматривать как граф вычислений. Это выражение:

y = relu(x @ W + b)

Это можно представить в виде графа связанных операций:

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

Слияние операторов и нагрузка на память

Возьмём модель, содержащую такую цепочку:

MatMul
   ↓
Add
   ↓
ReLU

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

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

Промежуточные представления

Далее рассматривается концепция, знакомая по курсам по классическим компиляторам — промежуточное представление (IR). Традиционный компилятор редко переходит сразу к этому этапу:

Source Code

вот сюда:

Machine Code

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

variables
functions
loops

Компилятор ИИ также должен уметь представлять такие элементы, как:

tensors
shapes
matrix operations
convolutions
attention
memory layouts

Именно это делает проекты вроде MLIR заметными. MLIR обеспечивает инфраструктуру для представления и преобразования вычислений на нескольких уровнях абстракции.

Постепенное снижение уровня абстракции

Компилятор постепенно переходит от высокоуровневой формы к формам, ближе к целевой аппаратуре. Представьте себе спуск по лестнице:

High-level ML operation
          ↓
Tensor operation
          ↓
Lower-level operation
          ↓
Hardware-oriented operation
          ↓
Machine instructions

На каждом уровне абстрактные идеи становятся немного более конкретными, пока компилятор не сгенерирует что-то, что может использоваться целевым оборудованием.

Основное оборудование

Теперь мы подходим к той части, о которой никто не думает при написании кода:

model(input)

Финальные вычисления могут выполняться на CPU, GPU, NPU, TPU или другом ускорителе ИИ. Ни один из них не обращает внимания на то, что в исходном коде было указано:

y = relu(x @ W + b)

Для работы с таким устройством необходимы исполняемые инструкции и данные, организованные соответствующим образом. Для GPU могут потребоваться код и модели выполнения, специфичные для него. CPU использует собственные инструкции и технологии векторизации. Ускорители могут требовать совершенно иного подхода к обработке данных. Поэтому одна и та же высокоуровневая модель может компилироваться по-разному в зависимости от целевого устройства.

Так где же здесь место Python?

Python не исчез, и он по-прежнему очень полезен. Он служит интерфейсом для определения моделей, управления процессом выполнения, загрузки данных и создания приложений. Однако он не обязательно выполняет сам все числовые операции. Упрощённая картина:

Это меняет представление о искусственном интеллекте. Если смотреть только с точки зрения модели, всё выглядит просто:

Data → Model → Prediction

Добавьте компилятор, и появится ещё один уровень между ними:

Model
  ↓
Compiler
  ↓
Hardware

Модель определяет, какие вычисления необходимо выполнить. Компилятор разрабатывает способы их представления и оптимизации. Выполнение осуществляется аппаратным обеспечением.

Почему требуется дополнительный слой

Если уже существуют фреймворки, зачем добавлять слой компилятора? Потому что аппаратное обеспечение для ИИ становится всё более специализированным. Существуют процессоры, графические процессоры, TPU, NPU, ускорители инференса и собственные чипы, каждый из которых имеет свои особенности. Модель, которая хорошо работает на одной архитектуре, может требовать совершенно иного настройки на другой. Компилятор служит связующим звеном между этими двумя сторонами. Одной из сторон является вопрос о том, какие вычисления хочет выполнить модель:

WHAT
The ML model wants to compute

Другой аспект — это то, как конкретное устройство должно делать это эффективно:

HOW
A specific piece of hardware should execute it efficiently

Этот мост постоянно приобретает всё большее значение.

Знакомая философия в новой области

Любой, кто знаком с традиционными компиляторами, распознаёт базовую схему обработки:

Source Code
    ↓
Lexer
    ↓
Parser
    ↓
AST
    ↓
IR
    ↓
Optimisation
    ↓
Machine Code

Машинное обучение применяет аналогичную философию к моделям:

ML Model
    ↓
ML Representation
    ↓
ML IR
    ↓
Optimisation
    ↓
Lowering
    ↓
Hardware Code

Домен отличается, а проблемы гораздо сложнее, однако подход остается прежним: представление, преобразование, оптимизация, снижение затрат, выполнение.

Естественной следующей идеей является создание небольшого компилятора для ИИ, чтобы понять, как он работает. Он не сможет конкурировать с LLVM, XLA или MLIR, а цель будет заключаться лишь в том, чтобы создать что-то достаточно простое для понимания, возможно, начав с простого языка машинного обучения:

Simple ML Language
        ↓
       Parser
        ↓
      ML AST
        ↓
       ML IR
        ↓
   Simple Optimiser
        ↓
    Lowering
        ↓
   LLVM / Backend

Стек в одной строке

Искусственный интеллект обычно представлен через очень простой интерфейс:

model(input)

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