Merge pull request #65 from andrea-fasoli/int8_smoothquant

feat: Support for int8 smoothquant

Author: chichun-charlie-liu

fms_mo/aiu_addons/i8i8/i8i8_aiu_adapter.py

@@ -103,6 +103,11 @@ def _add_defaults_and_concat(
     "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(
@@ -121,3 +126,12 @@ def _add_defaults_and_concat(
 serialization.register_adapter(
     "roberta", "fms_mo", ["hf_to_fms_names", "weight_fusion", "int8_qparams_aiu"]
 )
+serialization.register_adapter(
+    "roberta_question_answering",
+    "fms_mo",
+    [
+        "hf_to_fms_names",
+        "weight_fusion",
+        "int8_qparams_aiu",
+    ],
+)

fms_mo/aiu_addons/i8i8/i8i8_aiu_linear.py

@@ -15,6 +15,7 @@
 
 # Standard
 from dataclasses import dataclass
+from functools import partial
 from typing import Any, Mapping, Optional
 
 # Third Party
@@ -201,7 +202,7 @@ def get_int8_aiu_linear(
     out_features: int,
     bias: bool,
     linear_config: Optional[Mapping[str, Any]] = None,
-    use_smoothquant: bool = True,
+    use_smoothquant: bool = False,
 ) -> torch.nn.Module:
     """Retrieve a W8A8 Linear module"""
 
@@ -281,4 +282,8 @@ def shard_int8_aiu_linear(
 
 
 register_linear_type_to_module_map("int8_aiu", get_int8_aiu_linear)
+register_linear_type_to_module_map(
+    "int8_smoothquant_aiu",
+    partial(get_int8_aiu_linear, use_smoothquant=True),
+)
 register_linear_type_to_sharding_map("int8_aiu", shard_int8_aiu_linear)

fms_mo/aiu_addons/i8i8/i8i8_aiu_op.py

@@ -84,7 +84,7 @@ def i8i8_aiu(
         x_dq = quant_dequant_activ(x, a_cv, a_cvn, sq, activ_quant_type)
         w_dq = dequant_weights(weight, w_cv, sq, weight_quant_type)
 
-        return F.linear(x_dq.to(dtype), w_dq.to(dtype), bias)
+        return F.linear(x_dq.to(dtype), w_dq.to(dtype), bias.to(dtype))
 
     @torch.library.impl_abstract(op_namespace_id)
     def i8i8_aiu_abstract(