使用docker搭建分布式pytorch+cuda训练模型

# 基于具有CUDA支持的PyTorch镜像
FROM pytorch/pytorch:1.8.0-cuda11.1-cudnn8-devel
 
# 安装NCCL
RUN apt-get update && \
    apt-get install -y --no-install-recommends libnccl2=2.15.0-1+cuda11.1 libnccl-dev=2.15.0-1+cuda11.1 && \
    apt-get clean && \
    ldconfig
 
# 安装其他依赖
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
    git \
    vim \
    && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*
 
# 设置环境变量
ENV NCCL_VERSION 2.15.0-1+cuda11.1
ENV NCCL_HOME /usr/lib/x86_64-linux-gnu
ENV LD_LIBRARY_PATH $NCCL_HOME:$LD_LIBRARY_PATH
ENV PATH $NCCL_HOME/bin:$PATH
ENV NCCL_SOCKET_IFNAME eth0
 
# 复制NCCL库
COPY --from=nvidia/nccl-2:2.15.0-1+cuda11.1 /nccl /usr/lib/x86_64-linux-gnu
 
# 设置工作目录
WORKDIR /workspace

这段代码示例修正了原始代码中的一些问题，并提供了一个更为合理的NCCL安装方法。它使用了apt-get进行安装，并通过COPY --from指令复制了NCCL的二进制文件，这样做可以避免在构建过程中直接下载NCCL，从而加快构建速度并减少构建的复杂性。