added expectation test for roberta squad v2 qa

Signed-off-by: Joshua Rosenkranz <jmrosenk@us.ibm.com>

Author: JRosenkranz

tests/models/test_model_expectations.py

@@ -8,39 +8,71 @@
 )
 import os
 
+if "HF_HOME" not in os.environ:
+    os.environ["HF_HOME"] = "/tmp/models/hf_cache"
+
 model_dir = os.environ.get("FMS_TESTING_MODEL_DIR", "/tmp/models")
 LLAMA_194M = f"{model_dir}/llama-194m"
 GRANITE_7B_BASE = f"{model_dir}/granite-7b-base"
 GRANITE_8B_CODE_BASE = f"{model_dir}/granite-8b-code-base"
 GRANITE_3_8B_CODE_BASE = f"{model_dir}/granite-3-8b-base"
 
 models = [LLAMA_194M, GRANITE_7B_BASE, GRANITE_8B_CODE_BASE, GRANITE_3_8B_CODE_BASE]
+mini_models = {LLAMA_194M, GRANITE_7B_BASE, GRANITE_8B_CODE_BASE, GRANITE_3_8B_CODE_BASE}
 
 class AIUModelFixtureMixin(ModelFixtureMixin):
 
     @pytest.fixture(scope="class", autouse=True)
     def uninitialized_model(self, model_id):
+        if model_id in mini_models:
+            get_model_kwargs = {"architecture": "hf_configured", "nlayers": 3}
+        else:
+            get_model_kwargs = {"architecture": "hf_pretrained"}
+
         aiu_model = get_model(
-            "hf_configured",
-            model_id,
+            variant=model_id,
             device_type="cpu",
             unfuse_weights=True,
-            nlayers=3
+            **get_model_kwargs
         )
         torch.compile(aiu_model, backend="sendnn")
         return aiu_model
-    
+
+class TestAIUModels(
+    ModelConsistencyTestSuite,
+    AIUModelFixtureMixin,
+):
+
+    # x is the main parameter for this model which is the input tensor
+    _get_signature_params = ["x"]
+
     @pytest.fixture(scope="class", autouse=True, params=models)
     def model_id(self, request):
         return request.param
 
