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

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

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

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

*ссылка на видео*

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

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

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

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

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

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

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

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

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

Существуют также неизбежные проблемы с производительностью, даже если для библиотеки используется более быстрый скомпилированный язык реализации. Одной из основных проблем является отсутствие “слияния” – то есть вызов множества скомпилированных функций подряд приводит к большим накладным расходам, поскольку данные преобразуются в форматы python и из них, а стоимость многократного переключения с Python на C и обратно должна быть уплачена.

Поэтому вместо этого мы должны написать специальные “объединенные” версии распространенных комбинаций функций (например, линейный слой, за которым следует исправленный линейный слой в нейронной сети) и вызвать эти объединенные версии из Python. Это означает, что нужно реализовать и запомнить гораздо больше библиотечных функций, и вам не повезет, если вы будете делать что-то даже слегка нестандартное, потому что для вас не будет объединенной версии.

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

Есть несколько неуклюжих способов написания параллельного кода, который использует более одного ядра, но они либо должны работать с совершенно отдельной памятью (и имеют много накладных расходов для запуска), либо им приходится обращаться к памяти по очереди (страшная “глобальная блокировка компилятора”, которая часто затрудняет параллельное выполнение). код на самом деле медленнее, чем однопоточный код!)

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

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

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

...а вот и внимание к себе:

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 Фабриса Белларда и ggml Георгия Герганова оба используют C и в результате могут в полной мере использовать преимущества производительности полностью скомпилированных языков

Добро пожаловать в Mojo.

Крис Латтнер отвечает за создание многих проектов, на которые мы все полагаемся сегодня – даже несмотря на то, что мы, возможно, даже не слышали обо всем, что он создал! В рамках своей докторской диссертации он приступил к разработке 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 может создать оптимизированный машинный код для реализации вашей функции. Кроме того, если вы используете структуру вместо класса, ваши атрибуты будут плотно упакованы в память, так что их можно будет использовать даже в структурах данных, не гоняясь за указателями. Это те функции, которые позволяют таким языкам, как C, быть такими быстрыми, и теперь они доступны и программистам на Python – просто изучив немного нового синтаксиса.

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

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

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

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

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

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

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

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

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

Следующим шагом было создание минимального питонического способа прямого вызова MLIR. Это была совсем несложная работа, но это было все, что требовалось, чтобы затем создать 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, написанную с нуля, составляет около 100 тыс.

Это означает, что Mojo - это гораздо больше, чем просто язык для приложений AI /ML. На самом деле это версия 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, позволяющий мне использовать один язык для написания всего, начиная с бэкэнда программы, заканчивая архитектурой модели и установщиком, а также позволяющий мне отлаживать и профилировать мой код непосредственно на языке, на котором я работаю. о чем я и написал.

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