Merge branch 'main' into dependabot/pip/accelerate-gte-0.20.3-and-neq-0.34-and-lt-1.9

Signed-off-by: chichun-charlie-liu <[email protected]>

Author: chichun-charlie-liu

README.md

@@ -98,6 +98,26 @@ cd fms-model-optimizer
 pip install -e .
 ```
 
+#### Optional Dependencies
+The following optional dependencies are available:
+- `fp8`: `llmcompressor` package for fp8 quantization
+- `gptq`: `GPTQModel` package for W4A16 quantization
+- `mx`: `microxcaling` package for MX quantization
+- `opt`: Shortcut for `fp8`, `gptq`, and `mx` installs
+- `torchvision`: `torch` package for image recognition training and inference
+- `visualize`: Dependencies for visualizing models and performance data
+- `test`: Dependencies needed for unit testing
+- `dev`: Dependencies needed for development
+
+To install an optional dependency, modify the `pip install` commands above with a list of these names enclosed in brackets.  The example below installs `llm-compressor` and `torchvision` with FMS Model Optimizer:
+
+```shell
+pip install fms-model-optimizer[fp8,torchvision]
+
+pip install -e .[fp8,torchvision]
+```
+If you have already installed FMS Model Optimizer, then only the optional packages will be installed.
+
 ### Try It Out!
 
 To help you get up and running as quickly as possible with the FMS Model Optimizer framework, check out the following resources which demonstrate how to use the framework with different quantization techniques:

fms_mo/calib.py

@@ -574,13 +574,19 @@ def qmodel_calib(
                     f"Qmodel calibration (clip_val analysis) in progress: {i}/{Nbatch}"
                 )
 
-        if "perCh" not in qcfg["qw_mode"]:
-            cv_sum_dict = {"layer": [], "value": []}
-            for k, v in tempmodel.state_dict().items():
-                if "clip" in k:
-                    cv_sum_dict["layer"].append(k)
-                    cv_sum_dict["value"].append(v.item())
-            logger.info(f"Observed clipvals: \n{ pd.DataFrame(cv_sum_dict) }")
+        cv_sum_dict = {"layer": [], "value": []}
+        for k, v in tempmodel.state_dict().items():
+            if "clip" not in k:
+                continue
+
+            if v.numel() > 1:
+                k = k + "*"
+                v = v.mean()
+            cv_sum_dict["layer"].append(k)
+            cv_sum_dict["value"].append(v.item())
+        logger.info(
+            f"Observed clipvals: ('*' if it's a vector) \n{ pd.DataFrame(cv_sum_dict) }"
+        )
 
     # Step 3: extract new clip_vals, params and buffers, then remove handles if needed
     temp_new_clipvals = {

fms_mo/dq.py

@@ -36,6 +36,9 @@
 
 # Local
 from fms_mo import qconfig_init, qmodel_prep
+from fms_mo.custom_ext_kernels.utils import (
+    lower_qmodel_triton,  # pylint: disable=unused-import
+)
 from fms_mo.fx.utils import model_size_Wb
 from fms_mo.quant.ptq import (
     calibration_llm_1GPU_v2,
@@ -256,6 +259,15 @@ def run_dq(model_args, data_args, opt_args, fms_mo_args):
         model.save_pretrained(opt_args.output_dir, use_safetensors=True)
         tokenizer.save_pretrained(opt_args.output_dir)
 
+    if fms_mo_args.aiu_sim_triton:
+        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", 1024),
+        )
+
     if fms_mo_args.eval_ppl:
         path_test = Path(data_args.test_data_path)
         arrow_files = list(path_test.glob("*.arrow"))

fms_mo/fx/dynamo_utils.py

@@ -29,6 +29,7 @@
     get_target_op_from_mod_or_str,
     get_target_op_from_node,
 )
+from fms_mo.utils.import_utils import available_packages
 
 logger = logging.getLogger(__name__)
 
@@ -1133,7 +1134,6 @@ def cus_backend_model_analyzer(
     from functools import partial
 
     # Third Party
-    from torchvision.models import VisionTransformer
     from transformers import PreTrainedModel
 
     if issubclass(type(model), torch.nn.Module):
@@ -1145,7 +1145,16 @@ def cus_backend_model_analyzer(
         model_to_be_traced = model
         model_param_size = 999
 
-    is_transformers = issubclass(type(model), (PreTrainedModel, VisionTransformer))
+    transformer_model_classes = (PreTrainedModel,)
+
+    if available_packages["torchvision"]:
+        # Third Party
+        # pylint: disable = import-error
+        from torchvision.models import VisionTransformer
+
+        transformer_model_classes += (VisionTransformer,)
+
+    is_transformers = issubclass(type(model), transformer_model_classes)
     if model_param_size > 1:
         # Standard
         import sys
@@ -1188,11 +1197,13 @@ def call_seq_hook(mod, *_args, **_kwargs):
 
         # only add last layer
         qcfg["qskip_layer_name"] += [qcfg["mod_call_seq"][-1]]
-        # unless it's a ViT, skip first Conv as well
-        if issubclass(type(model), VisionTransformer) and isinstance(
-            model.get_submodule(qcfg["mod_call_seq"][0]), torch.nn.Conv2d
-        ):
-            qcfg["qskip_layer_name"] += [qcfg["mod_call_seq"][0]]
+
+        if available_packages["torchvision"]:
+            # unless it's a ViT, skip first Conv as well
+            if issubclass(type(model), VisionTransformer) and isinstance(
+                model.get_submodule(qcfg["mod_call_seq"][0]), torch.nn.Conv2d
+            ):
+                qcfg["qskip_layer_name"] += [qcfg["mod_call_seq"][0]]
 
     with torch.no_grad():
         model_opt = torch.compile(
@@ -1271,21 +1282,23 @@ def qbmm_auto_check(_mod, *_args, **_kwargs):
     # c) identify RPN/FPN
     # TODO this hack only works for torchvision models. will use find_rpn_fpn_gm()
 
-    # Third Party
-    from torchvision.models.detection.rpn import RegionProposalNetwork
-    from torchvision.ops import FeaturePyramidNetwork
-
-    rpnfpn_prefix = []
-    rpnfpn_convs = []
-    for n, m in model.named_modules():
-        if isinstance(m, (FeaturePyramidNetwork, RegionProposalNetwork)):
-            rpnfpn_prefix.append(n)
-        if isinstance(m, torch.nn.Conv2d) and any(
-            n.startswith(p) for p in rpnfpn_prefix
-        ):
-            rpnfpn_convs.append(n)
-            if n not in qcfg["qskip_layer_name"]:
-                qcfg["qskip_layer_name"].append(n)
+    if available_packages["torchvision"]:
+        # Third Party
+        # pylint: disable = import-error
+        from torchvision.models.detection.rpn import RegionProposalNetwork
+        from torchvision.ops import FeaturePyramidNetwork
+
+        rpnfpn_prefix = []
+        rpnfpn_convs = []
+        for n, m in model.named_modules():
+            if isinstance(m, (FeaturePyramidNetwork, RegionProposalNetwork)):
+                rpnfpn_prefix.append(n)
+            if isinstance(m, torch.nn.Conv2d) and any(
+                n.startswith(p) for p in rpnfpn_prefix
+            ):
+                rpnfpn_convs.append(n)
+                if n not in qcfg["qskip_layer_name"]:
+                    qcfg["qskip_layer_name"].append(n)
 
     if qcfg["N_backend_called"] > 1:
         logger.warning(

fms_mo/modules/bmm.py

@@ -192,7 +192,7 @@ def forward(self, m1, m2):
             torch.Tensor: Output tensor after quantized bmm.
         """
         # pylint: disable = access-member-before-definition
-        if self.calib_counter:
+        if self.calib_counter > 0:
             with torch.no_grad():
                 qm1 = self.quantize_calib_m1(m1)
                 qm2 = self.quantize_calib_m2(m2)

fms_mo/modules/conv.py

@@ -270,7 +270,7 @@ def forward(self, x):
             torch.Tensor: Output tensor of shape (batch_size, out_channels, out_height, out_width).
         """
         # pylint: disable = access-member-before-definition
-        if self.calib_counter:
+        if self.calib_counter > 0:
             with torch.no_grad():
                 qinput = self.quantize_calib_feature(x)
                 qweight = self.quantize_calib_weight(self.weight)