Fix Baichuan2-7B-Chat (#1987)

2b981012 · firebook · GitHub · 6ccc0bff · 2b981012
Unverified Commit 2b981012 authored Dec 09, 2023 by firebook Committed by GitHub Dec 08, 2023
Show whitespace changes
Inline Side-by-side

Showing with 6 additions and 2 deletions

vllm/model_executor/models/baichuan.py vllm/model_executor/models/baichuan.py +6 -2

No files found.
--- a/vllm/model_executor/models/baichuan.py
+++ b/vllm/model_executor/models/baichuan.py
@@ -366,11 +366,15 @@ class BaiChuanBaseForCausalLM(nn.Module):
                weight_loader(param, loaded_weight)
-class BaichuanForCausalLM(BaiChuanBaseForCausalLM):  # baichuan 13b
+class BaichuanForCausalLM(BaiChuanBaseForCausalLM
+                          ):  # baichuan 13b, baichuan2 13b, baichuan2 7b
    def __init__(self,
                 config,
                 linear_method: Optional[LinearMethodBase] = None):
+        if config.hidden_size == 4096:  # baichuan2 7b
+            super().__init__(config, "ROPE", linear_method)
+        else:  # baichuan 13b, baichuan2 13b
            super().__init__(config, "ALIBI", linear_method)