AppleシリコンでのMLXの導入

AppleシリコンでのMLXの導入

MLXは、Appleシリコン上で数値計算や機械学習を行うための柔軟で効率的な配列フレームワークです。ユニファイドメモリ、遅延計算、関数変換など、MLXの基本的な機能を解説します。また、SwiftとPythonのAPIを使用して、各種のAppleプラットフォーム上で機械学習モデルを構築、高速化するためのより高度な手法についても紹介します。

関連する章
- 0:00 - イントロダクション
- 1:15 - MLXの概要
- 4:21 - 主な機能
- 10:15 - MLXの高速化
- 17:30 - MLX Swift
リソース
関連ビデオ

WWDC25
- MLX：Appleシリコンでの大規模言語モデルの実行
こんにちは Awniです本セッションでは MLXの概要を説明します MLXはAppleシリコン向けに開発されたオープンソースの配列フレームワークです幅広い用途に対応できる柔軟なツールで基本的な数値計算から Appleデバイス上で最大規模の最先端機械学習モデルまで実行できます大規模な言語モデルでテキストを生成する場合や最新モデルで画像オーディオまたはビデオを生成する場合は MLXが最適です機械学習モデルを Mac上で直接トレーニング、微調整カスタマイズするのにも使えますまずMLXとは何かどこで役立つのか少し詳しく説明しますまたPythonでMLXを使用するための基本も説明しますインストールと基本的な配列操作などですその後他のフレームワークと異なる顕著な主要機能をいくつか説明します次にAppleシリコン上で機械学習のワークロードを可能な限り高速化するMLXのツールをいくつか説明します
最後のセクションで MLX Swiftをどのように使い始めるか簡単に説明しますではMLXの概要を詳しく見ていきます MLXはAppleシリコンでのスピードを目的に一から設計され CPU上での実行や GPU上での高速化が可能です MLXはさまざまなアプリに使用でき小規模な数値計算から大規模な機械学習まで対応します使いやすく柔軟な設計でスピードと効率を損なうことはありません
MLXにはNumPyに準拠した主要なAPIが備わっています多くの場合ドロップイン代替品として使用してほとんどの数値計算を高速化できます MLXには機械学習に必要なツールもすべて揃っています自動微分や高レベルのライブラリなどですこれらの高レベルのAPIは PyTorchやJAXに似ていますこうしたフレームワークを学んだ方には MLXは馴染みやすく始めるのも簡単でしょう MLXを使用するとデバイス上で直接より高度な機械学習を行うことができますたとえばMLXは人気アプリLM Studioで使用されており自分のMacで直接大規模言語モデルを使ってテキストを生成できます
MLX上に構築された MLX LMパッケージを使いテキストを生成し最大数千億のパラメータサイズの言語モデルを微調整できますこれに関する詳細を知る上で役立つセッションは「MLX：Appleシリコンでの大規模言語モデルの実行」です MLXはフル機能のPython APIを備えラピッドプロトタイプに役立ちます SwiftのAPIもありますニューラルネットワークを構築し最適化する高レベルのパッケージがすべて含まれています MLXには C++とCのAPIもありこれらの言語で使用して最新の機械学習モデルを Appleシリコンで実行できますたとえばMac、iPhone、iPad Vision Proなどです MLXソフトウェアはすべてオープンソースで寛容なMITライセンスの下で提供されており主要ソフトウェアは GitHubで公開されており PythonとSwift APIで構築された数例とパッケージも含まれています MLXにはHugging Faceに活発なモデルクリエイターコミュニティもあります最新モデルの多くはすでに Hugging Face OrganizationのMLXコミュニティにあり新しいモデルが毎日アップされています PythonでMLXを使用開始するには PyPiからのインストールが最も簡単ですターミナルで1行pip3 install mlxを実行するだけです MLXは簡単に使えます配列の計算の実行を開始するには Pythonファイルを開いて MLXをインポートするだけです次にいくつかの配列を作成しそれらで基本的な操作を行えますたとえばここでは2つの整数配列を追加します
また形状やデータ型など配列に関する情報も簡単に調べることができます前に述べたようにMLX Python APIは NumPyに似ています通常操作の名前と署名は同じで同じように動作します NumPyのようなフレームワークを学んでいる場合は MLXは馴染みやすく始めるのも簡単です MLXが何でどのように役立つのかいくつかの基本がわかったところで主な機能について見ていきましょうユニファイドメモリや遅延評価関数変換、ニューラルネットワークの構築と最適化のための高レベルのパッケージなどです MLXは高性能なAppleシリコンを活用できるようにデザインされていますこれにはユニファイドメモリに固有の新しいプログラミングモデルが含まれます
機械学習用に広く使われているシステムには独立したメモリを持つ別個のGPUがあります一方Appleシリコンにはユニファイドメモリアーキテクチャがあります CPUとGPUが同一の物理メモリを共有しているということです
ユニファイドメモリを使用するには MLXは従来のフレームワークで慣れているものとは異なります従来のフレームワークでは計算はデータに従います配列がCPUメモリにある場合計算はCPUで行われます配列がGPUメモリにある場合計算はGPUで行われます
MLXでは配列はユニファイドメモリに割り当てられますサポートされているデバイスのどれで使うときでもどこにもコピーする必要はありません
デバイスに対して操作を実行するには操作自体にデバイスを指定しますたとえばここではGPUに aとbを加算し CPUでaとbを乗算しますこの操作は並列実行することもできますまたMLXは依存関係が存在する場合自動的に管理します MLXのもう一つの重要な特徴は遅延計算です MLXを可能な限り効率的にするために特に非常に大規模な計算の場合 MLXには遅延実行エンジンがあります
2つの配列に対する加算のような操作が呼び出されると実際の計算は行われません代わりにここに示すような計算グラフが作成されます
配列Cはまだ計算されていません実際に使用する必要がある場合にのみ計算されますたとえばCを印刷したり MLXからCを変換して Pythonリストに変換するときは計算が強制されます mx.evalでグラフの評価を明示的に強制できます遅延計算にはいくつかの利点があります計算グラフの構築と実行を切り離すことで MLXは結果を計算する前にグラフ上で変換と最適化を行えますまた遅延計算への支払いは使用分だけです関数変換はMLXのもう一つの重要な機能です高速配列フレームワークから機械学習モデルのトレーニングと性能向上に役立つはるかに強力なツールへと進化しています
関数変換は関数を入力として受け取り新しい関数を返す変換です
MLXにはいくつかの関数変換がありますそれらは通常2つのカテゴリのいずれかに分類されます自動微分の変換は計算グラフを最適化するための変換ですたとえば関数変換を使用して自動的に MLXの任意の関数の勾配を計算できます
入力の正弦を計算する基本的な関数があるとしますこの関数の勾配を取るには mx.grad関数変換を使用できます mx.gradの出力は新しい関数で配列で呼び出すと導関数が得られます関数変換は任意に構成可能です 2次導関数を取るには mx.gradの出力に mx.gradを使用するだけです
結果は別の関数であり配列で呼び出すと微分の2次導関数が得られます
MLXにはニューラルネットワークの構築とトレーニングを容易にできる高レベルのパッケージも2つあります
1つはmlx.nnでニューラルネットワークの構築に使用するモジュラーライブラリです 2つ目は一般的な最適化アルゴリズムのライブラリであるmlx.optimizersですこの2つはスタンドアロンでも使用できますが互いにシームレスに統合することも可能です
mlx.nnパッケージはニューラルネットワークの構築に必要なすべてを備えています nn.moduleはすべてのレイヤーとコンテナが継承するプライマリ基底クラスですパラメータへのアクセスやパラメータの読み込み保存などに役立つ方法を表示します
nnライブラリにはnn.Linearのような標準の既製層のセットがありますが nn.Moduleから継承することで独自の層を構築することもできます
一般的に使用される損失関数と初期化方法も nn.lossesとnn.init サブパッケージに含まれています
mlx.nnを使って単純な多層ニューラルネットワークを構築する方法を紹介します
最初のステップはnn.Moduleを継承するカスタムクラスを作成することです今回は単純な1つの隠れ層のニューラルネットワークを使用しますモジュールの初期化メソッド内に線形層を作成します次に呼び出し関数を実装しますこれは指定された入力に対するモジュールの出力を計算します呼び出し関数は最初の線形層を呼び出し ReLU活性化関数を適用し次に2番目の線形層を呼び出します MLXはAppleシリコンの統一メモリアーキテクチャ向けに設計最適化されていますが簡単に始められるように上位のニューラルネットワークパッケージは PyTorchのような一般的に使用される機械学習フレームワークと類似した設計となっています
MLXモデルの実装をPyTorchの同様の実装と比較してみます 2つはほぼ同じですが出力を計算する関数にわずかな違いが2つだけあります PyTorchでモデルを構築したことがある方は MLXへの切り替えは非常に簡単なはずです MLXの主な機能のほとんどを見てきたので機械学習ワークロードを可能な限り高速かつ効率的に作成する方法について話しますまず関数をコンパイルして高速化する方法を説明します次にmx.fastサブパッケージについて説明しますこのサブパッケージには一般的な機械学習操作の高速実装が標準で用意されており独自のMetalカーネルを追加するためのAPIも提供されていますその後メモリ使用量の削減とモデル実行の高速化を実現する量子化の使い方を紹介します最後にMLXを使って多数のマシン間で計算を配信する方法を紹介します MLXのほとんどすべての現実的な計算は配列に対していくつかの操作を実行する関数で構成されますそのような関数を高速化する簡単な方法は mx.compile関数変換を使用することですいくつかの要素ごとの操作を行う関数があるとしますたとえばこのGELU活性化関数などですこの関数の計算グラフにはいくつかのノードが含まれていますこれらの各ノードは内部で別のGPUカーネルの起動に対応します
グラフのコンパイルではこれらの個別のカーネルをすべて単一の融合カーネルに統合しますこれによりメモリ帯域幅とグラフ実行のオーバーヘッドを節約しはるかに効率的な計算ができます
mx.compileの使用は簡単で関数にmx.compile関数変換をデコレーターとして適用するだけですコンパイルは多くの場合うまく機能しますが機械学習の一般的な操作などのより複雑な操作では mx.fastサブパッケージを使用するとより高速に実行できる可能性がありますたとえばTransformerモデルの主な構成要素の多くは mx.fastにある操作を使用しますたとえば位置エンコーディング正規化層スケールドット積アテンションなどです mx.fastの操作はより特殊化されていますがトレーニングと推論の両方で可能な限り高速になるように高度に調整されていますまた高度に構成可能で特定の計算の多くのバリエーションを高速化できますたとえばスケールドット積アテンションの操作は入力に任意のマスクを受け取れますマスクは加算マスクブール型マスクマスクタイプを示す文字列にできます mx.fastの操作の1つである RMS normを詳しく見てみましょう RMS normは最新のTransformerベースの多くの大規模言語モデルで使われています MLXの操作を使用した基本的な実装では大規模な計算グラフが生成されます
代わりに実装全体を1回の mx.fast.rms_normへの呼び出しに置き換えることができますコードは単純になり計算グラフにノードは1つだけとなり計算自体がはるかに高速になります
MLXには関数の最適化のためによりカスタマイズされた実装が必要な場合でそれがmx.fastにないときカスタムMetalカーネルを追加するための APIがありますカスタムカーネルを書き込めば MLXが残りのすべてを処理しますジャストインタイムコンパイルと実行が含まれますこれらのカーネルはMetalで書かれますがこれはApple GPUで関数を実行するために Appleが開発した言語とAPIですカーネルを構築するには Metalコードのソース文字列と入力と出力に関する情報を渡します
カーネルを呼び出すにはグリッドサイズと出力の形状とタイプを指定します MLXはカーネル呼び出しを他の操作と同じように扱います計算グラフにノードを作成し遅延評価されます
機械学習のワークロードを高速化するツールキットの中でもう1つが量子化です大規模なモデルが高速であるためには大量のメモリとメモリ帯域幅が必要です多くの場合トレーニングで推論と同等または同程度の品質を得るためには推論に実際に必要な精度よりもはるかに高い精度が必要になります精度を下げることでより大きなモデルをメモリに収め実行を加速できます
モデルが32ビット浮動小数点精度の場合最初のステップとして bfloat16またはfloat16を使用しメモリ使用量を半分に減らすことができます 16ビットが多すぎると MLXには配列をさらに小さく量子化しその中で操作を行うためのルーチンが組み込まれていますたとえば要素あたり4ビットを使用して量子化するとメモリ要件をさらに削減できます
行列を量子化するには mx.quantizeを使用します要素ごと使用するビット数とグループサイズを指定しますグループサイズにより量子化行列において共有されるスケール値とバイアス値を持つ要素の数が決定されますビットが小さくグループサイズが大きいほど結果は小さく高速になります MLXにはビットとグループサイズにさまざまな選択肢があり可能な限りの柔軟性を実現します
定量化されていない入力ベクトルまたは行列を定量化された行列と乗算するには mx.quantized_matmulを使用します
mx.dequantizeを使用して元の入力の近似値を回復します
MLX.mnにはモジュールを1つのコマンドで量子化する便利なユーティリティもありますたとえば埋め込み層に続いて複数の線形層が積み重なったモデルがあるとします nn.quantizeを使用してモデルを量子化できます quantizeコマンドにはオプションのコールバックを受け取る機能がありユーザーはどの層を量子化するか各レイヤーでどのような精度を使用するかをより細かく制御できます
大規模言語モデルでテキストを生成する際量子化はメモリ使用量を大幅に削減し毎秒生成されるトークンを増やします場合によっては 1台のマシンでは不十分ですたとえば1台のマシンのメモリに収まらない大規模なモデルを使用してテキストを生成する場合があるかもしれません大規模なデータセットでモデルを微調整または評価する場合もありますどちらも並列化が容易で複数のマシンを使用するとはるかに高速になります
MLXは設定なしですぐに使え複数のマシンで任意の計算を実行しますマシンはEthernetまたは Thunderbolt接続できます
mx.distributed サブパッケージを使用すると複数のマシン間で計算を配信できます MX distributedは主に一連の通信操作です
たとえばall_Sumはすべてのマシンの入力配列を加算します all_Sumの出力は入力の合計でこれは各マシンで同じです複数のマシンで配列を合計する方法を詳しく見てみましょう mx.distributed.initを使って分散バックエンドを初期化しますこのステップは任意ですがコミュニケーショングループへのアクセスが必要なときに呼び出すものですコミュニケーショングループにはプロセスの総数や現在のプロセスインデックスなど有用な情報が添付されています
次に各プロセスに単一の値を含む配列を作成し mx.distributed.all_sumを呼び出してすべてのプロセスにわたる配列を合算します
MLXには複数のマシンでMLXプログラムを実行するための便利なランチャがあります 4台のマシンでプログラムを実行するには4つのホストIPアドレスで mlx.launchを使用しますここまではすべてPythonでMLXを使用することを説明しました多くの場合Pythonの使いやすさと柔軟性が好まれるかもしれませんが Swiftの方がよい場合もありますそのためMLXはSwiftでのフル機能のAPIを提供しています Metal上に構築され macOS、iOS、iPadOS visionOSなどすべてにわたってうまく機能します SwiftでMLXを始めるにはみなさんのXcodeプロジェクトにパッケージのまま追加するだけですプロジェクトをクリックし次にパッケージの依存関係タブのプラス記号をクリックします次にMLX Swift GitHub リポジトリのURLを入力しパッケージを追加するボタンをクリックします MLX Swiftで構築を開始するのに必要なのはこれだけです PythonとSwiftの間の移行をできるだけ簡単にするために 2つの言語間のAPIは意図的に似せていますこれまでのPythonコードスニペットと MLX Swiftのものを並べて比較したものです配列の作成それらに対する操作の実行配列に関するメタデータの検査は PythonとSwiftでほぼ同じですこれまで見てきた PythonでMLXを使用するときのすべての主要機能とともに MLXの高速化の部分で話した最適化は MLX Swiftにも同様に適用されます MLXの多くの主要な機能の概要を説明してきましたフレームワークの詳細は MLX Webサイトをご覧くださいこれにはドキュメントや例などへのリンクがあります Python APIとSwift APIの両方にサンプルリポジトリがありここに一般的な機械学習のユースケースの例が多くありますこれには言語モデルのトレーニングと生成画像生成音声認識などが含まれますこれらの例はMLXについて学びを深めるのに最適で独自のプロジェクトでMLXを使用して構築するための良い出発点となりますありがとうございましたみなさんがMLXで何を構築するか楽しみです

