Mojo — это новый язык программирования, основанный на Python, который устраняет имеющиеся у него проблемы производительности и развёртывания.

Об авторе: Джереми Говард (Jeremy Howard) — Data Scientist, исследователь, разработчик, преподаватель и предприниматель. Джереми является одним из основателей исследовательского института fast.ai, занимающегося тем, чтобы сделать глубокое обучение более доступным, а также он является почётным профессором Университета Квинсленда. Ранее Джереми был выдающимся научным сотрудником в Университете Сан‑Франциско, где он был основателем Инициативы Уиклоу «Искусственный интеллект в медицинских исследованиях».

Джереми был генеральным директором‑основателем Enlitic, которая была первой компанией, применившей глубокое обучение в медицине, и два года подряд была выбрана в качестве одной из 50 самых умных компаний мира по версии MIT Tech Review. Он был президентом и главным научным сотрудником платформы по обработке и анализу данных Kaggle, где 2 года подряд занимал первое место в международных соревнованиях по машинному обучению. Он был генеральным директором‑основателем двух успешных австралийских стартапов (FastMail и Optimal Decisions Group, позднее приобретённых Lexis‑Nexis). До этого он 8 лет проработал в управленческом консалтинге в McKinsey & Co и AT Kearney. Джереми инвестировал, наставлял и консультировал многие стартапы, а также участвовал во многих проектах с открытым исходным кодом.

Я помню, как впервые использовал Visual Basic v. 1.0. Тогда это была программа для DOS. До него написание программ было чрезвычайно сложным делом, и мне никогда не удавалось добиться значительного прогресса, выходящего за рамки самых простых приложений. Но с помощью VB я нарисовал кнопку на экране, набрал всего одну строку кода, которую хотел запустить при нажатии на неё, и теперь у меня было готовое приложение, которое я теперь мог запустить. Это был такой удивительный опыт, что я никогда не забуду это чувство.

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

Написание кода на Mojo, новом языке программирования от Modular, — это второй раз в моей жизни, когда у меня возникло точно такое же чувство. Вот как это выглядит:

Почему бы просто не использовать Python?

Прежде чем я объясню, почему я так взволнован Mojo, я должен сказать несколько слов о Python.

Python — это язык, который я использовал почти во всех своих работах за последние несколько лет. Это прекрасный язык. У него есть элегантное ядро, на котором строится все остальное. Такой подход означает, что Python может (и делает) абсолютно все. Но у него есть и обратная сторона — производительность.

Несколько процентов тут или там не имеет значения. Но Python во много тысяч раз медленнее, чем C++ подобные языки. Это делает непрактичным использование Python для чувствительных к производительности частей кода — внутренних циклов, где производительность имеет решающее значение.

Однако у Python есть один козырь: он может обращаться к коду, написанному на быстрых языках. Таким образом, программисты Python учатся избегать его использования для реализации критичных к производительности участков, вместо этого используя Python-оболочки для кода C, FORTRAN, Rust и т. д. Такие библиотеки, как Numpy и PyTorch, предоставляют интерфейсы Python для высокопроизводительного кода, позволяя программистам Python чувствовать себя как в родной среде, даже если они используют высокооптимизированные числовые библиотеки.

Сегодня почти все модели ИИ разрабатываются на Python, благодаря гибкому и элегантному языку программирования, фантастическим инструментам и экосистеме, а также высокопроизводительным скомпилированным библиотекам.

Но у такого «двуязычного» подхода есть серьёзные недостатки. Например, модели ИИ часто приходится преобразовывать из Python в более быструю реализацию, такую как ONNX или TorchScript. Но эти подходы к развёртыванию не могут поддерживать все функции Python, поэтому программистам Python приходится учиться использовать подмножество языка, соответствующее их цели развёртывания. Очень сложно профилировать или отлаживать версию кода для развёртывания, и нет никакой гарантии, что она будет работать идентично версии Python.

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

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

Существуют также неизбежные проблемы с производительностью, даже если для библиотеки используется более быстрый язык (с компилируемой реализацией). Одной из основных проблем является отсутствие «слияния», то есть вызов множества скомпилированных функций подряд приводит к большим накладным расходам. Поскольку данные преобразуются из формата Python в другой и обратно, то за многократное переключение с Python на C приходиться платить. Поэтому вместо этого приходится писать специальные «сведённые» версии обычных комбинаций функций (таких как линейный слой, за которым следует выпрямленный линейный слой в нейронной сети) и вызывать эти «сведённые» версии из Python. Это означает, что вам придётся реализовывать и запоминать гораздо больше библиотечных функций, и вам не повезло, если вы делаете что-то хотя бы немного нестандартное, потому что для вас не будет сведённой версии.

