Fly.io 上で Livebook から GPU を動かす【Fly GPUs】 #Elixir

はじめに

Fly.io で GPU が使えるようになりました

ざっくり Fly.io の紹介、使い方はこちら

一言で言えば「安く簡単に Web アプリを構築・公開できるクラウドサービス」です

今までは GPU が使えないため、 AI モデルを高速に動かす、といったことができませんでした

ついに GPU が使えるようになったため、まずは使ってみました

Fly GPUs の価格

L40S: $2.50/hr per GPU
A100 40G PCIe: $2.50/hr per GPU
A100 80G SXM: $3.50/hr per GPU

参考記事

ほぼこの記事の通りで実行できました

一部コマンドの修正やエラー対応が必要だったので備忘として記事を残します

実行手順

Fly.io の新規アプリケーション作成

新規アプリを作成します

fly apps create <アプリケーション名>

参考記事ではリージョンを指定していましたが、この時点で --region は指定できません

アプリケーション設定ファイルの作成

fly.toml を以下の内容で作成します

app = <アプリケーション名>
primary_region = 'ord'

[build]
  image = 'ghcr.io/livebook-dev/livebook:latest-cuda12.1'

[env]
  ELIXIR_ERL_OPTIONS = '-proto_dist inet6_tcp +sssdio 128'
  LIVEBOOK_DATA_PATH = '/data'
  LIVEBOOK_HOME = '/data'
  LIVEBOOK_IP = '::'
  LIVEBOOK_ROOT_PATH = '/data'
  PORT = '8080'

[http_service]
  internal_port = 8080
  force_https = true
  auto_stop_machines = false
  auto_start_machines = false
  min_machines_running = 1
  processes = ['app']

[[vm]]
  size = 'a100-40gb'
  cpu_kind = 'performance'
  cpus = 8
  memory_mb = 32768

primary_region = 'ord' でシカゴリージョンを指定しています

また、 size = 'a100-40gb' で GPU モデルを指定しています

GPU モデル毎に使用できるリージョンが決まっており、今回は a100-40gb を使うため、シカゴリージョンを指定します

l40s: ord
a100-40gb: ord
a100-80gb: ams, iad, mia, sjc, syd

cpus や memory_mb は必要に応じて変更してください

以下の箇所は永続的に Livebook を起動しておくための設定です

  auto_stop_machines = false
  auto_start_machines = false
  min_machines_running = 1

デプロイ

Fly.io にアプリケーションをデプロイします

fly deploy --ha=false

--ha=false でクラスタリングしないようにしています（最小構成とするため）

エラー対応

上記の手順でデプロイ（ fly deploy ）を実行した時点で以下のエラーが発生しました

✖ Failed: error creating a new machine: failed to launch VM: Your
organization does not have enough trust to create GPU machines. Please
contact billing@fly.io to remove the restriction.

和訳

新しいマシンの作成に失敗しました。
エラー内容: VMの起動に失敗しました。お使いの組織はGPUマシンを作成するための十分な信頼を得ていません。
制限を解除してもらうために、billing@fly.io に連絡してください。

私の個人アカウントは無償範囲内でしか使っていなかったため、支払い実績がなく、 GPU マシンを立てることができませんでした

エラーメッセージにある通り、 billing@fly.io 宛にメールで問い合わせると数日後にメールの返信が来ました

You're all set.

和訳

準備が完了しました

やったぜ

デプロイ結果の確認

デプロイコマンドが正常すると、 Fly.io でマシンの状態が確認できます

確かに GPU が付いています

Livebook で GPU 確認

起動したアプリケーションの URL にアクセスします

認証トークンの入力を求められるので、実行ログ（ Monitoring メニュー、もしくは fly logs で確認可能）から起動時の認証トークンをコピー&ペーストします

新しいノートブックを作成し、 GPU が使えることを確認しましょう

セットアップ

Nx　、 EXLA 、 Torchx 、 evision などをインストールします

各環境変数で CUDA のバージョンを指定しています（2024年2月現在で使用可能だった最新版）

Torchx は GitHub 上では cu121 が書かれていましたが、エラーになりました

Mix.install(
  [
    {:benchee, "~> 1.1"},
    {:nx, "~> 0.6"},
    {:exla, "~> 0.6"},
    {:torchx, "~> 0.6"},
    {:evision, "~> 0.1"},
    {:kino, "~> 0.12"}
  ],
  system_env: [
    {"XLA_TARGET", "cuda120"},
    {"EXLA_TARGET", "cuda"},
    {"LIBTORCH_TARGET", "cu118"},
    {"EVISION_ENABLE_CUDA", "true"},
    {"EVISION_ENABLE_CONTRIB", "true"},
    {"EVISION_CUDA_VERSION", "121"}
  ]
)

各種バックエンドでのテンソル作成

以前、 Google Colab で Livebook を動かしたときと同じ手順でテンソルを確認します

バイナリバックエンド

tensor = Nx.tensor([1.0, 2.0, 3.0], backend: Nx.BinaryBackend)
Nx.add(tensor, tensor)

実行結果は以下のようになります

#Nx.Tensor<
  f32[3]
  [2.0, 4.0, 6.0]
>

バイナリバックエンドは GPU を使いません

EXLA バックエンド

tensor = Nx.tensor([1.0, 2.0, 3.0], backend: {EXLA.Backend, device_id: 0})
Nx.add(tensor, tensor)

実行結果は以下のようになります

#Nx.Tensor<
  f32[3]
  EXLA.Backend<cuda:0, 0.967388242.3511812113.20823>
  [2.0, 4.0, 6.0]
>

EXLA.Backend<cuda:0... というように、 GPU が使われていることがわかります

Torchx バックエンド

tensor = Nx.tensor([1.0, 2.0, 3.0], backend: {Torchx.Backend, device: :cuda})
Nx.add(tensor, tensor)

実行結果は以下のようになります

#Nx.Tensor<
  f32[3]
  Torchx.Backend(cuda)
  [2.0, 4.0, 6.0]
>

こちらも (cuda) で GPU を使っていることがわかります

evision バックエンド

tensor = Nx.tensor([1.0, 2.0, 3.0], backend: Evision.Backend)
Nx.add(tensor, tensor)

実行結果は以下のようになります

#Nx.Tensor<
  f32[3]
  Evision.Backend
  [2.0, 4.0, 6.0]
>

evision は特に GPU を使っていることが分かりません

ベンチマーク

200 x 200 の行列を足し算する処理でベンチマークします

bench = fn backend ->
  tensor = Nx.iota({200, 200}, type: {:f, 64}, backend: backend)
  Nx.add(tensor, tensor)
end

Benchee.run(%{
  "binary" => fn -> bench.(Nx.BinaryBackend) end,
  "exla" => fn -> bench.({EXLA.Backend, device_id: 0}) end,
  "torchx_cpu" => fn -> bench.({Torchx.Backend, device: :cpu}) end,
  "torchx" => fn -> bench.({Torchx.Backend, device: :cuda}) end,
  "evision" => fn -> bench.(Evision.Backend) end
})

実行結果は以下のようになりました

Name                 ips        average  deviation         median         99th %
torchx            9.27 K      107.86 μs    ±23.96%      103.90 μs      191.13 μs
exla              3.83 K      260.81 μs    ±24.47%      251.39 μs      407.82 μs
torchx_cpu        2.24 K      446.79 μs    ±33.37%      431.90 μs      834.56 μs
evision           0.93 K     1077.83 μs    ±17.67%     1014.58 μs     1555.76 μs
binary          0.0684 K    14628.29 μs     ±9.43%    14235.81 μs    18655.33 μs

Comparison: 
torchx            9.27 K
exla              3.83 K - 2.42x slower +152.96 μs
torchx_cpu        2.24 K - 4.14x slower +338.94 μs
evision           0.93 K - 9.99x slower +969.97 μs
binary          0.0684 K - 135.63x slower +14520.44 μs

バックエンド	実行時間（最速との比）
Torchx	最速
EXLA	2.42
Torchx(CPU)	4.14
Evision	9.99
Binary	135.63

GPU を使うことで高速化できています

注意

GPU はそれなりに高価なので、使わなくなったらアプリケーションを削除しておきましょう

まとめ

Fly.io で GPU を使うことができました

より応用的な実行例は別途記事を書きます