Bump to optimum 1.17 - Adapt to optimum exporter refactoring

.github/workflows/test_inf1_export.yml

@@ -1,4 +1,4 @@
-name: Optimum neuron / Test INF1 export
+name: Optimum neuron / Test INF1 partial export
 
 on:
   push:
@@ -18,7 +18,7 @@ concurrency:
 
 jobs:
   do-the-job:
-    name: Run INF1 tests
+    name: Run INF1 export tests
     runs-on: [self-hosted, 4-aws-inf1, 24-cpu, ci]
     env:
       AWS_REGION: us-east-1
@@ -46,4 +46,5 @@ jobs:
       - name: Run export tests
         run: |
           source aws_neuron_venv_pytorch/bin/activate
+          export MAX_EXPORT_TEST_COMBINATIONS=1
           HF_TOKEN=${{ secrets.HF_TOKEN_OPTIMUM_NEURON_CI }} pytest -m is_inferentia_test tests/exporters

.github/workflows/test_inf1_full_export.yml

@@ -0,0 +1,47 @@
+name: Optimum neuron / Test INF1 full export
+
+on:
+  push:
+    branches: [ main ]
+    paths:
+      - "optimum/exporters/neuron/*.py"
+  pull_request:
+    branches: [ main ]
+    paths:
+      - "optimum/exporters/neuron/*.py"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+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
+    steps:
+      - name: Check AMI
+        run: dpkg -l | grep neuron
+      - name: Checkout
+        uses: actions/checkout@v2
+      - name: Install system packages
+        run: |
+          sudo apt install python3.8-venv -y
+      - name: Install python packages
+        run: |
+          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
+          python -m pip install .[neuron,tests]
+          python -m pip uninstall optimum -y
+          python -m pip install optimum
+      - name: Run CLI tests
+        run: |
+          source aws_neuron_venv_pytorch/bin/activate
+          HF_TOKEN=${{ secrets.HF_TOKEN_OPTIMUM_NEURON_CI }} pytest -m is_inferentia_test tests/cli
+      - name: Run export tests
+        run: |
+          source aws_neuron_venv_pytorch/bin/activate
+          HF_TOKEN=${{ secrets.HF_TOKEN_OPTIMUM_NEURON_CI }} pytest -m is_inferentia_test tests/exporters

.github/workflows/test_inf1_inference.yml

@@ -1,4 +1,4 @@
-name: Optimum neuron / Test INF1 inference & pipelines
+name: Optimum neuron / Test INF1 inference
 
 on:
   push:
@@ -43,7 +43,3 @@ jobs:
         run: |
           source aws_neuron_venv_pytorch/bin/activate
           HF_TOKEN=${{ secrets.HF_TOKEN_OPTIMUM_NEURON_CI }} pytest -m is_inferentia_test tests/inference
-      - name: Run pipelines tests
-        run: |
-          source aws_neuron_venv_pytorch/bin/activate
-          HF_TOKEN=${{ secrets.HF_TOKEN_OPTIMUM_NEURON_CI }} pytest -m is_inferentia_test tests/pipelines

.github/workflows/test_inf1_pipelines.yml

@@ -0,0 +1,43 @@
+name: Optimum neuron / Test INF1 pipelines
+
+on:
+  push:
+    branches: [ main ]
+    paths:
+      - "optimum/neuron/pipelines/**.py"
+  pull_request:
+    branches: [ main ]
+    paths:
+      - "optimum/neuron/pipelines/**.py"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  do-the-job:
+    name: Run INF1 tests
+    runs-on: [self-hosted, 4-aws-inf1, 24-cpu, ci]
+    env:
+      AWS_REGION: us-east-1
+    steps:
+      - name: Check AMI
+        run: dpkg -l | grep neuron
+      - name: Checkout
+        uses: actions/checkout@v2
+      - name: Install system packages
+        run: |
+          sudo apt install python3.8-venv -y
+      - name: Install python packages
+        run: |
+          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
+          python -m pip install .[neuron,tests]
+          python -m pip uninstall optimum -y
+          python -m pip install optimum
+      - name: Run pipelines tests
+        run: |
+          source aws_neuron_venv_pytorch/bin/activate
+          HF_TOKEN=${{ secrets.HF_TOKEN_OPTIMUM_NEURON_CI }} pytest -m is_inferentia_test tests/pipelines

.github/workflows/test_inf2_export.yml

@@ -1,4 +1,4 @@
-name: Optimum neuron / Test INF2 export
+name: Optimum neuron / Test INF2 partial export
 
 on:
   push:
@@ -18,7 +18,7 @@ concurrency:
 
 jobs:
   do-the-job:
-    name: Run INF2 tests
+    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
@@ -38,4 +38,5 @@ jobs:
       - name: Run exporters tests
         run: |
           source aws_neuron_venv_pytorch/bin/activate