import mlx.core as mx

# Make an array
a = mx.array([1, 2, 3])

# Make another array
b = mx.array([4, 5, 6])

# Do an operation
c = a + b

# Access information about the array
shape = c.shape
dtype = c.dtype

print(f"Result c: {c}")
print(f"Shape: {shape}")
print(f"Data type: {dtype}")

5:31 - Unified memory

import mlx.core as mx

a = mx.array([1, 2, 3])
b = mx.array([4, 5, 6])

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

print(f"c computed on the GPU: {c}")
print(f"d computed on the CPU: {d}")

6:20 - Lazy computation

import mlx.core as mx

# Make an array
a = mx.array([1, 2, 3])

# Make another array
b = mx.array([4, 5, 6])

# Do an operation
c = a + b

# Evaluates c before printing it
print(c)

# Also evaluates c
c_list = c.tolist()

# Also evaluates c
mx.eval(c)

print(f"Evaluate c by converting to list: {c_list}")
print(f"Evaluate c using print: {c}")
print(f"Evaluate c using mx.eval(): {c}")

7:32 - Function transformation

import mlx.core as mx

def sin(x):
    return mx.sin(x)

dfdx = mx.grad(sin)

def sin(x):
    return mx.sin(x)

d2fdx2 = mx.grad(mx.grad(mx.sin))

