Apple открывает исходный код Apple Silicon-Optimized Machine Learning Framework MLX

MLX от Apple сочетает в себе знакомые API, преобразования компонуемых функций и ленивые вычисления для создания среды машинного обучения, вдохновленной NumPy и PyTorch, оптимизированной для Apple Silicon. Платформа, реализованная на Python и C++, призвана предоставить удобное и эффективное решение для обучения и развертывания моделей машинного обучения на Apple Silicon.

По словам Apple, MLX разработан исследователями машинного обучения для исследователей машинного обучения и выпущен под лицензией MIT, чтобы им было легче расширять и улучшать его. Он поддерживает обучение языковой модели трансформера, крупномасштабную генерацию текста с помощью Mistral, генерацию изображений с помощью Stable Diffusion и распознавание речи с помощью Whisper.

MLX предлагает низкоуровневый API Python на основе NumPy, а также полнофункциональный API C++, который точно его отражает. Кроме того, он предоставляет API более высокого уровня, который можно использовать для создания более сложных моделей с помощью PyTorch API.

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

Отличительной особенностью MLX является использование Apple Silicon. единая память, что отличает его от других платформ машинного обучения, говорит Apple. Короче говоря, это означает, что массивы находятся в общей памяти, и операции с массивами могут выполняться как на ЦП, так и на графическом процессоре без передачи данных через память. Например, когда вы создаете массив, вам не нужно указывать его местоположение, поскольку он находится в единой памяти, а вы сами выбираете, запускать ли для него преобразование на CPU или GPU при его выполнении:

a = mx.random.normal((100,))
b = mx.random.normal((100,))

mx.add(a, b, stream=mx.cpu)
mx.add(a, b, stream=mx.gpu)

MLX можно использовать на любых процессорах Apple Silicon, включая M1, и использовать встроенный графический процессор, поэтому исследователи могут выбрать оборудование, которое лучше всего подходит для их нужд.

Репозиторий MLX включает несколько примеров его использования с различными моделями, включая BERT, Llama, Mistral, Stable Diffusion и другие. В каждом примере перечислены зависимости, необходимые для requirements.txt файл и предоставляет готовые к использованию инструменты CLI. Например, чтобы создавать изображения с помощью Stable Diffusion, вы сначала устанавливаете все необходимые зависимости, а затем запускаете txt2image.py команда:

pip install -r requirements.txt
python txt2image.py "A photo of an astronaut riding a horse on Mars." --n_images 4 --n_rows 2

Apple еще не опубликовала публичные тесты, поэтому мы не знаем, как они себя чувствуют по сравнению с PyTorch/MPS или Георгия Герганова Вызов.cpp.

Однако пример Stable Diffusion включает сравнение производительности UNet в Stable Diffusion при запуске через PyTorch и MLX. Это показывает, что MLX достигает пропускной способности примерно на 40 % выше, чем PyTorch при размере пакета 16 и примерно на 15 % выше при сравнении оптимальных размеров пакетов.

PyTorch, однако, работает лучше для пакетов меньшего размера: пропускная способность на ~50% выше для пакетов размером 1 и на ~10% выше для пакетов размером 4. По мнению Apple, преимущество PyTorch в этих случаях можно объяснить скоростью компиляции, когда модели не загружаются в память, а ядра графов MPS PyTorch не кэшируются.

Если вы хотите попробовать MLX, вы можете посетить его Инструкция по началу работы или полная документация.