Нам также приходится иметь дело с отсутствием эффективной параллельной обработки в Python. В настоящее время у всех нас есть компьютеры с большим количеством ядер, но Python обычно использует только одно за раз. Есть несколько неуклюжих способов написания параллельного кода, который использует более одного ядра, но они либо должны работать с совершенно отдельной памятью (и требуют больших накладных расходов для запуска), либо им приходится обращаться к общей памяти по очереди (т.н. страшная «глобальная блокировка интерпретатора», которая часто делает многопоточный код фактически медленнее, чем однопоточный код!)

Такие библиотеки, как PyTorch, разрабатывают все более изощрённые способы решения этих проблем с производительностью, а недавно выпущенный PyTorch 2 даже включает функцию compile(), которая использует сложную бэкенд компиляцию для создания высокопроизводительной реализации кода Python. Однако такая функциональность не может творить чудеса: существуют фундаментальные ограничения на то, что возможно с Python, в зависимости от того, как устроен сам язык.

Вы можете себе представить, что на практике существует лишь небольшое количество строительных блоков для моделей ИИ, и поэтому на самом деле не имеет значения, нужно ли нам реализовывать каждый из них на C. Кроме того, в целом это довольно простые алгоритмы, верно? Например, модели transformers почти полностью реализованы несколькими слоями из двух компонентов, многослойных персептронов (MLP) и слоя внимания (self-attention layer), которые могут быть исполнены всего несколькими строками Python с PyTorch. Вот реализация MLP:

nn.Sequential(nn.Linear(ni,nh), nn.GELU(), nn.LayerNorm(nh), nn.Linear(nh,ni))

… а вот слой внимания (self-attention layer):

def forward(self, x):
    x = self.qkv(self.norm(x))
    x = rearrange(x, 'n s (h d) -> (n h) s d', h=self.nheads)
    q,k,v = torch.chunk(x, 3, dim=-1)
    s = (q@k.transpose(1,2))/self.scale
    x = s.softmax(dim=-1)@v
    x = rearrange(x, '(n h) s d -> n s (h d)', h=self.nheads)
    return self.proj(x)

Но это скрывает тот факт, что реальные реализации этих операций гораздо сложнее. Например, посмотрите на эту оптимизированную для памяти реализацию «flash attention» в CUDA C. Она также скрывает тот факт, что эти общие подходы к построению моделей оставляют «за бортом» огромную часть производительности. Например, подходы с «разреженным блоком» могут значительно улучшить скорость и использование памяти. Исследователи работают над настройками почти каждой части распространённых архитектур и придумывают новые архитектуры (и оптимизаторы SGD, методы расширения данных и т.д.) — мы даже близко не подошли к созданию какой-либо аккуратно обёрнутой системы, которую все будут использовать вечно.

На практике большая часть самого быстрого кода, используемого сегодня для языковых моделей, пишется на C и C++. Например, TextSynth Фабриса Белларда (Fabrice Bellard) и ggml Георгия Герганова (Georgi Gerganov) используют C, и в результате они могут в полной мере воспользоваться преимуществами производительности полностью компилируемых языков.

Введение в Mojo

Крис Латтнер (Chris Lattner) участвовал в создании многих продуктов, на которые мы все сегодня полагаемся, хотя мы, возможно, даже не слышали обо всём, что он создал! В рамках своей PhD диссертации он начал разработку LLVM, которая коренным образом изменила способ создания компиляторов и сегодня составляет основу многих наиболее широко используемых языковых экосистем в мире. Затем он запустил Clang, компилятор C и C++, который находится поверх LLVM и используется большинством наиболее значительных разработчиков программного обеспечения в мире (включая обеспечение основы для критически важного в части производительности кода Google). LLVM включает в себя «промежуточное представление» (IR), специальный язык, предназначенный для чтения и записи машинами (а не для людей), что позволило огромному сообществу разработчиков программного обеспечения совместно работать над улучшением функциональности языка программирования на более широком спектре аппаратных средств.

Однако Крис видел, что C и C++ не в полной мере используют возможности LLVM, поэтому, работая в Apple, он разработал новый язык под названием Swift, который он описывает как «синтаксический сахар для LLVM». Swift стал одним из наиболее широко используемых языков программирования в мире, в частности потому, что сегодня это основной способ создания приложений iOS для iPhone, iPad, macOS и Apple TV.

