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
フォーラムでも同じ現象になっている人がいるので修正を待っています。
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が使えるようになった当時から要望があったと思い、私もこれを期待していました。
すでにここまで使えるようになっているのはすごいですね。
今後のアップデートに期待!
カテゴリー