Add qconfig save during saving for AIU

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

Author: andrea-fasoli

fms_mo/utils/aiu_utils.py

@@ -17,6 +17,7 @@
 import logging
 
 # Third Party
+from fms_mo.utils.qconfig_utils import qconfig_save
 from transformers.modeling_utils import PreTrainedModel
 import torch
 
@@ -217,7 +218,7 @@ def convert_sd_for_aiu(
 
 def save_sd_for_aiu(
     model: PreTrainedModel,
-    output_dir: str,
+    output_dir: str = "./",
     savename: str = "qmodel_state_dict.pt",
     verbose: bool = False,
 ) -> None:
@@ -226,3 +227,31 @@ def save_sd_for_aiu(
     converted_sd = convert_sd_for_aiu(model, verbose)
     torch.save(converted_sd, Path(output_dir) / savename)
     logger.info("Model saved.")
+
+
+def save_for_aiu(
+    model: PreTrainedModel,
+    qcfg: dict,
+    output_dir: str = "./",
+    file_name: str = "qmodel.pt",
+    cfg_name: str = "qcfg.json",
+    recipe: str | None = None,
+    verbose: bool = False,
+) -> None:
+    """Save quantized model and configuration in the format request by the AIU.
+    The checkpoint saving is customized for AIU compatibility.
+    The general qconfig_save function is used to save the quantization configuration.
+    """
+
+    save_sd_for_aiu(model, output_dir, file_name, verbose)
+
+    # define specific keys needed when reloading model for AIU
+    qcfg["keys_to_save"] = [
+        "qa_mode",
+        "qw_mode",
+        "smoothq",
+        "scale_layers",
+        "qskip_layer_name",
+        "qskip_large_mag_layers",
+    ]
+    qconfig_save(qcfg, recipe=recipe, minimal=True, fname=Path(output_dir) / cfg_name)

fms_mo/utils/qconfig_utils.py

@@ -538,9 +538,9 @@ def add_wanted_defaults_to_config(config: dict, minimal: bool = True) -> None:
 
 def qconfig_save(
     qcfg: dict,
-    recipe: str = None,
+    recipe: str | None = None,
     minimal: bool = True,
-    fname="qcfg.json",
+    fname: str = "qcfg.json",
 ) -> None:
     """
     Try to save qcfg into a JSON file (or use .pt format if something really can't be text-only).
@@ -550,8 +550,8 @@ def qconfig_save(
     Args:
         qcfg (dict): Quantized config.
         recipe (str, optional): String name for a save recipe. Defaults to None.
-        minimal (bool, optional): Save a minimal quantized config. Defaults to True.
-        fname (str, optional): File name to save quantized config. Defaults to "qcfg.json".
+        minimal (bool): Save a minimal quantized config. Defaults to True.
+        fname (str): File name to save quantized config. Defaults to "qcfg.json".
     """
 
     # First check in qcfg for added save list
@@ -598,6 +598,7 @@ def qconfig_save(
         warnings.warn(message, UserWarning)
     with open(fname, "w", encoding="utf-8") as outfile:
         json.dump(temp_qcfg, outfile, indent=4)
+    logger.info(f"Quantization configuration saved to {fname}")
 
 
 def qconfig_load(fname: str = "qcfg.json") -> dict: