Merge branch 'main' into support_resize_token_embeddings_for_tp

huggingface · Sep 17, 2024 · 9a6c01e · 9a6c01e
2 parents d2498e6 + 3b381c0
commit 9a6c01e
Show file tree

Hide file tree

Showing 65 changed files with 1,771 additions and 305 deletions.
diff --git a/.github/workflows/inference_cache.yml → .github/workflows/inference_cache_llm.yml b/.github/workflows/inference_cache.yml → .github/workflows/inference_cache_llm.yml
@@ -1,4 +1,4 @@
-name: Optimum neuron inference cache builder
+name: Optimum neuron LLM inference cache builder
 
 on:
   workflow_dispatch:
@@ -12,7 +12,8 @@ concurrency:
 jobs:
   cache:
     name: Create optimum-neuron inference cache
-    runs-on: [self-hosted, 12-aws-inf2, 192-cpu, ci] # run the job on the newly created runner
+    runs-on:
+      group: aws-inf2-48xlarge
     env:
       AWS_REGION: us-east-1
     strategy:
@@ -38,13 +39,13 @@ jobs:
           EOF
           wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -
           sudo apt-get update -y
-          sudo apt-get install aws-neuronx-tools=2.17.1.0 aws-neuronx-runtime-lib=2.20.22.0-1b3ca6425 aws-neuronx-collectives=2.20.22.0-c101c322e  -y
+          sudo apt-get install aws-neuronx-tools=2.18.3.0 aws-neuronx-runtime-lib=2.21.41.0-fb1705f5f aws-neuronx-collectives=2.21.46.0-69b77134b  -y
           export PATH=/opt/aws/neuron/bin:$PATH
       - name: Checkout
         uses: actions/checkout@v4
       - name: Install python and create venv
         run: |
-          sudo apt install python3.8-venv python3-dev -y
+          sudo apt install python3-venv python3-dev -y
           python3 -m venv aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           python -m pip install -U pip

diff --git a/.github/workflows/inference_cache_stable_diffusion.yml b/.github/workflows/inference_cache_stable_diffusion.yml
@@ -0,0 +1,52 @@
+name: Optimum neuron SD inference cache builder
+
+on:
+  workflow_dispatch:
+  schedule:
+    # Schedule the workflow to run every Saturday at midnight UTC
+    - cron: '0 0 * * 6'
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+
+jobs:
+  cache:
+    name: Create optimum-neuron inference cache
+    runs-on:
+      group: aws-inf2-8xlarge
+    env:
+      AWS_REGION: us-east-1
+    strategy:
+      fail-fast: false
+      matrix:
+        config: [stable-diffusion]
+    steps:
+      - name: Install Neuron runtime
+        run: |
+          . /etc/os-release
+          sudo tee /etc/apt/sources.list.d/neuron.list > /dev/null <<EOF
+          deb https://apt.repos.neuron.amazonaws.com ${VERSION_CODENAME} main
+          EOF
+          wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -
+          sudo apt-get update -y
+          sudo apt-get install aws-neuronx-tools=2.18.3.0 aws-neuronx-runtime-lib=2.21.41.0-fb1705f5f aws-neuronx-collectives=2.21.46.0-69b77134b  -y
+          export PATH=/opt/aws/neuron/bin:$PATH
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Install python and create venv
+        run: |
+          sudo apt install python3-venv python3-dev -y
+          python3 -m venv aws_neuron_venv_pytorch
+          source aws_neuron_venv_pytorch/bin/activate
+          python -m pip install -U pip
+          python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com
+      - name: Install optimum neuron
+        run: |
+          source aws_neuron_venv_pytorch/bin/activate
+          python -m pip install .[neuronx,diffusers]
+      - name: Create cache for ${{matrix.config}} models
+        run: |
+          source aws_neuron_venv_pytorch/bin/activate
+          config_prefix_url=https://huggingface.co/aws-neuron/optimum-neuron-cache/raw/main/inference-cache-config
+          HF_TOKEN=${{secrets.HF_TOKEN_OPTIMUM_NEURON_CACHE}} \
+            python tools/auto_fill_inference_cache.py --config_file ${config_prefix_url}/${{matrix.config}}.json
diff --git a/.github/workflows/test_inf1_export.yml b/.github/workflows/test_inf1_export.yml
@@ -10,9 +10,8 @@ concurrency:
 jobs:
   do-the-job:
     name: Run INF1 export tests
-    runs-on: [self-hosted, 4-aws-inf1, 24-cpu, ci]
-    env:
-      AWS_REGION: us-east-1
+    runs-on:
+      group: aws-inf1-6xlarge
     steps:
       - name: Install Neuron runtime
         run: |
@@ -28,7 +27,7 @@ jobs:
         uses: actions/checkout@v2
       - name: Install system packages
         run: |
