Fix pipelines

.github/workflows/test_inf2.yml

@@ -41,6 +41,7 @@ jobs:
           python3 -m venv aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           python -m pip install -U pip
+          pip install --upgrade setuptools==69.5.1
           python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com
           python -m pip install .[neuronx,tests]
       - name: Run cache tests

.github/workflows/test_inf2_export.yml

@@ -41,6 +41,7 @@ jobs:
           python3 -m venv aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           python -m pip install -U pip
+          pip install --upgrade setuptools==69.5.1
           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

.github/workflows/test_inf2_full_export.yml

@@ -39,6 +39,7 @@ jobs:
           python3 -m venv aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           python -m pip install -U pip
+          pip install --upgrade setuptools==69.5.1
           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

.github/workflows/test_inf2_inference.yml

@@ -41,6 +41,7 @@ jobs:
           python3 -m venv aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           python -m pip install -U pip
+          pip install --upgrade setuptools==69.5.1
           python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com
           python -m pip install .[neuronx,tests]
       - name: Run inference tests

.github/workflows/test_trainium_common.yml

@@ -42,6 +42,7 @@ jobs:
           python3 -m venv aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           python -m pip install -U pip
+          pip install --upgrade setuptools==69.5.1
           python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com
           python -m pip install .[neuronx,tests]
       - name: Collect tests on Neuron Cores

.github/workflows/test_trainium_distributed.yml

@@ -44,6 +44,7 @@ jobs:
           python3 -m venv aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           python -m pip install -U pip
+          pip install --upgrade setuptools==69.5.1
           python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com
           python -m pip install .[neuronx,tests]
       - name: Collect tests on Neuron Cores

.github/workflows/test_trainium_examples.yml

@@ -2,8 +2,8 @@ name: Optimum Neuron - Test Example Scripts
 
 on:
   workflow_dispatch:
-    inputs: 
-      coverage: 
+    inputs:
+      coverage:
         description: Coverage
         type: choice
         options:
@@ -12,7 +12,7 @@ on:
           - middle
           - low
         required: true
-      model_size: 
+      model_size:
         description: Size of models
         type: choice
         options:
@@ -53,6 +53,7 @@ jobs:
           python3 -m venv aws_neuron_venv_pytorch
           source aws_neuron_venv_pytorch/bin/activate
           python -m pip install -U pip
+          pip install --upgrade setuptools==69.5.1
           python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com
           python -m pip install .[neuronx,tests]
       - name: Collect example tests on Neuron Cores

docs/source/package_reference/modeling.mdx

@@ -18,12 +18,12 @@ limitations under the License.
 
 ## Generic model classes
 
-### NeuronBaseModel
+### NeuronTracedModel
 
-The `NeuronBaseModel` class is available for instantiating a base Neuron model without a specific head.
+The `NeuronTracedModel` class is available for instantiating a base Neuron model without a specific head.
 It is used as the base class for all tasks but text generation.
 
-[[autodoc]] modeling_base.NeuronBaseModel
+[[autodoc]] modeling_traced.NeuronTracedModel
 
 ### NeuronDecoderModel
 
@@ -104,4 +104,4 @@ The following Neuron model classes are available for natural language processing
 
 ### NeuronStableDiffusionXLInpaintPipeline
 [[autodoc]] modeling_diffusion.NeuronStableDiffusionXLInpaintPipeline
-    - __call__
+    - __call__

optimum/neuron/__init__.py

@@ -29,7 +29,7 @@
     "hf_argparser": ["NeuronHfArgumentParser"],
     "trainers": ["NeuronTrainer", "Seq2SeqNeuronTrainer"],
     "training_args": ["NeuronTrainingArguments", "Seq2SeqNeuronTrainingArguments"],
-    "modeling_base": ["NeuronBaseModel"],
+    "modeling_traced": ["NeuronTracedModel"],
     "modeling": [
         "NeuronModelForFeatureExtraction",
         "NeuronModelForSentenceTransformers",
@@ -73,7 +73,6 @@
         NeuronModelForSequenceClassification,
         NeuronModelForTokenClassification,
     )
-    from .modeling_base import NeuronBaseModel
     from .modeling_decoder import NeuronDecoderModel
     from .modeling_diffusion import (
         NeuronLatentConsistencyModelPipeline,
@@ -85,6 +84,7 @@
         NeuronStableDiffusionXLPipeline,
     )
     from .modeling_seq2seq import NeuronModelForSeq2SeqLM
