feat: fast model inference

Signed-off-by: Omobayode Fagbohungbe <[email protected]>

Author: bayo-ibm

fms_mo/dq.py

@@ -50,6 +50,8 @@
 from fms_mo.utils.dq_utils import config_quantize_smooth_layers
 from fms_mo.utils.eval_utils import Evaluator, eval_llm_1GPU
 from fms_mo.utils.utils import patch_torch_bmm, prepare_input
+from fms_mo.utils.dq_inf import load_fp8_vllm, save_vllm_fp8
+from accelerate import load_checkpoint_and_dispatch
 
 logger = logging.getLogger(__name__)
 
@@ -134,7 +136,11 @@ def run_dq(model_args, data_args, opt_args, fms_mo_args):
     logger.info(f"Initialized model is: \n {model}")
     logger.info(f"Model is at {model.device} after intialization")
     logger.info(f"Tokenizer is {tokenizer}, block size is {block_size}")
-    qcfg = qconfig_init(recipe="dq", args=fms_mo_args)
+    
+    if not fms_mo_args.inference or fms_mo_args.vllm_fp8_load:
+        qcfg = qconfig_init(recipe="dq", args=fms_mo_args)
+    else:
+        qcfg = qconfig_init(recipe=opt_args.output_dir+"/qcfg")
 
     model_size = model_size_Wb(model, unit="GB")
     gpu_mem_util_per = model_size / total_gpu_memory
@@ -190,7 +196,7 @@ def run_dq(model_args, data_args, opt_args, fms_mo_args):
     )
 
     # For loading or creating smoothquant scale. Sometimes we may include scales in ckpt as well.
-    if qcfg["smoothq"]:
+    if not fms_mo_args.inference and qcfg["smoothq"] :
         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)
@@ -224,53 +230,76 @@ def run_dq(model_args, data_args, opt_args, fms_mo_args):
             use_layer_name_pattern_matching=use_layer_name_pattern_matching,
             use_dynamo=use_dynamo,
             dev=dev,
+            mode=fms_mo_args.inference,
             save_fname="dq",
+            folder=opt_args.output_dir,
         )
         logger.info(f"Quantized model {model}")
         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 not fms_mo_args.inference:
+        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"]:
+                calibration_llm_1GPU_v2(qcfg, model, dq_dataloader)
+            else:
+                model.to("cuda")
+                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)
+
+        if opt_args.save_ckpt_for_aiu:
+            logger.info(
+                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_for_vllm:
+            logger.info(
+                f"Saving model processed for vLLM and tokenizer to {opt_args.output_dir}"
+            )
+            save_vllm_fp8(model,qcfg,tokenizer,opt_args.output_dir)
+        elif opt_args.save_ckpt:
+            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.aiu_sim_triton:
+            # NOTE plz apply correct HW settings here, defaults are not real HW params
+            lower_qmodel_triton(
+                model,
+                use_dyn_max_act=-1 if qcfg["qa_mode"] == "pertokenmax" else False,
+                max_acc_bits=qcfg.get("max_acc_bits", 32),
+                num_lsb_to_truncate=qcfg.get("lsb_trun_bits", 0),
+                chunk_size=qcfg.get("chunk_size", 32),  # 1024
+                clamp_acc_to_dl16=fms_mo_args.aiu_sim_triton == "fp8",
+                # layer_to_exclude=["lm_head",]
+            )
+    else:
+        if fms_mo_args.vllm_fp8_load:
+            logger.info("loading llmcompressor fp8 model saved_checkpoint")
+            model = load_fp8_vllm( model=model, checkpoint=opt_args.output_dir)
 
-    if qcfg["qmodel_calibration_new"] > 0:
-        logger.info("Starting to calibrate activation clip_val")
-        if qcfg["large_model"]:
-            calibration_llm_1GPU_v2(qcfg, model, dq_dataloader)
         else:
-            model.to("cuda")
-            pbar = tqdm(
-                dq_dataloader,
-                desc=" calibration after applying smoothq scale and before inference",
-                total=qcfg["qmodel_calibration_new"],
+            logger.info("loading dq fms_mo fp8 model saved_checkpoint")
+            model = load_checkpoint_and_dispatch( 
+                model,
+                checkpoint=opt_args.output_dir,
+                device_map=None,
+                no_split_module_classes=['Block']
             )
-            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)
-
-    if opt_args.save_ckpt_for_aiu:
-        logger.info(
-            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:
-        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.aiu_sim_triton:
-        # NOTE plz apply correct HW settings here, defaults are not real HW params
-        lower_qmodel_triton(
-            model,
-            use_dyn_max_act=-1 if qcfg["qa_mode"] == "pertokenmax" else False,
-            max_acc_bits=qcfg.get("max_acc_bits", 32),
-            num_lsb_to_truncate=qcfg.get("lsb_trun_bits", 0),
-            chunk_size=qcfg.get("chunk_size", 32),  # 1024
-            clamp_acc_to_dl16=fms_mo_args.aiu_sim_triton == "fp8",
-            # layer_to_exclude=["lm_head",]
-        )
 
     if fms_mo_args.eval_ppl:
         path_test = Path(data_args.test_data_path)

fms_mo/prep.py

@@ -570,7 +570,42 @@ 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: nn.Module, qcfg: dict):
+    """Go through all model.named_modules(), try to create an equivalent Qbmm layer to replace each of
+    the existing linear Bmm layers.
 
+    Args:
+        model (nn.Module): input model to be "prepared"
+        qcfg (dict): quant config
+
+    Returns: updated model is returned with the Qbmm added
+        
+    """
+
+    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,
     dloader,
@@ -582,7 +617,9 @@ def qmodel_prep(
     Qcali=False,
     dev=None,
     use_dynamo=False,
+    mode=False,
     verbose=False,
+    folder=None,
     **kwargs,
 ):
     """Prepare a given PyTorch model for quantization process through three parts:
@@ -657,7 +694,14 @@ def qmodel_prep(
     Returns:
         nn.Module: quantized model ready for further PTQ/QAT
     """
+    if mode:
+        
+        if qcfg.get("QBmm"): 
+            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
@@ -907,7 +951,7 @@ def qmodel_prep(
             model, device_ids=DPorDDPdevices
         )
 
-    qconfig_save(qcfg, fname="qcfg.json")
+    qconfig_save(qcfg, fname=folder+"/qcfg.json")
     qcfg["tb_writer"] = tb_writer
 
     logger.info(f"--- Quantized model --- \n{model}\n")

fms_mo/recipes/quant.json

@@ -0,0 +1,44 @@
+{
+"quantization_config": {
+    "config_groups": {
+      "group_0": {
+        "input_activations": {
+          "actorder": null,
+          "block_structure": null,
+          "dynamic": true,
+          "group_size": null,
+          "num_bits": 8,
+          "observer": null,
+          "observer_kwargs": {},
+          "strategy": "token",
+          "symmetric": true,
+          "type": "float"
+        },
+        "output_activations": null,
+        "targets": [
+          "Linear"
+        ],
+        "weights": {
+          "actorder": null,
+          "block_structure": null,
+          "dynamic": false,
+          "group_size": null,
+          "num_bits": 8,
+          "observer": "minmax",
+          "observer_kwargs": {},
+          "strategy": "channel",
+          "symmetric": true,
+          "type": "float"
+        }
+      }
+    },
+    "format": "float-quantized",
+    "global_compression_ratio": null,
+    "ignore": [
+      "lm_head"
+    ],
+    "kv_cache_scheme": null,
+    "quant_method": "compressed-tensors",
+    "quantization_status": "compressed"
+  }
+}

fms_mo/training_args.py

@@ -160,6 +160,10 @@ class OptArguments(TypeChecker):
         default=False,
         metadata={"help": "Prepare and save AIU-compliant checkpoint."},
     )
