fix: use weights from base_layer

huggingface · Jun 28, 2024 · c4feb98 · c4feb98
1 parent fb98ab2
commit c4feb98
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)