WSL2 Docker 启用 GPU CUDA 支持

  • WSL2 下 Docker 的 GPU CUDA 支持

  • 资料来源:

    https://zhuanlan.zhihu.com/p/408403790
    https://forums.developer.nvidia.com/t/guide-to-run-cuda-wsl-docker-with-latest-versions-21382-windows-build-470-14-nvidia/178365/13

  • 更新

    1
    2
    2021.10.06 初始
    2022.03.05 更新一点驱动内容

导语

WSL2 的 GPU 透传已经内测很久了,因为需要进入预览版通道,主力机一直没敢上,直到昨天 win11 发布,看更新才知道 win11 正式版已经带了 wsl2 的 gpu 透传了,还顺带了 WSLg.

WSLg 的体验等到下一篇博文,这一篇是 wsl2 下 cuda 配置.

正文

前置:

  • 安装启用 wsl2 ,内核版本 5.4
  • docker for windows,4.0 版本
  • 目前是 win11 正式版,无法升级 win11 也可等 10 月中旬的 win10 21h2. win10 21h2 早就是 win10 默认版本了.

安装:

  • CUDA on Windows Subsystem for Linux (WSL) 下载适用于 wsl 的显卡驱动,一般人手上的应该都是 GEFORCE 而不是 QUADRO,点击 GEFORCE DRIVER. 到 NV 官网下载正式版驱动 51x.x 版本 带 WHQL 标记即可.
  • 更新驱动并重启.这一步软件会自动安装 CUDA DirectML DirectX

完毕 (???),😂,确实到这就够了.

启动 docker 镜像时别忘了 --gpus all,其他什么都不需要加.

验证

参考 Guide to run CUDA + WSL + Docker with latest versions (21382 Windows build + 470.14 Nvidia)

运行

1
docker run --gpus all nvcr.io/nvidia/k8s/cuda-sample:nbody -gpu -benchmark

得到

1
2
3
4
5
6
7
8
9
10
> Windowed mode
> Simulation data stored in video memory
> Single precision floating point simulation
> 1 Devices used for simulation
GPU Device 0: "Turing" with compute capability 7.5

> Compute 7.5 CUDA device: [NVIDIA GeForce RTX 2060]
30720 bodies, total time for 10 iterations: 65.960 ms
= 143.074 billion interactions per second
= 2861.488 single-precision GFLOP/s at 20 flops per interaction

5步搭建wsl2+cuda+docker解决windows深度学习开发问题 提到了 docker for windows 下只有 cuda-sample:nbody 能跑,其他的 cuda torch tf之类的镜像都检测不到 gpu. 这个 bug 似乎是已经修复了.

运行同样的镜像

1
docker run --name mytorch --gpus all -it pytorch/pytorch:1.9.0-cuda10.2-cudnn7-runtime bash

进入 py 交互模式,得到的结果完全正常.

1
2
3
4
5
6
7
8
9
10
Python 3.7.10 (default, Feb 26 2021, 18:47:35)
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.zeros([1]).cuda()
tensor([0.], device='cuda:0')
>>> torch.version.cuda
'10.2'
>>> torch.backends.cudnn.version()
7605

测试 tensorflow

1
docker run --gpus all -it tensorflow/tensorflow:2.2.3-gpu-py3-jupyter bash

之后直接 vscode 进入容器运行

1
2
3
4
5
6
7
8
9
10
#!/usr/bin/python3

import tensorflow as tf

a = tf.test.is_built_with_cuda() # 判断CUDA是否可以用
b = tf.test.is_gpu_available(cuda_only=False, min_cuda_compute_capability=None) # GPU 支持
c = tf.test.gpu_device_name() # GPU 设备
print(a)
print(b)
print(c)

gpu 能正常检测到,显存等都正常.

尾巴

升级 win11 还算正常,搜索栏还有闪退 bug.

看资料 wsl2 下 gpu 训练性能和 linux 原生接近,远比 windows 下要好.还未测试.