WWDC Quick Look 💓 By SwiftGGTeam
中文 EN 日本語
Get started with MLX for Apple silicon

Get started with MLX for Apple silicon

機械学習
元の動画を見る

ハイライト

MLX は Apple が自社開発したオープンソースの配列フレームワークで、Apple Silicon のユニファイドメモリアーキテクチャに最適化されています。核心的なコンセプトは、CPU と GPU がメモリを共有し、データを手動で移動させる必要がないという点です。

主要内容

Mac で PyTorch を動かしたことがある方なら、きっとこのような悩みを経験したことがあるでしょう。CPU と GPU は同一チップ上にあるのに、モデルのテンソルは二つの「仮想」メモリ間を行き来する必要があります。.to("mps") のたびにコピーが発生し、大規模モデルの推論時には VRAM が不足しがちです。Apple が WWDC 2025 で発表した MLX は、まさにこの問題を解決するために生まれました。

MLX の核心的な変更点は、「計算がデータを追いかける」という従来の考え方を根本から覆すことです。従来のフレームワークでは、配列が CPU メモリ上にあれば CPU でしか計算できず、GPU メモリ上にあれば GPU でしか計算できませんでした。MLX は配列を Apple Silicon のユニファイドメモリに配置し、特定のデバイスに属さないようにします。各オペレーションの呼び出し時に stream=mx.cpu または stream=mx.gpu で実行場所を指定します。同じデータに対して CPU と GPU が並列に操作でき、フレームワークが依存関係を自動的に管理します。この基盤を中心に、MLX は NumPy スタイルの API、PyTorch スタイルの nn.Module、JAX スタイルの関数変換（mx.grad、mx.compile）を単一のパッケージに統合し、Mac、iPhone、iPad、Vision Pro 上で同じコードを実行できるようにしています。さらに Python、Swift、C++、C の 4 つの API を網羅しています。

詳細

インストールと基本的な配列操作（03:48）。pip3 install mlx で一行でインストールでき、API は NumPy とほぼ一致しています。

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}")

ポイント：

  • import mlx.core as mx：MLX のコア配列ライブラリ。命名規則は NumPy を模倣しています。
  • mx.array([...])：Python リストから配列を作成。データはユニファイドメモリに格納され、device の概念がありません。
  • a + b：演算子オーバーロード。返される c は未評価の「遅延」配列です（後述）。
  • c.shape、c.dtype：NumPy と同名同義で、NumPy の経験があればそのまま移行できます。

ユニファイドメモリプログラミング（05:31）。これは MLX が他のフレームワークと最も異なる点です。stream パラメータは配列ではなくオペレーションに紐付けられます。

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}")

ポイント：

  • a と b は同じデータであり、「CPU か GPU か」という区別がありません。
  • mx.add(a, b, stream=mx.gpu)：加算を GPU 上で実行します。
  • mx.multiply(a, b, stream=mx.cpu)：同じ入力に対して、乗算を CPU 上で実行します。
  • 二つのオペレーションに依存関係がなければ、並列実行が可能です。依存関係があれば、MLX が自動的に順序付けします。

遅延評価（06:20）。MLX はオペレーションを呼び出す際に計算グラフを構築するだけで、実際の計算は行いません。

import mlx.core as mx

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

# Evaluates c before printing it
print(c)

# Also evaluates c
c_list = c.tolist()

# Also evaluates c
mx.eval(c)

ポイント：

  • c = a + b の時点では計算は発生せず、グラフに加算ノードを追加するだけです。
  • print(c)、.tolist()、mx.eval(c) のいずれかが評価をトリガーします。結果を使わなければ、完全に計算されません。
  • これにより MLX はグラフ最適化の余地を持ち、複数のオペレーションを単一の kernel に融合できます。

関数変換：自動微分（07:32）。mx.grad は関数を新しい関数に変換します。

import mlx.core as mx

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

dfdx = mx.grad(sin)

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

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

ポイント：

  • mx.grad(sin) は微分後の関数 dfdx を返します。これ自体も関数です。
  • 関数変換は任意にネスト可能です。mx.grad(mx.grad(...)) で二階微分を直接取得できます。
  • このスタイルは JAX に由来し、モデルの訓練や勾配計算に便利です。

mlx.nn でニューラルネットワークを構築する（09:16）：

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

ポイント：

  • nn.Module はすべてのレイヤーとコンテナの基底クラスで、パラメータへのアクセス、読み込み、保存などのメソッドを提供します。
  • nn.Linear、nn.relu などの標準レイヤーは、名前とシグネチャが PyTorch の慣習に準じています。
  • MLX は PyTorch の forward の代わりに __call__ を使用する点に注意してください。これは PyTorch との唯一の大きな違いの一つです。

mx.compile で kernel を融合する（11:35）。デコレータを追加するだけで、複数の element-wise オペレーションが単一の GPU kernel に融合され、重複するメモリ往復と kernel 起動のオーバーヘッドを削減できます。

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

