deepops/mpi-nccl-test

cuda:11.6.0-ubuntu20.04 镜像文档

镜像概述

本镜像基于Ubuntu 20.04操作系统构建，集成了CUDA 11.6.0、TensorFlow 2.7.0、NCCL 2.11.4-1+cuda11.6、OpenMPI 4.1.3和MOFED 5.5-1.0.3.2等关键组件，旨在提供一个开箱即用的GPU加速深度学习与高性能计算环境。

核心功能与特性

核心组件版本

• CUDA: 11.6.0 - 提供GPU计算能力支持，为深度学习和科学计算任务提供硬件加速
• TensorFlow: 2.7.0 - 主流深度学习框架，支持神经网络构建、训练与推理
• NCCL: 2.11.4-1+cuda11.6 - NVIDIA Collective Communications Library，优化多GPU间通信效率
• OpenMPI: 4.1.3 - 高性能消息传递接口，支持分布式计算任务的并行执行
• MOFED: 5.5-1.0.3.2 - Mellanox OpenFabrics Enterprise Distribution，提供高性能网络适配器驱动与工具

特性

• 基于Ubuntu 20.04稳定版，软件生态完善
• 预配置GPU加速环境，无需手动安装CUDA及相关依赖
• 集成分布式训练所需的通信库（NCCL+OpenMPI）
• 支持高性能网络通信（MOFED），适用于多节点集群环境

使用场景与适用范围

• 深度学习模型开发与训练（特别是基于TensorFlow的项目）
• 多GPU分布式训练任务
• 高性能计算（HPC）应用
• 需要GPU加速的科学计算与数据分析
• 多节点集群环境下的并行计算任务

使用方法与配置说明

前置要求

• 主机需配备NVIDIA GPU，且已安装NVIDIA驱动（版本需兼容CUDA 11.6.0）
• 已安装Docker和nvidia-docker2（或nvidia-container-toolkit）

基本使用（docker run）

bash
docker run --gpus all -it --rm cuda:11.6.0-ubuntu20.04 /bin/bash

参数说明

• --gpus all: 允许容器访问所有GPU设备（也可指定具体GPU，如--gpus "device=0,1"）
• -it: 交互式终端模式
• --rm: 容器退出后自动删除

挂载数据卷（持久化存储）

bash
docker run --gpus all -it --rm -v /host/path:/container/path cuda:11.6.0-ubuntu20.04 /bin/bash

验证环境

容器启动后，可通过以下命令验证关键组件是否正常工作：

bash
# 验证CUDA版本
nvcc --version

# 验证TensorFlow
python -c "import tensorflow as tf; print('TensorFlow version:', tf.__version__); print('GPU available:', tf.test.is_gpu_available())"

# 验证NCCL
nccl-tests (需额外安装nccl-tests包)

# 验证OpenMPI
mpirun --version

# 验证MOFED
ibv_devinfo

分布式训练示例

使用OpenMPI启动分布式TensorFlow训练任务：

bash
mpirun -np 2 --allow-run-as-root python your_training_script.py

注意事项

• 确保主机NVIDIA驱动版本与CUDA 11.6.0兼容（推荐驱动版本≥450.80.02）
• 多节点分布式训练需配置网络互通，并确保MOFED驱动与硬件适配
• 如需扩展镜像功能，可基于此镜像构建新的Dockerfile，添加额外依赖包

维护与更新

• 本镜像基于官方CUDA基础镜像构建，核心组件版本固定
• 如需更新组件版本，建议基于官方最新CUDA镜像重新构建