К сожалению, контроль Apple над Swift означает, что у него не было времени проявить себя за пределами замкнутого мира Apple. Некоторое время Крис возглавлял команду в Google, чтобы попытаться вывести Swift из зоны комфорта Apple, чтобы он стал заменой Python в разработке моделей ИИ. Я был очень взволнован этим проектом, но, к сожалению, он не получил необходимой поддержки ни от Apple, ни от Google, и в итоге не увенчался успехом.

Тем не менее, Крис, работая в Google, разработал ещё один проект, который стал чрезвычайно успешным: MLIR. MLIR — это замена IR от LLVM для современной эпохи многоядерных вычислений и рабочих нагрузок искусственного интеллекта. Это крайне важно для полного использования возможностей аппаратного обеспечения, такого как графические процессоры, TPU и векторные модули, которые все чаще добавляются в процессоры серверного класса.

Итак, если Swift являлся «синтаксическим сахаром для LLVM», то что же является «синтаксическим сахаром для MLIR»? Ответ — Mojo! Mojo – это совершенно новый язык, разработанный для того, чтобы в полной мере использовать все преимущества MLIR. А ещё Mojo — это Python.

Постойте! Что?

Хорошо, позвольте мне объяснить. Может быть, лучше сказать, что Mojo — это Python++. Когда он будет завершён, он станет строгим надмножеством языка Python. Но он также обладает дополнительной функциональностью, так что мы можем писать высокопроизводительный код, использующий преимущества современных ускорителей.

Mojo кажется мне более прагматичным подходом, чем Swift. В то время как Swift был совершенно новым языком, содержащим всевозможные интересные функции, основанные на последних исследованиях в области дизайна языков программирования, Mojo по своей сути является просто Python. Это кажется разумным не только потому, что Python уже хорошо понятен миллионам программистов, но и потому, что после десятилетий использования его возможности и ограничения теперь хорошо изучены. Полагаться на новейшие исследования в области языков программирования довольно круто, но это также и потенциально опасная спекуляция, потому что вы никогда по-настоящему не знаете, как все обернётся. (Признаюсь, лично я, например, часто путался в мощной, но причудливой системе типов Swift, а иногда даже умудрялся запутать компилятор Swift и взрывать его полностью!)

Ключевой трюк в Mojo заключается в том, что вы как разработчик можете в любое время перейти в более быстрый «режим», используя «fn» вместо «def» для создания своей функции. В этом режиме вы должны точно объявить тип каждой переменной, и в результате Mojo сможет создать оптимизированный машинный код для реализации вашей функции. Более того, если вы используете «struct» вместо «class», ваши атрибуты будут плотно упакованы в память, так что их можно будет использовать даже в структурах данных, не гоняясь за указателями. Это те функции, которые позволяют таким языкам, как C, быть такими быстрыми, и теперь они доступны и для программистов Python — просто изучив немного нового синтаксиса.

Как это возможно?

На данный момент за десятилетия были предприняты сотни попыток создать языки программирования, которые были бы лаконичными, гибкими, быстрыми, практичными и простыми в использовании, но без особого успеха. Но каким-то образом Modular, кажется, сделал это. Как это могло произойти? Мы можем предложить несколько гипотез:

1. Mojo на самом деле не достигла этих результатов, и за шикарной демонстрацией скрывается разочаровывающая реальная производительность, или

2. Modular — это огромная компания, в которой сотни разработчиков работают годами, тратя много времени на то, чтобы добиться чего-то, чего раньше не удавалось.

Ни то, ни другое не является правдой. Демонстрация, на самом деле, была создана всего за несколько дней до того, как я записал видео в начале статьи. Два примера, которые мы привели выше [matmul (матричное умножение) и Мандельброт], не были тщательно отобраны после десятков сравнений; а скорее, это были единственные вещи, которые мы пробовали для демонстрации, и они сработали с первого раза! Несмотря на то, что на этой ранней стадии отсутствует множество функций (Mojo ещё даже не выпущен для широкой публики, кроме онлайн-песочницы), демоверсия, которую вы видите, действительно работает именно так, как вы её видите. И действительно, вы можете запустить её в песочнице теперь самостоятельно.

