From 44d0b4c772ea6e2a2bad8e5f84bd159863142e4e Mon Sep 17 00:00:00 2001
From: Alexandre Marques <alexandre@neuralmagic.com>
Date: Fri, 13 Oct 2023 17:22:45 -0400
Subject: [PATCH] Add down_proj smoothing

---
 src/sparseml/experimental/sparsegpt/llama2.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/sparseml/experimental/sparsegpt/llama2.py b/src/sparseml/experimental/sparsegpt/llama2.py
index 9c99f487ab9..9beed5dcf40 100644
--- a/src/sparseml/experimental/sparsegpt/llama2.py
+++ b/src/sparseml/experimental/sparsegpt/llama2.py
@@ -38,6 +38,10 @@
         "module_to_balance": ["gate_proj", "up_proj"],
         "module_to_merge": "post_attention_layernorm",
     },
+    {
+        "module_to_balance": ["down_proj"],
+        "module_to_merge": "up_proj",
+    },
 ]