【2024年10月】Google Colaboratory で Livebook を起動し、 Nx を EXLA や Evision の最新版 GPU バックエンドで動かす

## はじめに

Nx は Elixir で行列演算を実行するためのモジュールです

https://github.com/elixir-nx/nx/tree/main/nx#readme

Nx は機械学習や画像処理、データ解析などに利用されています

Nx は EXLA や Torchx 、 Evision などをバックエンドとして利用することで、より高速に動作させることができます

特に GPU 環境ではバックエンドによる高速化の効果が高く、機械学習を実行する場合は必須と言えます

以前記事にも書きましたが、 Google Colaboratory (以下、 Colab )はブラウザから気軽に GPU を使える環境を提供しており、 Livebook を起動することも可能です

Colab 上の Livebook から EXLA バックエンドなどを使用すれば、誰でも GPU が基本無料で使えるわけです

https://qiita.com/RyoWakabayashi/items/113b94866780c7646af1

ただし、Colab の環境は定期的に更新されており、 Nx や EXLA も更新されています

2024年10月現在では少し工夫しないと各種バックエンドを使えなかったので、その解決方法を記事に残しておきます

実装したノートブックはこちら

https://github.com/RyoWakabayashi/elixir-learning/blob/main/colab/Livebook.ipynb

https://github.com/RyoWakabayashi/elixir-learning/blob/main/colab/bench.livemd

## Google Colab の2024年10月最新版

2024年10月時点では、 Google Colab の GPU ランタイム環境は以下のようになっています

- OS: Ubuntu 22.04.3
- CUDA: 12.2.140
- cuDNN: 8.9.6

この CUDA が 12 で cuDNN が 8 というのが厄介でした

## Nx の2024年10月最新版

Nx は 2024年9月26日に 0.9.0 がリリースされました

https://github.com/elixir-nx/nx/blob/main/nx/CHANGELOG.md

本記事では 0.9.0 を使用します

## EXLA の2024年10月最新版

Nx の更新に伴って、 EXLA も 0.9.0 がリリースされています

https://github.com/elixir-nx/nx/blob/main/exla/CHANGELOG.md

EXLA 用のコンパイル済バイナリとして XLA というモジュールを使用します

https://github.com/elixir-nx/xla

XLA の最新版は 0.8.0 です

https://github.com/elixir-nx/xla/releases

0.8.0 のリリースバイナリは CUDA 12 、 cuDNN 9.1 以上の組み合わせ用にコンパイルされています

この組み合わせでない場合、自分の環境内でコンパイルしなければなりません

そして、 Colab は CUDA 12 と cuDNN 8.6 なので、コンパイルする必要があります

### コンパイルする場合の手順（ボツ案）

コンパイルのためには bazel （ビルドツール）が必要です

https://bazel.build/

しかも、バージョンは 6.5 でないといけません（7 以降は使えない）