Modular — довольно небольшой стартап, которому всего год, и только часть компании работает над языком Mojo. Разработка Mojo началась совсем недавно. Это небольшая команда, работающая в течение короткого времени, так как же они так много добились?

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

В его основе лежит MLIR, который уже разрабатывался в течение многих лет, первоначально начатый Крисом Латтнером в Google. Он осознал, что потребуется для основ ядра «языка программирования эпохи ИИ», и сосредоточился на их создании. MLIR был ключевым элементом. Подобно тому, как за последнее десятилетие LLVM значительно облегчил разработку новых мощных языков программирования (таких как Rust, Julia и Swift, которые в свою очередь основаны на LLVM), MLIR предоставляет ещё более мощное ядро для языков, построенных на нём.

Ещё одним ключевым фактором, способствующим быстрому развитию Mojo, является решение использовать синтаксис Python. Разработка и итерация синтаксиса — одна из самых подверженных ошибкам, сложных и противоречивых частей разработки языка. Просто передав это на аутсорсинг существующему языку (который также является наиболее широко используемым языком сегодня), вся эта часть исчезает! Относительно небольшое количество новых элементов синтаксиса, необходимых поверх Python, в значительной степени подходит вполне естественно, поскольку база уже имеется.

Следующим шагом было создание минимального Pythonic-способа прямого вызова MLIR. Это была совсем несложная работа, но этого было достаточно, чтобы затем создать весь Mojo поверх него самого — и работать непосредственно в Mojo для всего остального. Это означало, что разработчики Mojo могли прибегнуть к «догфудингу» (прим., практике использования сотрудниками компании собственных продуктов и сервисов) при написании Mojo почти с самого начала. Каждый раз, когда они обнаруживали, что что-то работает не совсем хорошо при разработке Mojo, они могли добавить необходимую функцию в сам Mojo, чтобы облегчить им разработку следующей части Mojo!

Это очень похоже на язык Julia, разработанный на минимальном LISP-подобном ядре, которое предоставляет элементы языка Julia, а затем эти элементы привязываются к базовым операциям LLVM. Почти всё в Julia построено поверх этого, используя Julia.

Я не могу описать все маленькие (и большие!) идеи, воплощённые в дизайне и реализации Mojo — это плоды десятилетней работы Криса и его команды над компилятором и проектированием (языков программирования), включающие в себя все приёмы и с трудом приобретённый опыт — но то, что я могу описать, так это удивительный результат, который я увидел своими глазами.

Команда Modular внутри объявила, что они решили запустить Mojo с видео, включая демонстрацию, и назначили дату всего через несколько недель. Но в то время Mojo был самым примитивным языком. Не было пригодного для использования ядра ноутбука, синтаксис Python практически не применялся, и ничего не было оптимизировано. Я не мог понять, как они надеялись реализовать всё это за считанные недели, не говоря уже о том, чтобы это вообще хоть как-то сделать! То, что я увидел за это время, поразило меня. Каждый день или два внедрялись новые языковые функции, и как только их было достаточно, чтобы попробовать запустить алгоритмы, как правило, они сразу же достигали самого высокого уровня производительности или приближались к нему! Я понял, что все основы уже заложены и что они были специально разработаны для создания того, что сейчас находится в стадии разработки. Так что неудивительно, что всё заработало, и заработало хорошо — в конце концов, таков был план с самого начала!

Это повод с оптимизмом смотреть в будущее Mojo. Хотя это только начало проекта, моё предположение, основанное на том, что я наблюдал за последние несколько недель, заключается в том, что он будет развиваться быстрее и дальше, чем большинство из нас ожидает ...

Развёртывание

Я оставил один из моментов, который меня больше всего волнует, на потом: развёртывание. В настоящее время, если вы хотите дать другу свою крутую программу на Python, вам придётся сказать ему, чтобы он сначала установил Python! Или вы могли бы предоставить ему огромный файл, включающий в себя весь Python и библиотеки, которые вы используете, упакованные вместе, которые будут извлечены и загружены при запуске вашей программы.

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

Сравните это с развёртыванием статически скомпилированного приложения на C: вы буквально можете просто скомпилировать программу, которая сразу же становится доступной для прямой загрузки. Её размер может составлять всего 100 КБ или около того, и она будет просто запускаться и работать быстро.

Существует также подход, используемый Go, который не может генерировать такие небольшие приложения, как это делает C, а вместо этого включает «среду выполнения» в каждое упакованное приложение. Этот подход представляет собой компромисс между Python и C и обеспечивает более простое развёртывание, чем Python, но по-прежнему увеличивает размер двоичного файла на десятки мегабайтов.

Как компилируемый язык, история развёртывания Mojo в основном такая же, как и у C. Например, программа, включающая версию matmul (матричное умножение), написанную на Mojo с нуля, составляет около 100 КБ.

Это означает, что Mojo — это гораздо больше, чем просто язык для приложений в области искусственного интеллекта / машинного обучения. На самом деле это версия Python, которая позволяет нам писать быстрые, небольшие, легко развёртываемые приложения, использующие все доступные ядра процессора и ускорители!

Альтернативы Mojo

Mojo — не единственная попытка решить проблему производительности и развёртывания Python. С точки зрения языков, Julia, пожалуй, самая сильная из существующих альтернатив. Он обладает многими преимуществами Mojo, и с его помощью уже построено множество отличных проектов. Ребята из Julia были достаточно любезны, чтобы пригласить меня выступить с основным докладом на их недавней конференции, и я воспользовался этой возможностью, чтобы описать то, что, по моему мнению, было текущими недостатками (и возможностями) Julia:

Как обсуждалось в этом видео, самая большая проблема Julia связана с большим временем выполнения, которое, в свою очередь, связано с решением использовать сборку мусора в языке. Кроме того, подход с несколькими отправками, используемый в Julia, является довольно необычным выбором, который открывает много дверей для создания крутых вещей в языке, но также может усложнить жизнь разработчикам. (Я настолько в восторге от этого подхода, что создал его версию на Python, и в результате я также хорошо осведомлён о его ограничениях!)

В Python наиболее известным текущим решением, вероятно, является Jax, который эффективно создаёт доменно-специфический язык (DSL), используя Python. Результатом этого языка является XLA – компилятор машинного обучения, который предшествует MLIR (и который, я полагаю, постепенно переносится на MLIR). Jax наследует ограничения как Python (например, язык не имеет способа представления структур, или выделения памяти напрямую, или создания быстрых циклов), так и XLA (который в значительной степени ограничен специфическими концепциями машинного обучения и в первую очередь ориентирован на TPU), но имеет огромный плюс в том, что для него не требуется новый язык или новый компилятор.

Как обсуждалось ранее, есть также новый компилятор PyTorch, а также Tensorflow, способный генерировать XLA-код. Лично я нахожу использование Python таким образом в конечном счёте неудовлетворительным. На самом деле я не могу использовать всю мощь Python, но должен использовать подмножество, совместимое с серверной частью, на которую я ориентируюсь. Я не могу легко отлаживать и профилировать скомпилированный код, и происходит так много «волшебства», что трудно даже понять, что на самом деле в итоге выполняется. В итоге я даже не получаю автономный двоичный файл, а вместо этого вынужден использовать специальные среды выполнения и иметь дело со сложными API. (Я здесь не одинок — все, кого я знаю, кто использовал PyTorch или Tensorflow для нацеливания на периферийные устройства или оптимизации обслуживающей инфраструктуры, описали это как одну из самых сложных и разочаровывающих задач, которые они когда-либо решали! И я даже не уверен, что знаю кого-либо, кто на самом деле выполнил любую из этих задач, используя Jax.)

Ещё одно интересное направление для Python — Numba и Cython. Я большой поклонник этих проектов и использовал их как в преподавании, так и при разработке программного обеспечения. Numba использует специальный декоратор для компиляции функции Python в оптимизированный машинный код с использованием LLVM. Cython в этом схож, но также предоставляет Python-подобный язык, который обладает некоторыми особенностями Mojo, и преобразует этот диалект Python в C, который затем компилируется. Ни один из них не решает проблему развёртывания, но они могут во многом помочь с проблемой производительности.

Ни один из них не может быть нацелен на ряд ускорителей с универсальным кроссплатформенным кодом, хотя Numba предоставляет очень полезный способ написания CUDA-кода (и, таким образом, позволяет нацеливаться на графические процессоры NVIDIA).

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

Но я бы предпочёл использовать язык, такой же элегантный, как Python, и такой же быстрый, как написанный экспертами C, позволяющий мне использовать один язык для написания всего, начиная с сервера приложений, заканчивая архитектурой модели и установщиком, и позволяющий мне отлаживать и профилировать мой код непосредственно на том языке, на котором я его написал.

А вам бы понравился такой язык?

Джереми Говард (Jeremy Howard)

Основатель исследовательского института fast.ai