Use AWS Neuron SDK 2.25

.github/actions/install_neuronx_runtime/action.yml

@@ -12,5 +12,5 @@ runs:
           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.24.54.0 aws-neuronx-runtime-lib=2.26.42.0-2ff3b5c7d aws-neuronx-collectives=2.26.43.0-47cc904ea -y
+          sudo apt-get install aws-neuronx-tools=2.25.145.0 aws-neuronx-runtime-lib=2.27.23.0-8deec4dbf aws-neuronx-collectives=2.27.34.0-ec8cd5e8b -y
           export PATH=/opt/aws/neuron/bin:$PATH

infrastructure/ami/hcl2-files/variables.pkr.hcl

@@ -10,7 +10,7 @@ variable "instance_type" {
 }
 
 variable "source_ami" {
-  default     = "ami-0ffd183ece0ca0475"
+  default     = "ami-0aea5bf206fcd4a09"
   description = "Base Image"
   type        = string
   /*

optimum/neuron/version.py

@@ -14,4 +14,4 @@
 
 __version__ = "0.3.1.dev0"
 
-__sdk_version__ = "2.24.0"
+__sdk_version__ = "2.25.0"

pyproject.toml

@@ -98,12 +98,12 @@ neuron = [
 ]
 neuronx = [
     "wheel",
-    "neuronx-cc==2.19.8089.0",
-    "torch-neuronx==2.7.0.2.8.6734+ac864f72",
+    "neuronx-cc==2.20.9961.0+0acef03a",
+    "torch-neuronx==2.7.0.2.9.9357+08e1f40d",
     "torch==2.7.0.*",
     "torchvision==0.22.*",
-    "neuronx_distributed==0.13.14393",
-    "libneuronxla==2.2.4410.0",
+    "neuronx_distributed==0.14.18461+9ac233f2",
+    "libneuronxla==2.2.8201.0+f46ac1ef",
 ]
 diffusers = [
     "diffusers>=0.31.0, <=0.34.0",

tests/cli/test_export_cli.py

@@ -22,6 +22,8 @@
 from optimum.neuron.utils.testing_utils import is_inferentia_test, requires_neuronx
 from optimum.utils import logging
 
+from ..helpers import skip_for_sdk
+
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
@@ -200,6 +202,7 @@ def test_pixart(self):
                     check=True,
                 )
 
+    @skip_for_sdk(["2.25.0"])
     @requires_neuronx
     def test_flux_tp2(self):
         model_ids = ["hf-internal-testing/tiny-flux-pipe-gated-silu"]

tests/helpers.py

@@ -0,0 +1,18 @@
+import functools
+
+import pytest
+
+from optimum.neuron.version import __sdk_version__
+
+
+def skip_for_sdk(sdk_versions):
+    def skip_for_sdk_decorator(test):
+        @functools.wraps(test)
+        def test_wrapper(*args, **kwargs):
+            if __sdk_version__ in sdk_versions:
+                pytest.skip(f"Fails with Neuron SDK {__sdk_version__}")
+            test(*args, **kwargs)
+
+        return test_wrapper
+
+    return skip_for_sdk_decorator

tests/inference/test_flux.py

@@ -27,7 +27,10 @@
 from optimum.neuron.utils.testing_utils import is_inferentia_test, requires_neuronx
 from optimum.utils.testing_utils import require_diffusers
 
+from ..helpers import skip_for_sdk
 
+
+@skip_for_sdk(["2.25.0"])
 @pytest.mark.order(0)
 @is_inferentia_test
 @requires_neuronx
@@ -48,6 +51,7 @@ def test_flux_txt2img(neuron_flux_tp2_path):
     assert isinstance(image, PIL.Image.Image)
 
 
+@skip_for_sdk(["2.25.0"])
 @pytest.mark.order(0)
 @is_inferentia_test
 @requires_neuronx