+    from .modeling_traced import NeuronTracedModel
     from .pipelines import pipeline
     from .trainers import NeuronTrainer, Seq2SeqNeuronTrainer
     from .training_args import NeuronTrainingArguments, Seq2SeqNeuronTrainingArguments

optimum/neuron/modeling.py

@@ -43,8 +43,8 @@
 from transformers.utils import ModelOutput
 
 from .generation import TokenSelector
-from .modeling_base import NeuronBaseModel
 from .modeling_decoder import NeuronDecoderModel
+from .modeling_traced import NeuronTracedModel
 
 
 if TYPE_CHECKING:
@@ -61,13 +61,13 @@
 _TOKENIZER_FOR_DOC = "AutoTokenizer"
 
 NEURON_MODEL_START_DOCSTRING = r"""
-    This model inherits from [`~neuron.modeling.NeuronBaseModel`]. Check the superclass documentation for the generic methods the
+    This model inherits from [`~neuron.modeling.NeuronTracedModel`]. Check the superclass documentation for the generic methods the
     library implements for all its model (such as downloading or saving)
 
     Args:
         config (`transformers.PretrainedConfig`): [PretrainedConfig](https://huggingface.co/docs/transformers/main_classes/configuration#transformers.PretrainedConfig) is the Model configuration class with all the parameters of the model.
             Initializing with a config file does not load the weights associated with the model, only the
-            configuration. Check out the [`optimum.neuron.modeling.NeuronBaseModel.from_pretrained`] method to load the model weights.
+            configuration. Check out the [`optimum.neuron.modeling.NeuronTracedModel.from_pretrained`] method to load the model weights.
         model (`torch.jit._script.ScriptModule`): [torch.jit._script.ScriptModule](https://pytorch.org/docs/stable/generated/torch.jit.ScriptModule.html) is the TorchScript module with embedded NEFF(Neuron Executable File Format) compiled by neuron(x) compiler.
 """
 
@@ -125,7 +125,7 @@
     """,
     NEURON_MODEL_START_DOCSTRING,
 )
-class NeuronModelForFeatureExtraction(NeuronBaseModel):
+class NeuronModelForFeatureExtraction(NeuronTracedModel):
     """
     Feature Extraction model on Neuron devices.
     """
@@ -198,7 +198,7 @@ def forward(
     """,
     NEURON_MODEL_START_DOCSTRING,
 )
-class NeuronModelForSentenceTransformers(NeuronBaseModel):
+class NeuronModelForSentenceTransformers(NeuronTracedModel):
     """
     Sentence Transformers model on Neuron devices.
     """
@@ -283,7 +283,7 @@ def forward(
     """,
     NEURON_MODEL_START_DOCSTRING,
 )
-class NeuronModelForMaskedLM(NeuronBaseModel):
+class NeuronModelForMaskedLM(NeuronTracedModel):
     """
     Masked language model for on Neuron devices.
     """
@@ -353,7 +353,7 @@ def forward(
     """,
     NEURON_MODEL_START_DOCSTRING,
 )
-class NeuronModelForQuestionAnswering(NeuronBaseModel):
+class NeuronModelForQuestionAnswering(NeuronTracedModel):
     """
     Question Answering model on Neuron devices.
     """
@@ -422,7 +422,7 @@ def forward(
     """,
     NEURON_MODEL_START_DOCSTRING,
 )
-class NeuronModelForSequenceClassification(NeuronBaseModel):
+class NeuronModelForSequenceClassification(NeuronTracedModel):
     """
     Sequence Classification model on Neuron devices.
     """
@@ -490,7 +490,7 @@ def forward(
     """,
     NEURON_MODEL_START_DOCSTRING,
 )
-class NeuronModelForTokenClassification(NeuronBaseModel):
+class NeuronModelForTokenClassification(NeuronTracedModel):
     """
     Token Classification model on Neuron devices.
     """
@@ -571,7 +571,7 @@ def forward(
     """,
     NEURON_MODEL_START_DOCSTRING,
 )
-class NeuronModelForMultipleChoice(NeuronBaseModel):
+class NeuronModelForMultipleChoice(NeuronTracedModel):
     """
     Multiple choice model on Neuron devices.
     """