-class TestAIUModels(
+    def test_model_unfused(self, model, signature):
+        pytest.skip("All AIU models are already unfused")
+
+
+ROBERTA_SQUAD_v2 = "deepset/roberta-base-squad2"
+tuple_output_models = [ROBERTA_SQUAD_v2]
+
+class TestAIUModelsTupleOutput(
     ModelConsistencyTestSuite,
     AIUModelFixtureMixin,
 ):
-
+    
     # x is the main parameter for this model which is the input tensor
     _get_signature_params = ["x"]
 
+    @pytest.fixture(scope="class", autouse=True, params=tuple_output_models)
+    def model_id(self, request):
+        return request.param
+    
+    @staticmethod
+    def _get_signature_logits_getter_fn(f_out) -> torch.Tensor:
+        return torch.cat([f_out[0], f_out[1]], dim=-1)
+    
     def test_model_unfused(self, model, signature):
-        pytest.skip("All AIU models are already unfused")
+        pytest.skip("All AIU models are already unfused")
+    

tests/resources/expectations/models.test_model_expectations.TestAIUModelsTupleOutput.roberta-base-squad2.test_model_output

@@ -0,0 +1 @@
+9.834766387939453e-07,3.5762786865234375e-07,8.940696716308594e-07,6.258487701416016e-07,8.344650268554688e-07,1.1324882507324219e-06,6.556510925292969e-07,1.2516975402832031e-06,1.6391277313232422e-06,0.0,2.384185791015625e-07,1.1324882507324219e-06,4.172325134277344e-07,9.238719940185547e-07,4.76837158203125e-07,1.1622905731201172e-06,0.2104383111000061,0.2104375958442688,0.21043795347213745,0.21043813228607178,0.21043753623962402,0.21043819189071655,0.2104378342628479,0.21043813228607178,0.21043860912322998,0.21043741703033447,0.21043741703033447,0.21043819189071655,0.21043717861175537,0.21043848991394043,0.21043795347213745,0.21043837070465088

tests/resources/expectations/models.test_model_expectations.TestAIUModelsTupleOutput.roberta-base-squad2.test_model_weight_keys

@@ -0,0 +1 @@
+base_model.embedding.weight,base_model.enc_norm.bias,base_model.enc_norm.weight,base_model.layers.0.attn.dense.bias,base_model.layers.0.attn.dense.weight,base_model.layers.0.attn.in_proj.qkv_fused.bias,base_model.layers.0.attn.in_proj.qkv_fused.weight,base_model.layers.0.ff_ln.bias,base_model.layers.0.ff_ln.weight,base_model.layers.0.ff_sub_layer.w1.bias,base_model.layers.0.ff_sub_layer.w1.weight,base_model.layers.0.ff_sub_layer.w2.bias,base_model.layers.0.ff_sub_layer.w2.weight,base_model.layers.0.ln.bias,base_model.layers.0.ln.weight,base_model.layers.1.attn.dense.bias,base_model.layers.1.attn.dense.weight,base_model.layers.1.attn.in_proj.qkv_fused.bias,base_model.layers.1.attn.in_proj.qkv_fused.weight,base_model.layers.1.ff_ln.bias,base_model.layers.1.ff_ln.weight,base_model.layers.1.ff_sub_layer.w1.bias,base_model.layers.1.ff_sub_layer.w1.weight,base_model.layers.1.ff_sub_layer.w2.bias,base_model.layers.1.ff_sub_layer.w2.weight,base_model.layers.1.ln.bias,base_model.layers.1.ln.weight,base_model.layers.10.attn.dense.bias,base_model.layers.10.attn.dense.weight,base_model.layers.10.attn.in_proj.qkv_fused.bias,base_model.layers.10.attn.in_proj.qkv_fused.weight,base_model.layers.10.ff_ln.bias,base_model.layers.10.ff_ln.weight,base_model.layers.10.ff_sub_layer.w1.bias,base_model.layers.10.ff_sub_layer.w1.weight,base_model.layers.10.ff_sub_layer.w2.bias,base_model.layers.10.ff_sub_layer.w2.weight,base_model.layers.10.ln.bias,base_model.layers.10.ln.weight,base_model.layers.11.attn.dense.bias,base_model.layers.11.attn.dense.weight,base_model.layers.11.attn.in_proj.qkv_fused.bias,base_model.layers.11.attn.in_proj.qkv_fused.weight,base_model.layers.11.ff_ln.bias,base_model.layers.11.ff_ln.weight,base_model.layers.11.ff_sub_layer.w1.bias,base_model.layers.11.ff_sub_layer.w1.weight,base_model.layers.11.ff_sub_layer.w2.bias,base_model.layers.11.ff_sub_layer.w2.weight,base_model.layers.11.ln.bias,base_model.layers.11.ln.weight,base_model.layers.2.attn.dense.bias,base_model.layers.2.attn.dense.weight,base_model.layers.2.attn.in_proj.qkv_fused.bias,base_model.layers.2.attn.in_proj.qkv_fused.weight,base_model.layers.2.ff_ln.bias,base_model.layers.2.ff_ln.weight,base_model.layers.2.ff_sub_layer.w1.bias,base_model.layers.2.ff_sub_layer.w1.weight,base_model.layers.2.ff_sub_layer.w2.bias,base_model.layers.2.ff_sub_layer.w2.weight,base_model.layers.2.ln.bias,base_model.layers.2.ln.weight,base_model.layers.3.attn.dense.bias,base_model.layers.3.attn.dense.weight,base_model.layers.3.attn.in_proj.qkv_fused.bias,base_model.layers.3.attn.in_proj.qkv_fused.weight,base_model.layers.3.ff_ln.bias,base_model.layers.3.ff_ln.weight,base_model.layers.3.ff_sub_layer.w1.bias,base_model.layers.3.ff_sub_layer.w1.weight,base_model.layers.3.ff_sub_layer.w2.bias,base_model.layers.3.ff_sub_layer.w2.weight,base_model.layers.3.ln.bias,base_model.layers.3.ln.weight,base_model.layers.4.attn.dense.bias,base_model.layers.4.attn.dense.weight,base_model.layers.4.attn.in_proj.qkv_fused.bias,base_model.layers.4.attn.in_proj.qkv_fused.weight,base_model.layers.4.ff_ln.bias,base_model.layers.4.ff_ln.weight,base_model.layers.4.ff_sub_layer.w1.bias,base_model.layers.4.ff_sub_layer.w1.weight,base_model.layers.4.ff_sub_layer.w2.bias,base_model.layers.4.ff_sub_layer.w2.weight,base_model.layers.4.ln.bias,base_model.layers.4.ln.weight,base_model.layers.5.attn.dense.bias,base_model.layers.5.attn.dense.weight,base_model.layers.5.attn.in_proj.qkv_fused.bias,base_model.layers.5.attn.in_proj.qkv_fused.weight,base_model.layers.5.ff_ln.bias,base_model.layers.5.ff_ln.weight,base_model.layers.5.ff_sub_layer.w1.bias,base_model.layers.5.ff_sub_layer.w1.weight,base_model.layers.5.ff_sub_layer.w2.bias,base_model.layers.5.ff_sub_layer.w2.weight,base_model.layers.5.ln.bias,base_model.layers.5.ln.weight,base_model.layers.6.attn.dense.bias,base_model.layers.6.attn.dense.weight,base_model.layers.6.attn.in_proj.qkv_fused.bias,base_model.layers.6.attn.in_proj.qkv_fused.weight,base_model.layers.6.ff_ln.bias,base_model.layers.6.ff_ln.weight,base_model.layers.6.ff_sub_layer.w1.bias,base_model.layers.6.ff_sub_layer.w1.weight,base_model.layers.6.ff_sub_layer.w2.bias,base_model.layers.6.ff_sub_layer.w2.weight,base_model.layers.6.ln.bias,base_model.layers.6.ln.weight,base_model.layers.7.attn.dense.bias,base_model.layers.7.attn.dense.weight,base_model.layers.7.attn.in_proj.qkv_fused.bias,base_model.layers.7.attn.in_proj.qkv_fused.weight,base_model.layers.7.ff_ln.bias,base_model.layers.7.ff_ln.weight,base_model.layers.7.ff_sub_layer.w1.bias,base_model.layers.7.ff_sub_layer.w1.weight,base_model.layers.7.ff_sub_layer.w2.bias,base_model.layers.7.ff_sub_layer.w2.weight,base_model.layers.7.ln.bias,base_model.layers.7.ln.weight,base_model.layers.8.attn.dense.bias,base_model.layers.8.attn.dense.weight,base_model.layers.8.attn.in_proj.qkv_fused.bias,base_model.layers.8.attn.in_proj.qkv_fused.weight,base_model.layers.8.ff_ln.bias,base_model.layers.8.ff_ln.weight,base_model.layers.8.ff_sub_layer.w1.bias,base_model.layers.8.ff_sub_layer.w1.weight,base_model.layers.8.ff_sub_layer.w2.bias,base_model.layers.8.ff_sub_layer.w2.weight,base_model.layers.8.ln.bias,base_model.layers.8.ln.weight,base_model.layers.9.attn.dense.bias,base_model.layers.9.attn.dense.weight,base_model.layers.9.attn.in_proj.qkv_fused.bias,base_model.layers.9.attn.in_proj.qkv_fused.weight,base_model.layers.9.ff_ln.bias,base_model.layers.9.ff_ln.weight,base_model.layers.9.ff_sub_layer.w1.bias,base_model.layers.9.ff_sub_layer.w1.weight,base_model.layers.9.ff_sub_layer.w2.bias,base_model.layers.9.ff_sub_layer.w2.weight,base_model.layers.9.ln.bias,base_model.layers.9.ln.weight,base_model.position_embedding.weight,qa_head.bias,qa_head.weight