MLXについて調べる

README に書いてあること

MLX is an array framework for machine learning on Apple silicon, brought to you by Apple machine learning research.
MLXは、アップルの機械学習研究によってもたらされた、アップルのシリコン上での機械学習のための配列フレームワークです。

Some key features of MLX include:
MLXの主な特徴は以下の通り：

• Familiar APIs: MLX has a Python API that closely follows NumPy. MLX also has a fully featured C++ API, which closely mirrors the Python API. MLX has higher-level packages like mlx.nn and mlx.optimizers with APIs that closely follow PyTorch to simplify building more complex models.
  おなじみのAPI： MLXは、NumPyに忠実なPython APIを持っています。MLXはまた、Python APIを忠実に反映した完全な機能を備えたC++ APIも備えている。MLXにはmlx.nnやmlx.optimizersのような高レベルのパッケージがあり、PyTorchに忠実なAPIで、より複雑なモデルの構築を簡素化します。

• Composable function transformations: MLX supports composable function transformations for automatic differentiation, automatic vectorization, and computation graph optimization.
  コンポーザブル関数変換 MLXは、自動微分、自動ベクトル化、計算グラフの最適化などのための合成可能な関数変換をサポートしています。

• Lazy computation: Computations in MLX are lazy. Arrays are only materialized when needed.
  遅延計算 MLXの計算は遅延します．配列は必要なときだけ実体化されます．

• Dynamic graph construction: Computation graphs in MLX are constructed dynamically. Changing the shapes of function arguments does not trigger slow compilations, and debugging is simple and intuitive.
  動的なグラフ構築 MLXの計算グラフは動的に構築されます．関数の引数の形が変わってもコンパイルに時間がかかることはありません．

• Multi-device: Operations can run on any of the supported devices (currently the CPU and the GPU).
  マルチデバイス： サポートされているデバイス（現在はCPUとGPU）のいずれでも演算を実行できます。

• Unified memory: A notable difference from MLX and other frameworks is the unified memory model. Arrays in MLX live in shared memory. Operations on MLX arrays can be performed on any of the supported device types without transferring data.
  ユニファイド・メモリ： MLXと他のフレームワークとの顕著な違いは、ユニファイド・メモリ・モデルである。MLXのアレイは共有メモリ上に存在する。MLXの配列に対する操作は、データを転送することなく、サポートされているどのタイプのデバイスでも実行できる。

MLX is designed by machine learning researchers for machine learning researchers. The framework is intended to be user-friendly, but still efficient to train and deploy models. The design of the framework itself is also conceptually simple. We intend to make it easy for researchers to extend and improve MLX with the goal of quickly exploring new ideas.
MLXは、機械学習研究者による機械学習研究者のためのフレームワークです。このフレームワークは、ユーザーフレンドリーでありながら、モデルの学習とデプロイを効率的に行うことを目的としている。フレームワーク自体の設計も概念的にシンプルである。我々は、研究者がMLXを簡単に拡張・改良できるようにすることで、新しいアイデアを素早く探求することを目標としている。

The design of MLX is inspired by frameworks like NumPy, PyTorch, Jax, and ArrayFire.
MLXのデザインは、NumPy、PyTorch、Jax、ArrayFireのようなフレームワークにインスパイアされている。

Core MLとの関係

Apple Silicon用に最適化されたMLフレームワークということで、Apple Silicon Macで爆速で学習できて、そしてCore MLモデルを直接出力できたりするのかなと思ったのだが、READMEにはそんなことはまったく書かれていない。

リポジトリ内で検索してみると、

https://github.com/ml-explore/mlx/issues/241

CoreML + MLX?

Just wondering can MLX perform the conversions of standard models to coreml, didn't see any examples, and just wondered how the 2 fit together.

→ コメントなしで自分でクローズ

https://github.com/ml-explore/mlx/issues/93

→ コメントなし

https://github.com/ml-explore/mlx/issues/15

Swift bindings? #15

→ 結構議論が盛り上がっている。後で読む。

というわけで少なくとも現状ではCore MLモデルのエクスポートはサポートされてなさそうだ。

"coremltools"で検索してみる

Core MLモデルを出力するのであれば、必ずcoremltoolsを利用しているはず。

というわけでリポジトリ内検索してみたところ、次のissueだけが出てきた：

ANE support · Issue #18 · ml-explore/mlx

The top level readme mentions that current device support is limited to COU and GPU, is ANE support in the works?
（トップレベルのReadmeによると、現在のデバイスのサポートはCOUとGPUに限られているとのことだが、ANEのサポートは準備中なのか？）

リアクションが多くついている回答を以下に引用する

回答 1

