fix: use the base layers weight in mistral rocm (#2155)

2024-07-02 05:56:25 -04:00 · 2024-07-02 05:56:25 -04:00 · b966bc0d35
parent 5d97e0c4a3
commit b966bc0d35
1 changed files with 3 additions and 1 deletions
--- a/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py
+++ b/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py
@ -315,7 +315,9 @@ class MistralMLP(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)