Nx のバックエンドとして EMLX を使用し、 EXLA や TorchX と比較する

Last updated at 2024-12-25Posted at 2024-12-25

はじめに

MLX は Apple Silicon を利用した機械学習ライブラリです

EMLX は Nx のバックエンドとして MLX を使用します

EMLX を使うことで、 MacBook 上での機械学習が高速化できるはず

ということで、ベンチマークしてみました

実行環境

端末: MacBook Pro 13インチ、M2、2022
チップ: Apple M2
コアの総数: 8（パフォーマンス: 4、効率性: 4）
メモリ: 24 GB
Elixir: 1.18.1
Erlang: 27.2

Livebook の起動

MacBook 上で起動するため、 Livebook をソースコードから実行します

git clone -b v0.14.5 https://github.com/livebook-dev/livebook.git
cd livebook
mix setup
ELIXIR_ERL_OPTIONS="-epmd_module Elixir.Livebook.EPMD" mix phx.server

セットアップ

新しいノートブックを開き、セットアップセルで必要なモジュールをインストールします

Mix.install([
  {:nx, "~> 0.9"},
  {:exla, "~> 0.9"},
  {:torchx, "~> 0.9"},
  {:emlx, github: "elixir-nx/emlx"},
  {:kino_benchee, "~> 0.1"}
])

KinoBenchee はベンチマークのためのモジュールです

ベンチマーク

ベンチマーク用のモジュールを定義します

defmodule MyBenchmark do
  def add(backend) do
    tensor = Nx.iota({200, 200}, type: {:f, 64}, backend: backend)
    Nx.add(tensor, tensor)
  end

  def dot(backend) do
    tensor = Nx.iota({200, 200}, type: {:f, 64}, backend: backend)
    Nx.dot(tensor, tensor)
  end

  def run_add() do
    Benchee.run(
      %{
        "binary" => fn -> add(Nx.BinaryBackend) end,
        "exla" => fn -> add(EXLA.Backend) end,
        "torchx" => fn -> add(Torchx.Backend) end,
        "emlx" => fn -> add(EMLX.Backend) end
      },
      memory_time: 2,
      reduction_time: 2
    )
  end

  def run_dot() do
    Benchee.run(
      %{
        "binary" => fn -> dot(Nx.BinaryBackend) end,
        "exla" => fn -> dot(EXLA.Backend) end,
        "torchx" => fn -> dot(Torchx.Backend) end,
        "emlx" => fn -> dot(EMLX.Backend) end
      },
      memory_time: 2,
      reduction_time: 2
    )
  end
end