-          sudo apt install python3.8-venv python3-dev -y
+          sudo apt install python3-venv python3-dev -y
       - name: Install python packages
         run: |
           python3 -m venv aws_neuron_venv_pytorch

diff --git a/.github/workflows/test_inf1_full_export.yml b/.github/workflows/test_inf1_full_export.yml
@@ -10,9 +10,8 @@ concurrency:
 jobs:
   do-the-job:
     name: Run INF1 full export tests
-    runs-on: [self-hosted, 4-aws-inf1, 24-cpu, ci]
-    env:
-      AWS_REGION: us-east-1
+    runs-on:
+      group: aws-inf1-6xlarge
     steps:
       - name: Install Neuron runtime
         run: |
@@ -28,7 +27,7 @@ jobs:
         uses: actions/checkout@v2
       - name: Install system packages
         run: |
-          sudo apt install python3.8-venv python3-dev -y
+          sudo apt install python3-venv python3-dev -y
       - name: Install python packages
         run: |
           python3 -m venv aws_neuron_venv_pytorch

diff --git a/.github/workflows/test_inf1_inference.yml b/.github/workflows/test_inf1_inference.yml
@@ -10,9 +10,8 @@ concurrency:
 jobs:
   do-the-job:
     name: Run INF1 tests
-    runs-on: [self-hosted, 4-aws-inf1, 24-cpu, ci]
-    env:
-      AWS_REGION: us-east-1
+    runs-on:
+      group: aws-inf1-6xlarge
     steps:
       - name: Install Neuron runtime
         run: |
@@ -28,7 +27,7 @@ jobs:
         uses: actions/checkout@v2
       - name: Install system packages
         run: |
-          sudo apt install python3.8-venv python3-dev -y
+          sudo apt install python3-venv python3-dev -y
       - name: Install python packages
         run: |
           python3 -m venv aws_neuron_venv_pytorch

diff --git a/.github/workflows/test_inf1_pipelines.yml b/.github/workflows/test_inf1_pipelines.yml
@@ -10,9 +10,8 @@ concurrency:
 jobs:
   do-the-job:
     name: Run INF1 tests
-    runs-on: [self-hosted, 4-aws-inf1, 24-cpu, ci]
-    env:
-      AWS_REGION: us-east-1
+    runs-on:
+      group: aws-inf1-6xlarge
     steps:
       - name: Install Neuron runtime
         run: |
@@ -28,7 +27,7 @@ jobs:
         uses: actions/checkout@v2
       - name: Install system packages
         run: |
-          sudo apt install python3.8-venv python3-dev -y
+          sudo apt install python3-venv python3-dev -y
       - name: Install python packages
         run: |
           python3 -m venv aws_neuron_venv_pytorch

diff --git a/.github/workflows/test_inf2.yml b/.github/workflows/test_inf2.yml
@@ -6,11 +6,13 @@ on:
     paths:
       - "setup.py"
       - "optimum/**.py"
+      - ".github/workflows/test_inf2.yml"
   pull_request:
     branches: [ main ]
     paths:
       - "setup.py"
       - "optimum/**.py"
+      - ".github/workflows/test_inf2.yml"
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
@@ -19,9 +21,8 @@ concurrency:
 jobs:
   do-the-job:
     name: Run INF2 tests
-    runs-on: [self-hosted, 1-aws-inf2, 32-cpu, ci] # run the job on the newly created runner
-    env:
-      AWS_REGION: us-east-1
+    runs-on:
+      group: aws-inf2-8xlarge
     steps:
       - name: Install Neuron runtime
         run: |
@@ -37,7 +38,7 @@ jobs:
         uses: actions/checkout@v2
       - name: Install python dependencies
         run: |
-          sudo apt install python3.8-venv python3-dev -y
+          sudo apt install python3-venv python3-dev -y
           python3 -m venv aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           python -m pip install -U pip

diff --git a/.github/workflows/test_inf2_export.yml b/.github/workflows/test_inf2_export.yml
@@ -6,11 +6,13 @@ on:
     paths:
       - "setup.py"
       - "optimum/**.py"
+      - ".github/workflows/test_inf2_export.yml"
   pull_request:
     branches: [ main ]
     paths:
       - "setup.py"
       - "optimum/**.py"
+      - ".github/workflows/test_inf2_export.yml"
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
@@ -19,9 +21,8 @@ concurrency:
 jobs:
   do-the-job:
     name: Run INF2 export tests
-    runs-on: [self-hosted, 1-aws-inf2, 32-cpu, ci] # run the job on the newly created runner
-    env:
-      AWS_REGION: us-east-1
+    runs-on:
+      group: aws-inf2-8xlarge
     steps:
       - name: Install Neuron runtime
         run: |
@@ -37,7 +38,7 @@ jobs:
         uses: actions/checkout@v2
       - name: Install python dependencies
         run: |
-          sudo apt install python3.8-venv python3-dev -y
+          sudo apt install python3-venv python3-dev -y
           python3 -m venv aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           python -m pip install -U pip

