fix: use weights from base_layer (#2141)

huggingface · Jul 1, 2024 · 25f57e2 · 25f57e2
1 parent b4552f9
commit 25f57e2
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py b/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py
@@ -309,7 +309,9 @@ def forward(self, hidden_states, adapter_data):
                 dtype=hidden_states.dtype,
                 device="cuda",
             )
-            _custom_C.LLMM_Silu(self.gate_up_proj.linear.weight, hidden_states, out, 8)
+            _custom_C.LLMM_Silu(
+                self.gate_up_proj.base_layer.linear.weight, hidden_states, out, 8
+            )
             return self.down_proj(out, adapter_data)
         else:
             gate_up_states = self.gate_up_proj(hidden_states, adapter_data)