[Bugfix] Fix KeyError on loading GPT-NeoX (vllm-project#3925)

Temirulan · Apr 9, 2024 · 13d3c36 · 13d3c36
1 parent 15a1d71
commit 13d3c36
Showing 1 changed file with 5 additions and 0 deletions.
diff --git 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: