Soyukke.Dev

2021-04-05

Flux.jl(CUDA.jl)をcuda on wsl2で動かす

Windows Insider Previewではwsl2からGPUを使えるようになっています。

pytorchやTensorflowを使ったデモに関する記事は存在しますが、Flux.jl、CUDA.jlを動かしている記事がないので書きます。

私の環境ではCUDA.jlでSegmentation Faultが発生してしまいまともに使うことができませんでしたが、原因を調査して動かすことができたので共有します。

Windowsで動かしていた学習コードを動かすことができることを確認できました。

環境 (2021/03/31)

  • Windows 10 Home Insider Preview 21343.1000

  • Geforce Game Ready ドライバー 470.14

  • Geforce RTX 3070

ドライバー470.14の問題

これは一部のユーザだけの問題なのかnvmlが動きません。 そのため、wsl2からnvidia-smiを実行しても以下のエラーが表示されます。

NVIDIA-SMI has failed because it couldn't communicate with the NVIDIA driver. Make sure that the latest NVIDIA driver is installed and running.

Failed to properly shut down NVML: Driver Not Loaded

フォーラムでも同じ現象になっている人がいるので修正を待っています。

https://forums.developer.nvidia.com/t/docker-and-nvidia-smi-not-working-with-clean-install-on-driver-470-14-and-insider-preview-build-21343-ubuntu-20-04/173483

nvml自体は最近サポートされているようなので、動くのが正常だと思います。

pytorchを使ったcnnの学習デモを実行してみましたが、正常に動いていることを確認したのでnvml以外は問題ないようです。

JuliaのCUDA.jlがnvmlを利用しているのでsegmentation faultとなり、まともに使うことができないのです。

例えば以下のコードを実行するとJuliaが落ちます。

NVMLの関数を実行するために初期化しているのですが、何らかの理由で初期化できないのでdriver_versionなどが動きません。

using CUDA
# これでsegumentation fault
NVML.driver_version()

ちなみに、一部の演算の前でバージョンチェックを行っている実装があるので、演算の種類によっては落ちたり落ちなかったりします。

julia> x = rand(3)
3-element Vector{Float64}:
 0.6989937796473533
 0.2125750077973394
 0.6686848527807518

julia> y = cu(x)
3-element CuArray{Float32, 1}:
 0.6989938
 0.212575
 0.66868484

julia> y' * y
0.9809199f0

julia> y * y'
ERROR: 
signal (11): Segmentation fault

問題の回避

解決はできませんが、回避はできます。 方法は単純でNVML.driver_version()を実行しないようにします。

https://github.com/JuliaGPU/CUDA.jl/compare/v2.6.2...Soyukke:runnable-wsl-cuda

応急処置的な方法ですが、これで通常通り使用することができました。 テストの実行などはしていないのですべて期待通りに動いているかは未確認です。

Flux.jlを実行する

Flux.jlに依存するCUDA.jlを修正したCUDA.jlに変更した後に、小さいコードでGPU実行してみました。

using Flux

model = Chain(
    Dense(3, 3, relu),
    GRU(3, 3),
    Dense(3, 3)
) |> gpu

data = [(gpu(rand(3, 3)),)]
opt = ADAM()
ps = params(model)

loss(x) = sum(x)

Flux.train!(loss, ps, data, opt)

モデル定義から勾配計算まで問題なしに実行することができました。

おわり

nvmlの問題はありましたが、GPUを使った学習自体には影響がないようでまだ開発版とはいえ十分実用に耐えうると感じました。

wslでのGPUサポート自体はwslが使えるようになった当時から要望があったと思い、私もこれを期待していました。

すでにここまで使えるようになっているのはすごいですね。

今後のアップデートに期待！

カテゴリー