From e83bfa09e8b9b3911fd98e5f7a63fb16e2b9b184 Mon Sep 17 00:00:00 2001 From: Junichi Sato Date: Wed, 10 Apr 2024 04:11:31 +0900 Subject: [PATCH] [Bugfix] Fix KeyError on loading GPT-NeoX (#3925) --- vllm/model_executor/models/gpt_neox.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vllm/model_executor/models/gpt_neox.py b/vllm/model_executor/models/gpt_neox.py index 673900487cc96..a5b5d717d9846 100644 --- a/vllm/model_executor/models/gpt_neox.py +++ b/vllm/model_executor/models/gpt_neox.py @@ -274,6 +274,11 @@ def load_weights(self, if ("attention.bias" in name or "attention.masked_bias" in name or "rotary_emb.inv_freq" in name): continue + if ("rotary_emb.cos_cached" in name + or "rotary_emb.sin_cached" in name): + # Models trained using OpenRLHF may include + # these tensors in the checkpoint. Skip them. + continue param = params_dict[name] if "query_key_value" in name: