feat: modified fms_mo for inference

Signed-off-by: omobayode.fagbohungbe <[email protected]>

Author: bayo-ibm

fms_mo/__init__.py

@@ -19,7 +19,7 @@
 
 # Local
 from fms_mo.prep import qmodel_prep
-from fms_mo.utils.qconfig_utils import qconfig_init
+from fms_mo.utils.qconfig_utils import qconfig_init, qconfig_load
 
 VERSION_FALLBACK = "0.0.0"
 

fms_mo/dq.py

@@ -35,7 +35,7 @@
 import torch
 
 # Local
-from fms_mo import qconfig_init, qmodel_prep
+from fms_mo import qconfig_init, qmodel_prep, qconfig_load
 from fms_mo.fx.utils import model_size_Wb
 from fms_mo.quant.ptq import (
     calibration_llm_1GPU,
@@ -145,7 +145,7 @@ def run_dq(model_args, data_args, opt_args, fms_mo_args):
     ]
     qcfg["large_model"] = any(
         name in model_args.model_name_or_path for name in known_large_models
-    ) or (gpu_mem_util_per > 0.7)
+    ) or (gpu_mem_util_per > 0.1)
     dev = "cpu" if qcfg["large_model"] else "cuda"
     if model_args.device_map is None:
         model.to(dev)
@@ -154,6 +154,7 @@ def run_dq(model_args, data_args, opt_args, fms_mo_args):
         qcfg["model_type"] = model.config.model_type
 
     qcfg["model"] = model_args.model_name_or_path
+    qcfg["qskip_large_mag_layers"] = True
     # config layers to skip, smooth scale
     config_quantize_smooth_layers(qcfg)
 
@@ -174,6 +175,7 @@ def run_dq(model_args, data_args, opt_args, fms_mo_args):
     qcfg["model"] = model_args.model_name_or_path
     qcfg["smoothq"] = True
     qcfg["plotsvg"] = False
+    
 
     calibration_dataset = load_from_disk(data_args.training_data_path)
     calibration_dataset = calibration_dataset.with_format("torch")
@@ -184,62 +186,80 @@ def run_dq(model_args, data_args, opt_args, fms_mo_args):
         collate_fn=default_data_collator,
         batch_size=1,
     )
-
+    #print(fms_mo_args)
+    #ii
     # For loading or creating smoothquant scale. Sometimes we may include scales in ckpt as well.
-    scale_file = Path(f"./act_scales/{qcfg['model'].replace('/', '-')}.pt")
-    if qcfg.get("act_scale_path", None):
-        # user provided a scale file (or a dir)
-        scale_file_or_dir = Path(qcfg["act_scale_path"])
-        if scale_file_or_dir.is_dir():
-            scale_file = scale_file_or_dir / f"{qcfg['model'].replace('/', '-')}.pt"
-        elif scale_file_or_dir.is_file():
-            scale_file = scale_file_or_dir
+    if not fms_mo_args.inference:
+        scale_file = Path(f"./act_scales/{qcfg['model'].replace('/', '-')}.pt")
+        if qcfg.get("act_scale_path", None):
+            # user provided a scale file (or a dir)
+            scale_file_or_dir = Path(qcfg["act_scale_path"])
+            if scale_file_or_dir.is_dir():
+                scale_file = scale_file_or_dir / f"{qcfg['model'].replace('/', '-')}.pt"
+            elif scale_file_or_dir.is_file():
+                scale_file = scale_file_or_dir
 
-    if not scale_file.parent.exists():
-        scale_file.parent.mkdir(exist_ok=False)
+        if not scale_file.parent.exists():
+            scale_file.parent.mkdir(exist_ok=False)
 
-    if scale_file.exists():
-        act_scales = torch.load(scale_file, map_location=getattr(model, "device", dev))
-    else:
-        logger.info("Generate activation scales")
-        if qcfg["large_model"]:
-            act_scales = get_act_scales_1gpu(model, dq_dataloader, qcfg)
+        if scale_file.exists():
+            act_scales = torch.load(scale_file, map_location=getattr(model, "device", dev))
         else:
