feat: Added testing helpers for state dicts

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

Author: BrandonGroth

tests/models/test_model_utils.py

@@ -168,6 +168,7 @@ def load_json(file_path: str = "qcfg.json"):
     assert json_file is not None, f"JSON at {file_path} was not found"
     return json_file
 
+
 def save_json(data, file_path: str = "qcfg.json"):
     """
     Save data object to json file
@@ -179,6 +180,7 @@ def save_json(data, file_path: str = "qcfg.json"):
     with open(file_path, "w", encoding="utf-8") as outfile:
         json.dump(data, outfile, indent=4)
 
+
 def save_serialized_json(config: dict, file_path: str = "qcfg.json"):
     """
     Save qconfig by serializing it first
@@ -195,3 +197,44 @@ def save_serialized_json(config: dict, file_path: str = "qcfg.json"):
 
     serialize_config(config)  # Only remove stuff necessary to dump
     save_json(config, file_path)
+
+
+################################
+# General state dict functions #
+################################
+
+
+def load_state_dict(fname: str = "qmodel_for_aiu.pt") -> dict:
+    """
+    Load a model state dict .pt file.
+
+    Args:
+        fname (str, optional): File for state dict of model. Defaults to "qmodel_for_aiu.pt".
+
+    Returns:
+        dict: Model state dictionary
+    """
+    return torch.load(fname)
+
+
+def check_linear_dtypes(state_dict: dict, linear_names: list):
+    """
+    Checks a state dict for proper dtypes of linear names saved for AIU
+
+    Args:
+        state_dict (dict): Saved model state dict
+        linear_names (list): List of layer names that correspond to torch.nn.Linear layers
+    """
+    assert state_dict is not None
+
+    # Check all quantized linear layers are int8 and everything else is fp16
+    assert all(
+        v.dtype == torch.int8
+        for k, v in state_dict.items()
+        if any(n in k for n in linear_names) and k.endswith(".weight")
+    )
+    assert all(
+        v.dtype == torch.float16
+        for k, v in state_dict.items()
+        if all(n not in k for n in linear_names) or not k.endswith(".weight")
+    )