model: dbrx: fix missing embedding tensor, mix with output layer

This commit is contained in:
Pierrick HYMBERT 2024-04-06 23:27:29 +02:00
parent 9c7dedb0f3
commit fe8089871e

View file

@ -24,6 +24,7 @@ class TensorNameMap:
"backbone.embedding", # mamba
"backbone.embeddings", # mamba-hf
"transformer.in_out_embed", # Grok
"transformer.wte.weight", # dbrx
),
# Token type embeddings
@ -52,7 +53,7 @@ class TensorNameMap:
"output", # llama-pth bloom internlm2
"word_embeddings_for_head", # persimmon
"lm_head.linear", # phi2
"transformer.wte.weight", # dbrx
"lm_head.weight", # dbrx
),
# Output norm