Skip GPTNeo tests

huggingface · Jan 16, 2024 · 7bdad6a · 7bdad6a
1 parent bce46b5
commit 7bdad6a
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/tests/distributed/test_model_parallelization.py b/tests/distributed/test_model_parallelization.py
@@ -172,6 +172,13 @@ def _generate_supported_model_classes(
 
 MODEL_CLASSES_TO_IGNORE = [
     "BertForPreTraining",  # There is a compilation issue, and testing TP for BertForPretraining is not really important.
+    # TODO
+    # GPTNeo's attention mechanism is broken in transformers==4.36.2, this should be re-enabled once there is a release
+    # containing this PR: https://github.com/huggingface/transformers/pull/28533
+    "GPTNeoForSequenceClassification",
+    "GPTNeoForTokenClassification",
+    "GPTNeoForQuestionAnswering",
+    "GPTNeoForCausalLM",
 ]