Add entrypoint in docker file and update document

.github/workflows/unit-tests.yaml

@@ -17,7 +17,7 @@ jobs:
         - torch: "1.14"
           nvcr: 22.12-py3
           dir: torch1
-        # 2.1.0a0+fe05266f
+        # 2.1.0a0+32f93b1
         - torch: "2.1"
           nvcr: 23.10-py3
           dir: torch2
@@ -57,6 +57,10 @@ jobs:
         export LD_PRELOAD="/usr/local/lib/libmsamp_dist.so:/usr/local/lib/libnccl.so:${LD_PRELOAD}"
         cd ${{ matrix.dir }}/
         python3 setup.py test
+    - name: Clean repository
+      if: always()
+      run: |
+        rm -rf ${{ matrix.dir }}/
     # - name: Report coverage results
     #   run: |
     #     bash <(curl -s https://codecov.io/bash)

dockerfile/entrypoint.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+
+ldconfig
+
+exec "$@"

dockerfile/torch1.14-cuda11.8.dockerfile

@@ -57,3 +57,8 @@ RUN python3 -m pip install . && \
     make postinstall
 
 ENV LD_PRELOAD="/usr/local/lib/libmsamp_dist.so:/usr/local/lib/libnccl.so:${LD_PRELOAD}"
+
+# Set up entrypoint
+COPY dockerfile/entrypoint.sh /entrypoint.sh
+RUN chmod +x /entrypoint.sh
+ENTRYPOINT ["/entrypoint.sh"]

dockerfile/torch2.1-cuda12.2.dockerfile

@@ -1,11 +1,11 @@
 FROM nvcr.io/nvidia/pytorch:23.10-py3
 
 # Ubuntu: 22.04
-# Python: 3.8
+# Python: 3.10
 # CUDA: 12.2.0
 # cuDNN: 8.9.5
 # NCCL: v2.16.2-1 + FP8 Support
-# PyTorch: 2.1.0a0+fe05266f
+# PyTorch: 2.1.0a0+32f93b1
 
 LABEL maintainer="MS-AMP"
 
@@ -57,3 +57,8 @@ RUN python3 -m pip install . && \
     make postinstall
 
 ENV LD_PRELOAD="/usr/local/lib/libmsamp_dist.so:/usr/local/lib/libnccl.so:${LD_PRELOAD}"
+
+# Set up entrypoint
+COPY dockerfile/entrypoint.sh /entrypoint.sh
+RUN chmod +x /entrypoint.sh
+ENTRYPOINT ["/entrypoint.sh"]

docs/getting-started/installation.mdx

@@ -18,7 +18,7 @@ Here're the system requirements for MS-AMP.
 * CUDA version 11 or later (which can be checked by running `nvcc --version`).
 * PyTorch version 1.14 or later (which can be checked by running `python -c "import torch; print(torch.__version__)"`).
 
-You can try MS-AMP in two ways: Using Docker or installing from source:  
+You can try MS-AMP in two ways: Using Docker or installing from source.  
 
 * Using Docker is a convenient way to get started with MS-AMP. You can use the pre-built Docker image to quickly set up an environment for running MS-AMP.  
 * On the other hand, installing from source gives you more control over the installation process and allows you to customize the installation to your needs.
@@ -28,8 +28,8 @@ You can try MS-AMP in two ways: Using Docker or installing from source:
 You can try the latest MS-AMP Docker container with the following commands:
 
 ```bash
-sudo docker run -it -d --name=msampcu121 --privileged --net=host --ipc=host --gpus=all -v /:/hostroot ghcr.io/azure/msamp:main-cuda12.1 bash
-sudo docker exec -it msampcu121 bash
+sudo docker run -it -d --name=msampcu122 --privileged --net=host --ipc=host --gpus=all -v /:/hostroot ghcr.io/azure/msamp:main-cuda12.2 bash
+sudo docker exec -it msampcu122 bash
 ```
 
 MS-AMP is pre-installed in Docker container and you can verify it by running:
@@ -46,7 +46,7 @@ We strongly recommend using [PyTorch NGC Container](https://catalog.ngc.nvidia.c
 For example, to start PyTorch 2.1 container, run the following command:
 
 ```bash
-sudo docker run -it -d --name=msamp --privileged --net=host --ipc=host --gpus=all nvcr.io/nvidia/pytorch:23.04-py3 bash
+sudo docker run -it -d --name=msamp --privileged --net=host --ipc=host --gpus=all nvcr.io/nvidia/pytorch:23.10-py3 bash
 sudo docker exec -it msamp bash
 ```