fix merge conflicts

ani300 · ani300 · commit f4ec8360e010 · 2025-06-30T20:53:47.000Z
Signed-off-by: Antoni Viros i Martin &lt;aviros@ibm.com&gt;
diff --git a/README.md b/README.md
@@ -104,8 +104,10 @@ The following optional dependencies are available:
 - `gptq`: `GPTQModel` package for W4A16 quantization
 - `mx`: `microxcaling` package for MX quantization
 - `opt`: Shortcut for `fp8`, `gptq`, and `mx` installs
+- `aiu`: `ibm-fms` package for AIU model deployment
 - `torchvision`: `torch` package for image recognition training and inference
 - `triton`: `triton` package for matrix multiplication kernels
+- `examples`: Dependencies needed for examples
 - `visualize`: Dependencies for visualizing models and performance data
 - `test`: Dependencies needed for unit testing
 - `dev`: Dependencies needed for development
diff --git a/fms_mo/aiu_addons/__init__.py b/fms_mo/aiu_addons/__init__.py
@@ -21,8 +21,10 @@ def _infer_quantization_config(quant_config: dict) -> dict | None:
             and quant_config["config_groups"]["group_0"]["weights"]["num_bits"] == 8
         ):
             # First, import required FP8 linear classes from fms-mo
-            import fms_mo.aiu_addons.fp8.fp8_linear  # pylint: disable=unused-import
+            # Local
             import fms_mo.aiu_addons.fp8.fp8_adapter  # pylint: disable=unused-import
+            import fms_mo.aiu_addons.fp8.fp8_linear  # pylint: disable=unused-import
+
             # This is used by get_linear to decide whether a linear layer
             # will be quantized or not inside the model
             def fp8_linear_type(name: str) -> str:
