Adding unit tests for run_quant.py

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

Author: tharapalanivel

fms_mo/utils/dq_utils.py

@@ -38,7 +38,7 @@ def config_quantize_smooth_layers(qcfg):
         "granite-20b-code",
         "granite-20b-code",
     ]
-    if any(model in qcfg["model"] for model in llama_architecture):
+    if any(model in qcfg["model"] for model in llama_architecture) or any(model in qcfg["model_type"] for model in llama_architecture):
         qcfg["qlayer_name_pattern"] = ["model.layers."]
         qcfg["scale_layers"] = ["k_proj", "v_proj", "gate_proj", "up_proj"]
         qcfg["qskip_layer_name"] = []

tests/test_run_quant.py

@@ -0,0 +1,136 @@
+# Copyright The FMS Model Optimizer Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for run_quant.py"""
+
+# Standard
+import copy
+import json
+import os
+
+# Third Party
+import pytest
+
+# Local
+from fms_mo.run_quant import get_parser, parse_arguments, quantize
+from fms_mo.training_args import (
+    DataArguments,
+    FMSMOArguments,
+    FP8Arguments,
+    GPTQArguments,
+    ModelArguments,
+    OptArguments,
+)
+from tests.artifacts.testdata import MODEL_NAME, WIKITEXT_TOKENIZED_DATA_JSON
+
+MODEL_ARGS = ModelArguments(
+    model_name_or_path=MODEL_NAME, torch_dtype="float16"
+)
+DATA_ARGS = DataArguments(
+    training_data_path=WIKITEXT_TOKENIZED_DATA_JSON,
+)
+OPT_ARGS = OptArguments(quant_method="dq", output_dir="tmp")
+GPTQ_ARGS = GPTQArguments(
+    bits=4,
+    group_size=64,
+)
+FP8_ARGS = FP8Arguments()
+DQ_ARGS = FMSMOArguments(
+    nbits_w=8,
+    nbits_a=8,
+    nbits_kvcache=32,
+    qa_mode="fp8_e4m3_scale",
+    qw_mode="fp8_e4m3_scale",
+    qmodel_calibration_new=0,
+)
+
+
+def test_run_train_requires_output_dir():
+    """Check fails when output dir not provided."""
+    updated_output_dir_opt_args = copy.deepcopy(OPT_ARGS)
+    updated_output_dir_opt_args.output_dir = None
+    with pytest.raises(TypeError):
+        quantize(
+            model_args=MODEL_ARGS,
+            data_args=DATA_ARGS,
+            opt_args=updated_output_dir_opt_args,
+            fms_mo_args=DQ_ARGS,
+        )
+
+
+def test_run_train_fails_training_data_path_not_exist():
+    """Check fails when data path not found."""
+    updated_data_path_args = copy.deepcopy(DATA_ARGS)
+    updated_data_path_args.training_data_path = "fake/path"
+    with pytest.raises(FileNotFoundError):  # TPP Should this be FileNotFoundError or ValueError?
+        quantize(
+            model_args=MODEL_ARGS,
+            data_args=updated_data_path_args,
+            opt_args=OPT_ARGS,
+            fms_mo_args=DQ_ARGS,
+        )
+
+
+HAPPY_PATH_DUMMY_CONFIG_PATH = os.path.join(
+    os.path.dirname(__file__), "artifacts", "configs", "dummy_job_config.json"
+)
+
+
+# Note: job_config dict gets modified during process training args
+@pytest.fixture(name="job_config", scope="session")
+def fixture_job_config():
+    with open(HAPPY_PATH_DUMMY_CONFIG_PATH, "r", encoding="utf-8") as f:
+        dummy_job_config_dict = json.load(f)
+    return dummy_job_config_dict
+
+
+############################# Arg Parsing Tests #############################
+
+
+def test_parse_arguments(job_config):
+    parser = get_parser()
+    job_config_copy = copy.deepcopy(job_config)
+    (
+        model_args,
+        data_args,
+        opt_args,
+        fms_mo_args,
+        gptq_args,
+        fp8_args,
+    ) = parse_arguments(parser, job_config_copy)
+    assert str(model_args.torch_dtype) == "torch.bfloat16"
+    assert data_args.training_data_path == "data_train"
+    assert opt_args.output_dir == "models/Maykeye/TinyLLama-v0-GPTQ"
+    assert opt_args.quant_method == "gptq"
+
+
+def test_parse_arguments_defaults(job_config):
+    parser = get_parser()
+    job_config_defaults = copy.deepcopy(job_config)
+    assert "torch_dtype" not in job_config_defaults
+    assert "max_seq_length" not in job_config_defaults
+    assert "model_revision" not in job_config_defaults
+    assert "nbits_kvcache" not in job_config_defaults
+    (
+        model_args,
+        data_args,
+        opt_args,
+        fms_mo_args,
+        gptq_args,
+        fp8_args,
+    ) = parse_arguments(parser, job_config_defaults)
+    assert str(model_args.torch_dtype) == "torch.bfloat16"
+    assert model_args.model_revision == "main"
+    assert data_args.max_seq_length == 2048
+    assert fms_mo_args.nbits_kvcache == 32