fix: use weights from base_layer (#2141)

25f57e2e · drbh · GitHub · b4552f9d · 25f57e2e
Unverified Commit 25f57e2e authored Jul 01, 2024 by drbh Committed by GitHub Jul 01, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

server/text_generation_server/models/custom_modeling/flash_llama_modeling.py ...ion_server/models/custom_modeling/flash_llama_modeling.py +3 -1

No files found.
--- 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 @@ class LlamaMLP(nn.Module):
                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)