chore: Added license to ffn_tmp.py and other minor fixes

Signed-off-by: Brandon Groth <[email protected]>

Author: BrandonGroth

examples/MX/ffn_tmp.py

@@ -1,6 +1,17 @@
-# Third Party
+# Copyright The FMS Model Optimizer Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
-# from mx import Linear as Linear_mx  # Need to amend mx's Linear class
 # Third Party
 import numpy as np
 import torch

fms_mo/modules/linear.py

@@ -1964,77 +1964,6 @@ def extra_repr(self) -> str:
         "quantize_backprop": False,
     }
 
-    # class QLinearMX(mx.Linear):
-    #     """This is just a placeholder. Brandon is still working on it."""
-    #     @classmethod
-    #     def from_fms_mo(cls, fms_mo_qlinear, **kwargs):
-    #         """
-    #         Converts a QLinear module to QLinearMX.
-
-    #         Args:
-    #             cls: The class of the QLinearMX to be created.
-    #             fms_mo_qlinear: The QLinear module to be converted.
-    #             kwargs: Additional keyword arguments.
-
-    #         Returns:
-    #             A QLinearMX object initialized with the weights and biases from the
-    #                 QLinear module.
-    #         """
-    #         mx_supported_formats = {
-    #             "mx_fp8_e5m2",
-    #             "mx_fp8_e4m3",
-    #             "mx_fp4_e2m1",
-    #             "mx_fp4",
-    #             "mx_int8",
-    #             "mx_int4",
-    #             "mx_fp16",
-    #             "mx_float16",
-    #             "mx_bf16",
-    #             "mx_bfloat16",
-    #         }
-    #         assert (
-    #             fms_mo_qlinear.qa_mode in mx_supported_formats
-    #             and fms_mo_qlinear.qw_mode in mx_supported_formats
-    #         ), "Please check MX quantization mode settings!"
-    #         a_elem_format = fms_mo_qlinear.qa_mode.removeprefix("mx_")
-    #         w_elem_format = fms_mo_qlinear.qw_mode.removeprefix("mx_")
-
-    #         block_size = kwargs.pop("block_size")
-    #         mx_supported_block_sizes = {8, 16, 32, 64, 128}
-    #         assert (
-    #             block_size in mx_supported_block_sizes
-    #         ), "Please check MX block size setting!"
-
-    #         target_device = kwargs.get(
-    #             "target_device", next(fms_mo_qlinear.parameters()).device
-    #         )
-    #         use_ptq = fms_mo_qlinear
-
-    #         mx_specs = {
-    #             "a_elem_format": a_elem_format,
-    #             "w_elem_format": w_elem_format,
-    #             "block_size": block_size,
-    #             "bfloat": 16,
-    #             "custom_cuda": True,
-    #             # For quantization-aware finetuning, do backward pass in FP32
-    #             "quantize_backprop": False,
-    #         }
-
-    #         # Create mx.Linear class from QLinear
-    #         qlinear_mx = cls(
-    #             in_features=fms_mo_qlinear.in_features,
-    #             out_features=fms_mo_qlinear.out_features,
-    #             bias=isinstance(fms_mo_qlinear.bias, torch.Tensor),
-    #             mx_specs=fms_mo_qlinear.qcfg["mx_specs"],
-    #             name=None,
-    #         )
-
-    #     def extra_repr(self) -> str:
-    #         return (
-    #             f"in={self.in_features}, out={self.out_features}, bias={self.bias is not None}, "
-    #             f"mx_spec={self.mx_spec}"
-    #         )
-
     class QLinearMX(torch.nn.Linear):
         """Modified from mx.linear class. Only mildly changed init() and add extra_repr.
         1. Add **kwargs to receive extra (unused) params passed from qmodel_prep

fms_mo/utils/qconfig_utils.py

@@ -1182,18 +1182,18 @@ def check_config(config, model_dtype=None):
                 )
 
         # If mapping is defined, check for MX  classes
-        # Local
-        from fms_mo.modules.bmm import QBmmMX
-        from fms_mo.modules.linear import QLinearMX
-
-        mapping = config.get("mapping", None)
+        if available_packages["mx"]:
+            # Local
+            from fms_mo.modules.bmm import QBmmMX
+            from fms_mo.modules.linear import QLinearMX
 
-        # partial was used to init this mapping --> use .func pointer
-        if mapping is not None:
-            if not mapping[nn.Linear].func is QLinearMX:
-                raise ValueError("MX mapping for nn.Linear is not QLinearMX")
+            mapping = config.get("mapping", None)
 
-            if mapping["matmul_or_bmm"].func is QBmmMX:
-                raise ValueError("MX mapping for matmul_or_bmm is not QBmmMX")
+            # partial was used to init this mapping --> use .func pointer
+            if mapping is not None:
+                if not mapping[nn.Linear].func is QLinearMX:
+                    raise ValueError("MX mapping for nn.Linear is not QLinearMX")
 
+                if mapping["matmul_or_bmm"].func is QBmmMX:
+                    raise ValueError("MX mapping for matmul_or_bmm is not QBmmMX")
     # End mx_specs checks

fms_mo/utils/torchscript_utils.py

@@ -28,7 +28,6 @@
 import torch
 
 # Local
-# from fms_mo.modules import QBmm
 from fms_mo.quant.quantizers import transformers_prepare_input
 from fms_mo.utils.import_utils import available_packages
 from fms_mo.utils.utils import move_to, patch_torch_bmm, prepare_data_4_fwd

pyproject.toml

@@ -39,7 +39,6 @@ dependencies = [
 "safetensors",
 "ibm-fms>=0.0.8",
 "pkginfo>1.10",
-# "mx @ git+https://github.com/microsoft/microxcaling.git"
 ]
 
 [project.optional-dependencies]

tests/models/test_mx.py

@@ -4,16 +4,20 @@
 
 # Local
 from fms_mo import qmodel_prep
-from fms_mo.modules.bmm import QBmmMX
-from fms_mo.modules.linear import QLinearMX
 from fms_mo.utils.import_utils import available_packages
 from fms_mo.utils.qconfig_utils import check_config, set_mx_specs
 from tests.models.test_model_utils import delete_config, qmodule_error
 
-mx_qmodules = [
-    QLinearMX,
-    QBmmMX,
-]
+if available_packages["mx"]:
+    # Local
+    # pylint: disable=ungrouped-imports
+    from fms_mo.modules.bmm import QBmmMX
+    from fms_mo.modules.linear import QLinearMX
+
+    mx_qmodules = [
+        QLinearMX,
+        QBmmMX,
+    ]
 
 @pytest.mark.skipif(
     not available_packages["mx"],
@@ -92,7 +96,7 @@ def test_config_mx_error(
 
 @pytest.mark.skipif(
     not torch.cuda.is_available()
-    and not available_packages["mx"],
+    or not available_packages["mx"],
     reason="Skipped because CUDA or MX library was not available",
 )
 def test_residualMLP(