chore: align with AWS Neuron SDK 2.15.1 (#325)

* chore: align with AWS Neuron SDK 2.15.1

* chore: run all tests when setup.py is modified

.github/workflows/check_code_quality.yml

@@ -4,13 +4,15 @@ on:
   push:
     branches: [ main ]
     paths:
+      - "setup.py"
       - "optimum/neuron/**.py"
       - "tests/**.py"
       - "examples/**.py"
 
   pull_request:
     branches: [ main ]
     paths:
+      - "setup.py"
       - "optimum/neuron/**.py"
       - "tests/**.py"
       - "examples/**.py"

.github/workflows/test_common.yml

@@ -6,14 +6,16 @@ on:
   push:
     branches: [ main ]
     paths:
+      - "setup.py"
       - "optimum/neuron/**.py"
       - "tests/**.py"
       - "examples/**.py"
       - "tools/**.py"
-      
+
   pull_request:
     branches: [ main ]
     paths:
+      - "setup.py"
       - "optimum/neuron/**.py"
       - "tests/**.py"
       - "examples/**.py"

.github/workflows/test_inf1.yml

@@ -4,10 +4,12 @@ on:
   push:
     branches: [ main ]
     paths:
+      - "setup.py"
       - "optimum/**.py"
   pull_request:
     branches: [ main ]
     paths:
+      - "setup.py"
       - "optimum/**.py"
 
 concurrency:
@@ -26,7 +28,7 @@ jobs:
       - name: Install python dependencies
         run: |
           sudo apt install python3.8-venv -y
-          python3 -m venv aws_neuron_venv_pytorch 
+          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

.github/workflows/test_inf2.yml

@@ -4,10 +4,12 @@ on:
   push:
     branches: [ main ]
     paths:
+      - "setup.py"
       - "optimum/**.py"
   pull_request:
     branches: [ main ]
     paths:
+      - "setup.py"
       - "optimum/**.py"
 
 concurrency:

.github/workflows/test_trainium_common.yml

@@ -4,10 +4,12 @@ on:
   push:
     branches: [ main ]
     paths:
+      - "setup.py"
       - "optimum/**.py"
   pull_request:
     branches: [ main ]
     paths:
+      - "setup.py"
       - "optimum/**.py"
 
 concurrency:
@@ -35,5 +37,5 @@ jobs:
       - name: Run tests on Neuron cores
         run: |
           HF_TOKEN_OPTIMUM_NEURON_CI=${{ secrets.HF_TOKEN_OPTIMUM_NEURON_CI }} USE_VENV="false" pytest -m "is_trainium_test" $TESTS_TO_IGNORE_FLAGS tests
-      - name: Run staging tests on Neuron cores 
+      - name: Run staging tests on Neuron cores
         run: HUGGINGFACE_CO_STAGING=1 pytest -m "is_trainium_test and is_staging_test" $TESTS_TO_IGNORE_FLAGS tests -s

.github/workflows/test_trainium_distributed.yml

@@ -4,10 +4,12 @@ on:
   push:
     branches: [ main ]
     paths:
+      - "setup.py"
       - "optimum/neuron/distributed/**.py"
   pull_request:
     branches: [ main ]
     paths:
+      - "setup.py"
       - "optimum/neuron/distributed/**.py"
 
 concurrency:
@@ -17,7 +19,7 @@ concurrency:
 
 jobs:
   optimum-neuron-tests:
-    runs-on: [self-hosted, 16-aws-trn, 128-cpu, ci] 
+    runs-on: [self-hosted, 16-aws-trn, 128-cpu, ci]
     env:
       AWS_REGION: us-east-1
     steps:
@@ -32,4 +34,4 @@ jobs:
       - name: Run tests on Neuron cores
         run: |
           HF_TOKEN_OPTIMUM_NEURON_CI=${{ secrets.HF_TOKEN_OPTIMUM_NEURON_CI }} pytest -m "is_trainium_test" tests/distributed/
-  
+

setup.py

@@ -54,7 +54,7 @@
     "neuronx": [
         "wheel",
         "neuronx-cc==2.11.0.34",
-        "torch-neuronx==1.13.1.1.12.0",
+        "torch-neuronx==1.13.1.1.12.1",
         "transformers-neuronx==0.8.268",
         "torch==1.13.1.*",
         "torchvision==0.14.*",

text-generation-inference/Dockerfile

@@ -100,7 +100,7 @@ ENV PATH="/opt/bin/:/opt/aws/neuron/bin:${PATH}"
 
 RUN pip3 install \
     neuronx-cc==2.11.0.34 \
-    torch-neuronx==1.13.1.1.12.0 \
+    torch-neuronx==1.13.1.1.12.1 \
     transformers-neuronx==0.8.268 \
     --extra-index-url=https://pip.repos.neuron.amazonaws.com