[RFC] basic model builder infra

examples/llama2_inference/llama2.ipynb

@@ -7,19 +7,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
-    "import numpy as np\n",
-    "from shark_turbine.dynamo.passes import turbine_cpu_pass_pipeline\n",
-    "from torch._export import dynamic_dim\n",
-    "from torch._export.constraints import constrain_as_size, constrain_as_value\n",
     "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
     "import torch\n",
     "from torch.utils import _pytree as pytree\n",
-    "from torch.fx import GraphModule\n",
-    "from torch._export.constraints import constrain_as_size, constrain_as_value\n",
     "from shark_turbine.aot import *\n",
-    "from iree.compiler.ir import Context\n",
-    "from iree import runtime as ireert\n"
+    "from iree.compiler.ir import Context\n"
    ]
   },
   {

python/turbine_models/model_builder.py

@@ -0,0 +1,64 @@
+from transformers import AutoModel, AutoTokenizer, AutoConfig
+import torch
+import shark_turbine.aot as aot
+
+
+class HFTransformerBuilder:
+    """
+    A model builder that uses Hugging Face's transformers library to build a PyTorch model.
+
+    Args:
+        example_input (torch.Tensor): An example input tensor to the model.
+        hf_id (str): The Hugging Face model ID.
+        auto_model (AutoModel): The AutoModel class to use for loading the model.
+        auto_tokenizer (AutoTokenizer): The AutoTokenizer class to use for loading the tokenizer.
+        auto_config (AutoConfig): The AutoConfig class to use for loading the model configuration.
+    """
+
+    def __init__(
+        self,
+        example_input: torch.Tensor,
+        hf_id: str,
+        auto_model: AutoModel = AutoModel,
+        auto_tokenizer: AutoTokenizer = None,
+        auto_config: AutoConfig = None,
+        hf_auth_token=None,
+    ) -> None:
+        self.example_input = example_input
+        self.hf_id = hf_id
+        self.auto_model = auto_model
+        self.auto_tokenizer = auto_tokenizer
+        self.auto_config = auto_config
+        self.hf_auth_token = hf_auth_token
+        self.model = None
+        self.tokenizer = None
+        self.build_model()
+
+    def build_model(self) -> None:
+        """
+        Builds a PyTorch model using Hugging Face's transformers library.
+        """
+        # TODO: check cloud storage for existing ir
+        self.model = self.auto_model.from_pretrained(
+            self.hf_id, auth_token=self.hf_auth_token, config=self.auto_config
+        )
+        if self.auto_tokenizer is not None:
+            self.tokenizer = self.auto_tokenizer.from_pretrained(
+                self.hf_id, auth_token=self.hf_auth_token
+            )
+        else:
+            self.tokenizer = None
+
+    def get_compiled_module(self, save_to: str = None) -> aot.CompiledModule:
+        """
+        Compiles the PyTorch model into a compiled module using SHARK-Turbine's AOT compiler.
+
+        Args:
+            save_to (str): one of: input (Torch IR) or import (linalg).
+
+        Returns:
+            aot.CompiledModule: The compiled module binary.
+        """
+        module = aot.export(self.model, self.example_input)
+        compiled_binary = module.compile(save_to=save_to)
+        return compiled_binary

tests/examples/aot_mlp_test.py

@@ -17,15 +17,13 @@ def _run(local_path: str):
     path = REPO_DIR / local_path
     subprocess.check_call([sys.executable, str(path)])
 
-
 class AOTMLPTest(unittest.TestCase):
     def testMLPExportSimple(self):
         _run("examples/aot_mlp/mlp_export_simple.py")
 
     def testMLPExportSimple(self):
         _run("examples/aot_mlp/mlp_export_dynamic.py")
 
-
 if __name__ == "__main__":
     logging.basicConfig(level=logging.DEBUG)
     unittest.main()

tests/models/hf_dict.py

@@ -0,0 +1,15 @@
+import torch
+
+class ModelData():
+    def __init__(self, input_shape, torch_dtype: torch.dtype, xfail: bool = True):
+        self.xfail = xfail
+        self.input_shape = input_shape
+        self.dtype = torch_dtype
+
+model_dict = {
+    'distilgpt2': ModelData(input_shape=(1, 1), torch_dtype=torch.int64),
+    'gpt2': ModelData(input_shape=(1, 1), torch_dtype=torch.int64),
+    'gpt2-medium': ModelData(input_shape=(1, 1), torch_dtype=torch.int64),
+    'bert-base-uncased': ModelData(input_shape=(1, 1), torch_dtype=torch.int64),
+    'bert-large-uncased': ModelData(input_shape=(1, 1), torch_dtype=torch.int64),
+}

tests/models/transformer_builder_tests.py

@@ -0,0 +1,23 @@
+import unittest
+from turbine_models.model_builder import HFTransformerBuilder
+from hf_dict import model_dict
+import torch
+
+class TestHFTransformerBuilder(unittest.TestCase):
+    pass
+
+def create_test(model_name, model_data):
+    def test(self):
+        example_input = torch.ones(*model_data.input_shape, dtype=model_data.torch_dtype)
+        builder = HFTransformerBuilder(example_input, model_name, auto_tokenizer=None)
+        compiled_module = builder.get_compiled_module()
+        self.assertIsNotNone(compiled_module)
+    return test
+
+for model_name, model_data in model_dict.items():
+    test_method = create_test(model_name, model_data)
+    test_method = unittest.expectedFailure(test_method) if model_data.xfail else test_method
+    setattr(TestHFTransformerBuilder, f'test_{model_name}', test_method)
+
+if __name__ == '__main__':
+    unittest.main()

turbine-models-requirements.txt

@@ -0,0 +1,4 @@
+protobuf
+sentencepiece
+shark_turbine
+transformers @ git+https://github.com/huggingface/transformers.git@7d8ff3629b2725ec43ace99c1a6e87ac1978d433