leafspark
/
AutoGGUF
mirror of https://github.com/leafspark/AutoGGUF
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

refactor: adapt gguf library to project 

- remove comments
- remove argparse help text
This commit is contained in:
BuildTools 2024-08-16 19:58:29 -07:00
parent f7f9a457ea
commit a7e8bf673e
No known key found for this signature in database
GPG Key ID: 3270C066C15D530B
3 changed files with 1007 additions and 909 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
src/convert_hf_to_gguf.py
View File

@ -4403,83 +4403,81 @@ def __torch_function__(cls, func, types, args=(), kwargs=None):
def parse_args() -> argparse.Namespace: def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(description="")
description="Convert a huggingface model to a GGML compatible file"
)
parser.add_argument( parser.add_argument(
"--vocab-only", "--vocab-only",
action="store_true", action="store_true",
help="extract only the vocab", help="",
) )
parser.add_argument( parser.add_argument(
"--outfile", "--outfile",
type=Path, type=Path,
help="path to write to; default: based on input. {ftype} will be replaced by the outtype.", help="",
) )
parser.add_argument( parser.add_argument(
"--outtype", "--outtype",
type=str, type=str,
choices=["f32", "f16", "bf16", "q8_0", "auto"], choices=["f32", "f16", "bf16", "q8_0", "auto"],
default="f16", default="f16",
help="output format - use f32 for float32, f16 for float16, bf16 for bfloat16, q8_0 for Q8_0, auto for the highest-fidelity 16-bit float type depending on the first loaded tensor type", help="",
) )
parser.add_argument( parser.add_argument(
"--bigendian", "--bigendian",
action="store_true", action="store_true",
help="model is executed on big endian machine", help="",
) )
parser.add_argument( parser.add_argument(
"model", "model",
type=Path, type=Path,
help="directory containing model file", help="",
) )
parser.add_argument( parser.add_argument(
"--use-temp-file", "--use-temp-file",
action="store_true", action="store_true",
help="use the tempfile library while processing (helpful when running out of memory, process killed)", help="",
) )
parser.add_argument( parser.add_argument(
"--no-lazy", "--no-lazy",
action="store_true", action="store_true",
help="use more RAM by computing all outputs before writing (use in case lazy evaluation is broken)", help="",
) )
parser.add_argument( parser.add_argument(
"--model-name", "--model-name",
type=str, type=str,
default=None, default=None,
help="name of the model", help="",
) )
parser.add_argument( parser.add_argument(
"--verbose", "--verbose",
action="store_true", action="store_true",
help="increase output verbosity", help="",
) )
parser.add_argument( parser.add_argument(
"--split-max-tensors", "--split-max-tensors",
type=int, type=int,
default=0, default=0,
help="max tensors in each split", help="",
) )
parser.add_argument( parser.add_argument(
"--split-max-size", "--split-max-size",
type=str, type=str,
default="0", default="0",
help="max size per split N(M|G)", help="",
) )
parser.add_argument( parser.add_argument(
"--dry-run", "--dry-run",
action="store_true", action="store_true",
help="only print out a split plan and exit, without writing any new files", help="",
) )
parser.add_argument( parser.add_argument(
"--no-tensor-first-split", "--no-tensor-first-split",
action="store_true", action="store_true",
help="do not add tensors to the first split (disabled by default)", help="",
) )
parser.add_argument( parser.add_argument(
"--metadata", "--metadata",
type=Path, type=Path,
help="Specify the path for an authorship metadata override file", help="",
) )
return parser.parse_args() return parser.parse_args()

1137
src/gguf-py/gguf/constants.py
View File

File diff suppressed because it is too large Load Diff

747
src/gguf-py/gguf/tensor_mapping.py
View File

