added granite support; fixed adapters to ignore model_config

Signed-off-by: JOSHUA ROSENKRANZ <[email protected]>

Author: JRosenkranz

fms_mo/aiu_addons/__init__.py

@@ -22,7 +22,7 @@
 
 
 def _gptq_qweights_transpose_aiu(
-    input_sd: Mapping[str, torch.Tensor],
+    input_sd: Mapping[str, torch.Tensor], **kwargs
 ) -> Mapping[str, torch.Tensor]:
     new_sd = {}
     for name, param in input_sd.items():
@@ -41,6 +41,7 @@ def _gptq_qweights_transpose_aiu(
 serialization.register_adapter_step(
     "gpt_bigcode", "gptq_qweights_transpose_aiu", _gptq_qweights_transpose_aiu
 )
+serialization.register_adapter_step("granite", "gptq_qweights_transpose_aiu", _gptq_qweights_transpose_aiu)
 serialization.register_adapter(
     "llama",
     "hf_gptq_aiu",
@@ -57,3 +58,9 @@ def _gptq_qweights_transpose_aiu(
     "hf_gptq_aiu",
     ["hf_to_fms_names", "weight_fusion", "gptq_qweights_transpose_aiu"],
 )
+serialization.register_adapter(
+    "granite",
+    "hf_gptq_aiu",
+    ["hf_to_fms_names", "hf_to_fms_rope", "hf_gptq_fusion_check", "weight_fusion", "gptq_qweights_transpose_aiu"]
+)
+

fms_mo/aiu_addons/gptq/gptq_aiu_adapter.py

@@ -22,7 +22,7 @@
 
 
 def _int8_qparams_aiu(
-    input_sd: Mapping[str, torch.Tensor],
+    input_sd: Mapping[str, torch.Tensor], **kwargs
 ) -> Mapping[str, torch.Tensor]:
     new_sd = {}
     modules_seen = set()