Merge pull request #68 from andrea-fasoli/fix_int8_linear_type

fix: handle linear_type callable at int8 linear instantiation

Author: chichun-charlie-liu

fms_mo/aiu_addons/i8i8/i8i8_aiu_adapter.py

@@ -37,7 +37,7 @@ def _int8_qparams_aiu(
 
             param_type = "w" if is_weight else "a"
             new_name = f"{module_name}.{param_type}_{name_split[-1]}"
-        elif "smoothq" in name:
+        elif "smoothq" in name and "smoothquant" not in name:
             new_name = name.replace("smoothq", "smoothquant")
 
         new_sd[new_name] = param

fms_mo/aiu_addons/i8i8/i8i8_aiu_linear.py

@@ -16,7 +16,8 @@
 # Standard
 from dataclasses import dataclass
 from functools import partial
-from typing import Any, Mapping, Optional
+from typing import Any, Callable, Optional, Union
+import copy
 
 # Third Party
 from fms.modules.linear import (
@@ -197,16 +198,38 @@ def __repr__(self) -> str:
         )
 
 
+def update_from_partial(
+    linear_config: dict[Union[str, Callable], Any],
+) -> dict[Union[str, Callable], Any]:
+    """Update linear config parameters using those of partial callable"""
+
+    linear_config_updated = copy.deepcopy(linear_config)
+    for k, v in linear_config["linear_type"].keywords.items():
+        linear_config_updated[k] = v
+    return linear_config_updated
+
+
 def get_int8_aiu_linear(
     in_features: int,
     out_features: int,
     bias: bool,
-    linear_config: Optional[Mapping[str, Any]] = None,
+    linear_config: dict[Union[str, Callable], Any],
+    linear_type: Optional[str] = None,
     use_smoothquant: bool = False,
 ) -> torch.nn.Module:
     """Retrieve a W8A8 Linear module"""
 
-    int8_config = W8A8LinearConfig(**linear_config)
+    # Preprocess linear_config if its linear_type field is a callable
+    # (which would not initialize correctly the dataclass parameters).
+    # We don't want to alter the original linear_config though.
+    linear_config_for_dataclass: Optional[dict[Union[str, Callable], Any]] = None
+    if callable(linear_config["linear_type"]):
+        linear_config_for_dataclass = update_from_partial(linear_config)
+        linear_config_for_dataclass["linear_type"] = linear_type
+    if not linear_config_for_dataclass:
+        linear_config_for_dataclass = linear_config
+
+    int8_config = W8A8LinearConfig(**linear_config_for_dataclass)
     linear = W8A8LinearAIU(
         in_features=in_features,
         out_features=out_features,
@@ -281,9 +304,20 @@ def shard_int8_aiu_linear(
     # return unused_keys
 
 
-register_linear_type_to_module_map("int8_aiu", get_int8_aiu_linear)
+register_linear_type_to_module_map(
+    "int8_aiu",
+    partial(
+        get_int8_aiu_linear,
+        linear_type="int8_aiu",
+        use_smoothquant=False,
+    ),
+)
 register_linear_type_to_module_map(
     "int8_smoothquant_aiu",
-    partial(get_int8_aiu_linear, use_smoothquant=True),
+    partial(
+        get_int8_aiu_linear,
+        linear_type="int8_smoothquant_aiu",
+        use_smoothquant=True,
+    ),
 )
 register_linear_type_to_sharding_map("int8_aiu", shard_int8_aiu_linear)