Fix invalid type checking test

Signed-off-by: Thara Palanivel <[email protected]>

Author: tharapalanivel

fms_mo/training_args.py

@@ -18,14 +18,34 @@
 
 # Standard
 from dataclasses import dataclass, field
-from typing import List, Optional, Union
+from typing import Optional, Union
 
 # Third Party
 import torch
 
 
 @dataclass
-class ModelArguments:
+class TypeChecker:
+    def __post_init__(self):
+        for name, field_type in self.__annotations__.items():
+            val = self.__dict__[name]
+            invalid_val = False
+            if not field_type is list:
+                if not isinstance(val, field_type):
+                    invalid_val = True
+            else:
+                if not isinstance(val, list) or not all(isinstance(item, int) for item in val):
+                    invalid_val = True
+
+            if invalid_val:
+                current_type = type(val)
+                raise TypeError(
+                    f"The field `{name}` was assigned by `{current_type}` instead of `{field_type}`"
+                )
+
+
+@dataclass
+class ModelArguments(TypeChecker):
     """Dataclass for model related arguments."""
 
     model_name_or_path: Optional[str] = field(default="facebook/opt-125m")
@@ -60,24 +80,26 @@ class ModelArguments:
             )
         },
     )
-    device: str = field(
+    device: Optional[str] = field(
         default=None,
         metadata={
-            "help": ("`torch.device`: The device on which the module is (assuming that all the module parameters are on the same device).")
-        }
+            "help": (
+                "`torch.device`: The device on which the module is (assuming that all the module parameters are on the same device)."
+            )
+        },
     )
 
 
 @dataclass
-class DataArguments:
+class DataArguments(TypeChecker):
     """Dataclass for data related arguments."""
 
-    training_data_path: str = field(
+    training_data_path: Optional[str] = field(
         default=None,
         metadata={"help": "Path to the training data in JSON/JSONL format"},
     )
-    training_data_config: str = field(default=None)
-    test_data_path: str = field(
+    training_data_config: Optional[str] = field(default=None)
+    test_data_path: Optional[str] = field(
         default=None,
         metadata={"help": "Path to the test data in JSON/JSONL format"},
     )
@@ -86,7 +108,7 @@ class DataArguments:
 
 
 @dataclass
-class OptArguments:
+class OptArguments(TypeChecker):
     """Dataclass for optimization related arguments."""
 
     quant_method: str = field(
@@ -104,7 +126,7 @@ class OptArguments:
 
 
 @dataclass
-class FMSMOArguments:
+class FMSMOArguments(TypeChecker):
     """Dataclass arguments used by fms_mo native quantization functions."""
 
     nbits_w: int = field(default=32, metadata={"help": ("weight precision")})
@@ -139,7 +161,7 @@ class FMSMOArguments:
 
 
 @dataclass
-class GPTQArguments:
+class GPTQArguments(TypeChecker):
     """Dataclass for GPTQ related arguments that will be used by auto-gptq."""
 
     bits: int = field(default=4, metadata={"choices": [2, 3, 4, 8]})
@@ -157,9 +179,9 @@ class GPTQArguments:
 
 
 @dataclass
-class FP8Arguments:
+class FP8Arguments(TypeChecker):
     """Dataclass for FP8 related arguments that will be used by llm-compressor."""
 
     targets: str = field(default="Linear")
     scheme: str = field(default="FP8_DYNAMIC")
-    ignore: List[str] = field(default_factory=lambda: ["lm_head"])
+    ignore: list[str] = field(default_factory=lambda: ["lm_head"])

tests/build/test_launch_script.py

@@ -209,6 +209,10 @@ def test_config_parsing_error():
         assert pytest_wrapped_e.value.code == USER_ERROR_EXIT_CODE
         assert os.stat(tempdir + "/termination-log").st_size > 0
 
+        with open(tempdir + "/termination-log", "r", encoding="utf-8") as f:
+            contents = f.read()
+        assert contents == "Exception raised during optimization. This may be a problem with your input: The field `nbits_w` was assigned by `<class 'str'>` instead of `<class 'int'>`"
+
 
 def _validate_termination_files_when_quantization_succeeds(base_dir):
     # Check termination log and .complete files exist