Merge branch 'main' into unit_test_int8

Signed-off-by: chichun-charlie-liu <[email protected]>

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
@@ -200,8 +201,8 @@ def get_int8_aiu_linear(
     in_features: int,
     out_features: int,
     bias: bool,
-    linear_config: Mapping[str, Any],
-    use_smoothquant: bool = True,
+    linear_config: Optional[Mapping[str, Any]] = None,
+    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)