fix: Added dependency guards for aiu_addon files

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

Author: BrandonGroth

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],

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

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],

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,