Python version update (#1122)

* Switch default command for docker image

* Rename pythia paths docker file for clarity

* Update docker build to use python 3.10

* Update github workflows to use ubuntu 22.04 and python 3.10

* Bump pytorch library patch versions

* Add pytest-html for reasonably formatted test reports

* Fix build after torch and cuda version bump

* Fix apex install for newer version

1) This, empirically, works, as tested by running the build and kicking off training.
2) Apex documentation says it is incorrect syntax and deprecated.
3) It takes so long to compile that it is probably, all by itself, something that needs fixing.
4) I will probably pull the fused adamw out of apex.
5) It has been building for twenty minutes so I am going to go do something else.

* Fix pip version to ensure apex compilation remains good

* Fix unit test for evaluate

* Fix pip requirement

Prevents possible build issues with apex especially across divergent pip versions

* Update dockerfile to point to stripped-down apex repo

* Revert "Update dockerfile to point to stripped-down apex repo"

This reverts commit 40c7656.

* Update apex version in dockerfile

* Switch to downloading prebuilt apex wheel

* Clean up docker copy commands

* Have docker build conditionally get binaries or build apex

* Apply precommit

.github/workflows/pull_request.yml

@@ -4,12 +4,12 @@ on: [pull_request]
 
 jobs:
   pre-commit:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v2
       - uses: actions/setup-python@v4
         with:
-          python-version: 3.8
+          python-version: 3.10
           cache: "pip"
           cache-dependency-path: "**/requirements*.txt"
       # Need the right version of clang-format
@@ -24,7 +24,7 @@ jobs:
         uses: docker/build-push-action@v2
 
   update-documentation:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     steps:
       - uses: actions/checkout@v3
         with:

Dockerfile

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-FROM nvidia/cuda:11.7.1-devel-ubuntu20.04
+FROM nvidia/cuda:12.1.1-devel-ubuntu22.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 
@@ -21,20 +21,20 @@ LABEL org.opencontainers.image.version = "2.0"
 LABEL org.opencontainers.image.authors = "[email protected]"
 LABEL org.opencontainers.image.source = "https://www.github.com/eleutherai/gpt-neox"
 LABEL org.opencontainers.image.licenses = " Apache-2.0"
-LABEL org.opencontainers.image.base.name="docker.io/nvidia/cuda:11.7.1-devel-ubuntu20.04"
+LABEL org.opencontainers.image.base.name="docker.io/nvidia/cuda:12.1.1-devel-ubuntu22.04"
 
 #### System package (uses default Python 3 version in Ubuntu 20.04)
 RUN apt-get update -y && \
     apt-get install -y \
-    git python3.9 python3-dev libpython3-dev python3-pip sudo pdsh \
-    htop llvm-9-dev tmux zstd software-properties-common build-essential autotools-dev \
+    git python3-dev libpython3-dev python3-pip sudo pdsh \
+    htop tmux zstd software-properties-common build-essential autotools-dev \
     nfs-common pdsh cmake g++ gcc curl wget vim less unzip htop iftop iotop ca-certificates ssh \
     rsync iputils-ping net-tools libcupti-dev libmlx4-1 infiniband-diags ibutils ibverbs-utils \
     rdmacm-utils perftest rdma-core nano && \
     update-alternatives --install /usr/bin/python python /usr/bin/python3 1 && \
     update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 1 && \
-    pip install --upgrade pip && \
-    pip install gpustat
+    python -m pip install --upgrade pip && \
+    python -m pip install gpustat
 
 ### SSH
 RUN mkdir /var/run/sshd && \
@@ -88,24 +88,31 @@ RUN mkdir -p /home/mchorse/.ssh /job && \
     echo 'export LD_LIBRARY_PATH=/usr/local/lib:/usr/local/mpi/lib:/usr/local/mpi/lib64:$LD_LIBRARY_PATH' >> /home/mchorse/.bashrc
 
 #### Python packages
