Што на самай працоўны спосаб керуе нейральной сеткай: кампайляры AI пад Python

Кожны, хто мае апаратнае програмаванне і певныя знанні па аспектах машыннага навчання, можа ніколі не з’ўязаў гэтыя два направы. Лексеры, парсеры, структуры аб’ектаў і LLVM IR належаць да аднаго света, а нейронныя сеткі — да іншага. Яны зустрочаюцца ў кампайлярах для AI/ML, і першыя неспадзянкі стаюць такімі: мадэль глыбокага навчання на самай працэ практычна не працуе на Python.

Гэта звучыць дзівна, адколі такі код выглядае так, нібы Python выкананае ўсё:

output = model(input)

Адзначыце, Python сам не выкананаў мільярды умножэнняў матрыц. Пад спадом ведзецца значна дыягональнейшы працэс.

З формулы да інструкцыйяў апаратуры

Уявіце сабе маленькую сетку, якая праходзіць ад вхідных дадзеных праз умножэнне матрыц і даўжэй:

Input
  ↓
Matrix Multiplication
  ↓
ReLU
  ↓
Matrix Multiplication
  ↓
Output

У Python ягоў карэнны код можа выглядаць так:

y = relu(x @ W + b)

Гэта выглядае проста, але падумайце пра тое, чаго патрэбуе апарат. У CPU чыра GPU няма жаднага уявлення пра гэтую функцыю:

relu()

альбо гэтага аператара:

x @ W

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

Спрасцаваныяе зображэнне ланцоўкі, ад модэлю да працоўніка:

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-ы. Кампайляры AI следуюць таму ж прынцыпу, але іх прэдставленнія павінны адражаць ідэі, важныя для машыннага навчання. Калі звычны кампайляр думае пра:

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 не знікла, і яна яшчэ дужа корыстная. Яна ўжо ўсё тая інтэрфейс для задання модэляў, кантролю выканання, завантажэння дадзеных і стварэння прыкладнікаў. Яна проста не павінна обавязкова сама выканаваць кожную чысловую операцыю. Апроставаны ментальны модэль:

Гэта зміняе візуальны аспект AI. Калі аднацца толькі да стороны модэля, картына простая:

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, а мета будзе проста — стварыць ўсё тое, што ўстаўляецца зрозумець, можа, пачынając з простага мовы машыннага навчання:

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

Стакан пад адной лініяй

Штучны інтэлект зазвычай выклікаецца праз вельмі просты інтэрфейс:

model(input)

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