add env in Dockerfile

finetune/Dockerfile

@@ -18,6 +18,11 @@ COPY id_rsa.pub /root/.ssh/id_rsa.pub
 RUN <<EOT
 #!/bin/bash
 apt-get update
+# 更新包管理器，并安装 RDMA 和 InfiniBand 相关的库和工具
+# rdma-core: 包含 RDMA 驱动和库
+# ibverbs-utils: 提供用于检查和配置 InfiniBand 设备的工具
+# perftest: 提供用于测试 InfiniBand 性能的工具
+apt-get install -y rdma-core ibverbs-utils perftest
 apt-get install -y wget curl htop jq vim bash libaio-dev build-essential openssh-server python3 python3-pip bzip2
 apt-get install -y --no-install-recommends software-properties-common build-essential autotools-dev nfs-common pdsh cmake g++ gcc curl wget vim tmux emacs less unzip htop iftop iotop ca-certificates openssh-client openssh-server rsync iputils-ping net-tools sudo llvm-dev re2c
 add-apt-repository ppa:git-core/ppa -y
@@ -439,5 +444,17 @@ pip3 install peft tiktoken \
     zstandard -i https://pypi.org/simple/ --trusted-host pypi.org
 EOT
 
+ARG DEEPSPEED_TRAIN='/data/train_data'
+ENV DEEPSPEED_TRAIN=DEEPSPEED_TRAIN
+ARG DEEPSPEED_VALIDATION='/data/validation_data'
+ENV DEEPSPEED_VALIDATION=DEEPSPEED_VALIDATION
+ARG NCCL_SOCKET_IFNAME='eth0'
+
+# RUN echo 'export CUDA_HOME=/usr/local/cuda' >> ~/.bashrc && \
+#     echo 'export PATH=${CUDA_HOME}/bin:${PATH}' >> ~/.bashrc && \
+#     echo 'export CUTLASS_PATH=/opt/cutlass' >> ~/.bashrc && \
+#     echo 'export PATH=/opt/conda/bin:$PATH' >> ~/.bashrc && \
+#     echo "source activate ${CONDA_ENV_NAME}" > ~/.bashrc
+
 # CMD ["/usr/sbin/sshd", "-D"]
 CMD ["/bin/bash", "-c", "/usr/sbin/sshd -D & while true; do sleep 1000; done"]