+          export MAX_EXPORT_TEST_COMBINATIONS=1
           HF_TOKEN=${{ secrets.HF_TOKEN_OPTIMUM_NEURON_CI }} pytest -m is_inferentia_test tests/exporters

.github/workflows/test_inf2_full_export.yml

@@ -0,0 +1,39 @@
+name: Optimum neuron / Test INF2 full export
+
+on:
+  push:
+    branches: [ main ]
+    paths:
+      - "optimum/exporters/neuron/*.py"
+  pull_request:
+    branches: [ main ]
+    paths:
+      - "optimum/exporters/neuron/*.py"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+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
+    steps:
+      - name: Check AMI
+        run: dpkg -l | grep neuron
+      - name: Checkout
+        uses: actions/checkout@v2
+      - name: Install python dependencies
+        run: |
+          sudo apt install python3.8-venv -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
+          python -m pip install .[neuronx,tests]
+      - name: Run exporters tests
+        run: |
+          source aws_neuron_venv_pytorch/bin/activate
+          HF_TOKEN=${{ secrets.HF_TOKEN_OPTIMUM_NEURON_CI }} pytest -m is_inferentia_test tests/exporters

optimum/exporters/neuron/__main__.py

@@ -108,6 +108,7 @@ def infer_task(task: str, model_name_or_path: str) -> str:
     return task
 
 
+# This function is not applicable for diffusers / sentence transformers models
 def get_input_shapes_and_config_class(task: str, args: argparse.Namespace) -> Dict[str, int]:
     config = AutoConfig.from_pretrained(args.model)
 
@@ -116,7 +117,10 @@ def get_input_shapes_and_config_class(task: str, args: argparse.Namespace) -> Di
         model_type = model_type + "-encoder"
 
     neuron_config_constructor = TasksManager.get_exporter_config_constructor(
-        model_type=model_type, exporter="neuron", task=task
+        model_type=model_type,
+        exporter="neuron",
+        task=task,
+        library_name="transformers",
     )
     input_args = neuron_config_constructor.func.get_input_args_for_task(task)
     input_shapes = {name: getattr(args, name) for name in input_args}

optimum/exporters/neuron/model_configs.py

@@ -190,7 +190,9 @@ class DebertaV2NeuronConfig(ElectraNeuronConfig):
     pass
 
 
-@register_in_tasks_manager("sentence-transformers-transformer", *["feature-extraction", "sentence-similarity"])
+@register_in_tasks_manager(
+    "transformer", *["feature-extraction", "sentence-similarity"], library_name="sentence_transformers"
+)
 class SentenceTransformersTransformerNeuronConfig(TextEncoderNeuronConfig):
     NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
     CUSTOM_MODEL_WRAPPER = SentenceTransformersTransformerNeuronWrapper
@@ -268,7 +270,9 @@ def outputs(self) -> List[str]:
 
 # TODO: We should decouple clip text and vision, this would need fix on Optimum main. For the current workaround
 # users can pass dummy text inputs when encoding image, vice versa.
-@register_in_tasks_manager("sentence-transformers-clip", *["feature-extraction", "sentence-similarity"])
+@register_in_tasks_manager(
+    "clip", *["feature-extraction", "sentence-similarity"], library_name="sentence_transformers"
+)
 class SentenceTransformersCLIPNeuronConfig(CLIPNeuronConfig):
     CUSTOM_MODEL_WRAPPER = SentenceTransformersCLIPNeuronWrapper
     ATOL_FOR_VALIDATION = 1e-3
@@ -282,7 +286,7 @@ def patch_model_for_export(self, model, dummy_inputs):
         return self.CUSTOM_MODEL_WRAPPER(model, list(dummy_inputs.keys()))
 
 
-@register_in_tasks_manager("unet", *["semantic-segmentation"])
+@register_in_tasks_manager("unet", *["semantic-segmentation"], library_name="diffusers")
 class UNetNeuronConfig(VisionNeuronConfig):
     ATOL_FOR_VALIDATION = 1e-3
     INPUT_ARGS = ("batch_size", "sequence_length", "num_channels", "width", "height")
@@ -356,7 +360,7 @@ def is_sdxl(self, is_sdxl: bool):
         self._is_sdxl = is_sdxl
 
 
-@register_in_tasks_manager("vae-encoder", *["semantic-segmentation"])
+@register_in_tasks_manager("vae-encoder", *["semantic-segmentation"], library_name="diffusers")
 class VaeEncoderNeuronConfig(VisionNeuronConfig):
     ATOL_FOR_VALIDATION = 1e-3
     MODEL_TYPE = "vae-encoder"
