使用预装环境的docker镜像启动pytorch多gpu训练只需三步：宿主机配置nvidia驱动与nvidia-container-toolkit并验证；选用pytorch/pytorch等官方多卡镜像；启动时用--gpus指定gpu资源并在代码中通过dataparallel或ddp启用多卡。
直接用预装好环境的 Docker 镜像启动 PyTorch + 多 GPU，比手动配驱动、CUDA、PyTorch快得多。核心就三步：宿主机准备好 GPU 支持、选对镜像、启动时声明 GPU 资源。
宿主机必须先装好 NVIDIA 驱动和容器工具链
容器本身不自带驱动，它靠宿主机提供。所以第一步不是拉镜像，而是确认底层已就绪：
- 运行 nvidia-smi，能看到 GPU 型号和驱动版本（比如 535.129.03），说明驱动已装
- 检查驱动与目标 CUDA 版本是否兼容：CUDA 12.x 要求驱动 ≥ 525；CUDA 11.x 要求 ≥ 450
- 安装 nvidia-container-toolkit（不是旧版 nvidia-docker），并重启 docker 服务
- 验证工具链：运行 docker run --rm --gpus 1 nvidia/cuda:12.2-base nvidia-smi，能输出 GPU 信息即成功
选用支持多卡的 PyTorch 镜像
别从 scratch 写 Dockerfile，直接用社区维护好的镜像，省去版本冲突风险：
Docker Image Puller
下载
使用 Python 脚本下载 Docker 镜像并打包为 tar 文件。用户说“下载镜像”“拉取镜像”“pull xxx”“帮我下载 nginx 镜像”等时触发。支持 SOCKS5 代理直连（推荐）和国内镜像站加速两种模式。
- 官方 pytorch/pytorch:2.8-cuda12.1-cudnn8-runtime：生产稳定，含完整 CUDA 运行时
- Jupyter 场景用 jupyter/pytorch-notebook:cuda12：自带 Jupyter Lab 和常用数据科学库
- 开发调试选 pytorch/pytorch:2.8-cuda12.1-devel：含编译工具，适合需要自定义算子的场景
- 国内拉取慢？可加镜像前缀，如用 registry.cn-hangzhou.aliyuncs.com/pytorch/pytorch:2.8-cuda12.1-runtime
启动容器并启用全部 GPU
启动命令里明确告诉 Docker 要用哪些 GPU，否则默认只能访问 CPU：
- 用所有 GPU：docker run --gpus all -p 8888:8888 pytorch/pytorch:2.8-cuda12.1-runtime
- 指定使用前两块：--gpus '"device=0,1"'（注意引号嵌套）
- 限制显存用量（防 OOM）：--gpus '"device=0,1" --shm-size=8gb'
- 进容器后，在 Python 中运行 torch.cuda.device_count() 应返回对应数量，torch.cuda.is_available() 必须为 True
在容器内写代码调用多 GPU
PyTorch 本身不自动并行，需在代码中显式启用：
- 简单数据并行（适合单机多卡）：model = torch.nn.DataParallel(model).cuda()
- 推荐分布式训练（DDP）：torch.distributed.run --nproc_per_node=4 train.py，脚本内初始化 init_process_group(backend="nccl")
- 注意：DDP 要求每个进程绑定到单独 GPU，用 torch.cuda.set_device(rank) 显式指定
- 验证是否生效：训练时观察 nvidia-smi，各卡的 GPU-Util 应同步有负载，且显存占用均衡
