Bringup tt-torch models in forge (#1314)

## Summary 

- This PR addresses [Issue
#1321](#1321) by
porting models from tt-torch to tt-forge. The current compilation status
of these models are available in the logs.

- A list of skipped models along with the reasons for their exclusion
can be found in [Issue
#1339](#1339). Any
models that were missed in this PR will be included in [PR
#1337](#1337).

**Note:**

- The following test cases are passing end-to-end (E2E), so a push
marker & post processing steps has been added:
```
forge/test/models/pytorch/text/albert/test_albert.py::test_albert_sequence_classification_pytorch
forge/test/models/pytorch/text/perceiverio/test_perceiverio.py::test_perceiverio_masked_lm_pytorch
forge/test/models/pytorch/text/bert/test_bert.py::test_bert_question_answering_pytorch[phiyodr/bert-large-finetuned-squad2]
```
    

**Logs:**

-
[albert_qa_seq_cls.log](https://github.com/user-attachments/files/19051541/mar3_albert_ques_seq.log)
-
[beit.log](https://github.com/user-attachments/files/19006872/feb25_beit.log)
-
[bloom.log](https://github.com/user-attachments/files/19006874/feb25_bloom.log)
-
[glpn_kitti.log](https://github.com/user-attachments/files/19006875/feb25_glpn_kitti.log)
-
[mamba.log](https://github.com/user-attachments/files/19006877/feb25_mamba.log)
-
[mobilenetv2_torchvision.log](https://github.com/user-attachments/files/19006878/feb25_mobilenetv2_torchvision.log)
-
[qwen2_tkn_cls.log](https://github.com/user-attachments/files/19006879/feb25_qwen2_tkn_cls.log)
-
[rmbg.log](https://github.com/user-attachments/files/19006880/feb25_rmbg.log)
-
[unet_carvana.log](https://github.com/user-attachments/files/19006881/feb25_unet_carvana.log)
-
[yolos.log](https://github.com/user-attachments/files/19006882/feb25_yolos.log)
-
[bert_phiyodr.log](https://github.com/user-attachments/files/19051481/mar3_bert_ques_forge.log)
-
[gpt_seq_cls.log](https://github.com/user-attachments/files/19006885/feb26_gpt_seq_cls.log)
-
[mgp.log](https://github.com/user-attachments/files/19006886/feb26_mgp.log)
-
[mlp_mixer.log](https://github.com/user-attachments/files/19006887/feb26_mlp_mixer.log)
-
[mnist.log](https://github.com/user-attachments/files/19006888/feb26_mnist.log)
-
[mobilenetv3_ssd.log](https://github.com/user-attachments/files/19006889/feb26_mobilenetv3_ssd.log)
-
[regnet_torchvision.log](https://github.com/user-attachments/files/19006890/feb26_regnet_torchvision.log)
-
[swin_torchvision.log](https://github.com/user-attachments/files/19006891/feb26_swin_torchvision.log)
-
[vgg_torchvision.log](https://github.com/user-attachments/files/19006892/feb26_vgg_torchvision.log)
-
[vit_torchvision.log](https://github.com/user-attachments/files/19006893/feb26_vit_torchvision.log)
-
[perceiver_mlm.log](https://github.com/user-attachments/files/19051489/mar3_perceiverio_1.log)
-
[speecht5_tts.log](https://github.com/user-attachments/files/19006896/feb27_speech_to_text.log)

env/core_requirements.txt

@@ -51,3 +51,5 @@ pytorch_forecasting==1.0.0
 patool
 openpyxl==3.1.5
 GitPython==3.1.44
+kornia==0.8.0
+mlp-mixer-pytorch==0.2.0

forge/test/models/pytorch/text/albert/test_albert.py

@@ -5,8 +5,11 @@
 import torch
 from transformers import (
     AlbertForMaskedLM,
+    AlbertForQuestionAnswering,
+    AlbertForSequenceClassification,
     AlbertForTokenClassification,
     AlbertTokenizer,
+    AutoTokenizer,
 )
 
 import forge
@@ -162,3 +165,83 @@ def test_albert_token_classification_pytorch(record_forge_property, size, varian
 
     print(f"Context: {sample_text}")
     print(f"Answer: {predicted_tokens_classes}")
+
+
+@pytest.mark.nightly
+@pytest.mark.parametrize("variant", ["twmkn9/albert-base-v2-squad2"])
+def test_albert_question_answering_pytorch(record_forge_property, variant):
+
+    # Build Module Name
+    module_name = build_module_name(
+        framework=Framework.PYTORCH,
+        model="albert",
+        variant=variant,
+        task=Task.QA,
+        source=Source.HUGGINGFACE,
+    )
+
+    # Record Forge Property
+    record_forge_property("tags.model_name", module_name)
+
+    # Load Albert tokenizer and model from HuggingFace
+    tokenizer = download_model(AutoTokenizer.from_pretrained, variant)
+    framework_model = download_model(AlbertForQuestionAnswering.from_pretrained, variant, return_dict=False)
+    framework_model.eval()
+
+    # Load data sample
+    question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
+
+    # Data preprocessing
+    input_tokens = tokenizer(question, text, return_tensors="pt")
+    inputs = [input_tokens["input_ids"], input_tokens["attention_mask"]]
+
+    # Forge compile framework model
+    compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name)
+
+    # Model Verification
+    verify(inputs, framework_model, compiled_model)
+
+
+@pytest.mark.nightly
+@pytest.mark.push
+@pytest.mark.parametrize("variant", ["textattack/albert-base-v2-imdb"])
+def test_albert_sequence_classification_pytorch(record_forge_property, variant):
+
+    # Build Module Name
+    module_name = build_module_name(
+        framework=Framework.PYTORCH,
+        model="albert",
+        variant=variant,
+        task=Task.SEQUENCE_CLASSIFICATION,
+        source=Source.HUGGINGFACE,
+    )
+
+    # Record Forge Property
+    record_forge_property("tags.model_name", module_name)
+
+    # Load Albert tokenizer and model from HuggingFace
+    tokenizer = download_model(AlbertTokenizer.from_pretrained, variant)
+    framework_model = download_model(AlbertForSequenceClassification.from_pretrained, variant, return_dict=False)
+    framework_model.eval()
+
+    # Load data sample
+    input_text = "Hello, my dog is cute."
+
+    # Data preprocessing
+    input_tokens = tokenizer(input_text, return_tensors="pt")
+    inputs = [input_tokens["input_ids"], input_tokens["attention_mask"]]
+
+    # Forge compile framework model
+    compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name)
+
+    # Model Verification
+    verify(inputs, framework_model, compiled_model)
+
+    # Inference
+    co_out = compiled_model(*inputs)
+
+    # post processing
+    predicted_class_id = co_out[0].argmax().item()
+    predicted_category = framework_model.config.id2label[predicted_class_id]
+
+    print(f"predicted category: {predicted_category}")

forge/test/models/pytorch/text/bert/test_bert.py

@@ -91,13 +91,20 @@ def generate_model_bert_qa_hf_pytorch(variant):
         return_tensors="pt",
     )
 
-    return model, [input_tokens["input_ids"]], {}
+    return model, [input_tokens["input_ids"]], tokenizer
+
+
+variants = [
+    pytest.param("phiyodr/bert-large-finetuned-squad2", marks=[pytest.mark.push]),
+    pytest.param("bert-large-cased-whole-word-masking-finetuned-squad"),
+]
 
 
 @pytest.mark.nightly
-@pytest.mark.parametrize("variant", ["bert-large-cased-whole-word-masking-finetuned-squad"])
+@pytest.mark.parametrize("variant", variants)
 def test_bert_question_answering_pytorch(record_forge_property, variant):
-    pytest.skip("Skipping due to the current CI/CD pipeline limitations")
+    if variant == "bert-large-cased-whole-word-masking-finetuned-squad":
+        pytest.skip("Skipping due to the current CI/CD pipeline limitations")
 
     # Build Module Name
     module_name = build_module_name(
@@ -107,14 +114,30 @@ def test_bert_question_answering_pytorch(record_forge_property, variant):
     # Record Forge Property
     record_forge_property("tags.model_name", module_name)
 
-    framework_model, inputs, _ = generate_model_bert_qa_hf_pytorch(variant)
+    framework_model, inputs, tokenizer = generate_model_bert_qa_hf_pytorch(variant)
+    framework_model.eval()
 
     # Forge compile framework model
     compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name)
 
     # Model Verification
     verify(inputs, framework_model, compiled_model, verify_cfg=VerifyConfig(verify_values=False))
 
+    # Inference
+    output = compiled_model(*inputs)
+
+    # post processing
+    start_logits = output[0]
+    end_logits = output[1]
+
+    answer_start_index = start_logits.argmax()
+    answer_end_index = end_logits.argmax()
+
+    input_ids = inputs[0]
+    predict_answer_tokens = input_ids[0, answer_start_index : answer_end_index + 1]
+
+    print("predicted answer ", tokenizer.decode(predict_answer_tokens, skip_special_tokens=True))
+
 
 def generate_model_bert_seqcls_hf_pytorch(variant):
     # Load Bert tokenizer and model from HuggingFace

forge/test/models/pytorch/text/bloom/test_bloom.py

@@ -0,0 +1,64 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import pytest
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+import forge
+from forge.verify.verify import verify
+
+from test.models.utils import Framework, Source, Task, build_module_name
+from test.utils import download_model
+
+
+# Wrapper to get around past key values
+class Wrapper(torch.nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, input_ids, attention_mask):
+        output = self.model(input_ids, None, attention_mask)
+        return output
+
+
+@pytest.mark.nightly
+@pytest.mark.parametrize("variant", ["bigscience/bloom-1b1"])
+def test_bloom(record_forge_property, variant):
+
+    # Build Module Name
+    module_name = build_module_name(
+        framework=Framework.PYTORCH,
+        model="bloom",
+        variant=variant,
+        source=Source.HUGGINGFACE,
+        task=Task.CAUSAL_LM,
+    )
+
+    # Record Forge Property
+    record_forge_property("tags.model_name", module_name)
+
+    # Load tokenizer and model from HuggingFace
+    tokenizer = download_model(AutoTokenizer.from_pretrained, variant, padding_side="left")
+    model = download_model(AutoModelForCausalLM.from_pretrained, variant, use_cache=False, return_dict=False)
+    model.eval()
+    framework_model = Wrapper(model)
+
+    # Prepare input
+    test_input = "This is a sample text from "
+    input_tokens = tokenizer.encode_plus(
+        test_input,
+        return_tensors="pt",
+        max_length=32,
+        padding="max_length",
+        add_special_tokens=True,
+        truncation=True,
+    )
+    inputs = [input_tokens["input_ids"], input_tokens["attention_mask"]]
+
+    # Forge compile framework model
+    compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name)
+
+    # Model Verification
+    verify(inputs, framework_model, compiled_model)

forge/test/models/pytorch/text/gpt2/test_gpt2.py

@@ -3,7 +3,12 @@
 # SPDX-License-Identifier: Apache-2.0
 import pytest
 import torch
-from transformers import GPT2Config, GPT2LMHeadModel
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    GPT2Config,
+    GPT2LMHeadModel,
+)
 
 import forge
 from forge.verify.verify import verify
@@ -12,6 +17,16 @@
 from test.utils import download_model
 
 
+# Wrapper to get around past key values
+class Wrapper(torch.nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, input_ids, attention_mask):
+        return self.model(input_ids, None, attention_mask)
+
+
 @pytest.mark.nightly
 @pytest.mark.parametrize("variant", ["gpt2"])
 def test_gpt2_text_gen(record_forge_property, variant):
@@ -31,15 +46,6 @@ def test_gpt2_text_gen(record_forge_property, variant):
     config = GPT2Config(**config_dict)
     model = download_model(GPT2LMHeadModel.from_pretrained, variant, config=config)
 
-    # Wrapper to get around past key values
-    class Wrapper(torch.nn.Module):
-        def __init__(self, model):
-            super().__init__()
-            self.model = model
-
-        def forward(self, input_ids, attention_mask):
-            return self.model(input_ids, None, attention_mask)
-
     input_ids = torch.cat(
         [torch.randint(1, model.config.vocab_size, (1, 255)), torch.zeros(1, 1, dtype=torch.int64)], dim=-1
     ).to(torch.int64)
@@ -53,3 +59,37 @@ def forward(self, input_ids, attention_mask):
 
     # Model Verification
     verify(inputs, framework_model, compiled_model)
+
+
+@pytest.mark.nightly
+@pytest.mark.parametrize("variant", ["mnoukhov/gpt2-imdb-sentiment-classifier"])
+def test_gpt2_sequence_classification(record_forge_property, variant):
+
+    # Build Module Name
+    module_name = build_module_name(
+        framework=Framework.PYTORCH,
+        model="gpt2",
+        variant=variant,
+        task=Task.SEQUENCE_CLASSIFICATION,
+        source=Source.HUGGINGFACE,
+    )
+
+    # Record Forge Property
+    record_forge_property("tags.model_name", module_name)
+
+    # Load tokenizer and model from HuggingFace
+    tokenizer = download_model(AutoTokenizer.from_pretrained, variant, padding_side="left")
+    model = download_model(AutoModelForSequenceClassification.from_pretrained, variant, return_dict=False)
+    model.eval()
+    framework_model = Wrapper(model)
+
+    # Prepare input
+    test_input = "This is a sample text from "
+    input_tokens = tokenizer(test_input, return_tensors="pt")
+    inputs = [input_tokens["input_ids"], input_tokens["attention_mask"]]
+
+    # Forge compile framework model
+    compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name)
+
+    # Model Verification
+    verify(inputs, framework_model, compiled_model)

forge/test/models/pytorch/text/mamba/test_mamba.py

@@ -0,0 +1,64 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+# Reference: https://huggingface.co/state-spaces/mamba-2.8b-hf
+
+import pytest
+import torch
+from transformers import AutoTokenizer, MambaForCausalLM
+
+import forge
+from forge.verify.verify import verify
+
+from test.models.utils import Framework, Source, Task, build_module_name
+from test.utils import download_model
+
+
+# Wrapper to return only the output tensor, excluding cache or additional outputs
+class Wrapper(torch.nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, input_ids):
+        output = self.model(input_ids)
+        return output[0]
+
+
+variants = [
+    "state-spaces/mamba-790m-hf",
+    "state-spaces/mamba-2.8b-hf",
+    "state-spaces/mamba-1.4b-hf",
+    "state-spaces/mamba-370m-hf",
+]
+
+
+@pytest.mark.nightly
+@pytest.mark.parametrize("variant", variants)
+def test_mamba(record_forge_property, variant):
+    if variant != "state-spaces/mamba-790m-hf":
+        pytest.skip("Skipping this variant; only testing the base model (mamba-790m-hf) for now.")
+
+    # Build Module Name
+    module_name = build_module_name(
+        framework=Framework.PYTORCH, model="mamba", variant=variant, task=Task.CAUSAL_LM, source=Source.HUGGINGFACE
+    )
+
+    # Record Forge Property
+    record_forge_property("tags.model_name", module_name)
+
+    # Load tokenizer and model from HuggingFace
+    tokenizer = download_model(AutoTokenizer.from_pretrained, variant)
+    model = download_model(MambaForCausalLM.from_pretrained, variant)
+    model.eval()
+    framework_model = Wrapper(model)
+
+    # Prepare input
+    prompt = "Hey how are you doing?"
+    inputs = [tokenizer(prompt, return_tensors="pt")["input_ids"]]
+
+    # Forge compile framework model
+    compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name)
+
+    # Model Verification
+    verify(inputs, framework_model, compiled_model)