@@ -392,7 +396,7 @@ def generate_dummy_inputs(self, return_tuple: bool = False, **kwargs):
             return dummy_inputs
 
 
-@register_in_tasks_manager("vae-decoder", *["semantic-segmentation"])
+@register_in_tasks_manager("vae-decoder", *["semantic-segmentation"], library_name="diffusers")
 class VaeDecoderNeuronConfig(VisionNeuronConfig):
     ATOL_FOR_VALIDATION = 1e-3
     MODEL_TYPE = "vae-decoder"

optimum/exporters/neuron/utils.py

@@ -168,7 +168,10 @@ def get_stable_diffusion_models_for_export(
     if DIFFUSION_MODEL_TEXT_ENCODER_NAME in models_for_export:
         text_encoder = models_for_export[DIFFUSION_MODEL_TEXT_ENCODER_NAME]
         text_encoder_config_constructor = TasksManager.get_exporter_config_constructor(
-            model=text_encoder, exporter="neuron", task="feature-extraction"
+            model=text_encoder,
+            exporter="neuron",
+            task="feature-extraction",
+            library_name="transformers",
         )
         text_encoder_neuron_config = text_encoder_config_constructor(
             text_encoder.config,
@@ -185,6 +188,7 @@ def get_stable_diffusion_models_for_export(
             exporter="neuron",
             task="feature-extraction",
             model_type="clip-text-with-projection",
+            library_name="transformers",
         )
         text_encoder_neuron_config_2 = text_encoder_config_constructor_2(
             text_encoder_2.config,
@@ -197,7 +201,11 @@ def get_stable_diffusion_models_for_export(
     # U-NET
     unet = models_for_export[DIFFUSION_MODEL_UNET_NAME]
     unet_neuron_config_constructor = TasksManager.get_exporter_config_constructor(
-        model=unet, exporter="neuron", task="semantic-segmentation", model_type="unet"
+        model=unet,
+        exporter="neuron",
+        task="semantic-segmentation",
+        model_type="unet",
+        library_name="diffusers",
     )
     unet_neuron_config = unet_neuron_config_constructor(
         unet.config,
@@ -216,6 +224,7 @@ def get_stable_diffusion_models_for_export(
         exporter="neuron",
         task="semantic-segmentation",
         model_type="vae-encoder",
+        library_name="diffusers",
     )
     vae_encoder_neuron_config = vae_encoder_config_constructor(
         vae_encoder.config,
@@ -232,6 +241,7 @@ def get_stable_diffusion_models_for_export(
         exporter="neuron",
         task="semantic-segmentation",
         model_type="vae-decoder",
+        library_name="diffusers",
     )
     vae_decoder_neuron_config = vae_decoder_config_constructor(
         vae_decoder.config,
@@ -382,7 +392,10 @@ def get_encoder_decoder_models_for_export(
     # Encoder
     model_type = getattr(model.config, "model_type") + "-encoder"
     encoder_config_constructor = TasksManager.get_exporter_config_constructor(
-        exporter="neuron", model_type=model_type, task=task
+        exporter="neuron",
+        model_type=model_type,
+        task=task,
+        library_name="transformers",
     )
     check_mandatory_input_shapes(encoder_config_constructor, task, input_shapes)
     encoder_neuron_config = encoder_config_constructor(
@@ -396,7 +409,10 @@ def get_encoder_decoder_models_for_export(
     # Decoder
     model_type = getattr(model.config, "model_type") + "-decoder"
     decoder_config_constructor = TasksManager.get_exporter_config_constructor(
-        exporter="neuron", model_type=model_type, task=task
+        exporter="neuron",
+        model_type=model_type,
+        task=task,
+        library_name="transformers",
     )
     check_mandatory_input_shapes(encoder_config_constructor, task, input_shapes)
     decoder_neuron_config = decoder_config_constructor(

setup.py

@@ -15,7 +15,7 @@
 INSTALL_REQUIRES = [
     "transformers == 4.36.2",
     "accelerate == 0.23.0",
-    "optimum >= 1.16.2",
+    "optimum ~= 1.17.0",
     "huggingface_hub >= 0.20.1",
     "numpy>=1.22.2, <=1.25.2",
     "protobuf<4",

tests/exporters/exporters_utils.py

@@ -44,8 +44,8 @@
 }
 
 SENTENCE_TRANSFORMERS_MODELS = {
-    "sentence-transformers-transformer": "sentence-transformers/all-MiniLM-L6-v2",
-    "sentence-transformers-clip": "sentence-transformers/clip-ViT-B-32",
+    "transformer": "sentence-transformers/all-MiniLM-L6-v2",
+    "clip": "sentence-transformers/clip-ViT-B-32",
 }
 
 WEIGHTS_NEFF_SEPARATION_UNSUPPORTED_ARCH = ["camembert", "roberta"]