ポイント：

  • gelu 内には 4 つのオペレーション（erf、除算、加算、乗法）があり、本来は 4 回の kernel 起動に対応します。
  • @mx.compile により、MLX はこれらのノードを単一の kernel にコンパイルします。
  • 繰り返し呼び出される小さな関数（活性化関数、正規化）で最も効果が顕著です。

mx.fast 高性能オペレーション（12:32）。Transformer の主要なオペレーション（RMS Norm、SDPA、RoPE）には、手書きの最適化版が用意されており、直接利用すると自分で組み立てるよりも高速です。カスタム Metal kernel（13:30）では、Metal シェーダーのソースコードを計算グラフに直接組み込むことができます。

import mlx.core as mx

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,
)

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]

ポイント：

  • source は生の Metal コードで、thread_position_in_grid は標準的な Metal 組み込み変数です。
  • mx.fast.metal_kernel(...) はこのソースコードを呼び出し可能なオブジェクトにラップします。
  • 呼び出し時に grid と threadgroup でスレッド分割を指定し、output_shapes/output_dtypes で出力テンソルを記述します。
  • MLX がカバーしていないオペレーションを手書きで極限まで最適化しつつ、遅延グラフスケジューリングの恩恵を受けられます。

4-bit 量子化（14:41）。mx.quantize は重みを 4-bit にパックし、mx.quantized_matmul と組み合わせて推論を行います。モデルレベルでは nn.quantize(model, bits=4, group_size=32) の一行で、すべての Linear/Embedding を置き換えられます。マルチマシン分散処理（16:50）は、mx.distributed.init() でグループを取得し、mx.distributed.all_sum などの集合通信プリミティブを使用します。mlx.launch --hosts ip1,ip2,ip3,ip4 my_script.py と組み合わせることで、イーサネットまたは Thunderbolt で接続された複数台の Mac 上で同じコードを実行できます。

MLX Swift（18:20）。同じ API が Swift にも移植されており、iOS/macOS アプリケーションへの統合が容易です。

import MLX

let a = MLXArray([1, 2, 3])
let b = MLXArray([1, 2, 3])
let c = a + b

let shape = c.shape
let dtype = c.dtype

ポイント：

  • MLXArray は Python の mx.array に対応し、構築方法も一致しています。
  • 演算子オーバーロード、プロパティ名（shape、dtype）も一致しており、言語間の移行が容易です。
  • iOS、iPadOS、visionOS 上で直接ビルドでき、追加のランタイムに依存しません。

重要ポイント

  • やるべきこと：Mac 上で既存の PyTorch 推論を MLX に移行し、ユニファイドメモリのレイテンシー改善を検証する。

    • なぜ価値があるか：MPS バックエンドのコピー/同期オーバーヘッドは、大規模モデルで無視できない規模になります。MLX には「データ移動」というステップがなく、4-bit 量子化もネイティブにサポートされています。
    • 始め方：7B 規模のモデルを一つ選び、MLX LM リポジトリで対応する重みを探します。nn.quantize(model, bits=4, group_size=32) で量子化し、token/s とピークメモリを比較します。

  • やるべきこと：訓練ループ内の活性化関数と正規化を @mx.compile でラップする。

    • なぜ価値があるか：これらの関数は繰り返し呼び出され、毎回複数の kernel を経由します。コンパイルによる融合後に削減されるのは kernel 起動のオーバーヘッドであり、これが最も効果的に加速できます。
    • 始め方：まず mx.fast.rms_norm、mx.fast.scaled_dot_product_attention などの既存実装で手書き版を置き換えます。残りの独自の小さな関数には @mx.compile デコレータを追加します。

  • やるべきこと：2〜3 台の Mac mini を Thunderbolt で接続し、ローカルの分散クラスタを構築する。

    • なぜ価値があるか：単一マシンの VRAM を最大限に使うのではなく、水平スケーリングによりより大きなモデルを実行でき、MLX が提供する mlx.launch は起動の複雑さを一行のコマンドに押し込めます。
    • 始め方：各マシンに MLX をインストールし、mx.distributed.init() でグループを取得します。最小限の all_sum demo を作成して接続を検証し、既存の訓練スクリプトを mlx.launch --hosts に接続します。

  • やるべきこと：iOS/visionOS アプリに MLX Swift を組み込み、オンデバイス推論を行う。

    • なぜ価値があるか：Foundation Models とは異なり、MLX ではモデルの重みとグラフ構造を自分で持ち込めます。独自に訓練したモデルやコミュニティモデルに最も適しています。
    • 始め方：mlx-swift-examples リポジトリから最小限の推論サンプルを clone し、モデルの重みを自分の Hugging Face MLX community 版に置き換えます。動作確認後、SwiftUI に統合します。

関連セッション

コメント

GitHub Issues · utterances