diff --git a/.github/workflows/test_inf2_full_export.yml b/.github/workflows/test_inf2_full_export.yml
@@ -5,10 +5,12 @@ on:
     branches: [ main ]
     paths:
       - "optimum/exporters/neuron/*.py"
+      - ".github/workflows/test_inf2_full_export.yml"
   pull_request:
     branches: [ main ]
     paths:
       - "optimum/exporters/neuron/*.py"
+      - ".github/workflows/test_inf2_full_export.yml"
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
@@ -17,9 +19,8 @@ concurrency:
 jobs:
   do-the-job:
     name: Run INF2 full export tests
-    runs-on: [self-hosted, 1-aws-inf2, 32-cpu, ci] # run the job on the newly created runner
-    env:
-      AWS_REGION: us-east-1
+    runs-on:
+      group: aws-inf2-8xlarge
     steps:
       - name: Install Neuron runtime
         run: |
@@ -35,7 +36,7 @@ jobs:
         uses: actions/checkout@v2
       - name: Install python dependencies
         run: |
-          sudo apt install python3.8-venv python3-dev -y
+          sudo apt install python3-venv python3-dev -y
           python3 -m venv aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           python -m pip install -U pip

diff --git a/.github/workflows/test_inf2_inference.yml b/.github/workflows/test_inf2_inference.yml
@@ -6,11 +6,13 @@ on:
     paths:
       - "setup.py"
       - "optimum/**.py"
+      - ".github/workflows/test_inf2_inference.yml"
   pull_request:
     branches: [ main ]
     paths:
       - "setup.py"
       - "optimum/**.py"
+      - ".github/workflows/test_inf2_inference.yml"
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
@@ -19,9 +21,8 @@ concurrency:
 jobs:
   do-the-job:
     name: Run INF2 tests
-    runs-on: [self-hosted, 1-aws-inf2, 32-cpu, ci] # run the job on the newly created runner
-    env:
-      AWS_REGION: us-east-1
+    runs-on:
+      group: aws-inf2-8xlarge
     steps:
       - name: Install Neuron runtime
         run: |
@@ -37,7 +38,7 @@ jobs:
         uses: actions/checkout@v2
       - name: Install python dependencies
         run: |
-          sudo apt install python3.8-venv python3-dev -y
+          sudo apt install python3-venv python3-dev -y
           python3 -m venv aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           python -m pip install -U pip

diff --git a/.github/workflows/test_inf2_tgi.yml b/.github/workflows/test_inf2_tgi.yml
@@ -7,12 +7,14 @@ on:
       - "setup.py"
       - "optimum/**.py"
       - "text-generation-inference/**"
+      - ".github/workflows/test_inf2_tgi.yml"
   pull_request:
     branches: [ main ]
     paths:
       - "setup.py"
       - "optimum/**.py"
       - "text-generation-inference/**"
+      - ".github/workflows/test_inf2_tgi.yml"
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
@@ -21,9 +23,8 @@ concurrency:
 jobs:
   do-the-job:
     name: Run TGI tests
-    runs-on: [self-hosted, 1-aws-inf2, 32-cpu, ci] # run the job on the newly created runner
-    env:
-      AWS_REGION: us-east-1
+    runs-on:
+      group: aws-inf2-8xlarge
     steps:
       - name: Install Neuron runtime
         run: |
@@ -39,7 +40,7 @@ jobs:
         uses: actions/checkout@v2
       - name: Install python and create venv
         run: |
-          sudo apt install python3.8-venv python3-dev -y
+          sudo apt install python3-venv python3-dev -y
           python3 -m venv aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           python -m pip install -U pip

diff --git a/.github/workflows/test_trainium_common.yml b/.github/workflows/test_trainium_common.yml
@@ -6,11 +6,13 @@ on:
     paths:
       - "setup.py"
       - "optimum/**.py"
+      - ".github/workflows/test_trainium_common.yml"
   pull_request:
     branches: [ main ]
     paths:
       - "setup.py"
       - "optimum/**.py"
+      - ".github/workflows/test_trainium_common.yml"
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
@@ -19,9 +21,9 @@ concurrency:
 jobs:
   optimum-neuron-tests:
     name: Run common tests on Trainium 1
-    runs-on: [self-hosted, 16-aws-trn, 128-cpu, ci]
+    runs-on:
+      group: aws-trn1-32xlarge
     env:
-      AWS_REGION: us-east-1
       TESTS_TO_IGNORE_FLAGS: --ignore tests/distributed/ --ignore tests/test_examples.py
     steps:
       - name: Install Neuron runtime
@@ -38,7 +40,7 @@ jobs:
         uses: actions/checkout@v2
       - name: Install python dependencies
         run: |
-          sudo apt install python3.8-venv python3-dev -y
+          sudo apt install python3-venv python3-dev -y
           python3 -m venv aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           python -m pip install -U pip