Update Dockerfile

Dockerfile

@@ -10,4 +10,21 @@ ENV DATASETS_CACHE=$HOME/.cache/huggingface/datasets \
 ENV TOKENIZERS_PARALLELISM=false \
   PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 \
   FLASH_ATTENTION=1 \
-  ACCELERATE_USE_DEEPSPEED=true
+  ACCELERATE_USE_DEEPSPEED=true
+
+ENV NCCL_TIMEOUT=1800000 \
+  NCCL_BLOCKING_WAIT=1 \
+  NCCL_ASYNC_ERROR_HANDLING=1 \
+  NCCL_DEBUG=WARN \
+  NCCL_P2P_DISABLE=1 \
+  NCCL_IB_DISABLE=1 \
+  NCCL_NET_GDR_LEVEL=0 \
+  NCCL_SOCKET_IFNAME=lo \
+  NCCL_TREE_THRESHOLD=0
+
+ENV TORCH_NCCL_BLOCKING_WAIT=1 \
+  TORCH_NCCL_ASYNC_ERROR_HANDLING=1 \
+  TORCH_DISTRIBUTED_DEBUG=DETAIL
+
+ENV CUDA_LAUNCH_BLOCKING=0 \
+  CUDNN_BENCHMARK=1