公式の紹介にある通り、最も簡単に bazel をインストールする手段として、 [asdf](https://asdf-vm.com/) が使用できます

Colab 上で以下のコードを実行すれば bazel 6.5.0 がインストールできます

```bash
!git clone https://github.com/asdf-vm/asdf.git ~/.asdf
```

```python
import os
os.environ['PATH'] = "/root/.asdf/shims:/root/.asdf/bin:" + os.environ['PATH']
```

```bash
!asdf plugin-add bazel
!asdf install bazel 6.5.0
!asdf global bazel 6.5.0
```

その後、後述する手順で Livebook を起動したら、セットアップセルで以下のコードを実行します

```elixir
Mix.install(
  [
    {:nx, "~> 0.9"},
    {:exla, "~> 0.9"}
  ],
  system_env: [
    {"XLA_BUILD", "true"},
    {"XLA_TARGET", "cuda"},
    {"EXLA_TARGET", "cuda"}
  ]
)
```

ビルドが開始されますが、恐ろしく時間がかかります

私は6時間程度経過した時点で諦めました

最終的な手段は後述します

## Torchx の2024年10月最新版

Torchx も 0.9.0 が最新版です

https://github.com/elixir-nx/nx/tree/main/torchx

Torchx は cuDNN のバージョン指定がなく、 CUDA 12 で使用可能です

環境変数 `LIBTORCH_TARGET` に `cu121` を指定しましょう

## Evision の2024年10月最新版

Evision は 0.2.10 が最新です

https://github.com/cocoa-xu/evision

コンパイル済のバイナリを利用したい場合、 CUDA 12 と cuDNN 9 の組み合わせが必要になります

https://github.com/cocoa-xu/evision/blob/dee1313201402e70228ed3138ca812f429980df8/do_release.sh#L11

Colab は CUDA 12 と cuDNN 8 の組み合わせなので、コンパイルが必要です

### Evision をコンパイルする場合の手順（ボツ案）

コンパイルする場合は以下のようなセットアップが必要です

```elixir
Mix.install(
  [
    {:nx, "~> 0.9"},
    {:evision, "~> 0.2"}
  ],
  system_env: [
    {"EVISION_PREFER_PRECOMPILED", "false"},
    {"EVISION_ENABLE_CUDA", "true"},
    {"EVISION_ENABLE_CONTRIB", "true"},
    {"EVISION_CUDA_VERSION", "12"},
    {"EVISION_CUDNN_VERSION", "8"}
  ]
)
```

こちらは4時間程度でコンパイルできるため、時間に余裕があるなら実行可能です

しかし、もちろんオススメはしません

## cuDNN 9 をインストールする

結局、 Livebook で cuDNN 8 用にコンパイルする場合は時間がかかりすぎるため、 Colab 上で cuDNN 9 をインストールすることにしました

cuDNN 9 のインストール方法はこちら

https://qiita.com/RyoWakabayashi/private/005f171568c53b15c2a1

## 最終的な実行手順

### 環境構築

CUDA 12 用の cuDNN 9 をインストールします

```bash
!apt-get -y install cudnn9-cuda-12
```

asdf をクローンしてきます

```bash
!git clone https://github.com/asdf-vm/asdf.git ~/.asdf
```

asdf にパスを通します

```python
import os
os.environ['PATH'] = "/root/.asdf/shims:/root/.asdf/bin:" + os.environ['PATH']
```

asdf で Erlang 27 をインストールします

```bash
!asdf plugin add erlang
!asdf install erlang 27.1
!asdf global erlang 27.1
```

asdf で Elxiir 1.17 をインストールします

```bash
!asdf plugin add elixir
!asdf install elixir 1.17.3-otp-27
!asdf global elixir 1.17.3-otp-27
```

Livebook をインストールし、コマンドラインから呼び出せるようにします

```bash
!mix local.hex --force
!mix local.rebar --force
!mix escript.install hex livebook --force
!asdf reshim elixir
```

[ngrok](https://ngrok.com/) CLI をインストールします

ngrok を使うことで、 Colab 上の localhost で起動した Livebook に外部からアクセスできます

:::note info
ngrok にアカウントを作成し、認証トークンを準備しておきましょう

![スクリーンショット 2024-10-06 9.12.04.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/6d22c3db-0917-2d0e-bb12-5327145a1af6.png)![スクリーンショット 2024-10-06 9.12.33.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/2a940705-290b-b7cd-e8e5-92c40eb0bad4.png)

:::

```bash
!curl -s https://ngrok-agent.s3.amazonaws.com/ngrok.asc \
	| sudo tee /etc/apt/trusted.gpg.d/ngrok.asc >/dev/null \
	&& echo "deb https://ngrok-agent.s3.amazonaws.com buster main" \
	| sudo tee /etc/apt/sources.list.d/ngrok.list \
	&& sudo apt update \
	&& sudo apt install ngrok
```

パスワード用テキスト入力を用意し、 ngrok の認証トークンを入力します

```python
from getpass import getpass

token = getpass()
```

![スクリーンショット 2024-10-06 9.11.56.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/9c881dfe-494c-34ee-47e9-18e0eaf11b34.png)

ngrok CLI に認証トークンを設定します

```bash
!ngrok config add-authtoken "$token"
```

### Livebook の起動

ngrok を起動し、 Colab 上の localohost 8888 ポートを外部に公開します

```python
get_ipython().system_raw('ngrok http 8888 &')
!sleep 5s
```

ngrok が外部に公開している URL を取得します

```bash
!curl -s http://localhost:4040/api/tunnels | python3 -c "import sys, json; print(json.load(sys.stdin)['tunnels'][0]['public_url'])"
```

Livebook を 8888 ポートで起動します

```bash
!livebook server --port 8888
```

ngrok の公開 URL にアクセスすると、 Livebook の認証画面が表示されます

![スクリーンショット 2024-10-06 9.12.49.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/a264c79b-6fec-8a9c-6873-5341cc9a7d4f.png)

Livebook 起動時に表示される URL 末尾の認証トークンを入力してください（下画像の赤枠部分）

![スクリーンショット 2024-10-06 9.12.33.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/bf70f1a3-0f80-f7cf-685e-96d47381f151.png)

### 各種バックエンドの動作確認

新しいノートブックを開き、セットアップセルで以下のコードを実行します

```elixir
Mix.install(
  [
    {:benchee, "~> 1.3"},
    {:nx, "~> 0.9"},
    {:exla, "~> 0.9"},
    {:torchx, "~> 0.9"},
    {:evision, "~> 0.2"},
    {:kino, "~> 0.14"}
  ],
  system_env: [
    {"XLA_TARGET", "cuda12"},
    {"EXLA_TARGET", "cuda"},
    {"LIBTORCH_TARGET", "cu121"},
    {"EVISION_ENABLE_CUDA", "true"},
    {"EVISION_ENABLE_CONTRIB", "true"},
    {"EVISION_CUDA_VERSION", "12"},
    {"EVISION_CUDNN_VERSION", "9"}
  ]
)
```

各種バックエンドでテンソルを作ってみます

#### バイナリバックエンド（デフォルト）

```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]
>
```

#### EXLA バックエンド（GPU）

```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.736664782.2207121422.121493>
  [2.0, 4.0, 6.0]
>
```

標準出力には、以下のように cuDNN のロードが表示されます

```
00:23:04.701 [info] Loaded cuDNN version 90400
```

#### Torchx バックエンド（GPU）

```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]
>
```

#### Torchx バックエンド（CPU）

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

実行結果

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

#### Evision バックエンド（GPU）

```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]
>
```

## まとめ

Colab 上の cuDNN を更新することで、最新版の EXLA バックエンドや Evision バックエンドを使用することができました

Elixir で GPU を使った機械学習が捗りますね