support qwen

Signed-off-by: changwangss <[email protected]>
huggingface · Oct 19, 2023 · 74912c0 · 74912c0
1 parent 8273e7f
commit 74912c0
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/optimum/intel/generation/modeling.py b/optimum/intel/generation/modeling.py
@@ -284,6 +284,12 @@ def forward(
                     if self.model_dtype is not None:
                         empty_tensor = empty_tensor.to(self.model_dtype)
                     past_key_values = tuple([empty_tensor] * num_layers)
+                elif self.config.model_type == "qwen":
+                    new_shape = [input_ids.shape[0], 0, num_key_value_heads, d_k]
+                    empty_tensor = torch.empty(size=new_shape)
+                    if self.model_dtype is not None:
+                        empty_tensor = empty_tensor.to(self.model_dtype)
+                    pkv = tuple(empty_tensor for _ in range(nb_pkv))
                 elif self.config.model_type != "bloom":
                     new_shape = [input_ids.shape[0], num_key_value_heads, 0, d_k]
                     empty_tensor = torch.empty(size=new_shape)