fix tests

optimum/neuron/modeling_seq2seq.py

@@ -462,14 +462,14 @@ def forward(
         decoder_hidden_states = None
 
         # Skip pkv which can't be copied from memory to buffer
-        if output_attentions and self.config.neuron.get("output_attentions"):
+        if output_attentions and self.configs["decoder"].neuron.get("output_attentions"):
             if self.config.is_encoder_decoder:
                 cross_attentions = outputs[-self.config.num_decoder_layers :]
                 cur_idx += self.config.num_decoder_layers
             decoder_attentions = outputs[-(self.config.num_decoder_layers + cur_idx) : -cur_idx]
             cur_idx += self.config.num_decoder_layers
 
-        if output_hidden_states and self.config.neuron.get("output_hidden_states"):
+        if output_hidden_states and self.configs["decoder"].neuron.get("output_hidden_states"):
             decoder_hidden_states = outputs[-(self.config.num_decoder_layers + 1 + cur_idx) : -cur_idx]
 
         decoder_outputs = ModelOutput(

optimum/neuron/version.py

@@ -12,6 +12,6 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-__version__ = "0.0.25.dev0"
+__version__ = "0.0.27.dev0"
 
 __sdk_version__ = "2.20.0"

tests/generation/test_generate.py

@@ -117,26 +117,6 @@ def test_seq2seq_generation_greedy_with_optional_outputs(neuron_seq2seq_greedy_p
     assert "decoder_hidden_states" in output
 
 
-@is_inferentia_test
-@requires_neuronx
-def test_seq2seq_generation_tp2(neuron_seq2seq_tp2_path):
-    model = NeuronModelForSeq2SeqLM.from_pretrained(neuron_seq2seq_tp2_path)
-    tokenizer = AutoTokenizer.from_pretrained(neuron_seq2seq_tp2_path)
-    inputs = tokenizer("translate English to German: Lets eat good food.", return_tensors="pt")
-
-    output = model.generate(
-        **inputs,
-        num_return_sequences=1,
-        max_length=20,
-        output_attentions=True,
-        output_hidden_states=True,
-        return_dict_in_generate=True,
-    )
-    assert "decoder_attentions" in output
-    assert "cross_attentions" in output
-    assert "decoder_hidden_states" in output
-
-
 @pytest.mark.skip("Makes pytest fail, to fix.")
 @pytest.mark.parametrize(
     "gen_kwargs",
@@ -180,10 +160,3 @@ def test_general_seq2seq_generation(export_seq2seq_id, export_seq2seq_model_clas
     model = export_seq2seq_model_class.from_pretrained(export_seq2seq_id)
     tokenizer = AutoTokenizer.from_pretrained(export_seq2seq_id)
     _test_model_generation_trn(model, tokenizer, 1, 10, **gen_kwargs)
-
-
-# Compulsory for multiprocessing tests, since we want children processes to be spawned only in the main program.
-# eg. tensor parallel tracing, `neuronx_distributed.parallel_model_trace` will spawn multiple processes to trace
-# and compile the model.
-if __name__ == "__main__":
-    pytest.main([__file__])

tests/generation/test_parallel.py

@@ -0,0 +1,48 @@
+# coding=utf-8
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytest
+from transformers import AutoTokenizer
+
+from optimum.neuron import NeuronModelForSeq2SeqLM
+from optimum.neuron.utils.testing_utils import is_inferentia_test, requires_neuronx
+
+
+@is_inferentia_test
+@requires_neuronx
+def test_seq2seq_generation_tp2(neuron_seq2seq_tp2_path):
+    model = NeuronModelForSeq2SeqLM.from_pretrained(neuron_seq2seq_tp2_path)
+    tokenizer = AutoTokenizer.from_pretrained(neuron_seq2seq_tp2_path)
+    inputs = tokenizer("translate English to German: Lets eat good food.", return_tensors="pt")
+
+    output = model.generate(
+        **inputs,
+        num_return_sequences=1,
+        max_length=20,
+        output_attentions=True,
+        output_hidden_states=True,
+        return_dict_in_generate=True,
+    )
+    assert "decoder_attentions" in output
+    assert "cross_attentions" in output
+    assert "decoder_hidden_states" in output
+
+
+# Compulsory for multiprocessing tests, since we want children processes to be spawned only in the main program.
+# eg. tensor parallel tracing, `neuronx_distributed.parallel_model_trace` will spawn multiple processes to trace
+# and compile the model.
+if __name__ == "__main__":
+    pytest.main([__file__])