diff --git a/vllm/model_executor/models/baichuan.py b/vllm/model_executor/models/baichuan.py index 54ed548ba8bc7..767230aeacc35 100644 --- a/vllm/model_executor/models/baichuan.py +++ b/vllm/model_executor/models/baichuan.py @@ -432,7 +432,9 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): class BaichuanForCausalLM(BaiChuanBaseForCausalLM): - """Baichuan 13B and Baichuan2 7B/13B.""" + """Baichuan 13B and Baichuan2 7B/13B. + NOTE: the class name has a lower case 'c'. + """ def __init__( self, @@ -450,7 +452,9 @@ def __init__( class BaiChuanForCausalLM(BaiChuanBaseForCausalLM): - """Baichuan 7B.""" + """Baichuan 7B. + NOTE: the class name has an upper case 'C'. + """ def __init__( self, diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 787c65743e894..db58414299070 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -26,8 +26,10 @@ "AquilaModel": ("llama", "LlamaForCausalLM"), "AquilaForCausalLM": ("llama", "LlamaForCausalLM"), # AquilaChat2 "ArcticForCausalLM": ("arctic", "ArcticForCausalLM"), - "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"), # baichuan-7b - "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"), # baichuan-13b + # baichuan-7b, upper case 'C' in the class name + "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"), + # baichuan-13b, lower case 'c' in the class name + "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"), "BloomForCausalLM": ("bloom", "BloomForCausalLM"), # ChatGLMModel supports multimodal "CohereForCausalLM": ("commandr", "CohereForCausalLM"),