@ -4,471 +4,569 @@
from .constants import MODEL_ARCH, MODEL_TENSOR, MODEL_TENSORS, TENSOR_NAMES from .constants import MODEL_ARCH, MODEL_TENSOR, MODEL_TENSORS, TENSOR_NAMES
class TensorNameMap: class TensorNameMap:
mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = { mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = {
# Token embeddings
MODEL_TENSOR.TOKEN_EMBD: ( MODEL_TENSOR.TOKEN_EMBD: (
"gpt_neox.embed_in", # gptneox "gpt_neox.embed_in",
"transformer.wte", # gpt2 gpt-j mpt refact qwen dbrx jais "transformer.wte",
"transformer.word_embeddings", # falcon "transformer.word_embeddings",
"word_embeddings", # bloom "word_embeddings",
"model.embed_tokens", # llama-hf "model.embed_tokens",
"tok_embeddings", # llama-pth "tok_embeddings",
"embeddings.word_embeddings", # bert nomic-bert "embeddings.word_embeddings",
"language_model.embedding.word_embeddings", # persimmon "language_model.embedding.word_embeddings",
"wte", # gpt2 "wte",
"transformer.embd.wte", # phi2 "transformer.embd.wte",
"model.tok_embeddings", # internlm2 "model.tok_embeddings",
"model.embedding", # mamba-qbert "model.embedding",
"backbone.embedding", # mamba "backbone.embedding",
"backbone.embeddings", # mamba-hf "backbone.embeddings",
"transformer.in_out_embed", # Grok "transformer.in_out_embed",
"embedding.word_embeddings", # chatglm "embedding.word_embeddings",
"transformer.token_embeddings", # openelm "transformer.token_embeddings",
"shared", # t5 "shared",
), ),
# Token type embeddings
MODEL_TENSOR.TOKEN_TYPES: ( MODEL_TENSOR.TOKEN_TYPES: (
"embeddings.token_type_embeddings", # bert nomic-bert "embeddings.token_type_embeddings",
), ),
# Normalization of token embeddings
MODEL_TENSOR.TOKEN_EMBD_NORM: ( MODEL_TENSOR.TOKEN_EMBD_NORM: (
"word_embeddings_layernorm", # bloom "word_embeddings_layernorm",
"embeddings.LayerNorm", # bert "embeddings.LayerNorm",
"emb_ln", # nomic-bert "emb_ln",
"transformer.norm", # openelm "transformer.norm",
), ),
# Position embeddings
MODEL_TENSOR.POS_EMBD: ( MODEL_TENSOR.POS_EMBD: (
"transformer.wpe", # gpt2 "transformer.wpe",
"embeddings.position_embeddings", # bert "embeddings.position_embeddings",
"wpe", # gpt2 "wpe",
), ),
# Output
MODEL_TENSOR.OUTPUT: ( MODEL_TENSOR.OUTPUT: (
"embed_out", # gptneox "embed_out",
"lm_head", # gpt2 mpt falcon llama-hf baichuan qwen mamba dbrx jais "lm_head",
"output", # llama-pth bloom internlm2 "output",
"word_embeddings_for_head", # persimmon "word_embeddings_for_head",
"lm_head.linear", # phi2 "lm_head.linear",
"output_layer", # chatglm "output_layer",
), ),
# Output norm
MODEL_TENSOR.OUTPUT_NORM: ( MODEL_TENSOR.OUTPUT_NORM: (
"gpt_neox.final_layer_norm", # gptneox "gpt_neox.final_layer_norm",
"transformer.ln_f", # gpt2 gpt-j falcon jais "transformer.ln_f",
"model.norm", # llama-hf baichuan internlm2 "model.norm",
"norm", # llama-pth "norm",
"transformer.norm_f", # mpt dbrx "transformer.norm_f",
"ln_f", # refact bloom qwen gpt2 "ln_f",
"language_model.encoder.final_layernorm", # persimmon "language_model.encoder.final_layernorm",
"model.final_layernorm", # persimmon "model.final_layernorm",
"lm_head.ln", # phi2 "lm_head.ln",
"model.norm_f", # mamba-qbert "model.norm_f",
"backbone.norm_f", # mamba "backbone.norm_f",
"transformer.rms_norm", # Grok "transformer.rms_norm",
"encoder.final_layernorm", # chatglm "encoder.final_layernorm",
"transformer.norm", # openelm "transformer.norm",
"model.norm",
), ),
# Rope frequencies
MODEL_TENSOR.ROPE_FREQS: ( MODEL_TENSOR.ROPE_FREQS: (
"rope.freqs", # llama-pth "rope.freqs",
"rotary_pos_emb.inv_freq", # chatglm "rotary_pos_emb.inv_freq",
), ),
} }
block_mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = { block_mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = {
# Attention norm
MODEL_TENSOR.ATTN_NORM: ( MODEL_TENSOR.ATTN_NORM: (
"gpt_neox.layers.{bid}.input_layernorm", # gptneox "gpt_neox.layers.{bid}.input_layernorm",
"transformer.h.{bid}.ln_1", # gpt2 gpt-j refact qwen jais "transformer.h.{bid}.ln_1",
"transformer.blocks.{bid}.norm_1", # mpt "transformer.blocks.{bid}.norm_1",
"transformer.h.{bid}.input_layernorm", # falcon7b "transformer.h.{bid}.input_layernorm",
"h.{bid}.input_layernorm", # bloom "h.{bid}.input_layernorm",
"transformer.h.{bid}.ln_mlp", # falcon40b "transformer.h.{bid}.ln_mlp",
"model.layers.{bid}.input_layernorm", # llama-hf "model.layers.{bid}.input_layernorm",
"layers.{bid}.attention_norm", # llama-pth "layers.{bid}.attention_norm",
"language_model.encoder.layers.{bid}.input_layernorm", # persimmon "language_model.encoder.layers.{bid}.input_layernorm",
"model.layers.{bid}.ln1", # yi "model.layers.{bid}.ln1",
"h.{bid}.ln_1", # gpt2 "h.{bid}.ln_1",
"transformer.h.{bid}.ln", # phi2 "transformer.h.{bid}.ln",
"model.layers.layers.{bid}.norm", # plamo "model.layers.layers.{bid}.norm",
"model.layers.{bid}.attention_norm", # internlm2 "model.layers.{bid}.attention_norm",
"model.layers.{bid}.norm", # mamba-qbert "model.layers.{bid}.norm",
"backbone.layers.{bid}.norm", # mamba "backbone.layers.{bid}.norm",
"transformer.decoder_layer.{bid}.rms_norm", # Grok "transformer.decoder_layer.{bid}.rms_norm",
"transformer.blocks.{bid}.norm_attn_norm.norm_1", # dbrx "transformer.blocks.{bid}.norm_attn_norm.norm_1",
"encoder.layers.{bid}.input_layernorm", # chatglm "encoder.layers.{bid}.input_layernorm",
"transformer.layers.{bid}.attn_norm", # openelm "transformer.layers.{bid}.attn_norm",
), ),
# Attention norm 2
MODEL_TENSOR.ATTN_NORM_2: ( MODEL_TENSOR.ATTN_NORM_2: (
"transformer.h.{bid}.ln_attn", # falcon40b "transformer.h.{bid}.ln_attn",
"encoder.layer.{bid}.layer_norm_1", # jina-v2-code "encoder.layer.{bid}.layer_norm_1",
), ),
# Attention query-key-value
MODEL_TENSOR.ATTN_QKV: ( MODEL_TENSOR.ATTN_QKV: (
"gpt_neox.layers.{bid}.attention.query_key_value", # gptneox "gpt_neox.layers.{bid}.attention.query_key_value",
"transformer.h.{bid}.attn.c_attn", # gpt2 qwen jais "transformer.h.{bid}.attn.c_attn",
"transformer.blocks.{bid}.attn.Wqkv", # mpt "transformer.blocks.{bid}.attn.Wqkv",
"transformer.blocks.{bid}.norm_attn_norm.attn.Wqkv", # dbrx "transformer.blocks.{bid}.norm_attn_norm.attn.Wqkv",
"transformer.h.{bid}.self_attention.query_key_value", # falcon "transformer.h.{bid}.self_attention.query_key_value",
"h.{bid}.self_attention.query_key_value", # bloom "h.{bid}.self_attention.query_key_value",
"language_model.encoder.layers.{bid}.self_attention.query_key_value", # persimmon "language_model.encoder.layers.{bid}.self_attention.query_key_value",
"model.layers.{bid}.self_attn.query_key_value", # persimmon "model.layers.{bid}.self_attn.query_key_value",
"h.{bid}.attn.c_attn", # gpt2 "h.{bid}.attn.c_attn",
"transformer.h.{bid}.mixer.Wqkv", # phi2 "transformer.h.{bid}.mixer.Wqkv",
"encoder.layers.{bid}.attn.Wqkv", # nomic-bert "encoder.layers.{bid}.attn.Wqkv",
"model.layers.{bid}.self_attn.qkv_proj", # phi3 "model.layers.{bid}.self_attn.qkv_proj",
"encoder.layers.{bid}.self_attention.query_key_value", # chatglm "encoder.layers.{bid}.self_attention.query_key_value",
"transformer.layers.{bid}.attn.qkv_proj", # openelm "transformer.layers.{bid}.attn.qkv_proj",
), ),
# Attention query
MODEL_TENSOR.ATTN_Q: ( MODEL_TENSOR.ATTN_Q: (
"model.layers.{bid}.self_attn.q_proj", # llama-hf "model.layers.{bid}.self_attn.q_proj",
"layers.{bid}.attention.wq", # llama-pth "layers.{bid}.attention.wq",
"encoder.layer.{bid}.attention.self.query", # bert "encoder.layer.{bid}.attention.self.query",
"transformer.h.{bid}.attn.q_proj", # gpt-j "transformer.h.{bid}.attn.q_proj",
"model.layers.layers.{bid}.self_attn.q_proj", # plamo "model.layers.layers.{bid}.self_attn.q_proj",
"model.layers.{bid}.attention.wq", # internlm2 "model.layers.{bid}.attention.wq",
"transformer.decoder_layer.{bid}.multi_head_attention.query", # Grok "transformer.decoder_layer.{bid}.multi_head_attention.query",
"transformer.h.{bid}.attn.attention.q_proj",
), ),
# Attention key
MODEL_TENSOR.ATTN_K: ( MODEL_TENSOR.ATTN_K: (
"model.layers.{bid}.self_attn.k_proj", # llama-hf "model.layers.{bid}.self_attn.k_proj",
"layers.{bid}.attention.wk", # llama-pth "layers.{bid}.attention.wk",
"encoder.layer.{bid}.attention.self.key", # bert "encoder.layer.{bid}.attention.self.key",
"transformer.h.{bid}.attn.k_proj", # gpt-j "transformer.h.{bid}.attn.k_proj",
"transformer.h.{bid}.attn.k", # refact "transformer.h.{bid}.attn.k",
"model.layers.layers.{bid}.self_attn.k_proj", # plamo "model.layers.layers.{bid}.self_attn.k_proj",
"model.layers.{bid}.attention.wk", # internlm2 "model.layers.{bid}.attention.wk",
"transformer.decoder_layer.{bid}.multi_head_attention.key", # Grok "transformer.decoder_layer.{bid}.multi_head_attention.key",
"transformer.h.{bid}.attn.attention.k_proj",
), ),
# Attention value
MODEL_TENSOR.ATTN_V: ( MODEL_TENSOR.ATTN_V: (
"model.layers.{bid}.self_attn.v_proj", # llama-hf "model.layers.{bid}.self_attn.v_proj",
"layers.{bid}.attention.wv", # llama-pth "layers.{bid}.attention.wv",
"encoder.layer.{bid}.attention.self.value", # bert "encoder.layer.{bid}.attention.self.value",
"transformer.h.{bid}.attn.v_proj", # gpt-j "transformer.h.{bid}.attn.v_proj",
"transformer.h.{bid}.attn.v", # refact "transformer.h.{bid}.attn.v",
"model.layers.layers.{bid}.self_attn.v_proj", # plamo "model.layers.layers.{bid}.self_attn.v_proj",
"model.layers.{bid}.attention.wv", # internlm2 "model.layers.{bid}.attention.wv",
"transformer.decoder_layer.{bid}.multi_head_attention.value", # Grok "transformer.decoder_layer.{bid}.multi_head_attention.value",
"transformer.h.{bid}.attn.attention.v_proj",
), ),
# Attention output
MODEL_TENSOR.ATTN_OUT: ( MODEL_TENSOR.ATTN_OUT: (
"gpt_neox.layers.{bid}.attention.dense", # gptneox "gpt_neox.layers.{bid}.attention.dense",
"transformer.h.{bid}.attn.c_proj", # gpt2 refact qwen jais "transformer.h.{bid}.attn.c_proj",
"transformer.blocks.{bid}.attn.out_proj", # mpt "transformer.blocks.{bid}.attn.out_proj",
"transformer.h.{bid}.self_attention.dense", # falcon "transformer.h.{bid}.self_attention.dense",
"h.{bid}.self_attention.dense", # bloom "h.{bid}.self_attention.dense",
"model.layers.{bid}.self_attn.o_proj", # llama-hf "model.layers.{bid}.self_attn.o_proj",
"layers.{bid}.attention.wo", # llama-pth "layers.{bid}.attention.wo",
"encoder.layer.{bid}.attention.output.dense", # bert "encoder.layer.{bid}.attention.output.dense",
"transformer.h.{bid}.attn.out_proj", # gpt-j "transformer.h.{bid}.attn.out_proj",
"language_model.encoder.layers.{bid}.self_attention.dense", # persimmon "language_model.encoder.layers.{bid}.self_attention.dense",
"model.layers.{bid}.self_attn.dense", # persimmon "model.layers.{bid}.self_attn.dense",
"h.{bid}.attn.c_proj", # gpt2 "h.{bid}.attn.c_proj",
"transformer.h.{bid}.mixer.out_proj", # phi2 "transformer.h.{bid}.mixer.out_proj",
"model.layers.layers.{bid}.self_attn.o_proj", # plamo "model.layers.layers.{bid}.self_attn.o_proj",
"model.layers.{bid}.attention.wo", # internlm2 "model.layers.{bid}.attention.wo",
"encoder.layers.{bid}.attn.out_proj", # nomic-bert "encoder.layers.{bid}.attn.out_proj",
"transformer.decoder_layer.{bid}.multi_head_attention.linear", # Grok "transformer.decoder_layer.{bid}.multi_head_attention.linear",
"transformer.blocks.{bid}.norm_attn_norm.attn.out_proj", # dbrx "transformer.blocks.{bid}.norm_attn_norm.attn.out_proj",
"encoder.layers.{bid}.self_attention.dense", # chatglm "encoder.layers.{bid}.self_attention.dense",
"transformer.layers.{bid}.attn.out_proj", # openelm "transformer.layers.{bid}.attn.out_proj",
"transformer.h.{bid}.attn.attention.out_proj",
), ),
# Attention output norm
MODEL_TENSOR.ATTN_OUT_NORM: ( MODEL_TENSOR.ATTN_OUT_NORM: (
"encoder.layer.{bid}.attention.output.LayerNorm", # bert "encoder.layer.{bid}.attention.output.LayerNorm",
"encoder.layers.{bid}.norm1", # nomic-bert "encoder.layers.{bid}.norm1",
"transformer.decoder_layer.{bid}.rms_norm_1", # Grok "transformer.decoder_layer.{bid}.rms_norm_1",
"transformer.blocks.{bid}.norm_attn_norm.norm_2", # dbrx "transformer.blocks.{bid}.norm_attn_norm.norm_2",
), ),
MODEL_TENSOR.ATTN_POST_NORM: ( MODEL_TENSOR.ATTN_POST_NORM: (
"model.layers.{bid}.post_attention_layernorm", # gemma2 "model.layers.{bid}.post_attention_layernorm",
), ),
# Rotary embeddings
MODEL_TENSOR.ATTN_ROT_EMBD: ( MODEL_TENSOR.ATTN_ROT_EMBD: (
"model.layers.{bid}.self_attn.rotary_emb.inv_freq", # llama-hf "model.layers.{bid}.self_attn.rotary_emb.inv_freq",
"layers.{bid}.attention.inner_attention.rope.freqs", # llama-pth "layers.{bid}.attention.inner_attention.rope.freqs",
"model.layers.layers.{bid}.self_attn.rotary_emb.inv_freq", # plamo "model.layers.layers.{bid}.self_attn.rotary_emb.inv_freq",
"transformer.h.{bid}.attn.rotary_emb.inv_freq", # codeshell "transformer.h.{bid}.attn.rotary_emb.inv_freq",
), ),
# Feed-forward norm
MODEL_TENSOR.FFN_NORM: ( MODEL_TENSOR.FFN_NORM: (
"gpt_neox.layers.{bid}.post_attention_layernorm", # gptneox "gpt_neox.layers.{bid}.post_attention_layernorm",
"transformer.h.{bid}.ln_2", # gpt2 refact qwen jais "transformer.h.{bid}.ln_2",
"h.{bid}.post_attention_layernorm", # bloom "h.{bid}.post_attention_layernorm",
"transformer.blocks.{bid}.norm_2", # mpt "transformer.blocks.{bid}.norm_2",
"model.layers.{bid}.post_attention_layernorm", # llama-hf "model.layers.{bid}.post_attention_layernorm",
"layers.{bid}.ffn_norm", # llama-pth "layers.{bid}.ffn_norm",
"language_model.encoder.layers.{bid}.post_attention_layernorm", # persimmon "language_model.encoder.layers.{bid}.post_attention_layernorm",
"model.layers.{bid}.ln2", # yi "model.layers.{bid}.ln2",
"h.{bid}.ln_2", # gpt2 "h.{bid}.ln_2",
"model.layers.{bid}.ffn_norm", # internlm2 "model.layers.{bid}.ffn_norm",
"transformer.decoder_layer.{bid}.rms_norm_2", # Grok "transformer.decoder_layer.{bid}.rms_norm_2",
"encoder.layers.{bid}.post_attention_layernorm", # chatglm "encoder.layers.{bid}.post_attention_layernorm",
"transformer.layers.{bid}.ffn_norm", # openelm "transformer.layers.{bid}.ffn_norm",
), ),
# Post feed-forward norm
MODEL_TENSOR.FFN_PRE_NORM: ( MODEL_TENSOR.FFN_PRE_NORM: (
"model.layers.{bid}.pre_feedforward_layernorm", # gemma2 "model.layers.{bid}.pre_feedforward_layernorm",
), ),
# Post feed-forward norm
MODEL_TENSOR.FFN_POST_NORM: ( MODEL_TENSOR.FFN_POST_NORM: (
"model.layers.{bid}.post_feedforward_layernorm", # gemma2 "model.layers.{bid}.post_feedforward_layernorm",
), ),
MODEL_TENSOR.FFN_GATE_INP: ( MODEL_TENSOR.FFN_GATE_INP: (
"layers.{bid}.feed_forward.gate", # mixtral "layers.{bid}.feed_forward.gate",
"model.layers.{bid}.block_sparse_moe.gate", # mixtral "model.layers.{bid}.block_sparse_moe.gate",
"model.layers.{bid}.mlp.gate", # qwen2moe "model.layers.{bid}.mlp.gate",
"transformer.decoder_layer.{bid}.router", # Grok "transformer.decoder_layer.{bid}.router",
"transformer.blocks.{bid}.ffn.router.layer", # dbrx "transformer.blocks.{bid}.ffn.router.layer",
), ),
MODEL_TENSOR.FFN_GATE_INP_SHEXP: ( MODEL_TENSOR.FFN_GATE_INP_SHEXP: (
"model.layers.{bid}.mlp.shared_expert_gate", # qwen2moe "model.layers.{bid}.mlp.shared_expert_gate",
), ),
# Feed-forward up
MODEL_TENSOR.FFN_UP: ( MODEL_TENSOR.FFN_UP: (
"gpt_neox.layers.{bid}.mlp.dense_h_to_4h", # gptneox "gpt_neox.layers.{bid}.mlp.dense_h_to_4h",
"transformer.h.{bid}.mlp.c_fc", # gpt2 jais "transformer.h.{bid}.mlp.c_fc",
"transformer.blocks.{bid}.ffn.up_proj", # mpt "transformer.blocks.{bid}.ffn.up_proj",
"transformer.h.{bid}.mlp.dense_h_to_4h", # falcon "transformer.h.{bid}.mlp.dense_h_to_4h",
"h.{bid}.mlp.dense_h_to_4h", # bloom "h.{bid}.mlp.dense_h_to_4h",
"model.layers.{bid}.mlp.up_proj", # llama-hf refact "model.layers.{bid}.mlp.up_proj",
"layers.{bid}.feed_forward.w3", # llama-pth "layers.{bid}.feed_forward.w3",
"encoder.layer.{bid}.intermediate.dense", # bert "encoder.layer.{bid}.intermediate.dense",
"transformer.h.{bid}.mlp.fc_in", # gpt-j "transformer.h.{bid}.mlp.fc_in",
"transformer.h.{bid}.mlp.linear_3", # refact "transformer.h.{bid}.mlp.linear_3",
"language_model.encoder.layers.{bid}.mlp.dense_h_to_4h", # persimmon "language_model.encoder.layers.{bid}.mlp.dense_h_to_4h",
"model.layers.{bid}.mlp.dense_h_to_4h", # persimmon "model.layers.{bid}.mlp.dense_h_to_4h",
"transformer.h.{bid}.mlp.w1", # qwen "transformer.h.{bid}.mlp.w1",
"h.{bid}.mlp.c_fc", # gpt2 "h.{bid}.mlp.c_fc",
"transformer.h.{bid}.mlp.fc1", # phi2 "transformer.h.{bid}.mlp.fc1",
"model.layers.{bid}.mlp.fc1", # phi2 "model.layers.{bid}.mlp.fc1",
"model.layers.{bid}.mlp.gate_up_proj", # phi3 "model.layers.{bid}.mlp.gate_up_proj",
"model.layers.layers.{bid}.mlp.up_proj", # plamo "model.layers.layers.{bid}.mlp.up_proj",
"model.layers.{bid}.feed_forward.w3", # internlm2 "model.layers.{bid}.feed_forward.w3",
"encoder.layers.{bid}.mlp.fc11", # nomic-bert "encoder.layers.{bid}.mlp.fc11",
"model.layers.{bid}.mlp.c_fc", # starcoder2 "model.layers.{bid}.mlp.c_fc",
"encoder.layer.{bid}.mlp.gated_layers_v", # jina-bert-v2 "encoder.layer.{bid}.mlp.gated_layers_v",
"model.layers.{bid}.residual_mlp.w3", # arctic "model.layers.{bid}.residual_mlp.w3",
"encoder.layers.{bid}.mlp.dense_h_to_4h", # chatglm "encoder.layers.{bid}.mlp.dense_h_to_4h",
"transformer.h.{bid}.mlp.c_fc_1",
), ),
MODEL_TENSOR.FFN_UP_EXP: ( MODEL_TENSOR.FFN_UP_EXP: (
"layers.{bid}.feed_forward.experts.w3", # mixtral (merged) "layers.{bid}.feed_forward.experts.w3",
"transformer.decoder_layer.{bid}.moe.linear_v", # Grok (merged) "transformer.decoder_layer.{bid}.moe.linear_v",
"transformer.blocks.{bid}.ffn.experts.mlp.v1", # dbrx "transformer.blocks.{bid}.ffn.experts.mlp.v1",
"model.layers.{bid}.mlp.experts.up_proj", # qwen2moe (merged) "model.layers.{bid}.mlp.experts.up_proj",
), ),
MODEL_TENSOR.FFN_UP_SHEXP: ( MODEL_TENSOR.FFN_UP_SHEXP: (
"model.layers.{bid}.mlp.shared_expert.up_proj", # qwen2moe "model.layers.{bid}.mlp.shared_expert.up_proj",
"model.layers.{bid}.mlp.shared_experts.up_proj", # deepseek2 "model.layers.{bid}.mlp.shared_experts.up_proj",
), ),
# AWQ-activation gate
MODEL_TENSOR.FFN_ACT: ("transformer.blocks.{bid}.ffn.act",), # mpt MODEL_TENSOR.FFN_ACT: (
# Feed-forward gate "transformer.blocks.{bid}.ffn.act",
),
MODEL_TENSOR.FFN_GATE: ( MODEL_TENSOR.FFN_GATE: (
"model.layers.{bid}.mlp.gate_proj", # llama-hf refact "model.layers.{bid}.mlp.gate_proj",
"layers.{bid}.feed_forward.w1", # llama-pth "layers.{bid}.feed_forward.w1",
"transformer.h.{bid}.mlp.w2", # qwen "transformer.h.{bid}.mlp.w2",
"transformer.h.{bid}.mlp.c_fc2", # jais "transformer.h.{bid}.mlp.c_fc2",
"model.layers.layers.{bid}.mlp.gate_proj", # plamo "model.layers.layers.{bid}.mlp.gate_proj",
"model.layers.{bid}.feed_forward.w1", # internlm2 "model.layers.{bid}.feed_forward.w1",
"encoder.layers.{bid}.mlp.fc12", # nomic-bert "encoder.layers.{bid}.mlp.fc12",
"encoder.layer.{bid}.mlp.gated_layers_w", # jina-bert-v2 "encoder.layer.{bid}.mlp.gated_layers_w",
"transformer.h.{bid}.mlp.linear_1", # refact "transformer.h.{bid}.mlp.linear_1",
"model.layers.{bid}.residual_mlp.w1", # arctic "model.layers.{bid}.residual_mlp.w1",
"transformer.h.{bid}.mlp.c_fc_0",
), ),
MODEL_TENSOR.FFN_GATE_EXP: ( MODEL_TENSOR.FFN_GATE_EXP: (
"layers.{bid}.feed_forward.experts.w1", # mixtral (merged) "layers.{bid}.feed_forward.experts.w1",
"transformer.decoder_layer.{bid}.moe.linear", # Grok (merged) "transformer.decoder_layer.{bid}.moe.linear",
"transformer.blocks.{bid}.ffn.experts.mlp.w1", # dbrx "transformer.blocks.{bid}.ffn.experts.mlp.w1",
"model.layers.{bid}.mlp.experts.gate_proj", # qwen2moe (merged) "model.layers.{bid}.mlp.experts.gate_proj",
), ),
MODEL_TENSOR.FFN_GATE_SHEXP: ( MODEL_TENSOR.FFN_GATE_SHEXP: (
"model.layers.{bid}.mlp.shared_expert.gate_proj", # qwen2moe "model.layers.{bid}.mlp.shared_expert.gate_proj",
"model.layers.{bid}.mlp.shared_experts.gate_proj", # deepseek2 "model.layers.{bid}.mlp.shared_experts.gate_proj",
), ),
# Feed-forward down
MODEL_TENSOR.FFN_DOWN: ( MODEL_TENSOR.FFN_DOWN: (
"gpt_neox.layers.{bid}.mlp.dense_4h_to_h", # gptneox "gpt_neox.layers.{bid}.mlp.dense_4h_to_h",
"transformer.h.{bid}.mlp.c_proj", # gpt2 refact qwen jais "transformer.h.{bid}.mlp.c_proj",
"transformer.blocks.{bid}.ffn.down_proj", # mpt "transformer.blocks.{bid}.ffn.down_proj",
"transformer.h.{bid}.mlp.dense_4h_to_h", # falcon "transformer.h.{bid}.mlp.dense_4h_to_h",
"h.{bid}.mlp.dense_4h_to_h", # bloom "h.{bid}.mlp.dense_4h_to_h",
"model.layers.{bid}.mlp.down_proj", # llama-hf "model.layers.{bid}.mlp.down_proj",
"layers.{bid}.feed_forward.w2", # llama-pth "layers.{bid}.feed_forward.w2",
"encoder.layer.{bid}.output.dense", # bert "encoder.layer.{bid}.output.dense",
"transformer.h.{bid}.mlp.fc_out", # gpt-j "transformer.h.{bid}.mlp.fc_out",
"language_model.encoder.layers.{bid}.mlp.dense_4h_to_h", # persimmon "language_model.encoder.layers.{bid}.mlp.dense_4h_to_h",
"model.layers.{bid}.mlp.dense_4h_to_h", # persimmon "model.layers.{bid}.mlp.dense_4h_to_h",
"h.{bid}.mlp.c_proj", # gpt2 "h.{bid}.mlp.c_proj",
"transformer.h.{bid}.mlp.fc2", # phi2 "transformer.h.{bid}.mlp.fc2",
"model.layers.{bid}.mlp.fc2", # phi2 "model.layers.{bid}.mlp.fc2",
"model.layers.layers.{bid}.mlp.down_proj", # plamo "model.layers.layers.{bid}.mlp.down_proj",
"model.layers.{bid}.feed_forward.w2", # internlm2 "model.layers.{bid}.feed_forward.w2",
"encoder.layers.{bid}.mlp.fc2", # nomic-bert "encoder.layers.{bid}.mlp.fc2",
"model.layers.{bid}.mlp.c_proj", # starcoder2 "model.layers.{bid}.mlp.c_proj",
"encoder.layer.{bid}.mlp.wo", # jina-bert-v2 "encoder.layer.{bid}.mlp.wo",
"transformer.layers.{bid}.ffn.proj_2", # openelm "transformer.layers.{bid}.ffn.proj_2",
"model.layers.{bid}.residual_mlp.w2", # arctic "model.layers.{bid}.residual_mlp.w2",
"encoder.layer.{bid}.mlp.down_layer", # jina-bert-v2 "encoder.layer.{bid}.mlp.down_layer",
"encoder.layers.{bid}.mlp.dense_4h_to_h", # chatglm "encoder.layers.{bid}.mlp.dense_4h_to_h",
"model.layers.h.{bid}.mlp.c_proj",
), ),
MODEL_TENSOR.FFN_DOWN_EXP: ( MODEL_TENSOR.FFN_DOWN_EXP: (
"layers.{bid}.feed_forward.experts.w2", # mixtral (merged) "layers.{bid}.feed_forward.experts.w2",
"transformer.decoder_layer.{bid}.moe.linear_1", # Grok (merged) "transformer.decoder_layer.{bid}.moe.linear_1",
"transformer.blocks.{bid}.ffn.experts.mlp.w2", # dbrx "transformer.blocks.{bid}.ffn.experts.mlp.w2",
"model.layers.{bid}.mlp.experts.down_proj", # qwen2moe (merged) "model.layers.{bid}.mlp.experts.down_proj",
), ),
MODEL_TENSOR.FFN_DOWN_SHEXP: ( MODEL_TENSOR.FFN_DOWN_SHEXP: (
"model.layers.{bid}.mlp.shared_expert.down_proj", # qwen2moe "model.layers.{bid}.mlp.shared_expert.down_proj",
"model.layers.{bid}.mlp.shared_experts.down_proj", # deepseek2 "model.layers.{bid}.mlp.shared_experts.down_proj",
), ),
MODEL_TENSOR.ATTN_Q_NORM: ( MODEL_TENSOR.ATTN_Q_NORM: (
"language_model.encoder.layers.{bid}.self_attention.q_layernorm", "language_model.encoder.layers.{bid}.self_attention.q_layernorm",
"model.layers.{bid}.self_attn.q_layernorm", # persimmon "model.layers.{bid}.self_attn.q_layernorm",
"model.layers.{bid}.self_attn.q_norm", # cohere "model.layers.{bid}.self_attn.q_norm",
"transformer.blocks.{bid}.attn.q_ln", # sea-lion "transformer.blocks.{bid}.attn.q_ln",
"encoder.layer.{bid}.attention.self.layer_norm_q", # jina-bert-v2 "encoder.layer.{bid}.attention.self.layer_norm_q",
"transformer.layers.{bid}.attn.q_norm", # openelm "transformer.layers.{bid}.attn.q_norm",
), ),
MODEL_TENSOR.ATTN_K_NORM: ( MODEL_TENSOR.ATTN_K_NORM: (
"language_model.encoder.layers.{bid}.self_attention.k_layernorm", "language_model.encoder.layers.{bid}.self_attention.k_layernorm",
"model.layers.{bid}.self_attn.k_layernorm", # persimmon "model.layers.{bid}.self_attn.k_layernorm",
"model.layers.{bid}.self_attn.k_norm", # cohere "model.layers.{bid}.self_attn.k_norm",
"transformer.blocks.{bid}.attn.k_ln", # sea-lion "transformer.blocks.{bid}.attn.k_ln",
"encoder.layer.{bid}.attention.self.layer_norm_k", # jina-bert-v2 "encoder.layer.{bid}.attention.self.layer_norm_k",
"transformer.layers.{bid}.attn.k_norm", # openelm "transformer.layers.{bid}.attn.k_norm",
), ),
MODEL_TENSOR.ROPE_FREQS: ( MODEL_TENSOR.ROPE_FREQS: (
"language_model.encoder.layers.{bid}.self_attention.rotary_emb.inv_freq", # persimmon "language_model.encoder.layers.{bid}.self_attention.rotary_emb.inv_freq",
), ),
MODEL_TENSOR.LAYER_OUT_NORM: ( MODEL_TENSOR.LAYER_OUT_NORM: (
"encoder.layer.{bid}.output.LayerNorm", # bert "encoder.layer.{bid}.output.LayerNorm",
"encoder.layers.{bid}.norm2", # nomic-bert "encoder.layers.{bid}.norm2",
"transformer.decoder_layer.{bid}.rms_norm_3", # Grok "transformer.decoder_layer.{bid}.rms_norm_3",
"encoder.layer.{bid}.mlp.layernorm", # jina-bert-v2 "encoder.layer.{bid}.mlp.layernorm",
"encoder.layer.{bid}.layer_norm_2", # jina-v2-code "encoder.layer.{bid}.layer_norm_2"
), ),
MODEL_TENSOR.SSM_IN: ( MODEL_TENSOR.SSM_IN: (
"model.layers.{bid}.in_proj", "model.layers.{bid}.in_proj",
"backbone.layers.{bid}.mixer.in_proj", "backbone.layers.{bid}.mixer.in_proj",
), ),
MODEL_TENSOR.SSM_CONV1D: ( MODEL_TENSOR.SSM_CONV1D: (
"model.layers.{bid}.conv1d", "model.layers.{bid}.conv1d",
"backbone.layers.{bid}.mixer.conv1d", "backbone.layers.{bid}.mixer.conv1d",
), ),
MODEL_TENSOR.SSM_X: ( MODEL_TENSOR.SSM_X: (
"model.layers.{bid}.x_proj", "model.layers.{bid}.x_proj",
"backbone.layers.{bid}.mixer.x_proj", "backbone.layers.{bid}.mixer.x_proj",
), ),
MODEL_TENSOR.SSM_DT: ( MODEL_TENSOR.SSM_DT: (
"model.layers.{bid}.dt_proj", "model.layers.{bid}.dt_proj",
"backbone.layers.{bid}.mixer.dt_proj", "backbone.layers.{bid}.mixer.dt_proj",
), ),
MODEL_TENSOR.SSM_A: ( MODEL_TENSOR.SSM_A: (
"model.layers.{bid}.A_log", "model.layers.{bid}.A_log",
"backbone.layers.{bid}.mixer.A_log", "backbone.layers.{bid}.mixer.A_log",
), ),
MODEL_TENSOR.SSM_D: ( MODEL_TENSOR.SSM_D: (
"model.layers.{bid}.D", "model.layers.{bid}.D",
"backbone.layers.{bid}.mixer.D", "backbone.layers.{bid}.mixer.D",
), ),
MODEL_TENSOR.SSM_OUT: ( MODEL_TENSOR.SSM_OUT: (
"model.layers.{bid}.out_proj", "model.layers.{bid}.out_proj",
"backbone.layers.{bid}.mixer.out_proj", "backbone.layers.{bid}.mixer.out_proj",
), ),
MODEL_TENSOR.ATTN_Q_A: ("model.layers.{bid}.self_attn.q_a_proj",), # deepseek2
MODEL_TENSOR.ATTN_Q_B: ("model.layers.{bid}.self_attn.q_b_proj",), # deepseek2 MODEL_TENSOR.ATTN_Q_A: (
"model.layers.{bid}.self_attn.q_a_proj",
),
MODEL_TENSOR.ATTN_Q_B: (
"model.layers.{bid}.self_attn.q_b_proj",
),
MODEL_TENSOR.ATTN_KV_A_MQA: ( MODEL_TENSOR.ATTN_KV_A_MQA: (
"model.layers.{bid}.self_attn.kv_a_proj_with_mqa", # deepseek2 "model.layers.{bid}.self_attn.kv_a_proj_with_mqa",
), ),
MODEL_TENSOR.ATTN_KV_B: ( MODEL_TENSOR.ATTN_KV_B: (
"model.layers.{bid}.self_attn.kv_b_proj", # deepseek2 "model.layers.{bid}.self_attn.kv_b_proj",
), ),
MODEL_TENSOR.ATTN_Q_A_NORM: ( MODEL_TENSOR.ATTN_Q_A_NORM: (
"model.layers.{bid}.self_attn.q_a_layernorm", # deepseek2 "model.layers.{bid}.self_attn.q_a_layernorm",
), ),
MODEL_TENSOR.ATTN_KV_A_NORM: ( MODEL_TENSOR.ATTN_KV_A_NORM: (
"model.layers.{bid}.self_attn.kv_a_layernorm", # deepseek2 "model.layers.{bid}.self_attn.kv_a_layernorm",
), ),
MODEL_TENSOR.ATTN_SUB_NORM: ( MODEL_TENSOR.ATTN_SUB_NORM: (
"model.layers.{bid}.self_attn.inner_attn_ln", # bitnet "model.layers.{bid}.self_attn.inner_attn_ln",
), ),
MODEL_TENSOR.FFN_SUB_NORM: ("model.layers.{bid}.mlp.ffn_layernorm",), # bitnet
MODEL_TENSOR.DEC_ATTN_NORM: ("decoder.block.{bid}.layer.0.layer_norm",), # t5 MODEL_TENSOR.FFN_SUB_NORM: (
MODEL_TENSOR.DEC_ATTN_Q: ("decoder.block.{bid}.layer.0.SelfAttention.q",), # t5 "model.layers.{bid}.mlp.ffn_layernorm",
MODEL_TENSOR.DEC_ATTN_K: ("decoder.block.{bid}.layer.0.SelfAttention.k",), # t5 ),
MODEL_TENSOR.DEC_ATTN_V: ("decoder.block.{bid}.layer.0.SelfAttention.v",), # t5
MODEL_TENSOR.DEC_ATTN_NORM: (
"decoder.block.{bid}.layer.0.layer_norm",
),
MODEL_TENSOR.DEC_ATTN_Q: (
"decoder.block.{bid}.layer.0.SelfAttention.q",
),
MODEL_TENSOR.DEC_ATTN_K: (
"decoder.block.{bid}.layer.0.SelfAttention.k",
),
MODEL_TENSOR.DEC_ATTN_V: (
"decoder.block.{bid}.layer.0.SelfAttention.v",
),
MODEL_TENSOR.DEC_ATTN_OUT: ( MODEL_TENSOR.DEC_ATTN_OUT: (
"decoder.block.{bid}.layer.0.SelfAttention.o", # t5 "decoder.block.{bid}.layer.0.SelfAttention.o",
), ),
MODEL_TENSOR.DEC_ATTN_REL_B: ( MODEL_TENSOR.DEC_ATTN_REL_B: (
"decoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias", # t5 "decoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias",
), ),
MODEL_TENSOR.DEC_CROSS_ATTN_NORM: ( MODEL_TENSOR.DEC_CROSS_ATTN_NORM: (
"decoder.block.{bid}.layer.1.layer_norm", # t5 "decoder.block.{bid}.layer.1.layer_norm",
), ),
MODEL_TENSOR.DEC_CROSS_ATTN_Q: ( MODEL_TENSOR.DEC_CROSS_ATTN_Q: (
"decoder.block.{bid}.layer.1.EncDecAttention.q", # t5 "decoder.block.{bid}.layer.1.EncDecAttention.q",
), ),
MODEL_TENSOR.DEC_CROSS_ATTN_K: ( MODEL_TENSOR.DEC_CROSS_ATTN_K: (
"decoder.block.{bid}.layer.1.EncDecAttention.k", # t5 "decoder.block.{bid}.layer.1.EncDecAttention.k",
), ),
MODEL_TENSOR.DEC_CROSS_ATTN_V: ( MODEL_TENSOR.DEC_CROSS_ATTN_V: (
"decoder.block.{bid}.layer.1.EncDecAttention.v", # t5 "decoder.block.{bid}.layer.1.EncDecAttention.v",
), ),
MODEL_TENSOR.DEC_CROSS_ATTN_OUT: ( MODEL_TENSOR.DEC_CROSS_ATTN_OUT: (
"decoder.block.{bid}.layer.1.EncDecAttention.o", # t5 "decoder.block.{bid}.layer.1.EncDecAttention.o",
), ),
MODEL_TENSOR.DEC_CROSS_ATTN_REL_B: ( MODEL_TENSOR.DEC_CROSS_ATTN_REL_B: (
"decoder.block.{bid}.layer.1.EncDecAttention.relative_attention_bias", # t5 "decoder.block.{bid}.layer.1.EncDecAttention.relative_attention_bias",
), ),
MODEL_TENSOR.DEC_FFN_NORM: ("decoder.block.{bid}.layer.2.layer_norm",), # t5
MODEL_TENSOR.DEC_FFN_NORM: (
"decoder.block.{bid}.layer.2.layer_norm",
),
MODEL_TENSOR.DEC_FFN_GATE: ( MODEL_TENSOR.DEC_FFN_GATE: (
"decoder.block.{bid}.layer.2.DenseReluDense.wi_0", # flan-t5 "decoder.block.{bid}.layer.2.DenseReluDense.wi_0",
), ),
MODEL_TENSOR.DEC_FFN_UP: ( MODEL_TENSOR.DEC_FFN_UP: (
"decoder.block.{bid}.layer.2.DenseReluDense.wi", # t5 "decoder.block.{bid}.layer.2.DenseReluDense.wi",
"decoder.block.{bid}.layer.2.DenseReluDense.wi_1", # flan-t5 "decoder.block.{bid}.layer.2.DenseReluDense.wi_1",
), ),
MODEL_TENSOR.DEC_FFN_DOWN: ( MODEL_TENSOR.DEC_FFN_DOWN: (
"decoder.block.{bid}.layer.2.DenseReluDense.wo", # t5 "decoder.block.{bid}.layer.2.DenseReluDense.wo",
), ),
MODEL_TENSOR.DEC_OUTPUT_NORM: ("decoder.final_layer_norm",), # t5
MODEL_TENSOR.ENC_ATTN_NORM: ("encoder.block.{bid}.layer.0.layer_norm",), # t5 MODEL_TENSOR.DEC_OUTPUT_NORM: (
MODEL_TENSOR.ENC_ATTN_Q: ("encoder.block.{bid}.layer.0.SelfAttention.q",), # t5 "decoder.final_layer_norm",
MODEL_TENSOR.ENC_ATTN_K: ("encoder.block.{bid}.layer.0.SelfAttention.k",), # t5 ),
MODEL_TENSOR.ENC_ATTN_V: ("encoder.block.{bid}.layer.0.SelfAttention.v",), # t5
MODEL_TENSOR.ENC_ATTN_NORM: (
"encoder.block.{bid}.layer.0.layer_norm",
),
MODEL_TENSOR.ENC_ATTN_Q: (
"encoder.block.{bid}.layer.0.SelfAttention.q",
),
MODEL_TENSOR.ENC_ATTN_K: (
"encoder.block.{bid}.layer.0.SelfAttention.k",
),
MODEL_TENSOR.ENC_ATTN_V: (
"encoder.block.{bid}.layer.0.SelfAttention.v",
),
MODEL_TENSOR.ENC_ATTN_OUT: ( MODEL_TENSOR.ENC_ATTN_OUT: (
"encoder.block.{bid}.layer.0.SelfAttention.o", # t5 "encoder.block.{bid}.layer.0.SelfAttention.o",
), ),
MODEL_TENSOR.ENC_ATTN_REL_B: ( MODEL_TENSOR.ENC_ATTN_REL_B: (
"encoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias", # t5 "encoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias",
), ),
MODEL_TENSOR.ENC_FFN_NORM: ("encoder.block.{bid}.layer.1.layer_norm",), # t5
MODEL_TENSOR.ENC_FFN_NORM: (
"encoder.block.{bid}.layer.1.layer_norm",
),
MODEL_TENSOR.ENC_FFN_GATE: ( MODEL_TENSOR.ENC_FFN_GATE: (
"encoder.block.{bid}.layer.1.DenseReluDense.wi_0", # flan-t5 "encoder.block.{bid}.layer.1.DenseReluDense.wi_0",
), ),
MODEL_TENSOR.ENC_FFN_UP: ( MODEL_TENSOR.ENC_FFN_UP: (
"encoder.block.{bid}.layer.1.DenseReluDense.wi", # t5 "encoder.block.{bid}.layer.1.DenseReluDense.wi",
"encoder.block.{bid}.layer.1.DenseReluDense.wi_1", # flan-t5 "encoder.block.{bid}.layer.1.DenseReluDense.wi_1",
), ),
MODEL_TENSOR.ENC_FFN_DOWN: ( MODEL_TENSOR.ENC_FFN_DOWN: (
"encoder.block.{bid}.layer.1.DenseReluDense.wo", # t5 "encoder.block.{bid}.layer.1.DenseReluDense.wo",
),
MODEL_TENSOR.ENC_OUTPUT_NORM: (
"encoder.final_layer_norm",
), ),
MODEL_TENSOR.ENC_OUTPUT_NORM: ("encoder.final_layer_norm",), # t5
} }
# architecture-specific block mappings
arch_block_mappings_cfg: dict[MODEL_ARCH, dict[MODEL_TENSOR, tuple[str, ...]]] = { arch_block_mappings_cfg: dict[MODEL_ARCH, dict[MODEL_TENSOR, tuple[str, ...]]] = {
MODEL_ARCH.ARCTIC: { MODEL_ARCH.ARCTIC: {
MODEL_TENSOR.FFN_NORM: ("model.layers.{bid}.residual_layernorm",), MODEL_TENSOR.FFN_NORM: (
MODEL_TENSOR.FFN_NORM_EXP: ("model.layers.{bid}.post_attention_layernorm",), "model.layers.{bid}.residual_layernorm",
),
MODEL_TENSOR.FFN_NORM_EXP: (
"model.layers.{bid}.post_attention_layernorm",
),
}, },
} }
@ -490,35 +588,31 @@ def __init__(self, arch: MODEL_ARCH, n_blocks: int):
if tensor not in MODEL_TENSORS[arch]: if tensor not in MODEL_TENSORS[arch]:
continue continue
tensor_name = TENSOR_NAMES[tensor].format(bid=bid) tensor_name = TENSOR_NAMES[tensor].format(bid = bid)
self.mapping[tensor_name] = (tensor, tensor_name) self.mapping[tensor_name] = (tensor, tensor_name)
for key in keys: for key in keys:
key = key.format(bid=bid) key = key.format(bid = bid)
self.mapping[key] = (tensor, tensor_name) self.mapping[key] = (tensor, tensor_name)
def get_type_and_name( def get_type_and_name(self, key: str, try_suffixes: Sequence[str] = ()) -> tuple[MODEL_TENSOR, str] | None:
self, key: str, try_suffixes: Sequence[str] = ()
) -> tuple[MODEL_TENSOR, str] | None:
result = self.mapping.get(key) result = self.mapping.get(key)
if result is not None: if result is not None:
return result return result
for suffix in try_suffixes: for suffix in try_suffixes:
if key.endswith(suffix): if key.endswith(suffix):
result = self.mapping.get(key[: -len(suffix)]) result = self.mapping.get(key[:-len(suffix)])
if result is not None: if result is not None:
return result[0], result[1] + suffix return result[0], result[1] + suffix
return None return None
def get_name(self, key: str, try_suffixes: Sequence[str] = ()) -> str | None: def get_name(self, key: str, try_suffixes: Sequence[str] = ()) -> str | None:
result = self.get_type_and_name(key, try_suffixes=try_suffixes) result = self.get_type_and_name(key, try_suffixes = try_suffixes)
if result is None: if result is None:
return None return None
return result[1] return result[1]
def get_type( def get_type(self, key: str, try_suffixes: Sequence[str] = ()) -> MODEL_TENSOR | None:
self, key: str, try_suffixes: Sequence[str] = () result = self.get_type_and_name(key, try_suffixes = try_suffixes)
) -> MODEL_TENSOR | None:
result = self.get_type_and_name(key, try_suffixes=try_suffixes)
if result is None: if result is None:
return None return None
return result[0] return result[0]
@ -535,6 +629,5 @@ def __contains__(self, key: str) -> bool:
def __repr__(self) -> str: def __repr__(self) -> str:
return repr(self.mapping) return repr(self.mapping)
def get_tensor_name_map(arch: MODEL_ARCH, n_blocks: int) -> TensorNameMap: def get_tensor_name_map(arch: MODEL_ARCH, n_blocks: int) -> TensorNameMap:
return TensorNameMap(arch, n_blocks) return TensorNameMap(arch, n_blocks)