-RUN pip install torch==1.13.0+cu117 torchvision==0.14.0+cu117 torchaudio==0.13.0 --extra-index-url https://download.pytorch.org/whl/cu117 && pip cache purge
-COPY requirements/requirements.txt .
-COPY requirements/requirements-wandb.txt .
-COPY requirements/requirements-onebitadam.txt .
-COPY requirements/requirements-sparseattention.txt .
-COPY requirements/requirements-flashattention.txt .
-RUN pip install -r requirements.txt && pip install -r requirements-onebitadam.txt
-RUN pip install -r requirements-sparseattention.txt
-RUN pip install -r requirements-flashattention.txt
-RUN pip install -r requirements-wandb.txt
-RUN pip install protobuf==3.20.*
-RUN pip cache purge
+RUN python -m pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
+COPY requirements/* ./
+RUN python -m pip install --no-cache-dir -r requirements.txt && pip install -r requirements-onebitadam.txt
+RUN python -m pip install -r requirements-sparseattention.txt
+RUN python -m pip install -r requirements-flashattention.txt
+RUN python -m pip install -r requirements-wandb.txt
+RUN python -m pip install protobuf==3.20.*
+RUN python -m pip cache purge
 
 ## Install APEX
-RUN pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" git+https://github.com/NVIDIA/apex.git@a651e2c24ecf97cbf367fd3f330df36760e1c597
+# Detect the architecture and install Apex accordingly
+RUN ARCH=$(uname -m) && \
+    if [ "$ARCH" = "x86_64" ]; then \
+        wget https://github.com/segyges/not-nvidia-apex/releases/download/jan-2024/apex-0.1-cp310-cp310-linux_x86_64.zip && \
+        unzip ./apex-0.1-cp310-cp310-linux_x86_64.zip && \
+        python -m pip install ./apex-0.1-cp310-cp310-linux_x86_64.whl; \
+    else \
+    # Install Apex directly from source for other architectures
+        python -m pip install -r requirements-apex-pip.txt && \
+        python -m pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings --global-option=--cpp_ext --config-settings --global-option=--cuda_ext git+https://github.com/NVIDIA/apex.git@141bbf1cf362d4ca4d94f4284393e91dda5105a5; \
+    fi
 
 COPY megatron/fused_kernels/ megatron/fused_kernels
-RUN python megatron/fused_kernels/setup.py install
+WORKDIR /megatron/fused_kernels
+RUN python setup.py install
 
 # Clear staging
 RUN mkdir -p /tmp && chmod 0777 /tmp

configs/neox_arguments.md

@@ -1024,7 +1024,7 @@ Text Generation arguments
 
 - **prompt_end**: str
 
-    Default = 
+    Default =
 
 
     a single prompt's end. Defaults to newline
@@ -1066,7 +1066,7 @@ Text Generation arguments
 
 - **eval_results_prefix**: str
 
-    Default = 
+    Default =
 
     prefix to which to save evaluation results - final fp will be {eval_results_prefix}_eval_results_yy-mm-dd-HH-MM.json
 
@@ -1810,7 +1810,7 @@ Args for deepspeed config
 
     Default = None
 
-    
+
 
 
 
@@ -2110,4 +2110,3 @@ Args for deepspeed runner (deepspeed.launcher.runner).
     Default = None
 
     Adds a `--account` to the DeepSpeed launch command. In DeeperSpeed this is passed on to the SlurmLauncher as well. Sometimes necessary for cluster rules, or so I've heard.
-

docker-compose-dockerhub.yml

@@ -1,7 +1,7 @@
 version: '3'
 services:
   gpt-neox:
-    command: nvidia-smi -q --loop=10
+    command: nvidia-smi dmon
     image: leogao2/gpt-neox:main
     shm_size: 1g
     ulimits:

docker-compose.yml

@@ -1,7 +1,7 @@
 version: '3'
 services:
   gpt-neox:
-    command: nvidia-smi -q --loop=10
+    command: nvidia-smi dmon
     image: gpt-neox
     build:
       context: .

requirements/requirements-apex-pip.txt

@@ -0,0 +1 @@
+pip==23.3.2

requirements/requirements-dev.txt

@@ -4,4 +4,5 @@ pre-commit>=2.17.0
 pytest>=6.2.3
 pytest-cov>=2.11.1
 pytest-forked>=1.3.0
+pytest-html==4.1.1
 pytest-xdist

requirements/requirements-sparseattention.txt

@@ -1 +1 @@
-triton==2.0.0.dev20221202
+triton==2.1.0

tests/unit/test_launcher_scripts.py

@@ -1,12 +1,13 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 
 import pytest
-from tools.datasets import preprocess_data
-import train
+
+import eval
 import generate
-import evaluate
+import train
 from megatron.neox_arguments import NeoXArgs
-from tests.common import simulate_deepy_env, save_random_model
+from tests.common import save_random_model, simulate_deepy_env
+from tools.datasets import preprocess_data
 
 
 @pytest.fixture(
@@ -90,7 +91,7 @@ def test_evaluate(monkeypatch, tmpdir, tmp_path):
         "eval_tasks": ["lambada"],  # ["lambada", "hellaswag", "piqa", "sciq"],
         "eval_results_prefix": sample_output_file,
     }
-    evaluate.main(input_args=deepspeed_main_args, overwrite_values=evaluate_args)
+    eval.main(input_args=deepspeed_main_args, overwrite_values=evaluate_args)
 
 
 @pytest.mark.skip(