Add TestInjectionPolicy inference unittest class for testing custom injection policies

tests/conftest.py

@@ -2,10 +2,13 @@
 
 import sys
 import pytest
+import os
 from os.path import abspath, dirname, join
 import torch
 import warnings
 
+os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'
+
 # allow having multiple repository checkouts and not needing to remember to rerun
 # 'pip install -e .[dev]' when switching between checkouts and running tests.
 git_repo_path = abspath(join(dirname(dirname(__file__)), "src"))

tests/unit/inference/test_inference.py

@@ -9,6 +9,7 @@
 from packaging import version as pkg_version
 from deepspeed.ops.op_builder import OpBuilder
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers.models.t5.modeling_t5 import T5Block
 from huggingface_hub import HfApi
 
 rocm_version = OpBuilder.installed_rocm_version()
@@ -46,15 +47,20 @@
     "facebook/opt-125m",        # 125m, 1.7B, ..., 175B variants have the same model architecture.
     "facebook/opt-350m",        # 350m applies layer norm after attnention layer which is different than other variants.
 ]
+_t5_models = [
+    "google/t5-v1_1-small",
+]
 _all_models = HfApi().list_models()
 
-test_models = set(_bert_models + _roberta_models + _gpt_models + _opt_models)
+test_models = set(_bert_models + _roberta_models + _gpt_models + _opt_models +
+                  _t5_models)
 test_tasks = [
     "fill-mask",
     "question-answering",
     "text-classification",
     "token-classification",
     "text-generation",
+    "text2text-generation",
 ]
 pytest.all_models = {
     task: [m.modelId for m in _all_models if m.pipeline_tag == task]
@@ -150,6 +156,8 @@ def query(model_w_task):
         return "My name is jean-baptiste and I live in montreal."
     elif task == "text-generation":
         return "DeepSpeed is the greatest"
+    elif task == "text2text-generation":
+        return "Is this review positive or negative? Review: this is the best cast iron skillet you will ever buy"
     else:
         NotImplementedError(f'query for task "{task}" is not implemented')
 
@@ -184,6 +192,11 @@ def text_generation_assert(x, y):
                                                           for res in y)
 
 
+def text2text_generation_assert(x, y):
+    return set(res["generated_text"] for res in x) == set(res["generated_text"]
+                                                          for res in y)
+
+
 @pytest.fixture
 def assert_fn(model_w_task):
     model, task = model_w_task
@@ -193,6 +206,7 @@ def assert_fn(model_w_task):
         "text-classification": text_classification_assert,
         "token-classification": token_classification_assert,
         "text-generation": text_generation_assert,
+        "text2text-generation": text2text_generation_assert,
     }
     assert_fn = assert_fn_dict.get(task, None)
     if assert_fn is None:
@@ -252,14 +266,28 @@ def test(
         torch.cuda.synchronize()
         bs_time = time.time() - start
 
-        pipe.model = deepspeed.init_inference(
-            pipe.model,
-            mp_size=1,
-            dtype=dtype,
-            replace_method="auto",
-            replace_with_kernel_inject=True,
-            enable_cuda_graph=enable_cuda_graph,
-        )
+        if "t5-v1_1-small" in model:
+            pipe.model = deepspeed.init_inference(
+                pipe.model,
+                mp_size=1,
+                dtype=dtype,
+                injection_policy={
+                    T5Block: ('SelfAttention.o',
+                              'EncDecAttention.o',
+                              'DenseReluDense.wo')
+                },
+                enable_cuda_graph=enable_cuda_graph,
+            )
+        else:
+            pipe.model = deepspeed.init_inference(
+                pipe.model,
+                mp_size=1,
+                dtype=dtype,
+                replace_method="auto",
+                replace_with_kernel_inject=True,
+                enable_cuda_graph=enable_cuda_graph,
+            )
+
         # Warm-up queries for perf measurement
         #for i in range(10):
         #    _ = pipe(query, **inf_kwargs)