# Computes the second derivative of sine at 1.0
d2fdx2(mx.array(1.0))

9:16 - Neural Networks in MLX

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim

class MLP(nn.Module):
    """A simple MLP."""

    def __init__(self, dim, h_dim):
        super().__init__()
        self.linear1 = nn.Linear(dim, h_dim)
        self.linear2 = nn.Linear(h_dim, dim)

    def __call__(self, x):
        x = self.linear1(x)
        x = nn.relu(x)
        x = self.linear2(x)
        return x

9:57 - MLX and PyTorch

# MLX version
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim

class MLP(nn.Module):
    """A simple MLP."""

    def __init__(self, dim, h_dim):
        super().__init__()
        self.linear1 = nn.Linear(dim, h_dim)
        self.linear2 = nn.Linear(h_dim, dim)

    def __call__(self, x):
        x = self.linear1(x)
        x = nn.relu(x)
        x = self.linear2(x)
        return x

# PyTorch version
import torch
import torch.nn as nn
import torch.optim as optim

class MLP(nn.Module):
    """A simple MLP."""

    def __init__(self, dim, h_dim):
        super().__init__()
        self.linear1 = nn.Linear(dim, h_dim)
        self.linear2 = nn.Linear(h_dim, dim)

    def forward(self, x):
        x = self.linear1(x)
        x = x.relu()
        x = self.linear2(x)
        return x