The ANE is mostly designed as an inference only device that supports only forward prop
The ANE has layer support implemented in hardware and can't be easily extended (?)
The ANE only supports half Float (Float 16) accelerated compute, everything is managed / converted to it by the runtime.
The ANE requires talking to the OS runtime for scheduling, and there's never a guarantee that you can be resident on the ANE, you can just request it, and hope you get a time slice
The API is only exposed via CoreML / Swift
Internally the API is Espresso (? IIRC from stack traces) / a C++ library which isn't public
CoreMLTools, the public pythonic way to create CoreML models has a CoreML runtime exposed, but requires CoreML model specs (protobufs) to run on the ANE
I doubt that Apple would let an open source project leak the internal tooling of the ANE Runtime (Espresso?)
ANEは推論専用デバイスとして設計されており、フォワード・プロップのみをサポートしている。
ANEはレイヤーサポートをハードウェアで実装しており、簡単に拡張できない（？）
ANEはハーフフロート（フロート16）アクセラレーション計算しかサポートしておらず、すべてはランタイムによって管理/変換されます。
ANEはスケジューリングのためにOSのランタイムと通信する必要があり、ANEに常駐できる保証はありません。
APIはCoreML / Swiftを介してのみ公開される。
内部的にはAPIはEspresso(? IRC from stack traces) / 公開されていないC++ライブラリです。
CoreMLToolsは、CoreMLモデルを作成するための公開されたpythonicな方法で、CoreMLランタイムを公開していますが、ANE上で実行するにはCoreMLモデルの仕様（protobufs）が必要です。
AppleがオープンソースプロジェクトにANEランタイム(Espresso?)の内部ツールをリークさせるとは思えない。

There's some ANE reverse engineering work that's sporadically happening, but I suspect this will be Metal / GPU for a while unless Apple exposes some cool new ways to publicly run arbitrary programs on the ANE (which would be dope)
散発的にANEのリバースエンジニアリングが行われているが、AppleがANE上で任意のプログラムを実行するクールな新しい方法を公開しない限り（それはドープなことだ）、これはしばらくMetal / GPUになるだろう。

回答 2

One thought which would be cool however, to get both MLX and ANE inference would be:
Implement MLX via swift runtime.
Expose custom layers to your CoreML model export
Implement those layers via MLX in a Swift native app
Use CoreML runtime to load the model, request ANE inference
You could also implement MLX preprocessing to get IOSurface backed memory buffers in half float easily which would grant your app the same unified memory access and avoid a ton of overhead of moving data to the ANE which is default path without IOSurface backed buffers.
しかし、MLXとANEの両方の推論を得るために、1つの考えがある：
MLXをswiftランタイムで実装する。
カスタムレイヤーをCoreMLモデルエクスポートに公開する。
SwiftネイティブアプリにMLX経由でレイヤーを実装する。
CoreMLランタイムを使用してモデルをロードし、ANE推論を要求する。
また、MLXの前処理を実装することで、IOSurfaceのメモリバッファを簡単に半分のフロートで取得することができ、同じように統一されたメモリアクセスをアプリに与えることができます。

In theory you'd get:
Fast MLX pre processing with unified memory
CoreML ANE acceleration on native layers
MLX layer ops with unified memory outputs
理論的には
ユニファイドメモリによる高速MLX前処理
ネイティブレイヤーでのCoreML ANEアクセラレーション
ユニファイドメモリ出力によるMLXレイヤーオペレーション

回答 3

@vade basically said it all already, but at the moment we don't have plans to support ANE in MLX given it is a closed source API.
基本的には@vadeがすでにすべてを語っている、
しかし、クローズドソースのAPIであるため、現時点ではMLXでANEをサポートする予定はありません。

If / when that changes we will be first in line to add it as a supported device.
もしそれが変更されれば、私たちは真っ先に対応デバイスとして追加するつもりです。

クイックスタートガイド

Quick Start Guide — MLX 0.0.7 documentation

Stable Diffusionサンプル

mlx-examples というサンプルリポジトリ内に、Stable Diffusionサンプルがある。

https://github.com/ml-explore/mlx-examples/tree/main/stable_diffusion

実行方法

1. git cloneする
2. cd mlx-examples/stable_diffusion
3. pip install -r requirements.txt

あとはサンプルコードを実行するだけ。

Text 2 Image

% python txt2image.py "A photo of an astronaut riding a horse on Mars." --n_images 1 --n_rows 1  

Image 2 Image

original.pngを用意した上で、

python image2image.py --strength 0.5 original.png 'A lit fireplace'

のように実行する。

モデルのダウンロード

サンプルを実行すると、初回のみモデルのダウンロードが行われる。

これぐらいのサイズのこういうファイルがダウンロードされている、というメモとしてコンソール出力を転載しておく：

Downloading (…)cheduler_config.json: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 346/346 [00:00<00:00, 1.41MB/s]
Downloading unet/config.json: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 911/911 [00:00<00:00, 3.31MB/s]
Downloading (…)ch_model.safetensors: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3.46G/3.46G [06:08<00:00, 9.41MB/s]
Downloading (…)_encoder/config.json: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 613/613 [00:00<00:00, 3.44MB/s]
Downloading model.safetensors: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1.36G/1.36G [02:24<00:00, 9.42MB/s]
Downloading vae/config.json: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 553/553 [00:00<00:00, 424kB/s]
Downloading (…)ch_model.safetensors: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 335M/335M [00:35<00:00, 9.39MB/s]
Downloading tokenizer/vocab.json: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1.06M/1.06M [00:00<00:00, 1.99MB/s]
Downloading tokenizer/merges.txt: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 525k/525k [00:00<00:00, 1.54MB/s]