Merge pull request #43 from tharapalanivel/opt_args

OptArguments

Author: chichun-charlie-liu

.spellcheck-en-custom.txt

@@ -26,10 +26,11 @@ eval
 fms
 fp
 FP
+FP8Arguments
 frac
 gptq
 GPTQ
-GPTQArgs
+GPTQArguments
 graphviz
 GPTQ
 hyperparameters

examples/FP8_QUANT/README.md

@@ -27,7 +27,7 @@ This is an example of mature FP8, which under the hood leverages some functional
 ## QuickStart
 This end-to-end example utilizes the common set of interfaces provided by `fms_mo` for easily applying multiple quantization algorithms with FP8 being the focus of this example. The steps involved are:
 
-1. **FP8 quantization through CLI**. Other arguments could be found here [FP8Args](../../fms_mo/training_args.py#L84).
+1. **FP8 quantization through CLI**. Other arguments could be found here [FP8Arguments](../../fms_mo/training_args.py#L84).
 
     ```bash
     python -m fms_mo.run_quant \
@@ -100,7 +100,7 @@ This end-to-end example utilizes the common set of interfaces provided by `fms_m
     tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
     ```
 
-2. Quantization setting is provided using `QuantizationModifier`, additional settings can be found in [FP8Args](../../fms_mo/training_args.py#L84).
+2. Quantization setting is provided using `QuantizationModifier`, additional settings can be found in [FP8Arguments](../../fms_mo/training_args.py#L84).
 
     ```python
     recipe = QuantizationModifier(

examples/GPTQ/README.md

@@ -32,7 +32,7 @@ This end-to-end example utilizes the common set of interfaces provided by `fms_m
 > - Tokenized data will be saved in `<path_to_save>_train` and `<path_to_save>_test`
 > - If you have trouble downloading Llama family of models from Hugging Face ([LLama models require access](https://www.llama.com/docs/getting-the-models/hugging-face/)), you can use `ibm-granite/granite-8b-code` instead
 
-2. **Quantize the model** using the data generated above, the following command will kick off the quantization job (by invoking `auto_gptq` under the hood.) Additional acceptable arguments can be found here in [GPTQArgs](../../fms_mo/training_args.py#L127).
+2. **Quantize the model** using the data generated above, the following command will kick off the quantization job (by invoking `auto_gptq` under the hood.) Additional acceptable arguments can be found here in [GPTQArguments](../../fms_mo/training_args.py#L127).
 
     ```bash
     python -m fms_mo.run_quant \

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

@@ -39,9 +39,10 @@
 from fms_mo.training_args import (
     DataArguments,
     FMSMOArguments,
-    FP8Args,
-    GPTQArgs,
+    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: GPTQArgs,
-    fp8_args: FP8Args,
-    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
 
@@ -64,23 +64,23 @@ def quantize(
             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 related arguments
         fms_mo_args (fms_mo.training_args.FMSMOArguments): Parameters to use for PTQ quantization
-        gptq_args (fms_mo.training_args.GPTQArgs): Parameters to use for GPTQ quantization
-        fp8_args (fms_mo.training_args.FP8Args): Parameters to use for FP8 quantization
-        quant_method (str): Quantization technique, options are gptq and fp8
-        output_dir (str) Output directory to write to
+        gptq_args (fms_mo.training_args.GPTQArguments): Parameters to use for GPTQ quantization
+        fp8_args (fms_mo.training_args.FP8Arguments): Parameters to use for FP8 quantization
     """
 
-    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,25 +89,26 @@ 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"
+            f"{opt_args.quant_method} is not a valid quantization technique option. \
+            Please choose from: gptq, fp8, dq"
         )
 
 
-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:
         model_args (fms_mo.training_args.ModelArguments): Model arguments to be used when loading
             the model
         data_args (fms_mo.training_args.DataArguments): Data arguments to be used when loading the
             tokenized dataset
-        gptq_args (fms_mo.training_args.GPTQArgs): Parameters to use for GPTQ quantization
-        output_dir (str) Output directory to write to
+        opt_args (fms_mo.training_args.OptArguments): Generic optimization related arguments
+        gptq_args (fms_mo.training_args.GPTQArguments): Parameters to use for GPTQ quantization
     """
 
     # Third Party
@@ -152,23 +153,25 @@ 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:
         model_args (fms_mo.training_args.ModelArguments): Model arguments to be used when loading
             the model
         data_args (fms_mo.training_args.DataArguments): Data arguments to be used when loading the
             tokenized dataset
-        fp8_args (fms_mo.training_args.FP8Args): Parameters to use for FP8 quantization
-        output_dir (str) Output directory to write to
+        opt_args (fms_mo.training_args.OptArguments): Generic optimization related arguments
+        fp8_args (fms_mo.training_args.FP8Arguments): Parameters to use for FP8 quantization
     """
 
     # Third Party
@@ -192,11 +195,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 +211,41 @@ def main():
         dataclass_types=(
             ModelArguments,
             DataArguments,
+            OptArguments,
             FMSMOArguments,
-            GPTQArgs,
-            FP8Args,
+            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."""
@@ -115,7 +133,7 @@ class FMSMOArguments:
 
 
 @dataclass
-class GPTQArgs:
+class GPTQArguments:
     """Dataclass for GPTQ related arguments that will be used by auto-gptq."""
 
     bits: int = field(default=4, metadata={"choices": [2, 3, 4, 8]})
@@ -133,7 +151,7 @@ class GPTQArgs:
 
 
 @dataclass
-class FP8Args:
+class FP8Arguments:
     """Dataclass for FP8 related arguments that will be used by llm-compressor."""
 
     targets: str = field(default="Linear")