Вокруг презентации ядра Jalapeño MLA на 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
Позвольте мне начать с того, что я не знаю, какой низкоуровневый код ИИ выдал для нашего ядра Jalapeño MLA, но я знаю, как ввести numpy для MLA.
В компиляторе XLA, над которым я ранее работал, мы объединяли эти операции numpy в кластеры, а затем использовали метапрограмму, называемую "эмиттером", чтобы понизить их до циклов, инструкций и низкоуровневых примитивов.
Когда компилятор/программа-эмиттер XLA делали это, мне было все равно, какой ассемблер получится на выходе. Для нашего стохастического оптимизатора ИИ концептуально занимает место метапрограммы-эмиттера – он одновременно понижает и оптимизирует, и мы можем попросить его оптимизировать все ближе и ближе к теоретическому пределу производительности.
Надеюсь, это проясняет, куда вписывается ИИ и как он аналогичен компоненту существующей системы оптимизирующего компилятора. Также полезно подумать: уровень, который мы считаем "ассемблерным кодом", теперь поднимается. Когда вы вводите обычный C++ и компилируете его с -O3 (самый высокий типичный уровень оптимизации), вы не ожидаете понимать полученный ассемблер, даже если вы понимали введенный C++. Мы делаем здесь аналогичную вещь, но с более высокоуровневой и более математической входной спецификацией.
Теперь ключевой вопрос заключается в том, как мы проверяем, что программа, полученная от ИИ, действительно эквивалентна высокоуровневому описанию / numpy. Этот механизм проверки устанавливает надежность процесса стохастической оптимизации ИИ. Я ожидаю, что будущий пост в блоге может углубиться в эту тему, но пока достаточно сказать, что тестирование семантической эквивалентности возможно, и мы это делаем. Программы для ускорителей особенно поддаются строгим, полным контрактам, которые мы можем проверить "точно ли они соответствуют тому, что делает оптимизированная программа ИИ", поскольку они довольно математичны и ориентированы на потоки данных в своем широком контексте.
Обратите внимание, что многие соответствующие методы в этой области были разработаны усилиями в подобласти синтеза программ. В то время как оптимизирующие компиляторы говорят: "вот программа с семантикой, сделай ее лучше, но с эквивалентной семантикой!", синтез программ говорит: "я считаю, что существует программа с такой семантикой, пожалуйста, попробуй найти лучшую из возможных". Синтез программ – более сложная задача, чем оптимизирующая компиляция, но он также менее фундаментально ограничен. Это, по сути, то, что делают люди, подобные Олли, когда они превосходят оптимизирующий компилятор, и это то, в чем ИИ теперь может помочь нам автоматизировать. ИИ может черпать "вдохновение" из исходной программы, но ему не нужно выполнять только незначительные локальные преобразования. Классические оптимизирующие компиляторы не увидят "о, вы написали пузырьковую сортировку" и, поняв контракт, не заменят ее на быструю сортировку, но и Олли, и ИИ способны на это. Это помещает нас в режим синтеза программ со стохастической оптимизацией, а не в классический режим оптимизирующего компилятора.
Все это сводится к тому, что вы можете начать с чего-то, что "не очень далеко от numpy", подождать 48 часов и получить оптимизированное ядро с той же семантикой, как мы показали в нашем выступлении на HotChips:
Как также отмечается на слайде, на нашей машине мы часто можем наблюдать, как производительность ИИ превосходит наших экспертов-людей даже на ядрах, которые мы считали довольно хорошо настроенными. Часто остается значительный достижимый процент просто из-за множества комбинаций/перестановок, которые, возможно, необходимо исследовать. Это часто непосильно утомительно для инженера по производительности-человека.
Резюме и Заключение
Компилятор, в конечном счете, это просто функция. Мы передаем нашу программу этой функции и получаем обратно улучшенную версию нашей программы. Мы ожидаем, что полученная программа и введенная программа будут иметь одинаковую семантику.
Традиционные оптимизирующие компиляторы улучшают программы с помощью правил потока данных и эвристик. Они полностью понятны в своем происхождении, но также могут быть более ограничены в том, какие преобразования они могут выполнять.
Напротив, ИИ, как стохастический оптимизатор, просто должен "хорошо подумать" и выдать что-то. Его преобразования не так фундаментально ограничены, что делает его более похожим на нашего эксперта-оптимизатора-человека. Нам нужны способы проверки того, что программы, которые он выдает, надежны и реализуют ту же семантику, которую мы ввели, и у нас они есть. И такой вид оптимизации ИИ особенно хорошо подходит для математических операций, которые имеют очень строгие контракты. Контракты устраняют необходимость понимать, что делает ядро строка за строкой.
Вот как мы получили ядро MLA, сгенерированное ИИ!
