Move generic opt params to OptArgs

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

Author: tharapalanivel

fms_mo/dq.py

@@ -51,7 +51,7 @@
 logger = logging.getLogger(__name__)
 
 
-def run_dq(model_args, data_args, fms_mo_args, output_dir):
+def run_dq(model_args, data_args, opt_args, fms_mo_args):
     """
     For direct quantization LLMs without optimization:
     Models are directly quantized into INT8 or FP8 precisions using
@@ -63,8 +63,9 @@ def run_dq(model_args, data_args, fms_mo_args, output_dir):
             the model
         data_args (fms_mo.training_args.DataArguments): Data arguments to be used when loading the
             tokenized dataset
+        opt_args (fms_mo.training_args.OptArguments): Generic optimization arguments to be used
+            during DQ
         fms_mo_args (fms_mo.training_args.FMSMOArguments): Parameters to use for DQ quantization
-        output_dir (str) Output directory to write to
     """
     # for attention or kv-cache quantization, need to use eager attention
     attn_bits = [
@@ -218,9 +219,9 @@ def run_dq(model_args, data_args, fms_mo_args, output_dir):
                 with patch_torch_bmm(qcfg):
                     model(**data_mb)
 
-    logger.info(f"Saving quantized model and tokenizer to {output_dir}")
-    model.save_pretrained(output_dir, use_safetensors=True)
-    tokenizer.save_pretrained(output_dir)
+    logger.info(f"Saving quantized model and tokenizer to {opt_args.output_dir}")
+    model.save_pretrained(opt_args.output_dir, use_safetensors=True)
+    tokenizer.save_pretrained(opt_args.output_dir)
 
     if fms_mo_args.eval_ppl:
         path_test = Path(data_args.test_data_path)

fms_mo/run_quant.py

@@ -42,6 +42,7 @@
     FP8Arguments,
     GPTQArguments,
     ModelArguments,
+    OptArguments,
 )
 from fms_mo.utils.import_utils import available_packages
 
@@ -51,11 +52,10 @@
 def quantize(
     model_args: ModelArguments,
     data_args: DataArguments,
-    fms_mo_args: FMSMOArguments,
-    gptq_args: GPTQArguments,
-    fp8_args: FP8Arguments,
-    quant_method: str,
-    output_dir: str,
+    opt_args: OptArguments,
+    fms_mo_args: FMSMOArguments = None,
+    gptq_args: GPTQArguments = None,
+    fp8_args: FP8Arguments = None,
 ):
     """Main entry point to quantize a given model with a set of specified hyperparameters
 
@@ -71,16 +71,17 @@ def quantize(
         output_dir (str) Output directory to write to
     """
 
-    logging.info(f"{fms_mo_args}\n{quant_method}\n")
-    if quant_method == "gptq":
+    logger.info(f"{fms_mo_args}\n{opt_args.quant_method}\n")
+
+    if opt_args.quant_method == "gptq":
         if not available_packages["auto_gptq"]:
             raise ImportError(
                 "Quantization method has been selected as gptq but unable to use external library, "
                 "auto_gptq module not found. For more instructions on installing the appropriate "
                 "package, see https://github.com/AutoGPTQ/AutoGPTQ?tab=readme-ov-file#installation"
             )
-        run_gptq(model_args, data_args, gptq_args, output_dir)
-    elif quant_method == "fp8":
+        run_gptq(model_args, data_args, opt_args, gptq_args)
+    elif opt_args.quant_method == "fp8":
         if not available_packages["llmcompressor"]:
             raise ImportError(
                 "Quantization method has been selected as fp8 but unable to use external library, "
@@ -89,16 +90,18 @@ def quantize(
                 "https://github.com/vllm-project/llm-compressor/tree/"
                 "main?tab=readme-ov-file#installation"
             )
-        run_fp8(model_args, data_args, fp8_args, output_dir)
-    elif quant_method == "dq":
-        run_dq(model_args, data_args, fms_mo_args, output_dir)
+        run_fp8(model_args, data_args, opt_args, fp8_args)
+    elif opt_args.quant_method == "dq":
+        run_dq(model_args, data_args, opt_args, fms_mo_args)
     else:
         raise ValueError(
-            "Not a valid quantization technique option. Please choose from: gptq, fp8, dq"
+            "{} is not a valid quantization technique option. Please choose from: gptq, fp8, dq".format(
+                opt_args.quant_method
+            )
         )
 
 
-def run_gptq(model_args, data_args, gptq_args, output_dir):
+def run_gptq(model_args, data_args, opt_args, gptq_args):
     """GPTQ quantizes a given model with a set of specified hyperparameters
 
     Args:
@@ -152,14 +155,16 @@ def run_gptq(model_args, data_args, gptq_args, output_dir):
         cache_examples_on_gpu=gptq_args.cache_examples_on_gpu,
     )
 
-    logger.info(f"Time to quantize model at {output_dir}: {time.time() - start_time}")
+    logger.info(
+        f"Time to quantize model at {opt_args.output_dir}: {time.time() - start_time}"
+    )
 
-    logger.info(f"Saving quantized model and tokenizer to {output_dir}")
-    model.save_quantized(output_dir, use_safetensors=True)
-    tokenizer.save_pretrained(output_dir)
+    logger.info(f"Saving quantized model and tokenizer to {opt_args.output_dir}")
+    model.save_quantized(opt_args.output_dir, use_safetensors=True)
+    tokenizer.save_pretrained(opt_args.output_dir)
 
 
-def run_fp8(model_args, data_args, fp8_args, output_dir):
+def run_fp8(model_args, data_args, opt_args, fp8_args):
     """FP8 quantizes a given model with a set of specified hyperparameters
 
     Args:
@@ -192,11 +197,13 @@ def run_fp8(model_args, data_args, fp8_args, output_dir):
         max_seq_length=data_args.max_seq_length,
         num_calibration_samples=data_args.num_calibration_samples,
     )
-    logger.info(f"Time to quantize model at {output_dir}: {time.time() - start_time}")
+    logger.info(
+        f"Time to quantize model at {opt_args.output_dir}: {time.time() - start_time}"
+    )
 
-    logger.info(f"Saving quantized model and tokenizer to {output_dir}")
-    model.save_pretrained(output_dir)
-    tokenizer.save_pretrained(output_dir)
+    logger.info(f"Saving quantized model and tokenizer to {opt_args.output_dir}")
+    model.save_pretrained(opt_args.output_dir)
+    tokenizer.save_pretrained(opt_args.output_dir)
 
 
 def main():
@@ -206,53 +213,41 @@ def main():
         dataclass_types=(
             ModelArguments,
             DataArguments,
+            OptArguments,
             FMSMOArguments,
             GPTQArguments,
             FP8Arguments,
         )
     )
 
-    parser.add_argument(
-        "--quant_method",
-        type=str.lower,
-        choices=["gptq", "fp8", None, "none", "dq"],
-        default="none",
-    )
-
-    parser.add_argument("--output_dir", type=str)
-
     (
         model_args,
         data_args,
+        opt_args,
         fms_mo_args,
         gptq_args,
         fp8_args,
-        additional,
         _,
     ) = parser.parse_args_into_dataclasses(return_remaining_strings=True)
-    quant_method = additional.quant_method
-    output_dir = additional.output_dir
 
     logger.debug(
-        "Input args parsed: \nmodel_args %s, data_args %s, fms_mo_args %s, "
-        "gptq_args %s, fp8_args %s, quant_method %s, output_dir %s",
+        "Input args parsed: \nmodel_args %s, data_args %s, opt_args %s, fms_mo_args %s, "
+        "gptq_args %s, fp8_args %s",
         model_args,
         data_args,
+        opt_args,
         fms_mo_args,
         gptq_args,
         fp8_args,
-        quant_method,
-        output_dir,
     )
 
     quantize(
         model_args=model_args,
         data_args=data_args,
+        opt_args=opt_args,
         fms_mo_args=fms_mo_args,
         gptq_args=gptq_args,
         fp8_args=fp8_args,
-        quant_method=quant_method,
-        output_dir=output_dir,
     )
 
 

fms_mo/training_args.py

@@ -18,18 +18,18 @@
 
 # Standard
 from dataclasses import dataclass, field
-from typing import List, Optional
+from typing import List, Optional, Union
+
+# Third Party
+import torch
 
 
 @dataclass
 class ModelArguments:
     """Dataclass for model related arguments."""
 
     model_name_or_path: Optional[str] = field(default="facebook/opt-125m")
-    torch_dtype: Optional[str] = field(
-        default=None,
-        metadata={"help": ["bfloat16", "float16", "float", "auto"]},
-    )
+    torch_dtype: Optional[Union[torch.dtype, str]] = torch.bfloat16
     use_fast_tokenizer: bool = field(
         default=True,
         metadata={
@@ -79,6 +79,24 @@ class DataArguments:
     num_calibration_samples: Optional[int] = field(default=512)
 
 
+@dataclass
+class OptArguments:
+    """Dataclass for optimization related arguments."""
+
+    quant_method: str = field(
+        metadata={"choices": ["gptq", "fp8", "dq"], "help": "Quantization technique"}
+    )
+    output_dir: str = field(
+        metadata={
+            "help": "Output directory to write quantized model artifacts and log files to"
+        }
+    )
+    log_level: str = field(
+        default="INFO",
+        metadata={"help": "The log level to adopt during optimization."},
+    )
+
+
 @dataclass
 class FMSMOArguments:
     """Dataclass arguments used by fms_mo native quantization functions."""