Unverified Commit e23a43ae authored by Junichi Sato's avatar Junichi Sato Committed by GitHub
Browse files

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

parent e7c7067b
...@@ -274,6 +274,11 @@ class GPTNeoXForCausalLM(nn.Module): ...@@ -274,6 +274,11 @@ class GPTNeoXForCausalLM(nn.Module):
if ("attention.bias" in name or "attention.masked_bias" in name if ("attention.bias" in name or "attention.masked_bias" in name
or "rotary_emb.inv_freq" in name): or "rotary_emb.inv_freq" in name):
continue 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] param = params_dict[name]
if "query_key_value" in name: if "query_key_value" in name:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment