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

## はじめに

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

https://fly.io/docs/gpus/

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

https://qiita.com/RyoWakabayashi/items/c8666c2797ff5bafe611

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

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

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

## Fly GPUs の価格

https://fly.io/docs/about/pricing/#gpus-and-fly-machines

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

## 参考記事

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

https://fly.io/phoenix-files/real-world-machine-learning-on-fly-gpu/

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

## 実行手順

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

新規アプリを作成します

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

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

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

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

```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 を起動しておくための設定です

```toml
  auto_stop_machines = false
  auto_start_machines = false
  min_machines_running = 1
```

### デプロイ

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

```sh
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 でマシンの状態が確認できます

![スクリーンショット 2024-02-21 21.45.34.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/3cfcd302-b734-9362-a82f-cd1755a01427.png)

確かに GPU が付いています

## Livebook で GPU 確認

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

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

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

### セットアップ

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

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

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

```elixir
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 を動かしたときと同じ手順でテンソルを確認します

https://qiita.com/RyoWakabayashi/items/48e52878bad57e468276

#### バイナリバックエンド

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

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

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

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

#### EXLA バックエンド

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

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

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

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


#### Torchx バックエンド

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

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

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

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

#### evision バックエンド

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

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

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

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

## ベンチマーク

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

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

```elixir
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 を使うことができました

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