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

はじめに

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

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

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

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

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

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

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

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

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

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 がリリースされました

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

EXLA の2024年10月最新版

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

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

XLA の最新版は 0.8.0 です

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

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

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

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

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

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

公式の紹介にある通り、最も簡単に bazel をインストールする手段として、 asdf が使用できます

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

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

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

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

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

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 が最新版です

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

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

Evision の2024年10月最新版

Evision は 0.2.10 が最新です

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

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

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

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

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 のインストール方法はこちら

最終的な実行手順

環境構築

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

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

asdf をクローンしてきます

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

asdf にパスを通します

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

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

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

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

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

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

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

ngrok CLI をインストールします

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

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

!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 の認証トークンを入力します

from getpass import getpass

token = getpass()

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

!ngrok config add-authtoken "$token"

Livebook の起動

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

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

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

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

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

!livebook server --port 8888

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

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

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

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

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

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

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

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

EXLA バックエンド（GPU）

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

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

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

Torchx バックエンド（GPU）

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

Torchx バックエンド（CPU）

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

実行結果

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

Evision バックエンド（GPU）

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

まとめ

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

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