-            act_scales = get_act_scales(model, dq_dataloader, qcfg)
-        torch.save(act_scales, scale_file)
+            logger.info("Generate activation scales")
+            if qcfg["large_model"]:
+                act_scales = get_act_scales_1gpu(model, dq_dataloader, qcfg)
+            else:
+                act_scales = get_act_scales(model, dq_dataloader, qcfg)
+            torch.save(act_scales, scale_file)
+    else:
+        import json
+        q_file = open('qcfg_llama.json', "r", encoding="utf-8")
+        a = json.load(q_file)
+        print(a)
+        qcfg.update(a)
+        print(qcfg)
+        
     qmodel_prep(
         model,
         dq_dataloader,
         qcfg,
         use_layer_name_pattern_matching=use_layer_name_pattern_matching,
         use_dynamo=use_dynamo,
         dev=dev,
+        mode=fms_mo_args.inference,
         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["qmodel_calibration_new"] > 0:
-        logger.info("Starting to calibrate activation clip_val")
-        if qcfg["large_model"]:
-            calibration_llm_1GPU(qcfg, model, dq_dataloader)
-        else:
-            model.to("cuda:0")
-            pbar = tqdm(
-                dq_dataloader,
-                desc=" calibration after applying smoothq scale and before inference",
-                total=qcfg["qmodel_calibration_new"],
-            )
-            for data_mb, _ in zip(pbar, range(qcfg["qmodel_calibration_new"])):
-                data_mb = prepare_input(model.device, data_mb)
-                with patch_torch_bmm(qcfg):
-                    model(**data_mb)
 
-    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 not fms_mo_args.inference:
+        logger.info("Starting to apply smooth scale")
+        dq_llm(model, act_scales, qcfg)
+        logger.info("Finished applying smooth scale")
+        logger.info("==" * 20)
+        if qcfg["qmodel_calibration_new"] > 0:
+            logger.info("Starting to calibrate activation clip_val")
+            if qcfg["large_model"]:
+                calibration_llm_1GPU(qcfg, model, dq_dataloader)
+            else:
+                model.to("cuda:0")
+                pbar = tqdm(
+                    dq_dataloader,
+                    desc=" calibration after applying smoothq scale and before inference",
+                    total=qcfg["qmodel_calibration_new"],
+                )
+                for data_mb, _ in zip(pbar, range(qcfg["qmodel_calibration_new"])):
+                    data_mb = prepare_input(model.device, data_mb)
+                    with patch_torch_bmm(qcfg):
+                        model(**data_mb)
+        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)
+    else:
+        pass
+        from accelerate import load_checkpoint_and_dispatch
+        model = load_checkpoint_and_dispatch( model, checkpoint=opt_args.output_dir, device_map=None, no_split_module_classes=['Block'])
+
+
 
     if fms_mo_args.eval_ppl:
         path_test = Path(data_args.test_data_path)
@@ -253,7 +273,7 @@ def run_dq(model_args, data_args, opt_args, fms_mo_args):
 
         logger.info(f"Model for evaluation: {model}")
         if qcfg["large_model"]:
-            eval_llm_1GPU(qcfg, model, test_dataset)
+            eval_llm_1GPU(qcfg, model.to('cpu'), test_dataset)
         else:
             model.to(torch.device("cuda:0"))
             n_samples = int(test_dataset.input_ids.shape[1] / block_size)

fms_mo/prep.py

@@ -535,6 +535,30 @@ def has_quantized_module(model):
     """Check if model is already quantized - do not want to quantize twice if so"""
     return any(isinstance(m, quantized_modules) for m in model.modules())
 
+def swap_qbmm(model, qcfg):
+    from fms_mo.modules import QBmm
+
+    qcfg["which2patch_contextmanager"] = qcfg["bmm_prep"][
+        "which2patch_contextmanager"
+    ]
+    isbmm = qcfg["which2patch_contextmanager"] == "torch.bmm"
+    for mod_name, line_nums in qcfg["bmm_prep"]["layers_with_bmm"].items():
+        mod_bmm_happened = model.get_submodule(mod_name)
+        for whichQBmm, ln in enumerate(line_nums, start=1):
+            nbits = qcfg[f"nbits_bmm{whichQBmm}"]
+            newQBmm = QBmm(
+                num_bits_m1=max(nbits, 8) if whichQBmm == 2 else nbits,
+                num_bits_m2=nbits,
+                qm1_mode=qcfg[f"bmm{whichQBmm}_qm1_mode"],
+                qm2_mode=qcfg[f"bmm{whichQBmm}_qm2_mode"],
+                m1_unidirectional=(whichQBmm == 2),
+                m1_bounded=(whichQBmm == 2),  # see Note 5
+                m2_unidirectional=False,
+                m2_bounded=False,
+                replaceBmm=isbmm,
+                qcfg=qcfg,
+            )
+            setattr(mod_bmm_happened, f"QBmm{ln}", newQBmm)
 
 def qmodel_prep(
     model,
@@ -548,6 +572,7 @@ def qmodel_prep(
     dev=None,
     use_dynamo=False,
     verbose=False,
+    mode=False,
     **kwargs,
 ):
     """Prepare a given PyTorch model for quantization process through three parts:
@@ -622,7 +647,15 @@ def qmodel_prep(
     Returns:
         nn.Module: quantized model ready for further PTQ/QAT
     """
+    if mode:
+        
+        if qcfg.get("QBmm"): 
+            pass
+            swap_qbmm(model,qcfg)
 
+        model = q_any_net_5(model, qcfg, verbose = False)
+        return model
+    
     sys.setrecursionlimit(4000)
 
     currDev = next(model.parameters()).device if dev is None else dev
@@ -869,7 +902,7 @@ def qmodel_prep(
             model, device_ids=DPorDDPdevices
         )
 
-    qconfig_save(qcfg, fname="qcfg.json")
+    qconfig_save(qcfg, fname="qcfg2.json")
     qcfg["tb_writer"] = tb_writer
 
     logger.info(f"--- Quantized model --- \n{model}\n")

fms_mo/training_args.py

@@ -173,6 +173,8 @@ class FMSMOArguments(TypeChecker):
         default=2048, metadata={"help": "input sequence length after tokenization"}
     )
     eval_ppl: bool = field(default=False)
+    inference: bool = field(default=False)
+
 
 
 @dataclass

fms_mo/utils/qconfig_utils.py

@@ -623,7 +623,7 @@ def qconfig_save(
 def qconfig_load(fname: str = "qcfg.json") -> dict:
     """Read config in json format, work together with qconfig_save"""
     config = get_recipe(fname)
-
+    
     if config:
         # Check that loaded file is a dict
         if not isinstance(config, dict):