From 61fdf72b423adc977d589504e5f288e1117302de Mon Sep 17 00:00:00 2001
From: Ooze <50599367+Ooooze@users.noreply.github.com>
Date: Tue, 18 Nov 2025 06:02:58 +0300
Subject: [PATCH] Add support for Granite MoE Hybrid in model.py by including
 down projections for shared MLP and MoE experts (#14)

---
 src/heretic/model.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/heretic/model.py b/src/heretic/model.py
index b01439e..b617052 100644
--- a/src/heretic/model.py
+++ b/src/heretic/model.py
@@ -141,6 +141,15 @@ class Model:
             # but thanks to PyTorch's broadcasting magic, it all just works anyway.
             try_add("mlp.down_proj", layer.mlp.experts.down_proj)
 
+        # Granite MoE Hybrid - attention layers with shared_mlp.
+        with suppress(Exception):
+            try_add("mlp.down_proj", layer.shared_mlp.output_linear.weight)
+
+        # Granite MoE Hybrid - MoE layers with experts.
+        with suppress(Exception):
+            for expert in layer.moe.experts:
+                try_add("mlp.down_proj", expert.output_linear.weight)
+
         # We need at least one MLP down-projection.
         assert matrices["mlp.down_proj"]