11:35 - Compiling MLX functions

import mlx.core as mx
import math

def gelu(x):
    return x * (1 + mx.erf(x / math.sqrt(2))) / 2

@mx.compile
def compiled_gelu(x):
    return x * (1 + mx.erf(x / math.sqrt(2))) / 2

x = mx.random.normal(shape=(4,))

out = gelu(x)
compiled_out = compiled_gelu(x)
print(f"gelu:          {out}")
print(f"compiled gelu: {compiled_out}")

12:32 - MLX Fast package

import mlx.core as mx
import time

def rms_norm(x, weight, eps=1e-5):
    y = x.astype(mx.float32)
    y = y * mx.rsqrt(mx.mean(
        mx.square(y), 
        axis=-1, 
        keepdims=True,
    ) + eps)
    return (weight * y).astype(x.dtype)

batch_size = 8192
feature_dim = 4096
iterations = 1000

x = mx.random.normal([batch_size, feature_dim])
weight = mx.ones(feature_dim)
bias = mx.zeros(feature_dim)

start_time = time.perf_counter()
for _ in range(iterations):
    y = rms_norm(x, weight, eps=1e-5)
    mx.eval(y)
rms_norm_time = time.perf_counter() - start_time
print(f"rms_norm execution: {gelu_time:0.4f} sec")

