Що насправді керує нейронною мережею: компілятори ШІ під Python

Будь-хто з досвідом роботи з компіляторами та базовими знаннями машинного навчання, можливо, ніколи не пов’язував ці дві галузі. Лексери, парсери, дерева синтаксичного аналізу та LLVM IR належать до одного світу, тоді як нейронні мережі — до іншого. Вони зустрічаються у компіляторах для штучного інтелекту/машинного навчання, і перше здивування полягає у тому, що модель глибокого навчання насправді не працює на 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 використовує інструкції CPU та векторизацію. Акселератор може потребувати зовсім іншого бекенду. Тож одну й ту саму модель високого рівня можна скомпілювати по-різному залежно від її цільової системи.

То де ж місце Python?

Python не зник, і він досі дуже корисний. Це інтерфейс для визначення моделей, керування виконанням, завантаження даних та створення додатків. Він просто не обов’язково виконує кожну числову операцію самостійно. Спрощена уявна модель:

Це змінює сприйняття штучного інтелекту. Якщо дивитися лише з боку моделі, ситуація проста:

Data → Model → Prediction

Додайте компілятор, і з’явиться ще один шар між ними:

Model
  ↓
Compiler
  ↓
Hardware

Модель визначає, які обчислення потрібно виконати. Компілятор знаходить спосіб їх представлення та оптимізації. Апаратне забезпечення виконує ці обчислення.

Чому потрібен ще один шар

Якщо вже існують фреймворки, навіщо додавати шар компілятора? Тому що апаратне забезпечення для ШІ стає все більш спеціалізованим. Існують CPU, GPU, 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 використовується лише як верхній шар.