diff --git a/fms_mo/aiu_addons/fp8/fp8_attn.py b/fms_mo/aiu_addons/fp8/fp8_attn.py
@@ -173,7 +173,7 @@ def _math_fp8_compute_op(
         )
 
     attn_weight = (
-        torch.ops.sendnn.scaled_bmm(
+        torch.ops.spyre.scaled_bmm(
             query,
             key_cache.transpose(-2, -1),
             q_scale,
diff --git a/fms_mo/aiu_addons/fp8/fp8_spyre_op.py b/fms_mo/aiu_addons/fp8/fp8_spyre_op.py
@@ -27,7 +27,7 @@
 
 
 @torch.library.custom_op("spyre::scaled_bmm", mutates_args=())
-def sendnn_scaled_bmm(
+def spyre_scaled_bmm(
     mat1: Tensor,
     mat2: Tensor,
     scale1: Tensor,
@@ -43,6 +43,8 @@ def sendnn_scaled_bmm(
     assert (
         mat1.shape[:-2] == mat2.shape[:-2]
     ), "batch dimensions must match for mat1 and mat2"
+    assert scale1.numel() == 1, "only per-tensor scales supported"
+    assert scale2.numel() == 1, "only per-tensor scales supported"
     mat1 = mat1.view(-1, *mat1.shape[-2:])
     mat2 = mat2.view(-1, *mat2.shape[-2:])
     out = torch.empty(
@@ -62,7 +64,7 @@ def sendnn_scaled_bmm(
     return out.view(*mat1.shape[:-2], mat1.shape[1], mat2.shape[2])
 
 
-@sendnn_scaled_bmm.register_fake
+@spyre_scaled_bmm.register_fake
 def _(
     mat1: Tensor,
     mat2: Tensor,
diff --git a/fms_mo/aiu_addons/fp8/fp8_utils.py b/fms_mo/aiu_addons/fp8/fp8_utils.py
@@ -72,7 +72,7 @@ def __init__(
         self._scaled = scaled
 
     def __tensor_flatten__(self):
-        ctx = {"scaled", self._scaled}
+        ctx = {"scaled": self._scaled}
         return ["_data", "_scale"], ctx
 
     @staticmethod
diff --git a/fms_mo/aiu_addons/gptq/gptq_aiu_adapter.py b/fms_mo/aiu_addons/gptq/gptq_aiu_adapter.py
@@ -17,9 +17,21 @@
 from typing import Mapping
 
 # Third Party
-from fms.utils import serialization
 import torch
 
+# Local
+from fms_mo.utils.import_utils import available_packages
+
+if not available_packages["fms"]:
+    raise ImportError(
+        "AIU functionality requires ibm-fms to be installed."
+        "See https://github.com/foundation-model-stack/foundation-model-stack for details."
+    )
+
+# Third Party
+# pylint: disable=import-error,wrong-import-position
+from fms.utils import serialization
+
 
 def _gptq_qweights_transpose_aiu(
     input_sd: Mapping[str, torch.Tensor],
diff --git a/fms_mo/aiu_addons/gptq/gptq_aiu_linear.py b/fms_mo/aiu_addons/gptq/gptq_aiu_linear.py
@@ -18,6 +18,19 @@
 import math
 
 # Third Party
+import torch
+
+# Local
+from fms_mo.utils.import_utils import available_packages
+
+if not available_packages["fms"]:
+    raise ImportError(
+        "AIU functionality requires ibm-fms to be installed."
+        "See https://github.com/foundation-model-stack/foundation-model-stack for details."
+    )
+
+# Third Party
+# pylint: disable=import-error,wrong-import-position,ungrouped-imports
 from fms.modules.linear import (
     LinearModuleShardingInfo,
     LinearParameterShardingInfo,
@@ -27,7 +40,6 @@
 )
 from fms.modules.tp import ShardType, TPModule
 from fms.utils.gptq import GPTQLinearConfig
-import torch
 
 # Local
 from fms_mo.aiu_addons.gptq.gptq_aiu_op import register_aiu_gptq_op
diff --git a/fms_mo/aiu_addons/i8i8/i8i8_aiu_adapter.py b/fms_mo/aiu_addons/i8i8/i8i8_aiu_adapter.py
@@ -17,9 +17,21 @@
 from typing import Mapping
 
 # Third Party
-from fms.utils import serialization
 import torch
 
+# Local
+from fms_mo.utils.import_utils import available_packages
+
+if not available_packages["fms"]:
+    raise ImportError(
+        "AIU functionality requires ibm-fms to be installed."
+        "See https://github.com/foundation-model-stack/foundation-model-stack for details."
+    )
+
+# Third Party
+# pylint: disable=import-error,wrong-import-position
+from fms.utils import serialization
+
 
 def _int8_qparams_aiu(
     input_sd: Mapping[str, torch.Tensor],
diff --git a/fms_mo/aiu_addons/i8i8/i8i8_aiu_linear.py b/fms_mo/aiu_addons/i8i8/i8i8_aiu_linear.py
@@ -19,7 +19,17 @@
 from typing import Any, Callable, Optional, Union
 import copy
 
+# Local
+from fms_mo.utils.import_utils import available_packages
+
+if not available_packages["fms"]:
+    raise ImportError(
+        "AIU functionality requires ibm-fms to be installed."
+        "See https://github.com/foundation-model-stack/foundation-model-stack for details."
+    )
+
 # Third Party
+# pylint: disable=import-error,wrong-import-position,ungrouped-imports
 from fms.modules.linear import (
     LinearModuleShardingInfo,
     LinearParameterShardingInfo,
diff --git a/fms_mo/quant/ptq.py b/fms_mo/quant/ptq.py
@@ -30,7 +30,6 @@
 import sys
 
 # Third Party
-from torch.utils.tensorboard import SummaryWriter
 from tqdm import tqdm
 import numpy as np
 import pandas as pd
@@ -1449,49 +1448,57 @@ def ptq_mod_optim_lm(_model, m, layers, qcfg, optim_mode="both", **kwargs):
                 # show loss on pbar
                 pbar2.set_description(pbar_desc + f"{PTQloss:.6f}")
 
-            if isinstance(qcfg["tb_writer"], SummaryWriter) and isOutput:
-                scalars2log = {}
-                hist2log = {}
+            if available_packages["tensorboard"]:
+                # Third Party
+                from torch.utils.tensorboard import SummaryWriter
 
-                for k, v in loss4plot.items():  # plot loss
-                    scalars2log[f"{mod_name}/PTQloss_{k}"] = v
-                for k, v in m.named_buffers():  # plot cv, delta, zp, alpha, and lr
-                    if any(kb in k for kb in ["delta", "zero_point", "clip_val"]):
-                        if len(v.shape) > 0 and v.shape[0] > 1:  # perCh
-                            hist2log[f"{mod_name}/{k}"] = v
-                        else:
-                            scalars2log[f"{mod_name}/{k}"] = v
-                for p, pname in zip(
-                    optim_a.param_groups[0]["params"], param_names[1]
-                ):  # cva
-                    scalars2log[f"{mod_name}/{pname}"] = p.item()
-                    scalars2log[f"{mod_name}/LR_cv_a"] = optim_a.param_groups[0]["lr"]
-                for p, pname in zip(
-                    optim_w.param_groups[0]["params"], param_names[0]
-                ):  # weights
-                    hist2log[f"{mod_name}/{pname}"] = p
-                    scalars2log[f"{mod_name}/LR_w"] = optim_w.param_groups[0]["lr"]
-                for p, pname in zip(
-                    optim_w.param_groups[1]["params"], param_names[2]
-                ):  # cvw
-                    if "alpha" in pname:
-                        hist2log[f"{mod_name}/{pname}"] = p
-                    else:
+                if isinstance(qcfg["tb_writer"], SummaryWriter) and isOutput:
+                    scalars2log = {}
+                    hist2log = {}
+
+                    for k, v in loss4plot.items():  # plot loss
+                        scalars2log[f"{mod_name}/PTQloss_{k}"] = v
+                    for k, v in m.named_buffers():  # plot cv, delta, zp, alpha, and lr
+                        if any(kb in k for kb in ["delta", "zero_point", "clip_val"]):
+                            if len(v.shape) > 0 and v.shape[0] > 1:  # perCh
+                                hist2log[f"{mod_name}/{k}"] = v
+                            else:
+                                scalars2log[f"{mod_name}/{k}"] = v
+                    for p, pname in zip(
+                        optim_a.param_groups[0]["params"], param_names[1]
+                    ):  # cva
                         scalars2log[f"{mod_name}/{pname}"] = p.item()
-                    scalars2log[f"{mod_name}/LR_cvw"] = optim_w.param_groups[1]["lr"]
-                if "adaround" in qcfg["qw_mode"]:
-                    scalars2log[f"{mod_name}/AdaR_beta"] = (
-                        loss_func.temp_decay.curr_beta
-                    )
-                    for lidx, l in enumerate(layers):
-                        if not hasattr(l, "quantize_m1"):
-                            hist2log[f"{mod_name}/W{lidx}"] = l.weight
+                        scalars2log[f"{mod_name}/LR_cv_a"] = optim_a.param_groups[0][
+                            "lr"
+                        ]
+                    for p, pname in zip(
+                        optim_w.param_groups[0]["params"], param_names[0]
+                    ):  # weights
+                        hist2log[f"{mod_name}/{pname}"] = p
+                        scalars2log[f"{mod_name}/LR_w"] = optim_w.param_groups[0]["lr"]
+                    for p, pname in zip(
+                        optim_w.param_groups[1]["params"], param_names[2]
+                    ):  # cvw
+                        if "alpha" in pname:
+                            hist2log[f"{mod_name}/{pname}"] = p
+                        else:
+                            scalars2log[f"{mod_name}/{pname}"] = p.item()
+                        scalars2log[f"{mod_name}/LR_cvw"] = optim_w.param_groups[1][
+                            "lr"
+                        ]
+                    if "adaround" in qcfg["qw_mode"]:
+                        scalars2log[f"{mod_name}/AdaR_beta"] = (
+                            loss_func.temp_decay.curr_beta
+                        )
+                        for lidx, l in enumerate(layers):
+                            if not hasattr(l, "quantize_m1"):
+                                hist2log[f"{mod_name}/W{lidx}"] = l.weight
 
-                # write every in one shot will mess up the folder, better write them one by one
-                for n, v in scalars2log.items():
-                    qcfg["tb_writer"].add_scalar(n, v, Niter)
-                for n, v in hist2log.items():
-                    qcfg["tb_writer"].add_histogram(n, v, Niter)
+                    # write every in one shot will mess up the folder, better write them one by one
+                    for n, v in scalars2log.items():
+                        qcfg["tb_writer"].add_scalar(n, v, Niter)
+                    for n, v in hist2log.items():
+                        qcfg["tb_writer"].add_histogram(n, v, Niter)
 
         for s in scheduler:
             s.step()  # we set up scheduler based on Nouterloop, not inner
diff --git a/fms_mo/run_quant.py b/fms_mo/run_quant.py
@@ -34,7 +34,6 @@
 
 # Third Party
 from datasets import load_from_disk
-from huggingface_hub.errors import HFValidationError
 from torch.cuda import OutOfMemoryError
 from transformers import AutoTokenizer
 import torch
@@ -353,12 +352,6 @@ def main():
         logger.error(traceback.format_exc())
         write_termination_log(f"Unable to load file: {e}")
         sys.exit(USER_ERROR_EXIT_CODE)
-    except HFValidationError as e:
-        logger.error(traceback.format_exc())
-        write_termination_log(
-            f"There may be a problem with loading the model. Exception: {e}"
-        )
-        sys.exit(USER_ERROR_EXIT_CODE)
     except (TypeError, ValueError, EnvironmentError) as e:
         logger.error(traceback.format_exc())
         write_termination_log(
diff --git a/fms_mo/utils/import_utils.py b/fms_mo/utils/import_utils.py
@@ -16,10 +16,18 @@
 Utils for storing what optional dependencies are available
 """
 
+# Standard
+import pkgutil
+import sys
+
 # Third Party
 from transformers.utils.import_utils import _is_package_available
 import torch
 
+all_available_modules = []
+for finder, name, ispkg in pkgutil.iter_modules(sys.path):
+    all_available_modules.append(name)
+
 optional_packages = [
     "gptqmodel",
     "gptqmodel_exllama_kernels",
@@ -32,12 +40,15 @@
     "fms",
     "triton",
     "torchvision",
+    "huggingface_hub",
     "torchao",
 ]
 
 available_packages = {}
 for package in optional_packages:
-    available_packages[package] = _is_package_available(package)
+    available_packages[package] = (
+        _is_package_available(package) or package in all_available_modules
+    )
 
 # cutlass is detected through torch.ops.cutlass_gemm
 available_packages["cutlass"] = hasattr(torch.ops, "cutlass_gemm") and hasattr(
diff --git a/pyproject.toml b/pyproject.toml
@@ -23,32 +23,28 @@ classifiers=[
 dynamic = ["version"]
 dependencies = [
 "numpy>=1.26.4,<2.3.0",
-"accelerate>=0.20.3,!=0.34,<1.7",
+"accelerate>=0.20.3,!=0.34,<1.9",
 "transformers>=4.45,<4.53",
-"torch>=2.2.0,<2.6", 
+"torch>=2.2.0,<2.6",
 "tqdm>=4.66.2,<5.0",
 "datasets>=3.0.0,<4.0",
-"ninja>=1.11.1.1,<2.0",
-"tensorboard",
-"notebook",
-"evaluate",
-"huggingface_hub",
 "pandas",
 "safetensors",
-"ibm-fms>=0.0.8",
 "pkginfo>1.10",
 "torchao"
 ]
 
 [project.optional-dependencies]
+examples = ["ninja>=1.11.1.1,<2.0", "evaluate", "huggingface_hub"]
 fp8 = ["llmcompressor", "torchao"]
 gptq = ["Cython", "gptqmodel>=1.7.3"]
 mx = ["microxcaling>=1.1"]
 opt = ["fms-model-optimizer[fp8, gptq, mx]"]
+aiu = ["ibm-fms>=0.0.8"]
 torchvision = ["torchvision>=0.17"]
 flash-attn = ["flash-attn>=2.5.3,<3.0"]
 triton = ["triton>=3.0,<3.4"]
-visualize = ["matplotlib", "graphviz", "pygraphviz"]
+visualize = ["matplotlib", "graphviz", "pygraphviz", "tensorboard", "notebook"]
 dev = ["pre-commit>=3.0.4,<5.0"]
 test = ["pytest", "pillow"]
 
diff --git a/tests/build/test_launch_script.py b/tests/build/test_launch_script.py
@@ -250,7 +250,7 @@ def _validate_quantization_output(base_dir, quant_method):
     assert os.path.exists(os.path.join(base_dir, "tokenizer.json")) is True
     assert os.path.exists(os.path.join(base_dir, "special_tokens_map.json")) is True
     assert os.path.exists(os.path.join(base_dir, "tokenizer_config.json")) is True
-    assert os.path.exists(os.path.join(base_dir, "tokenizer.model")) is True
+    # assert os.path.exists(os.path.join(base_dir, "tokenizer.model")) is True
 
     # Check quantized model files exist
     if quant_method == "gptq":

Original file line number	Diff line number	Diff line change
`@@ -173,7 +173,7 @@ def _math_fp8_compute_op(`
`173`	`173`	`)`
`174`	`174`
`175`	`175`	`attn_weight = (`
`176`		`- torch.ops.sendnn.scaled_bmm(`
	`176`	`+ torch.ops.spyre.scaled_bmm(`
`177`	`177`	`query,`
`178`	`178`	`key_cache.transpose(-2, -1),`
`179`	`179`	`q_scale,`