start_time = time.perf_counter()
for _ in range(iterations):
    mx.eval(mx.fast.rms_norm(x, weight, eps=1e-5))
fast_rms_norm_time = time.perf_counter() - start_time
print(f"mx.fast.rms_norm execution: {compiled_gelu_time:0.4f} sec")

print(f"mx.fast.rms_norm speedup: {rms_norm_time/fast_rms_norm_time:0.2f}x")

13:30 - Custom Metal kernel

import mlx.core as mx

# Build the kernel
source = """
    uint elem = thread_position_in_grid.x;
    out[elem] = metal::exp(inp[elem]);
"""
kernel = mx.fast.metal_kernel(
    name="myexp",
    input_names=["inp"],
    output_names=["out"],
    source=source,
)

# Call the kernel on a sample input
x = mx.array([1.0, 2.0, 3.0])
out = kernel(
    inputs=[x],
    grid=(x.size, 1, 1),
    threadgroup=(256, 1, 1),
    output_shapes=[x.shape],
    output_dtypes=[x.dtype],
)[0]
print(out)

14:41 - Quantization

import mlx.core as mx

x = mx.random.normal([1024])
weight = mx.random.normal([1024, 1024])

quantized_weight, scales, biases = mx.quantize(
        weight, bits=4, group_size=32,
)

