Fix Falcon weight mapping for H2O.ai checkpoints

During the safetensor conversion, duplicate weights are removed.
However, which of the duplicates gets removed, differs per checkpoint.
In some, like `h2oai/h2ogpt-oig-oasst1-falcon-40b`, the weight
`transformer.word_embeddings.weightSafetensor` gets removed. In others,
`lm_head.weight` gets removed. Long story long, we need to support both.

Originally, f018143 mapped `lm_head` to `word_embeddings`. Then ac736fd
switched this around. This commit merges them and allows for both.
This commit is contained in:
Vincent Brouwers 2023-08-30 09:50:49 +00:00
parent 7c2e0af2a6
commit e864b95656

View File

@ -54,7 +54,10 @@ class FlashRWSharded(FlashCausalLM):
device,
dtype,
process_group=self.process_group,
aliases={"lm_head.weight": ["transformer.word_embeddings.weight"]},
aliases={
"lm_head.weight": ["transformer.word_embeddings.weight"],
"transformer.word_embeddings.weight": ["lm_head.weight"],
},
)
config.quantize = quantize