Після презентації ядра MLA Jalapeño на HotChips та подальших коментарів від SemiAnalysis виникло багато обговорень. Як команда апаратного забезпечення OpenAI, ми ледь торкнулися цього маленького золотого самородка: того факту, що ШІ пише наші ядра, і що, коли він це робить, нам не потрібно розуміти, що робить ядро рядок за рядком. Ми явно залишили поза увагою: як таке взагалі можливо? Як правильно думати про це, порівняно з більш традиційним методом генерації коду? Чи є оптимізоване ядро таким же надійним, як і неоптимізоване?
Що стосується мого досвіду, я працював над компіляторами для прискорювачів понад десять років. Я започаткував XLA, яка є чудовою інфраструктурою компілятора з чудовою міжкомпанійною командою та зусиллями, що працюють над нею. Останні 2+ роки в OpenAI я намагаюся переосмислити, як мають працювати компілятори в епоху ШІ. Нові формулювання компіляторів спиратимуться на існуючі сильні сторони, але неможливо заперечувати, що в наборі інструментів з'явився потужний новий інструмент, який варто використовувати.
Це буде певна подорож, але я сподіваюся пролити світло на те, як ШІ використовується для автоматизації покращення комп'ютерних програм; тобто оптимізуючої компіляції. Я вважаю, що за допомогою ШІ ми можемо пережити те, що ми називаємо "компіляторами 2.0". ШІ принципово менш обмежений у тому, що він може запропонувати, і те, що він пропонує, є результатом навчання та контексту моделі, що призводить мене до класифікації його як "стохастичного оптимізатора" – це може створювати проблеми, але, як ми побачимо, також є джерелом великих переваг…
Значна частина академічних досліджень та промислового застосування вже рухається в цьому напрямку, швидко розкриваючи потенціал участі ШІ в царині оптимізуючих компіляторів, але ми перебуваємо на етапі, коли це потребує загального пояснення.
Передумови
Компілятори приймають програми та видають перекладені або покращені версії цих програм.
Програми, як на стороні входу, так і на стороні виходу, мають семантику, яка говорить нам, що означають програми, що вони можуть робити, і як міркувати про ці речі, які вони можуть робити.
Ті з нас, хто працює над компіляторами, думають про них як про чисті функції – вони приймають структуру даних і видають структуру даних, яка повинна мати відповідну семантику.
Іноді наші компілятори зосереджуються на "зниженні рівня" або "перекладі". Наприклад, вони можуть прийняти C і видати асемблер x86-64, який ми часто вважаємо "нижчим рівнем". Але часто вони роблять більше, ніж просто переклад, як частина цього процесу…
Наші компілятори, на практиці, зосереджуються на "оптимізації". Вони можуть прийняти структуру даних, яка представляє програму – на нашій мові "проміжне представлення" (IR) – і намагаються створити кращу версію цієї програми. Іноді "краще" означає, що вона виконується за меншу кількість циклів, іноді – що вона матиме менше непотрібного коду, іноді – спеціалізацію для речей, які ми можемо довести "мають бути правдою" про програму (часткове обчислення).
Тепер коротко розглянемо, що LLM спочатку були створені для перекладу людського тексту з однієї мови на іншу. Очевидно, що переклад входить до їхньої компетенції. І ми можемо бачити через використання LLM у повсякденних завданнях, що вони також можуть писати нові рішення та покращувати існуючі. Багато з нас, програмістів, також мають досвід прохання LLM "оптимізувати цей фрагмент коду", і вони чудово це роблять. (Однак нам потрібно знати, що вони правильно оптимізували код, до чого ми ще дійдемо!) Це просто для того, щоб підкреслити, що LLM мають можливості, які ми шукаємо в оптимізуючому компіляторі.
Оптимізація та Оптимальність
Оптимізуючі компілятори, не дивно, намагаються підвищити оптимальність програми, над якою вони працюють, за певним критерієм (зазвичай часом виконання). Це настільки важко зробити в загальному випадку для довільної програми, що існує теорема, яка називається теоремою повної зайнятості для інженерів-компіляторників. (Я дізнався про це лише після того, як вибрав бути інженером-компіляторником, але це все одно принесло мені полегшення!)
"Супероптимізатори" – це дивовижна маленька підгалузь оптимізуючих компіляторів. Уявіть, що є задана програма, і ми можемо сказати, що вона робить за допомогою семантики. Якою є найоптимальніша програма, яка має ту саму семантику? Ось що намагаються вирішити супероптимізатори, і це, по суті, проблема пошуку…
Уявіть, що я намагаюся знайти найкоротшу програму, яка має ту саму семантику, і в мене є спосіб запитати, чи має програма-кандидат таку саму семантику. Гіпотетично, я міг би перерахувати всі програми в порядку об'єктивності та вибрати найменшу, яка має ту саму семантику.
Однак перерахування всіх програм у порядку об'єктивності звучить досить нездійсненно. Одна з моїх улюблених академічних робіт, створена в 2013 році під назвою "STOKE" (Стохастична супероптимізація), запитала: "а що, якщо ми просто випадково змінюватимемо програми знову і знову, чи побачимо ми тоді врешті-решт найкращу програму?" Вони припустили, що за допомогою випадкового блукання (і з нашим старим другом машинного навчання, ланцюгом Маркова Монте-Карло / Метрополіс-Гастінгс), зрештою ви побачите цю оптимальну програму.
Монте-Карло модифікації зазвичай дурні (ви випадково обираєте зміну), але швидкі. LLM дуже розумні (багато токенів міркування), але порівняно повільні.
Що, якщо замість дурних/швидких модифікацій Монте-Карло ми доручимо LLM визначати напрямки, в яких слід рухати програми? Ми б отримали стохастичний оптимізатор, який був би дуже розумним, проводячи нашу програму через простір оптимізованих програм.
Інтуїція для Оптимізації
Давайте зробимо крок назад. Подумайте про людину, яку ви знаєте і яка найкраще уособлює "оптимізує фрагменти коду до краю". Для стислості назвемо її "Оптимізатор Оля". Оля, ймовірно, має інтуїтивне відчуття, які види змін коду можуть бути плідними. Оля, ймовірно, пробує деякі речі, щоб побачити, чи вони працюють, і якщо вони не спрацьовують, відкочує їх і пробує щось інше. Але вона має певну інтуїцію щодо того, які речі можливі і як вони можуть обійти компілятор.
Ці інтуїції, які має Оля, часто виходять за межі того, що роблять компілятори. Хоча сучасні оптимізуючі компілятори досить вражають своїми результатами, вони базуються на досить простих правилах і евристиках. Технічною мовою, вони базуються на ідеї локального перетворення потоку даних, яке виконується до фіксованої точки. Ми також впорядковуємо фази розгляду; тобто ми будуємо конвеєри компілятора, щоб розглянути A, а потім B, але не складену проблему AB. Планувальники та розподільники регістрів є відомим прикладом цього, багато докторських дисертацій було присвячено складеному планувальнику-розподільнику регістрів (щоб отримати переваги усунення впорядкування фаз), але їх було важко змусити працювати на практиці.
Ось чому досвід Олі є цінним. Часто Оля знає, як збалансувати кілька NP-повних проблем за допомогою евристик, які є індивідуальними для ситуації. Отже, є більше індивідуальної обізнаності та чутливості до контексту. Оля також здатна застосовувати методи, які оптимізуючі компілятори можуть не застосовувати прибутково, особливо в комбінації, такі речі, як винесення коду, створення спеціальних ABI або перетворень для увімкнення векторизації, або інші речі, які змушують нас бурчати "хотів би я, щоб компілятор мав спосіб просто зробити це…"
Тепер розглянемо, що ШІ, через свої механізми міркування, може діяти як міні-Оля. Він може не мати такої ж інтуїції щодо того, що принесе результат, але він має уявлення про те, що може бути прибутковим, і може робити багато, багато спроб.
За такого підходу, на відміну від роботи STOKE, ми не можемо гарантувати, що з часом, що прямує до нескінченності, ми побачимо оптимальну програму, але оскільки ШІ має "більш людські" механізми міркування, він може досягати значного, подібного до людського, прогресу за одиницю часу.
Пов'язуємо все разом: Ядро MLA
Дозвольте мені почати з того, що я не знаю, який низькорівневий код видав ШІ для нашого ядра MLA Jalapeño, але я знаю, як ввести numpy для MLA.
У компіляторі XLA, над яким я раніше працював, ми об'єднували ці операції numpy в групи, а потім використовували метапрограму, яка називається "емітер", щоб знизити їх до циклів, інструкцій та низькорівневих примітивів.
Коли компілятор / програма-емітер XLA робили це, мені не потрібно було піклуватися про те, який асемблер вийде на виході. Для нашого стохастичного оптимізатора ШІ концептуально займає місце метапрограми-емітера – він одночасно знижує рівень і оптимізує, і ми можемо попросити його оптимізувати далі і далі до межі продуктивності.
Сподіваюся, це прояснює, куди вбудовується ШІ і як він аналогічний компоненту в існуючій системі оптимізуючого компілятора. Також корисно подумати: який шар ми вважаємо "асемблерним кодом", тепер піднімається. Коли ви вводите звичайний C++ і компілюєте його з -O3 (найвищим типовим рівнем оптимізації), ви не очікуєте розуміти асемблер, який виходить, навіть якщо ви розуміли C++, який ви ввели. Ми робимо тут аналогічну річ, але з більш високорівневою та більш математичною вхідною специфікацією.
Тепер ключове питання: як ми перевіряємо, що програма, яку ми отримуємо від ШІ, дійсно еквівалентна високорівневому опису / numpy. Цей механізм перевірки встановлює надійність процесу стохастичної оптимізації ШІ. Я очікую, що майбутній допис у блозі може детальніше розповісти про це, але поки що достатньо сказати, що тестування на семантичну еквівалентність можливе, і ми це робимо. Програми для прискорювачів особливо піддаються сильним, повним контрактам, які ми можемо перевірити "чи є саме тим, що робить оптимізована програма ШІ", оскільки вони досить математичні та орієнтовані на потік даних у широкому контексті.
Зауважте, що багато відповідних методів у цій області були започатковані зусиллями в підгалузі синтезу програм. Тоді як оптимізуючі компілятори кажуть: "ось програма з семантикою, зроби її кращою, але з еквівалентною семантикою!", синтез програм каже: "я вірю, що існує програма з цією семантикою, будь ласка, спробуй знайти найкращу, яку зможеш". Синтез програм є складнішою проблемою, ніж оптимізуюча компіляція, але він також менш принципово обмежений. Це, по суті, те, що роблять люди, подібні до Олі, коли вони роблять краще, ніж оптимізуючий компілятор, і це те, що ШІ тепер може допомогти нам автоматизувати. ШІ може черпати "натхнення" з оригінальної програми, але йому не потрібно просто виконувати незначні локальні перетворення над нею. Класичні оптимізуючі компілятори не побачать "о, ви написали бульбашкове сортування" і, розуміючи контракт, не замінять його на швидке сортування, але і Оля, і ШІ здатні це зробити. Це ставить нас більше в режим синтезу програм зі стохастичною оптимізацією, ніж у класичний режим оптимізуючого компілятора.
Все це сходиться в тому, що ви можете почати з чогось, що "не дуже далеко від numpy", зачекати 48 годин і мати оптимізоване ядро з тією ж семантикою, як ми показали в нашій доповіді на HotChips:
Як також зазначено на слайді, на нашій машині ми часто можемо спостерігати, як ШІ підвищує продуктивність, обганяючи наших експертів-людей навіть на ядрах, які ми вважали добре налаштованими. Часто залишається ще значний досяжний відсоток просто через велику кількість комбінацій / перестановок, які можуть потребувати дослідження. Вони часто є нездійсненно нудними для інженера з продуктивності-людини.
Підсумок та Висновки
Компілятор, зрештою, це просто функція. Ми передаємо нашу програму цій функції, і отримуємо кращу версію нашої програми. Ми очікуємо, що програма, яку ми отримуємо, і програма, яку ми вводимо, мають однакову семантику.
Традиційні оптимізуючі компілятори роблять програми кращими за допомогою правил потоку даних та евристик. Вони повністю зрозумілі у своєму походженні, але також можуть бути більш обмеженими в тому, які зміни вони можуть зробити.
На противагу цьому, ШІ, як стохастичний оптимізатор, просто повинен "напружено думати" і видати щось. Його зміни не настільки принципово обмежені, що робить його більш схожим на нашого експерта-оптимізатора-людину. Нам потрібні способи перевірити, що програми, які він видає, є надійними та реалізують ту саму семантику, яку ми ввели, і ми маємо їх. І цей вид оптимізації ШІ особливо добре підходить для математичних операцій, які мають дуже сильні контракти. Контракти усувають необхідність розуміти, що робить ядро рядок за рядком.
Ось як ми отримали згенероване ШІ ядро MLA!