y = mx.quantized_matmul(
    x,
    quantized_weight,
    scales=scales,
    biases=biases,
    bits=4,
    group_size=32,
)

w_orig = mx.dequantize(
    quantized_weight,
    scales=scales,
    biases=biases,
    bits=4,
    group_size=32,
)

15:23 - Quantized models

import mlx.nn as nn

model = nn.Sequential(
    nn.Embedding(100, 32),
    nn.Linear(32, 32),
    nn.Linear(32, 32),
    nn.Linear(32, 1),
)

print(model)

nn.quantize(
    model,
    bits=4,
    group_size=32,
)

print(model)

16:50 - Distributed

import mlx.core as mx

group = mx.distributed.init()

world_size = group.size()
rank = group.rank()

x = mx.array([1.0])

x_sum = mx.distributed.all_sum(x)

print(x_sum)

17:20 - Distributed launcher

mlx.launch --hosts ip1, ip2, ip3, ip4 my_script.py

18:20 - MLX Swift

// Swift
import MLX

// Make an array
let a = MLXArray([1, 2, 3])

// Make another array
let b = MLXArray([1, 2, 3])

// Do an operation
let c = a + b

// Access information about the array
let shape = c.shape
let dtype = c.dtype

// Print results
print("a: \(a)")
print("b: \(b)")
print("c = a + b: \(c)")
print("shape: \(shape)")
print("dtype: \(dtype)")

関連する章

リソース

関連ビデオ

WWDC25