Merge pull request #92 from andrea-fasoli/int8_granite_addon

feat: int8 granite addon

Author: chichun-charlie-liu

fms_mo/aiu_addons/i8i8/i8i8_aiu_adapter.py

@@ -97,41 +97,22 @@ def _add_defaults_and_concat(
         )
 
 
-# registration of new adapter steps for each architecture
-serialization.register_adapter_step("llama", "int8_qparams_aiu", _int8_qparams_aiu)
-serialization.register_adapter_step(
-    "gpt_bigcode", "int8_qparams_aiu", _int8_qparams_aiu
-)
-serialization.register_adapter_step("roberta", "int8_qparams_aiu", _int8_qparams_aiu)
-serialization.register_adapter_step(
-    "roberta_question_answering",
-    "int8_qparams_aiu",
-    _int8_qparams_aiu,
-)
-
-# registration of multi-step adapter for each architecture
-serialization.register_adapter(
+# registration of new adapter step and adapter for each architecture
+for arch in [
     "llama",
-    "fms_mo",
-    [
-        "hf_to_fms_names",
-        "hf_to_fms_rope",
-        "weight_fusion",
-        "int8_qparams_aiu",
-    ],
-)
-serialization.register_adapter(
-    "gpt_bigcode", "fms_mo", ["hf_to_fms_names", "weight_fusion", "int8_qparams_aiu"]
-)
-serialization.register_adapter(
-    "roberta", "fms_mo", ["hf_to_fms_names", "weight_fusion", "int8_qparams_aiu"]
-)
-serialization.register_adapter(
+    "gpt_bigcode",
+    "granite",
+    "roberta",
     "roberta_question_answering",
-    "fms_mo",
-    [
-        "hf_to_fms_names",
-        "weight_fusion",
-        "int8_qparams_aiu",
-    ],
-)
+]:
+    serialization.register_adapter_step(arch, "int8_qparams_aiu", _int8_qparams_aiu)
+    if arch in ["llama", "granite"]:
+        steps_to_register = [
+            "hf_to_fms_names",
+            "hf_to_fms_rope",
+            "weight_fusion",
+            "int8_qparams_aiu",
+        ]
+    else:
+        steps_to_register = ["hf_to_fms_names", "weight_fusion", "int8_qparams_aiu"]
+    serialization.register_adapter(arch, "fms_mo", steps_to_register)

fms_mo/aiu_addons/i8i8/i8i8_aiu_linear.py

@@ -84,10 +84,10 @@ def __init__(
             "weight",
             torch.zeros(out_features, in_features, dtype=torch.int8),
         )
-        if bias:
-            self.register_buffer(
-                "bias", torch.zeros((out_features), dtype=torch.float16)
-            )
+
+        self.has_bias = bias
+        bias_size = out_features if self.has_bias else 1
+        self.register_buffer("bias", torch.zeros((bias_size), dtype=torch.float16))
 
         if config.weight_per_channel:
             w_clip_size = out_features
@@ -192,7 +192,7 @@ def __repr__(self) -> str:
         return (
             f"{self.__class__.__name__}"
             f"(in={self.in_features}, out={self.out_features}, "
-            f"bias={self.bias is not None}, wq={self.weight_quant_type}, "
+            f"bias={self.has_bias}, wq={self.weight_quant_type}, "
             f"aq={self.activ_quant_type}, smoothq={self.smoothquant}, "
             f"op={self.aiu_op})"
         )