Move aiu arguments to other dataclasses

Signed-off-by: Andrea Fasoli <[email protected]>

Author: andrea-fasoli

fms_mo/dq.py

@@ -51,7 +51,7 @@
 logger = logging.getLogger(__name__)
 
 
-def run_dq(model_args, data_args, opt_args, fms_mo_args, aiu_args = None):
+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
@@ -66,8 +66,6 @@ def run_dq(model_args, data_args, opt_args, fms_mo_args, aiu_args = None):
         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
-        aiu_args (fms_mo.training_args.AIUArguments): Parameters specific to AIU-compliant
-            checkpoint generation and saving
 
     NOTE:
         use dynamo tracing instead of torchscript by default. if torchscript is needed, change
@@ -175,7 +173,7 @@ def run_dq(model_args, data_args, opt_args, fms_mo_args, aiu_args = None):
 
     qcfg["seq_len"] = block_size
     qcfg["model"] = model_args.model_name_or_path
-    qcfg["smoothq"] = True
+    qcfg["smoothq"] = qcfg.get("smoothq_alpha", -1) >= 0
     qcfg["plotsvg"] = False
 
     calibration_dataset = load_from_disk(data_args.training_data_path)
@@ -224,10 +222,13 @@ def run_dq(model_args, data_args, opt_args, fms_mo_args, aiu_args = None):
         save_fname="dq",
     )
     logger.info(f"Quantized model {model}")
-    logger.info("Starting to apply smooth scale")
-    dq_llm(model, act_scales, qcfg)
-    logger.info("Finished applying smooth scale")
     logger.info("==" * 20)
+
+    if qcfg["smoothq"]:
+        logger.info("Starting to apply smooth scale")
+        dq_llm(model, act_scales, qcfg)
+        logger.info("Finished applying smooth scale")
+
     if qcfg["qmodel_calibration_new"] > 0:
         logger.info("Starting to calibrate activation clip_val")
         if qcfg["large_model"]:
@@ -244,9 +245,9 @@ def run_dq(model_args, data_args, opt_args, fms_mo_args, aiu_args = None):
                 with patch_torch_bmm(qcfg):
                     model(**data_mb)
 
-    if aiu_args is not None and aiu_args.save_ckpt_for_aiu:
+    if opt_args.save_ckpt_for_aiu:
         logger.info(
-            f"Saving model processed for AIU and tokenizer to {aiu_args.output_dir}"
+            f"Saving model processed for AIU and tokenizer to {opt_args.output_dir}"
         )
         save_for_aiu(model, qcfg, output_dir=opt_args.output_dir, verbose=True)
     elif opt_args.save_ckpt:

fms_mo/run_quant.py

@@ -43,7 +43,6 @@
 # Local
 from fms_mo.dq import run_dq
 from fms_mo.training_args import (
-    AIUArguments,
     DataArguments,
     FMSMOArguments,
     FP8Arguments,
@@ -68,7 +67,6 @@ def quantize(
     fms_mo_args: FMSMOArguments = None,
     gptq_args: GPTQArguments = None,
     fp8_args: FP8Arguments = None,
-    aiu_args: AIUArguments = None,
 ):
     """Main entry point to quantize a given model with a set of specified hyperparameters
 
@@ -107,7 +105,7 @@ def quantize(
             )
         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, aiu_args)
+        run_dq(model_args, data_args, opt_args, fms_mo_args)
     else:
         raise ValueError(
             f"{opt_args.quant_method} is not a valid quantization technique option. \
@@ -236,7 +234,6 @@ def get_parser():
             FMSMOArguments,
             GPTQArguments,
             FP8Arguments,
-            AIUArguments,
         )
     )
     return parser
@@ -273,7 +270,6 @@ def parse_arguments(parser, json_config=None):
             fms_mo_args,
             gptq_args,
             fp8_args,
-            aiu_args,
         ) = parser.parse_dict(json_config, allow_extra_keys=True)
     else:
         (
@@ -283,7 +279,6 @@ def parse_arguments(parser, json_config=None):
             fms_mo_args,
             gptq_args,
             fp8_args,
-            aiu_args,
             _,
         ) = parser.parse_args_into_dataclasses(return_remaining_strings=True)
 
@@ -298,7 +293,6 @@ def parse_arguments(parser, json_config=None):
         fms_mo_args,
         gptq_args,
         fp8_args,
-        aiu_args,
     )
 
 
@@ -317,15 +311,14 @@ def main():
             fms_mo_args,
             gptq_args,
             fp8_args,
-            aiu_args,
         ) = parse_arguments(parser, job_config)
 
         logger = set_log_level(opt_args.log_level, __name__)
 
         logger.debug(
             f"Input args parsed: \nmodel_args {model_args}, data_args {data_args}, "
             f"opt_args {opt_args}, fms_mo_args {fms_mo_args}, gptq_args {gptq_args}, "
-            f"fp8_args {fp8_args}, aiu_args {aiu_args}"
+            f"fp8_args {fp8_args}"
         )
     except Exception as e:  # pylint: disable=broad-except
         logger.error(traceback.format_exc())
@@ -345,7 +338,6 @@ def main():
             fms_mo_args=fms_mo_args,
             gptq_args=gptq_args,
             fp8_args=fp8_args,
-            aiu_args=aiu_args,
         )
     except (MemoryError, OutOfMemoryError) as e:
         logger.error(traceback.format_exc())

fms_mo/training_args.py

@@ -142,16 +142,6 @@ class OptArguments(TypeChecker):
         default=True,
         metadata={"help": "Save quantized checkpoint."},
     )
-
-
-@dataclass
-class AIUArguments(TypeChecker):
-    """Dataclass for AIU-related arguments. Only apply to Direct Quantization runs."""
-
-    recompute_narrow_weights: bool = field(
-        default=False,
-        metadata={"help": "Apply recomputation during checkpoint saving."},
-    )
     save_ckpt_for_aiu: bool = field(
         default=False,
         metadata={"help": "Prepare and save AIU-compliant checkpoint."},
@@ -191,6 +181,10 @@ class FMSMOArguments(TypeChecker):
         default=2048, metadata={"help": "input sequence length after tokenization"}
     )
     eval_ppl: bool = field(default=False)
+    recompute_narrow_weights: bool = field(
+        default=False,
+        metadata={"help": "Apply recomputation during checkpoint saving for AIU."},
+    )
 
 
 @dataclass

tests/test_run_quant.py

@@ -110,7 +110,6 @@ def test_parse_arguments(job_config):
         _,
         _,
         _,
-        _,
     ) = parse_arguments(parser, job_config_copy)
     assert str(model_args.torch_dtype) == "torch.bfloat16"
     assert data_args.training_data_path == "data_train"
@@ -133,7 +132,6 @@ def test_parse_arguments_defaults(job_config):
         fms_mo_args,
         _,
         _,
-        _,
     ) = parse_arguments(parser, job_config_defaults)
     assert str(model_args.torch_dtype) == "torch.bfloat16"
     assert model_args.model_revision == "main"