fix stateless llama testing (#600)

This commit fixes the stateless llama testing. Basically, with the setup
+ teardown approach of pytest unit testing, all the tests were sharing
the same model from hugging face. We were running the streaming llama
tests before the other tests. In these tests (run_torch_llm and
export_transformer_model with streaming_llm=True), we do
`enable_llama_pos_shift_attention(model)`, which changes the model we
are using. So, this was giving us inaccurate results by the time it came
to our base vmfb_comparison test. I created this issue to track and
provide more info on the error we are now seeing with torch 2.3 in
`test_ test_vmfb_comparison`:
#601. Also, marked it as
an expected failure for now. (changed runner because we are using
previous machine to repro issue for tinygrad folks which can lead to
instability and system hangs)

.github/workflows/test_models.yml

@@ -20,7 +20,7 @@ jobs:
     strategy:
       matrix:
         version: [3.11]
-        os: [nodai-amdgpu-w7900-x86-64]
+        os: [nodai-amdgpu-mi210-x86-64]
 
     runs-on: ${{matrix.os}}
     steps:

models/turbine_models/custom_models/llm_runner.py

@@ -1,6 +1,6 @@
 import argparse
 from turbine_models.model_runner import vmfbRunner
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from iree import runtime as ireert
 import torch
 import time
@@ -209,6 +209,18 @@ def run_torch_llm(
     model=None,
     tokenizer=None,
 ):
+    if model == None:
+        model = AutoModelForCausalLM.from_pretrained(
+            hf_model_name,
+            torch_dtype=torch.float,
+            token=hf_auth_token,
+        )
+    if tokenizer == None:
+        tokenizer = AutoTokenizer.from_pretrained(
+            hf_model_name,
+            use_fast=False,
+            token=hf_auth_token,
+        )
     if streaming_llm is True:
         enable_llama_pos_shift_attention(model)
 

models/turbine_models/tests/stateless_llama_test.py

@@ -76,6 +76,9 @@ def tearDownClass(cls):
         cls.tokenizer = None
         cls.mod = None
 
+    # See: https://github.com/nod-ai/SHARK-Turbine/issues/601
+    # Developed issues related to the pytorch 2.3 upgrade.
+    @unittest.expectedFailure
     def test_vmfb_comparison(self):
         """
         Test that the vmfb model produces the same output as the torch model
@@ -113,7 +116,7 @@ def test_vmfb_comparison(self):
             torch_str = llm_runner.run_torch_llm(
                 "Trelis/Llama-2-7b-chat-hf-function-calling-v2",
                 None,
-                self.DEFAULT_PROMPT,
+                DEFAULT_PROMPT,
                 model=self.mod,
                 tokenizer=self.tokenizer,
             )
@@ -152,7 +155,7 @@ def test_streaming_vmfb_comparison(self):
             target_triple="host",
             streaming_llm=True,
             vmfb_path="streaming_llama.vmfb",
-            mod=self.mod,
+            mod=None,
             tokenizer=self.tokenizer,
         )
 
@@ -169,7 +172,7 @@ def test_streaming_vmfb_comparison(self):
                 None,
                 DEFAULT_PROMPT,
                 streaming_llm=True,
-                model=self.mod,
+                model=None,
                 tokenizer=self.tokenizer,
             )
 
@@ -204,7 +207,7 @@ def test_rerotated_torch_comparison(self):
             None,
             DEFAULT_PROMPT,
             streaming_llm=True,
-            model=self.mod,
+            model=None,
             tokenizer=self.tokenizer,
         )
         check_output_string(torch_str, rotated_torch_str)