fix: Changed input_tiny to dataloader to work with qmodel_prep

Signed-off-by: Brandon Groth <[email protected]>

Author: BrandonGroth

tests/models/conftest.py

@@ -25,7 +25,6 @@
 from torchvision.io import read_image
 from torchvision.models import ResNet50_Weights, ViT_B_16_Weights, resnet50, vit_b_16
 from transformers import (
-    BatchEncoding,
     BertConfig,
     BertModel,
     BertTokenizer,
@@ -38,6 +37,7 @@
 import pytest
 import torch
 import torch.nn.functional as F
+from torch.utils.data import TensorDataset, DataLoader
 
 # Local
 # fms_mo imports
@@ -1327,25 +1327,27 @@ def model_residualMLP():
 size = (batch_size, max_position_embeddings)
 
 
-@pytest.fixture(scope="session")
-def input_tiny() -> BatchEncoding:
+@pytest.fixture(scope="function")
+def input_tiny() -> DataLoader:
     """
     Create a fake input for tiny models w/ fixed vocab_size and max_position_embeddings
 
     Returns:
-        BatchEncoding: Fake Encoding for a Tokenizer
+        DataLoader: Fake Encoding for a Tokenizer
     """
     # Random tokens and attention mask == 1
     random_tokens = torch.randint(low=0, high=vocab_size, size=size)
     attention_mask = torch.ones(size)
 
-    fake_tokenizer_output = BatchEncoding(
-        {
-            "input_ids": random_tokens,
-            "attention_mask": attention_mask,
-        }
+    dataset = TensorDataset(random_tokens, attention_mask)
+    # qmodel_prep expects dataloader batch=tuple(tensor, tensor)
+    # Without collate_fn, it returns batch=list(tensor,tensor)
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        collate_fn=lambda batch: tuple(torch.stack(samples) for samples in zip(*batch))
     )
-    return fake_tokenizer_output
 
 
 #############################
@@ -1398,6 +1400,7 @@ def model_tiny_bert(config_tiny_bert: BertConfig) -> BertModel:
         "nbits_w": 8,
         "qa_mode": "pertokenmax",
         "qw_mode": "max",
+        "qmodel_calibration": 1,
         "smoothq": False,
         "smoothq_scale_layers": [],
         "qskip_layer_name": [
@@ -1414,6 +1417,7 @@ def model_tiny_bert(config_tiny_bert: BertConfig) -> BertModel:
         "nbits_w": 8,
         "qa_mode": "maxsym",
         "qw_mode": "maxperCh",
+        "qmodel_calibration": 1,
         "smoothq": False,
         "smoothq_scale_layers": [],
         "qskip_layer_name": [
@@ -1512,6 +1516,7 @@ def model_tiny_llama(config_tiny_llama: LlamaConfig) -> LlamaModel:
         "nbits_w": 8,
         "qa_mode": "pertokenmax",
         "qw_mode": "max",
+        "qmodel_calibration": 1,
         "smoothq": False,
         "smoothq_scale_layers": [],
         "qskip_layer_name": [
@@ -1611,6 +1616,7 @@ def model_tiny_granite(config_tiny_granite: GraniteConfig) -> GraniteModel:
         "nbits_w": 8,
         "qa_mode": "pertokenmax",
         "qw_mode": "maxperCh",
+        "qmodel_calibration": 1,
         "smoothq": False,
         "smoothq_scale_layers": ["k_proj", "v_proj", "gate_proj", "up_proj"],
         "qskip_layer_name": ["lm_head"],