+    save_ckpt_for_vllm: bool = field(
+        default=False,
+        metadata={"help": "Prepare and save vllm-compliant checkpoint."},
+    )
 
 
 @dataclass
@@ -209,6 +213,9 @@ class FMSMOArguments(TypeChecker):
         default=False,
         metadata={"help": "Apply recomputation during checkpoint saving for AIU."},
     )
+    fp8_use_subnormal: bool = field(default=False)
+    inference: bool = field(default=False)
+    vllm_fp8_load: bool = field(default=False)
 
 
 @dataclass

fms_mo/utils/dq_inf.py

@@ -0,0 +1,89 @@
+import torch
+import fms_mo
+from fms_mo.quant.quantizers import to_fp8_scaled_perCh as fp8
+from huggingface_hub import save_torch_state_dict
+import json
+import os 
+import glob
+from fms_mo.utils.qconfig_utils import get_recipe
+from safetensors.torch import load_file, save_file
+from torch import nn
+
+def save_vllm_fp8(model: nn.Module, qcfg: dict, tokenizer = None, folder: str = None):   
+    """
+    Function to save fms_mo fp8 checkpoint in vllm fp8 format
+    """
+
+    st_dict={}
+ 
+    for k,v in model.state_dict().items():
+        if k[-11:] == "proj.weight":
+            weight, scale = fp8(v,emulate=False)
+            st_dict[k]= weight
+
+            if k[:-7] in qcfg["qskip_layer_name"]:
+                pass
+            else:
+                st_dict[k + "_scale"] = 1/scale 
+
+        elif k[-6:] == "weight":
+            st_dict[k]=v
+        else:
+            pass
+
+    config = model.config.to_dict()
+
+    #TO DO: To support multiple recipes, check qconfig arguments and update data loaded from quant.json
+    data = get_recipe('quant')
+    
+    config.update(data)
+
+    save_torch_state_dict(st_dict, folder)
+
+    tokenizer.save_pretrained(folder)
+
+    with open(folder+'/config.json', 'a') as f:
+        json.dump(config, f, indent=4)
+    
+
+
+def find_file_glob(pattern: str , search_path: str):
+    """
+    Finds files matching a pattern within a directory and its subdirectories.
+    """
+    # Use '**' for recursive search in modern Python versions (3.5+)
+    full_pattern = os.path.join(search_path, '**', pattern)
+    found_files = glob.glob(full_pattern, recursive=True)
+    return sorted(found_files)
+
+def load_fp8_vllm(model: nn.Module = None, checkpoint: str=None):
+    """
+    Function to help load vllm fp8 checkpoint into fms_mo
+    """   
+
+    merged_files_dict={}
+
+    files = find_file_glob('*.safetensors',checkpoint)
+
+    model_dict = model.state_dict()
+
+    for file in files:
+        merged_files_dict = load_file(file)
+
+        for k,v in merged_files_dict.items():
+
+            if k[-11:] == "proj.weight":
+                scale = merged_files_dict[k+ "_scale"].reshape(-1,1)
+                model_dict[k]= merged_files_dict[k].to(torch.float16) * scale
+
+            elif k[-6:] == "weight":
+                model_dict[k]=v
+
+            else:
+                pass
+
+    return model
+
+
+
+

fms_mo/utils/dq_utils.py

@@ -74,7 +74,7 @@ def config_quantize_smooth_layers(qcfg: dict):
             for llama_family, layers in large_mag_layers.items():
                 if llama_family in qcfg["model"]:
                     qcfg["qskip_layer_name"] += [
-                        f"model.layers.{i}.mlp.down_proj" for i in layers
+                        f"model.layers.{i}.mlp.down_projj" for i in layers
                     ]
                     break
     elif any(model in qcfg["model"] for model in granite_architecture) or any(