fix: corrected error associated with eval_llm_1GPU failing with granite-3-models

Signed-off-by: omobayode.fagbohungbe <[email protected]>

Author: bayo-ibm

fms_mo/quant/ptq.py

@@ -2140,14 +2140,22 @@ def get_blocks(model, model_type=None):
             None,
             "lm_head",
         ),
-        "granite": (
+        "granite_old": (
             "transformer.h",
             "transformer.wte",
             "transformer.wpe",
             None,
             "transformer.ln_f",
             "lm_head",
         ),
+        "granite": (
+            "model.layers",
+            "model.embed_tokens",
+            "model.rotary_emb",
+            None,
+            "model.norm",
+            "lm_head",
+        ),
         "llama": (
             "model.layers",
             "model.embed_tokens",

fms_mo/utils/eval_utils.py

@@ -107,6 +107,9 @@ def eval_llm_1GPU(qcfg, model, test_dataset, pre_cache_func=None, **kwargs):  #
         lm_head.to(dev)
         lm_logits = lm_head(hidden_states)
 
+        if model.config.model_type == "granite":
+            lm_logits /= model.config.logits_scaling
+
         # Shift so that tokens < n predict n
         shift_logits = lm_logits[:, :-1, :].contiguous().float()
         shift_labels = test_dataset.input_ids[:, (i * seq_len) : ((i + 1) * seq_len)][