| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523 |
- from __future__ import annotations
- from typing import Sequence
- from .constants import MODEL_ARCH, MODEL_TENSOR, MODEL_TENSORS, TENSOR_NAMES
- class TensorNameMap:
- mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = {
- # Token embeddings
- MODEL_TENSOR.TOKEN_EMBD: (
- "gpt_neox.embed_in", # gptneox
- "transformer.wte", # gpt2 gpt-j mpt refact qwen dbrx jais exaone
- "transformer.word_embeddings", # falcon
- "word_embeddings", # bloom
- "model.embed_tokens", # llama-hf nemotron olmoe olmo2 rwkv6qwen2 glm4-0414 plamo2 granite-hybrid
- "tok_embeddings", # llama-pth
- "embeddings.word_embeddings", # bert nomic-bert
- "language_model.embedding.word_embeddings", # persimmon
- "wte", # gpt2
- "transformer.embd.wte", # phi2
- "model.tok_embeddings", # internlm2
- "model.embedding", # mamba-qbert
- "backbone.embedding", # mamba
- "backbone.embeddings", # mamba-hf
- "transformer.in_out_embed", # Grok
- "embedding.word_embeddings", # chatglm
- "transformer.token_embeddings", # openelm
- "shared", # t5
- "rwkv.embeddings", # rwkv6
- "model.embeddings", # rwkv7
- "model.word_embeddings", # bailingmoe
- "language_model.model.embed_tokens", # llama4
- "encoder", # neobert
- "model.transformer.wte", # llada
- "embed_tokens", # qwen3-embedding
- ),
- # Token type embeddings
- MODEL_TENSOR.TOKEN_TYPES: (
- "embeddings.token_type_embeddings", # bert nomic-bert
- ),
- # Normalization of token embeddings
- MODEL_TENSOR.TOKEN_EMBD_NORM: (
- "word_embeddings_layernorm", # bloom
- "embeddings.LayerNorm", # bert
- "emb_ln", # nomic-bert
- "transformer.norm", # openelm
- "rwkv.blocks.0.pre_ln", # rwkv
- "rwkv.blocks.0.pre_ln", # rwkv6
- "model.pre_ln", # rwkv7
- "model.layers.0.pre_norm", # rwkv7
- "backbone.norm", # wavtokenizer
- "model.embedding_norm", # lfm2
- ),
- # Position embeddings
- MODEL_TENSOR.POS_EMBD: (
- "transformer.wpe", # gpt2
- "embeddings.position_embeddings", # bert
- "wpe", # gpt2
- ),
- # Output
- MODEL_TENSOR.OUTPUT: (
- "embed_out", # gptneox
- "lm_head", # gpt2 mpt falcon llama-hf baichuan qwen mamba dbrx jais nemotron exaone olmoe olmo2 phimoe plamo2
- "output", # llama-pth bloom internlm2
- "word_embeddings_for_head", # persimmon
- "lm_head.linear", # phi2
- "output_layer", # chatglm
- "head", # rwkv
- "head.out", # wavtokenizer
- "lm_head", # llama4
- "model.transformer.ff_out", # llada
- ),
- # Output norm
- MODEL_TENSOR.OUTPUT_NORM: (
- "gpt_neox.final_layer_norm", # gptneox
- "transformer.ln_f", # gpt2 gpt-j falcon jais exaone
- "model.norm", # llama-hf baichuan internlm2 olmoe olmo2 phimoe plamo2
- "norm", # llama-pth
- "transformer.norm_f", # mpt dbrx
- "ln_f", # refact bloom qwen gpt2
- "language_model.encoder.final_layernorm", # persimmon
- "model.final_layernorm", # persimmon
- "lm_head.ln", # phi2
- "model.norm_f", # mamba-qbert
- "backbone.norm_f", # mamba
- "transformer.rms_norm", # Grok
- "encoder.final_layernorm", # chatglm
- "transformer.norm", # openelm
- "model.norm", # nemotron
- "rwkv.ln_out", # rwkv6
- "model.ln_out", # rwkv7
- "backbone.final_layer_norm", # wavtokenizer
- "model.norm", # llama4
- "model.transformer.ln_f", # llada
- ),
- # Rope frequencies
- MODEL_TENSOR.ROPE_FREQS: (
- "rope.freqs", # llama-pth
- "rotary_pos_emb.inv_freq", # chatglm
- ),
- MODEL_TENSOR.ROPE_FACTORS_LONG: (),
- MODEL_TENSOR.ROPE_FACTORS_SHORT: (),
- MODEL_TENSOR.CONV1D: (
- "backbone.embed", # roberta
- ),
- }
- block_mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = {
- # Attention norm
- MODEL_TENSOR.ATTN_NORM: (
- "gpt_neox.layers.{bid}.input_layernorm", # gptneox
- "transformer.h.{bid}.ln_1", # gpt2 gpt-j refact qwen jais exaone
- "transformer.blocks.{bid}.norm_1", # mpt
- "transformer.h.{bid}.input_layernorm", # falcon7b
- "h.{bid}.input_layernorm", # bloom
- "transformer.h.{bid}.ln_mlp", # falcon40b
- "model.layers.{bid}.input_layernorm", # llama-hf nemotron olmoe phimoe granite-hybrid
- "layers.{bid}.attention_norm", # llama-pth
- "language_model.encoder.layers.{bid}.input_layernorm", # persimmon
- "model.layers.{bid}.ln1", # yi
- "h.{bid}.ln_1", # gpt2
- "transformer.h.{bid}.ln", # phi2
- "model.layers.layers.{bid}.norm", # plamo
- "model.layers.layers.{bid}.pre_mixer_norm", # plamo2
- "model.layers.{bid}.attention_norm", # internlm2
- "model.layers.{bid}.norm", # mamba-qbert
- "backbone.layers.{bid}.norm", # mamba
- "transformer.decoder_layer.{bid}.rms_norm", # Grok
- "transformer.blocks.{bid}.norm_attn_norm.norm_1", # dbrx
- "encoder.layers.{bid}.input_layernorm", # chatglm
- "transformer.layers.{bid}.attn_norm", # openelm
- "rwkv.blocks.{bid}.ln1", # rwkv6
- "model.layers.{bid}.ln1", # rwkv7
- "model.layers.{bid}.input_layernorm", # llama4
- "transformer_encoder.{bid}.attention_norm", # neobert
- "model.layers.{bid}.operator_norm", # lfm2
- "model.transformer.blocks.{bid}.attn_norm", # llada
- "layers.{bid}.input_layernorm", # qwen3-embedding
- ),
- # Attention norm 2
- MODEL_TENSOR.ATTN_NORM_2: (
- "transformer.h.{bid}.ln_attn", # falcon40b
- "encoder.layer.{bid}.layer_norm_1", # jina-v2-code
- "rwkv.blocks.{bid}.ln2", # rwkv6
- "model.layers.{bid}.ln2", # rwkv7
- ),
- # Attention query-key-value
- MODEL_TENSOR.ATTN_QKV: (
- "gpt_neox.layers.{bid}.attention.query_key_value", # gptneox
- "transformer.h.{bid}.attn.c_attn", # gpt2 qwen jais
- "transformer.blocks.{bid}.attn.Wqkv", # mpt
- "transformer.blocks.{bid}.norm_attn_norm.attn.Wqkv", # dbrx
- "transformer.h.{bid}.self_attention.query_key_value", # falcon
- "h.{bid}.self_attention.query_key_value", # bloom
- "language_model.encoder.layers.{bid}.self_attention.query_key_value", # persimmon
- "model.layers.{bid}.self_attn.query_key_value", # persimmon
- "h.{bid}.attn.c_attn", # gpt2
- "transformer.h.{bid}.mixer.Wqkv", # phi2
- "encoder.layers.{bid}.attn.Wqkv", # nomic-bert
- "encoder.layers.{bid}.mixer.Wqkv", # jina
- "model.layers.{bid}.self_attn.qkv_proj", # phi3
- "model.layers.layers.{bid}.mixer.qkv_proj", # plamo2
- "encoder.layers.{bid}.self_attention.query_key_value", # chatglm
- "transformer.layers.{bid}.attn.qkv_proj", # openelm
- "transformer_encoder.{bid}.qkv", # neobert
- ),
- # Attention query
- MODEL_TENSOR.ATTN_Q: (
- "model.layers.{bid}.self_attn.q_proj", # llama-hf nemotron olmoe olmo2 phimoe
- "model.layers.{bid}.self_attn.q_proj_no_perm", # llama-custom
- "layers.{bid}.attention.wq", # llama-pth
- "encoder.layer.{bid}.attention.self.query", # bert
- "transformer.layer.{bid}.attention.q_lin", # distillbert
- "transformer.h.{bid}.attn.q_proj", # gpt-j
- "model.layers.layers.{bid}.self_attn.q_proj", # plamo
- "model.layers.{bid}.attention.wq", # internlm2
- "transformer.decoder_layer.{bid}.multi_head_attention.query",# Grok
- "transformer.h.{bid}.attn.attention.q_proj", # exaone
- "model.layers.{bid}.self_attn.q_proj", # llama4
- "model.transformer.blocks.{bid}.q_proj", # llada
- "layers.{bid}.self_attn.q_proj", # qwen3-embedding
- "backbone.layers.{bid}.mixer.q_proj", # nemotron-h
- ),
- # Attention key
- MODEL_TENSOR.ATTN_K: (
- "model.layers.{bid}.self_attn.k_proj", # llama-hf nemotron olmoe olmo2 phimoe
- "model.layers.{bid}.self_attn.k_proj_no_perm", # llama-custom
- "layers.{bid}.attention.wk", # llama-pth
- "encoder.layer.{bid}.attention.self.key", # bert
- "transformer.layer.{bid}.attention.k_lin", # distillbert
- "transformer.h.{bid}.attn.k_proj", # gpt-j
- "transformer.h.{bid}.attn.k", # refact
- "model.layers.layers.{bid}.self_attn.k_proj", # plamo
- "model.layers.{bid}.attention.wk", # internlm2
- "transformer.decoder_layer.{bid}.multi_head_attention.key",# Grok
- "transformer.h.{bid}.attn.attention.k_proj", # exaone
- "model.layers.{bid}.self_attn.k_proj", # llama4
- "model.transformer.blocks.{bid}.k_proj", # llada
- "layers.{bid}.self_attn.k_proj", # qwen3-embedding
- "backbone.layers.{bid}.mixer.k_proj", # nemotron-h
- ),
- # Attention value
- MODEL_TENSOR.ATTN_V: (
- "model.layers.{bid}.self_attn.v_proj", # llama-hf nemotron olmoe olmo2 phimoe
- "layers.{bid}.attention.wv", # llama-pth
- "encoder.layer.{bid}.attention.self.value", # bert
- "transformer.layer.{bid}.attention.v_lin", # distillbert
- "transformer.h.{bid}.attn.v_proj", # gpt-j
- "transformer.h.{bid}.attn.v", # refact
- "model.layers.layers.{bid}.self_attn.v_proj", # plamo
- "model.layers.{bid}.attention.wv", # internlm2
- "transformer.decoder_layer.{bid}.multi_head_attention.value",# Grok
- "transformer.h.{bid}.attn.attention.v_proj", # exaone
- "model.layers.{bid}.self_attn.v_proj", # llama4
- "model.transformer.blocks.{bid}.v_proj", # llada
- "layers.{bid}.self_attn.v_proj", # qwen3-embedding
- "backbone.layers.{bid}.mixer.v_proj", # nemotron-h
- ),
- # Attention output
- MODEL_TENSOR.ATTN_OUT: (
- "gpt_neox.layers.{bid}.attention.dense", # gptneox
- "transformer.h.{bid}.attn.c_proj", # gpt2 refact qwen jais
- "transformer.blocks.{bid}.attn.out_proj", # mpt
- "transformer.h.{bid}.self_attention.dense", # falcon
- "h.{bid}.self_attention.dense", # bloom
- "model.layers.{bid}.self_attn.o_proj", # llama-hf nemotron olmoe olmo2 phimoe
- "model.layers.{bid}.self_attn.out_proj", # lfm2
- "model.layers.{bid}.self_attn.linear_attn", # deci
- "layers.{bid}.attention.wo", # llama-pth
- "encoder.layer.{bid}.attention.output.dense", # bert
- "transformer.layer.{bid}.attention.out_lin", # distillbert
- "transformer.h.{bid}.attn.out_proj", # gpt-j
- "language_model.encoder.layers.{bid}.self_attention.dense", # persimmon
- "model.layers.{bid}.self_attn.dense", # persimmon
- "h.{bid}.attn.c_proj", # gpt2
- "transformer.h.{bid}.mixer.out_proj", # phi2
- "model.layers.layers.{bid}.self_attn.o_proj", # plamo
- "model.layers.layers.{bid}.mixer.o_proj", # plamo2
- "model.layers.{bid}.attention.wo", # internlm2
- "encoder.layers.{bid}.attn.out_proj", # nomic-bert
- "encoder.layers.{bid}.mixer.out_proj", # jina
- "transformer.decoder_layer.{bid}.multi_head_attention.linear", # Grok
- "transformer.blocks.{bid}.norm_attn_norm.attn.out_proj", # dbrx
- "encoder.layers.{bid}.self_attention.dense", # chatglm
- "transformer.layers.{bid}.attn.out_proj", # openelm
- "transformer.h.{bid}.attn.attention.out_proj", # exaone
- "model.layers.{bid}.self_attn.o_proj", # llama4
- "transformer_encoder.{bid}.wo", # neobert
- "model.transformer.blocks.{bid}.attn_out", # llada
- "layers.{bid}.self_attn.o_proj", # qwen3-embedding
- "backbone.layers.{bid}.mixer.o_proj", # nemotron-h
- ),
- # Attention output norm
- MODEL_TENSOR.ATTN_OUT_NORM: (
- "encoder.layer.{bid}.attention.output.LayerNorm", # bert
- "transformer.layer.{bid}.sa_layer_norm", # distillbert
- "encoder.layers.{bid}.norm1", # nomic-bert
- "transformer.decoder_layer.{bid}.rms_norm_1", # Grok
- "transformer.blocks.{bid}.norm_attn_norm.norm_2", # dbrx
- ),
- MODEL_TENSOR.ATTN_POST_NORM: (
- "model.layers.{bid}.post_attention_layernorm", # gemma2 olmo2 # ge
- "model.layers.{bid}.post_self_attn_layernorm", # glm-4-0414
- "model.layers.layers.{bid}.post_mixer_norm.weight", # plamo2
- ),
- # Rotary embeddings
- MODEL_TENSOR.ATTN_ROT_EMBD: (
- "model.layers.{bid}.self_attn.rotary_emb.inv_freq", # llama-hf
- "layers.{bid}.attention.inner_attention.rope.freqs", # llama-pth
- "model.layers.layers.{bid}.self_attn.rotary_emb.inv_freq", # plamo
- "transformer.h.{bid}.attn.rotary_emb.inv_freq", # codeshell
- ),
- MODEL_TENSOR.ATTN_SINKS: (
- "model.layers.{bid}.self_attn.sinks", # openai-moe
- ),
- # Feed-forward norm
- MODEL_TENSOR.FFN_NORM: (
- "gpt_neox.layers.{bid}.post_attention_layernorm", # gptneox
- "transformer.h.{bid}.ln_2", # gpt2 refact qwen jais exaone
- "h.{bid}.post_attention_layernorm", # bloom
- "transformer.blocks.{bid}.norm_2", # mpt
- "model.layers.{bid}.post_attention_layernorm", # llama-hf nemotron olmoe phimoe
- "layers.{bid}.ffn_norm", # llama-pth
- "language_model.encoder.layers.{bid}.post_attention_layernorm", # persimmon
- "model.layers.{bid}.ln2", # yi
- "h.{bid}.ln_2", # gpt2
- "model.layers.{bid}.ffn_norm", # internlm2
- "transformer.decoder_layer.{bid}.rms_norm_2", # Grok
- "encoder.layers.{bid}.post_attention_layernorm", # chatglm
- "transformer.layers.{bid}.ffn_norm", # openelm
- "model.layers.{bid}.pre_ff_layernorm", # jamba granite-hybrid
- "model.layers.{bid}.pre_moe_layernorm", # mini-jamba
- "model.layers.{bid}.post_attention_layernorm", # llama4
- "transformer_encoder.{bid}.ffn_norm", # neobert
- "model.layers.layers.{bid}.pre_mlp_norm", # plamo2
- "model.transformer.blocks.{bid}.ff_norm", # llada
- "layers.{bid}.post_attention_layernorm", # qwen3-embedding
- ),
- # Post feed-forward norm
- MODEL_TENSOR.FFN_PRE_NORM: (
- "model.layers.{bid}.pre_feedforward_layernorm", # gemma2
- "model.layers.{bid}.pre_ff_layernorm.weight",
- ),
- # Post feed-forward norm
- MODEL_TENSOR.FFN_POST_NORM: (
- "model.layers.{bid}.post_feedforward_layernorm", # gemma2 olmo2
- "model.layers.{bid}.post_mlp_layernorm", # glm-4-0414
- "model.layers.layers.{bid}.post_mlp_norm.weight", # plamo2
- "model.layers.{bid}.feed_forward.up_proj",
- ),
- MODEL_TENSOR.FFN_GATE_INP: (
- "layers.{bid}.feed_forward.gate", # mixtral
- "model.layers.{bid}.block_sparse_moe.gate", # mixtral phimoe
- "model.layers.{bid}.mlp.gate", # qwen2moe olmoe
- "transformer.decoder_layer.{bid}.router", # Grok
- "transformer.blocks.{bid}.ffn.router.layer", # dbrx
- "model.layers.{bid}.block_sparse_moe.router.layer", # granitemoe
- "model.layers.{bid}.feed_forward.router", # llama4 jamba
- "encoder.layers.{bid}.mlp.router.layer", # nomic-bert-moe
- "model.layers.{bid}.mlp.router", # openai-moe
- "model.layers.{bid}.mlp.gate.wg", # hunyuan
- "model.layers.{bid}.block_sparse_moe.primary_router", # smallthinker
- ),
- MODEL_TENSOR.FFN_GATE_INP_SHEXP: (
- "model.layers.{bid}.mlp.shared_expert_gate", # qwen2moe
- ),
- MODEL_TENSOR.FFN_EXP_PROBS_B: (
- "model.layers.{bid}.mlp.gate.e_score_correction", # deepseek-v3 dots1
- "model.layers.{bid}.mlp.moe_statics.e_score_correction", # ernie4.5-moe
- ),
- # Feed-forward up
- MODEL_TENSOR.FFN_UP: (
- "gpt_neox.layers.{bid}.mlp.dense_h_to_4h", # gptneox
- "transformer.h.{bid}.mlp.c_fc", # gpt2 jais
- "transformer.blocks.{bid}.ffn.up_proj", # mpt
- "transformer.h.{bid}.mlp.dense_h_to_4h", # falcon
- "h.{bid}.mlp.dense_h_to_4h", # bloom
- "model.layers.{bid}.mlp.up_proj", # llama-hf refact nemotron olmo2
- "layers.{bid}.feed_forward.w3", # llama-pth
- "encoder.layer.{bid}.intermediate.dense", # bert
- "transformer.layer.{bid}.ffn.lin1", # distillbert
- "transformer.h.{bid}.mlp.fc_in", # gpt-j
- "transformer.h.{bid}.mlp.linear_3", # refact
- "language_model.encoder.layers.{bid}.mlp.dense_h_to_4h", # persimmon
- "model.layers.{bid}.mlp.dense_h_to_4h", # persimmon
- "transformer.h.{bid}.mlp.w1", # qwen
- "h.{bid}.mlp.c_fc", # gpt2
- "transformer.h.{bid}.mlp.fc1", # phi2
- "model.layers.{bid}.mlp.fc1", # phi2
- "model.layers.{bid}.mlp.gate_up_proj", # phi3 glm-4-0414
- "model.layers.layers.{bid}.mlp.up_proj", # plamo
- "model.layers.layers.{bid}.mlp.gate_up_proj", # plamo2
- "model.layers.{bid}.feed_forward.w3", # internlm2
- "encoder.layers.{bid}.mlp.fc11", # nomic-bert
- "encoder.layers.{bid}.mlp.fc1", # nomic-bert-moe
- "model.layers.{bid}.mlp.c_fc", # starcoder2
- "encoder.layer.{bid}.mlp.gated_layers_v", # jina-bert-v2 (split up/gate, no longer used)
- "encoder.layer.{bid}.mlp.gated_layers", # jina-bert-v2 (GEGLU)
- "encoder.layer.{bid}.mlp.up_gated_layer", # jina-v2-code (GEGLU)
- "model.layers.{bid}.residual_mlp.w3", # arctic
- "encoder.layers.{bid}.mlp.dense_h_to_4h", # chatglm
- "transformer.h.{bid}.mlp.c_fc_1", # exaone
- "model.layers.{bid}.feed_forward.up_proj", # llama4 jamba granite-hybrid
- "transformer_encoder.{bid}.ffn.w12", # neobert
- "model.layers.{bid}.block_sparse_moe.up", # smallthinker
- "model.transformer.blocks.{bid}.up_proj", # llada
- "layers.{bid}.mlp.up_proj", # qwen3-embedding
- "backbone.layers.{bid}.mixer.up_proj", # nemotron-h
- ),
- MODEL_TENSOR.FFN_UP_EXP: (
- "layers.{bid}.feed_forward.experts.w3", # mixtral (merged)
- "transformer.decoder_layer.{bid}.moe.linear_v", # Grok (merged)
- "transformer.blocks.{bid}.ffn.experts.mlp.v1", # dbrx
- "model.layers.{bid}.mlp.experts.up_proj", # qwen2moe olmoe (merged) ernie4.5-moe
- "model.layers.{bid}.block_sparse_moe.experts.w3", # phimoe (merged)
- "model.layers.{bid}.feed_forward.experts.up_proj", # llama4
- "encoder.layers.{bid}.mlp.experts.mlp.w1", # nomic-bert-moe
- "model.layers.{bid}.block_sparse_moe.experts.up", # smallthinker
- ),
- MODEL_TENSOR.FFN_UP_SHEXP: (
- "model.layers.{bid}.mlp.shared_expert.up_proj", # qwen2moe
- "model.layers.{bid}.mlp.shared_experts.up_proj", # deepseek deepseek2
- "model.layers.{bid}.feed_forward.shared_expert.up_proj", # llama4
- "model.layers.{bid}.feed_forward.down_proj",
- "model.layers.{bid}.mlp.shared_mlp.up_proj", # hunyuan
- ),
- # AWQ-activation gate
- MODEL_TENSOR.FFN_ACT: (
- "transformer.blocks.{bid}.ffn.act", # mpt
- ),
- # Feed-forward gate
- MODEL_TENSOR.FFN_GATE: (
- "model.layers.{bid}.mlp.gate_proj", # llama-hf refact olmo2
- "layers.{bid}.feed_forward.w1", # llama-pth
- "transformer.h.{bid}.mlp.w2", # qwen
- "transformer.h.{bid}.mlp.c_fc2", # jais
- "model.layers.layers.{bid}.mlp.gate_proj", # plamo
- "model.layers.{bid}.feed_forward.w1", # internlm2
- "encoder.layers.{bid}.mlp.fc12", # nomic-bert
- "encoder.layer.{bid}.mlp.gated_layers_w", # jina-bert-v2 (split up/gate, no longer used)
- "transformer.h.{bid}.mlp.linear_1", # refact
- "model.layers.{bid}.residual_mlp.w1", # arctic
- "transformer.h.{bid}.mlp.c_fc_0", # exaone
- "model.layers.{bid}.feed_forward.gate_proj", # llama4 jamba granite-hybrid
- "model.transformer.blocks.{bid}.ff_proj", # llada
- "layers.{bid}.mlp.gate_proj", # qwen3-embedding
- ),
- MODEL_TENSOR.FFN_GATE_EXP: (
- "layers.{bid}.feed_forward.experts.w1", # mixtral (merged)
- "transformer.decoder_layer.{bid}.moe.linear", # Grok (merged)
- "transformer.blocks.{bid}.ffn.experts.mlp.w1", # dbrx
- "model.layers.{bid}.mlp.experts.gate_proj", # qwen2moe olmoe (merged) ernie4.5-moe
- "model.layers.{bid}.block_sparse_moe.experts.w1", # phimoe (merged)
- "model.layers.{bid}.feed_forward.experts.gate_proj", # llama4
- "model.layers.{bid}.block_sparse_moe.experts.gate", # smallthinker
- ),
- MODEL_TENSOR.FFN_GATE_SHEXP: (
- "model.layers.{bid}.mlp.shared_expert.gate_proj", # qwen2moe
- "model.layers.{bid}.mlp.shared_experts.gate_proj", # deepseek deepseek2
- "model.layers.{bid}.feed_forward.shared_expert.gate_proj", # llama4
- "model.layers.{bid}.mlp.shared_mlp.gate_proj", # hunyuan
- ),
- # Feed-forward down
- MODEL_TENSOR.FFN_DOWN: (
- "gpt_neox.layers.{bid}.mlp.dense_4h_to_h", # gptneox
- "transformer.h.{bid}.mlp.c_proj", # gpt2 refact qwen jais
- "transformer.blocks.{bid}.ffn.down_proj", # mpt
- "transformer.h.{bid}.mlp.dense_4h_to_h", # falcon
- "h.{bid}.mlp.dense_4h_to_h", # bloom
- "model.layers.{bid}.mlp.down_proj", # llama-hf nemotron olmo2
- "layers.{bid}.feed_forward.w2", # llama-pth
- "encoder.layer.{bid}.output.dense", # bert
- "transformer.layer.{bid}.ffn.lin2", # distillbert
- "transformer.h.{bid}.mlp.fc_out", # gpt-j
- "language_model.encoder.layers.{bid}.mlp.dense_4h_to_h", # persimmon
- "model.layers.{bid}.mlp.dense_4h_to_h", # persimmon
- "h.{bid}.mlp.c_proj", # gpt2
- "transformer.h.{bid}.mlp.fc2", # phi2
- "model.layers.{bid}.mlp.fc2", # phi2
- "model.layers.layers.{bid}.mlp.down_proj", # plamo
- "model.layers.{bid}.feed_forward.w2", # internlm2
- "encoder.layers.{bid}.mlp.fc2", # nomic-bert
- "model.layers.{bid}.mlp.c_proj", # starcoder2
- "encoder.layer.{bid}.mlp.wo", # jina-bert-v2
- "transformer.layers.{bid}.ffn.proj_2", # openelm
- "model.layers.{bid}.residual_mlp.w2", # arctic
- "encoder.layer.{bid}.mlp.down_layer", # jina-bert-v2
- "encoder.layers.{bid}.mlp.dense_4h_to_h", # chatglm
- "model.layers.h.{bid}.mlp.c_proj", # exaone
- "model.layers.{bid}.feed_forward.down_proj", # llama4 jamba granite-hybrid
- "transformer_encoder.{bid}.ffn.w3", # neobert
- "model.layers.{bid}.block_sparse_moe.down", # smallthinker
- "model.transformer.blocks.{bid}.ff_out", # llada
- "layers.{bid}.mlp.down_proj", # qwen3-embedding
- "backbone.layers.{bid}.mixer.down_proj", # nemotron-h
- ),
- MODEL_TENSOR.FFN_DOWN_EXP: (
- "layers.{bid}.feed_forward.experts.w2", # mixtral (merged)
- "transformer.decoder_layer.{bid}.moe.linear_1", # Grok (merged)
- "transformer.blocks.{bid}.ffn.experts.mlp.w2", # dbrx
- "model.layers.{bid}.mlp.experts.down_proj", # qwen2moe olmoe (merged) ernie4.5-moe
- "model.layers.{bid}.block_sparse_moe.output_linear", # granitemoe
- "model.layers.{bid}.block_sparse_moe.experts.w2", # phimoe (merged)
- "model.layers.{bid}.feed_forward.experts.down_proj", # llama4
- "encoder.layers.{bid}.mlp.experts.mlp.w2", # nomic-bert-moe
- "model.layers.{bid}.block_sparse_moe.experts.down", # smallthinker
- ),
- MODEL_TENSOR.FFN_DOWN_SHEXP: (
- "model.layers.{bid}.mlp.shared_expert.down_proj", # qwen2moe
- "model.layers.{bid}.mlp.shared_experts.down_proj", # deepseek deepseek2
- "model.layers.{bid}.feed_forward.shared_expert.down_proj", # llama4
- "model.layers.{bid}.shared_mlp.output_linear", # granitemoe
- "model.layers.{bid}.mlp.shared_mlp.down_proj", # hunyuan
- ),
- MODEL_TENSOR.ATTN_Q_NORM: (
- "language_model.encoder.layers.{bid}.self_attention.q_layernorm",
- "model.layers.{bid}.self_attn.q_layernorm", # persimmon
- "model.layers.{bid}.self_attn.query_layernorm", # hunyuan
- "model.layers.{bid}.self_attn.q_norm", # cohere olmoe chameleon olmo2
- "transformer.blocks.{bid}.attn.q_ln", # sea-lion
- "encoder.layer.{bid}.attention.self.layer_norm_q", # jina-bert-v2
- "transformer.layers.{bid}.attn.q_norm", # openelm
- "model.layers.layers.{bid}.mixer.q", # plamo2
- "layers.{bid}.self_attn.q_norm", # qwen3-embedding
- ),
- MODEL_TENSOR.ATTN_K_NORM: (
- "language_model.encoder.layers.{bid}.self_attention.k_layernorm",
- "model.layers.{bid}.self_attn.k_layernorm", # persimmon
- "model.layers.{bid}.self_attn.key_layernorm", # hunyuan
- "model.layers.{bid}.self_attn.k_norm", # cohere olmoe chameleon olmo2
- "transformer.blocks.{bid}.attn.k_ln", # sea-lion
- "encoder.layer.{bid}.attention.self.layer_norm_k", # jina-bert-v2
- "transformer.layers.{bid}.attn.k_norm", # openelm
- "model.layers.layers.{bid}.mixer.k", # plamo2
- "layers.{bid}.self_attn.k_norm", # qwen3-embedding
- ),
- MODEL_TENSOR.ROPE_FREQS: (
- "language_model.encoder.layers.{bid}.self_attention.rotary_emb.inv_freq", # persimmon
- ),
- MODEL_TENSOR.LAYER_OUT_NORM: (
- "encoder.layer.{bid}.output.LayerNorm", # bert
- "transformer.layer.{bid}.output_layer_norm", # distillbert
- "encoder.layers.{bid}.norm2", # nomic-bert
- "transformer.decoder_layer.{bid}.rms_norm_3", # Grok
- "encoder.layer.{bid}.mlp.layernorm", # jina-bert-v2
- "encoder.layer.{bid}.layer_norm_2", # jina-v2-code
- ),
- MODEL_TENSOR.PER_LAYER_TOKEN_EMBD: (
- "model.embed_tokens_per_layer", # gemma3n
- ),
- MODEL_TENSOR.PER_LAYER_MODEL_PROJ: (
- "model.per_layer_model_projection", # gemma3n
- ),
- MODEL_TENSOR.PER_LAYER_PROJ_NORM: (
- "model.per_layer_projection_norm", # gemma3n
- ),
- MODEL_TENSOR.ALTUP_PROJ: (
- "model.altup_projections", # gemma3n
- ),
- MODEL_TENSOR.ALTUP_UNEMBD_PROJ: (
- "model.altup_unembed_projections", # gemma3n
- ),
- MODEL_TENSOR.PER_LAYER_INP_GATE: (
- "model.layers.{bid}.per_layer_input_gate", # gemma3n
- ),
- MODEL_TENSOR.PER_LAYER_PROJ: (
- "model.layers.{bid}.per_layer_projection", # gemma3n
- ),
- MODEL_TENSOR.PER_LAYER_POST_NORM: (
- "model.layers.{bid}.post_per_layer_input_norm", # gemma3n
- ),
- MODEL_TENSOR.ALTUP_CORRECT_COEF: (
- "model.layers.{bid}.altup.correction_coefs", # gemma3n
- ),
- MODEL_TENSOR.ALTUP_CORRECT_SCALE: (
- "model.layers.{bid}.altup.correct_output_scale", # gemma3n
- ),
- MODEL_TENSOR.ALTUP_PREDICT_COEF: (
- "model.layers.{bid}.altup.prediction_coefs", # gemma3n
- ),
- MODEL_TENSOR.ALTUP_ROUTER: (
- "model.layers.{bid}.altup.modality_router", # gemma3n
- ),
- MODEL_TENSOR.ALTUP_ROUTER_NORM: (
- "model.layers.{bid}.altup.router_norm", # gemma3n
- ),
- MODEL_TENSOR.LAUREL_L: (
- "model.layers.{bid}.laurel.linear_left", # gemma3n
- ),
- MODEL_TENSOR.LAUREL_R: (
- "model.layers.{bid}.laurel.linear_right", # gemma3n
- ),
- MODEL_TENSOR.LAUREL_POST_NORM: (
- "model.layers.{bid}.laurel.post_laurel_norm", # gemma3n
- ),
- MODEL_TENSOR.SSM_IN: (
- "model.layers.{bid}.in_proj", # mamba-hf
- "backbone.layers.{bid}.mixer.in_proj", # mamba
- "model.layers.{bid}.mamba.in_proj", # jamba falcon-h1 granite-hybrid
- "model.layers.layers.{bid}.mixer.in_proj", # plamo2
- ),
- MODEL_TENSOR.SSM_CONV1D: (
- "model.layers.{bid}.conv1d", # mamba-hf
- "backbone.layers.{bid}.mixer.conv1d", # mamba
- "model.layers.{bid}.mamba.conv1d", # jamba falcon-h1 granite-hybrid
- "model.layers.layers.{bid}.mixer.conv1d", # plamo2
- ),
- MODEL_TENSOR.SSM_X: (
- "model.layers.{bid}.x_proj", # mamba-hf
- "backbone.layers.{bid}.mixer.x_proj", # mamba
- "model.layers.{bid}.mamba.x_proj", # jamba
- "model.layers.layers.{bid}.mixer.bcdt_proj", # plamo2
- ),
- MODEL_TENSOR.SSM_DT: (
- "model.layers.{bid}.dt_proj", # mamba-hf
- "backbone.layers.{bid}.mixer.dt_proj", # mamba
- "model.layers.{bid}.mamba.dt_proj", # jamba falcon-h1 granite-hybrid
- "model.layers.layers.{bid}.mixer.dt_proj", # plamo2
- ),
- MODEL_TENSOR.SSM_DT_NORM: (
- "model.layers.layers.{bid}.mixer.dt_norm.weight", # plamo2
- "model.layers.{bid}.mamba.dt_layernorm", # jamba
- ),
- MODEL_TENSOR.SSM_A: (
- "model.layers.{bid}.A_log", # mamba-hf
- "backbone.layers.{bid}.mixer.A_log", # mamba
- "model.layers.{bid}.mamba.A_log", # jamba falcon-h1 granite-hybrid
- "model.layers.layers.{bid}.mixer.A_log", # plamo2
- ),
- MODEL_TENSOR.SSM_B_NORM: (
- "model.layers.{bid}.mamba.b_layernorm", # jamba
- "model.layers.{bid}.mamba.B_layernorm", # mini-jamba
- "model.layers.layers.{bid}.mixer.B_norm.weight", # plamo2
- ),
- MODEL_TENSOR.SSM_C_NORM: (
- "model.layers.{bid}.mamba.c_layernorm", # jamba
- "model.layers.{bid}.mamba.C_layernorm", # mini-jamba
- "model.layers.layers.{bid}.mixer.C_norm.weight", # plamo2
- ),
- MODEL_TENSOR.SSM_D: (
- "model.layers.{bid}.D", # mamba-hf
- "backbone.layers.{bid}.mixer.D", # mamba
- "model.layers.{bid}.mamba.D", # jamba falcon-h1 granite-hybrid
- "model.layers.layers.{bid}.mixer.D", # plamo2
- ),
- MODEL_TENSOR.SSM_NORM: (
- "model.layers.{bid}.mamba.norm", # falcon-h1 granite-hybrid
- "backbone.layers.{bid}.mixer.norm", # mamba2
- ),
- MODEL_TENSOR.SSM_OUT: (
- "model.layers.{bid}.out_proj", # mamba-hf
- "backbone.layers.{bid}.mixer.out_proj", # mamba
- "model.layers.{bid}.mamba.out_proj", # jamba falcon-h1 granite-hybrid
- "model.layers.layers.{bid}.mixer.out_proj", # plamo2
- ),
- MODEL_TENSOR.TIME_MIX_W0: (
- "model.layers.{bid}.attention.w0", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_W1: (
- "rwkv.blocks.{bid}.attention.time_maa_w1", # rwkv6
- "model.layers.{bid}.self_attn.time_maa_w1", # rwkv6qwen2
- "model.layers.{bid}.attention.w1", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_W2: (
- "rwkv.blocks.{bid}.attention.time_maa_w2", # rwkv6
- "model.layers.{bid}.self_attn.time_maa_w2", # rwkv6qwen2
- "model.layers.{bid}.attention.w2", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_A0: (
- "model.layers.{bid}.attention.a0", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_A1: (
- "model.layers.{bid}.attention.a1", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_A2: (
- "model.layers.{bid}.attention.a2", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_V0: (
- "model.layers.{bid}.attention.v0", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_V1: (
- "model.layers.{bid}.attention.v1", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_V2: (
- "model.layers.{bid}.attention.v2", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_G1: (
- "model.layers.{bid}.attention.g1", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_G2: (
- "model.layers.{bid}.attention.g2", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_K_K: (
- "model.layers.{bid}.attention.k_k", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_K_A: (
- "model.layers.{bid}.attention.k_a", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_R_K: (
- "model.layers.{bid}.attention.r_k", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_LERP_X: (
- "rwkv.blocks.{bid}.attention.time_maa_x", # rwkv6
- "model.layers.{bid}.self_attn.time_maa_x", # rwkv6qwen2
- ),
- MODEL_TENSOR.TIME_MIX_LERP_K: (
- "rwkv.blocks.{bid}.attention.time_maa_k", # rwkv6
- "model.layers.{bid}.self_attn.time_maa_k", # rwkv6qwen2
- ),
- MODEL_TENSOR.TIME_MIX_LERP_V: (
- "rwkv.blocks.{bid}.attention.time_maa_v", # rwkv6
- "model.layers.{bid}.self_attn.time_maa_v", # rwkv6qwen2
- ),
- MODEL_TENSOR.TIME_MIX_LERP_R: (
- "rwkv.blocks.{bid}.attention.time_maa_r", # rwkv6
- "model.layers.{bid}.self_attn.time_maa_r", # rwkv6qwen2
- ),
- MODEL_TENSOR.TIME_MIX_LERP_G: (
- "rwkv.blocks.{bid}.attention.time_maa_g", # rwkv6
- "model.layers.{bid}.self_attn.time_maa_g", # rwkv6qwen2
- ),
- MODEL_TENSOR.TIME_MIX_LERP_W: (
- "rwkv.blocks.{bid}.attention.time_maa_w", # rwkv6
- "model.layers.{bid}.self_attn.time_maa_w", # rwkv6qwen2
- ),
- MODEL_TENSOR.TIME_MIX_FIRST: (
- "rwkv.blocks.{bid}.attention.time_faaaa", # rwkv6
- ),
- MODEL_TENSOR.TIME_MIX_DECAY: (
- "rwkv.blocks.{bid}.attention.time_decay", # rwkv6
- "model.layers.{bid}.self_attn.time_decay", # rwkv6qwen2
- ),
- MODEL_TENSOR.TIME_MIX_DECAY_W1: (
- "rwkv.blocks.{bid}.attention.time_decay_w1", # rwkv6
- "model.layers.{bid}.self_attn.time_decay_w1", # rwkv6qwen2
- ),
- MODEL_TENSOR.TIME_MIX_DECAY_W2: (
- "rwkv.blocks.{bid}.attention.time_decay_w2", # rwkv6
- "model.layers.{bid}.self_attn.time_decay_w2", # rwkv6qwen2
- ),
- MODEL_TENSOR.TIME_MIX_KEY: (
- "rwkv.blocks.{bid}.attention.key", # rwkv6
- "model.layers.{bid}.self_attn.k_proj", # rwkv6qwen2
- "model.layers.{bid}.attention.key", # rwkv7
- "model.layers.{bid}.attention.k_proj", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_VALUE: (
- "rwkv.blocks.{bid}.attention.value", # rwkv6
- "model.layers.{bid}.self_attn.v_proj", # rwkv6qwen2
- "model.layers.{bid}.attention.value", # rwkv7
- "model.layers.{bid}.attention.v_proj", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_RECEPTANCE: (
- "rwkv.blocks.{bid}.attention.receptance", # rwkv6
- "model.layers.{bid}.self_attn.q_proj", # rwkv6qwen2
- "model.layers.{bid}.attention.receptance", # rwkv7
- "model.layers.{bid}.attention.r_proj", # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_GATE: (
- "rwkv.blocks.{bid}.attention.gate", # rwkv6
- "model.layers.{bid}.self_attn.gate", # rwkv6qwen2
- ),
- MODEL_TENSOR.TIME_MIX_LN: (
- "rwkv.blocks.{bid}.attention.ln_x", # rwkv6
- "model.layers.{bid}.attention.ln_x" # rwkv7
- ),
- MODEL_TENSOR.TIME_MIX_OUTPUT: (
- "rwkv.blocks.{bid}.attention.output", # rwkv6
- "model.layers.{bid}.self_attn.o_proj", # rwkv6qwen2
- "model.layers.{bid}.attention.output", # rwkv7
- "model.layers.{bid}.attention.o_proj", # rwkv7
- ),
- MODEL_TENSOR.CHANNEL_MIX_LERP_K: (
- "rwkv.blocks.{bid}.feed_forward.time_maa_k", # rwkv6
- "model.layers.{bid}.feed_forward.x_k", # rwkv7
- ),
- MODEL_TENSOR.CHANNEL_MIX_LERP_R: (
- "rwkv.blocks.{bid}.feed_forward.time_maa_r", # rwkv6
- ),
- MODEL_TENSOR.CHANNEL_MIX_KEY: (
- "rwkv.blocks.{bid}.feed_forward.key", # rwkv6
- "model.layers.{bid}.feed_forward.key", # rwkv7
- ),
- MODEL_TENSOR.CHANNEL_MIX_RECEPTANCE: (
- "rwkv.blocks.{bid}.feed_forward.receptance", # rwkv6
- ),
- MODEL_TENSOR.CHANNEL_MIX_VALUE: (
- "rwkv.blocks.{bid}.feed_forward.value", # rwkv6
- "model.layers.{bid}.feed_forward.value", # rwkv7
- ),
- 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_KV_A_MQA: (
- "model.layers.{bid}.self_attn.kv_a_proj_with_mqa", # deepseek2
- ),
- MODEL_TENSOR.ATTN_KV_B: (
- "model.layers.{bid}.self_attn.kv_b_proj", # deepseek2
- ),
- MODEL_TENSOR.ATTN_K_B: (
- "model.layers.{bid}.self_attn.k_b_proj", # deepseek2
- ),
- MODEL_TENSOR.ATTN_V_B: (
- "model.layers.{bid}.self_attn.v_b_proj", # deepseek2
- ),
- MODEL_TENSOR.ATTN_Q_A_NORM: (
- "model.layers.{bid}.self_attn.q_a_layernorm", # deepseek2
- ),
- MODEL_TENSOR.ATTN_KV_A_NORM: (
- "model.layers.{bid}.self_attn.kv_a_layernorm", # deepseek2
- ),
- MODEL_TENSOR.ATTN_SUB_NORM: (
- "model.layers.{bid}.self_attn.inner_attn_ln", # bitnet
- ),
- 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.DEC_ATTN_Q: (
- "decoder.block.{bid}.layer.0.SelfAttention.q", # t5
- ),
- 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_OUT: (
- "decoder.block.{bid}.layer.0.SelfAttention.o", # t5
- ),
- MODEL_TENSOR.DEC_ATTN_REL_B: (
- "decoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias", # t5
- ),
- MODEL_TENSOR.DEC_CROSS_ATTN_NORM: (
- "decoder.block.{bid}.layer.1.layer_norm", # t5
- ),
- MODEL_TENSOR.DEC_CROSS_ATTN_Q: (
- "decoder.block.{bid}.layer.1.EncDecAttention.q", # t5
- ),
- MODEL_TENSOR.DEC_CROSS_ATTN_K: (
- "decoder.block.{bid}.layer.1.EncDecAttention.k", # t5
- ),
- MODEL_TENSOR.DEC_CROSS_ATTN_V: (
- "decoder.block.{bid}.layer.1.EncDecAttention.v", # t5
- ),
- MODEL_TENSOR.DEC_CROSS_ATTN_OUT: (
- "decoder.block.{bid}.layer.1.EncDecAttention.o", # t5
- ),
- MODEL_TENSOR.DEC_CROSS_ATTN_REL_B: (
- "decoder.block.{bid}.layer.1.EncDecAttention.relative_attention_bias", # t5
- ),
- MODEL_TENSOR.DEC_FFN_NORM: (
- "decoder.block.{bid}.layer.2.layer_norm", # t5
- ),
- MODEL_TENSOR.DEC_FFN_GATE: (
- "decoder.block.{bid}.layer.2.DenseReluDense.wi_0", # flan-t5
- ),
- MODEL_TENSOR.DEC_FFN_UP: (
- "decoder.block.{bid}.layer.2.DenseReluDense.wi", # t5
- "decoder.block.{bid}.layer.2.DenseReluDense.wi_1", # flan-t5
- ),
- MODEL_TENSOR.DEC_FFN_DOWN: (
- "decoder.block.{bid}.layer.2.DenseReluDense.wo", # t5
- ),
- 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.ENC_ATTN_Q: (
- "encoder.block.{bid}.layer.0.SelfAttention.q", # t5
- ),
- 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_OUT: (
- "encoder.block.{bid}.layer.0.SelfAttention.o", # t5
- ),
- MODEL_TENSOR.ENC_ATTN_REL_B: (
- "encoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias", # t5
- ),
- MODEL_TENSOR.ENC_FFN_NORM: (
- "encoder.block.{bid}.layer.1.layer_norm", # t5
- ),
- MODEL_TENSOR.ENC_FFN_GATE: (
- "encoder.block.{bid}.layer.1.DenseReluDense.wi_0", # flan-t5
- ),
- MODEL_TENSOR.ENC_FFN_UP: (
- "encoder.block.{bid}.layer.1.DenseReluDense.wi", # t5
- "encoder.block.{bid}.layer.1.DenseReluDense.wi_1", # flan-t5
- ),
- MODEL_TENSOR.ENC_FFN_DOWN: (
- "encoder.block.{bid}.layer.1.DenseReluDense.wo", # t5
- ),
- ############################################################################
- # TODO: these do not belong to block_mappings_cfg - move them to mappings_cfg
- MODEL_TENSOR.ENC_OUTPUT_NORM: (
- "encoder.final_layer_norm", # t5
- "layer_norm", # neobert
- ),
- MODEL_TENSOR.CLS: (
- "classifier", # jina
- "classifier.dense", # roberta
- "pre_classifier", # distillbert
- "dense", # neobert
- ),
- MODEL_TENSOR.CLS_OUT: (
- "classifier.out_proj", # roberta
- ),
- #############################################################################
- MODEL_TENSOR.CONVNEXT_DW: (
- "backbone.convnext.{bid}.dwconv", # wavtokenizer
- ),
- MODEL_TENSOR.CONVNEXT_NORM: (
- "backbone.convnext.{bid}.norm", # wavtokenizer
- ),
- MODEL_TENSOR.CONVNEXT_PW1: (
- "backbone.convnext.{bid}.pwconv1", # wavtokenizer
- ),
- MODEL_TENSOR.CONVNEXT_PW2: (
- "backbone.convnext.{bid}.pwconv2", # wavtokenizer
- ),
- MODEL_TENSOR.CONVNEXT_GAMMA: (
- "backbone.convnext.{bid}.gamma", # wavtokenizer
- ),
- MODEL_TENSOR.POSNET_CONV1: (
- "backbone.posnet.{bid}.conv1", # wavtokenizer
- ),
- MODEL_TENSOR.POSNET_CONV2: (
- "backbone.posnet.{bid}.conv2", # wavtokenizer
- ),
- MODEL_TENSOR.POSNET_NORM: (
- "backbone.posnet.{bid}.norm", # wavtokenizer
- ),
- MODEL_TENSOR.POSNET_NORM1: (
- "backbone.posnet.{bid}.norm1", # wavtokenizer
- ),
- MODEL_TENSOR.POSNET_NORM2: (
- "backbone.posnet.{bid}.norm2", # wavtokenizer
- ),
- MODEL_TENSOR.POSNET_ATTN_NORM: (
- "backbone.posnet.{bid}.norm", # wavtokenizer
- ),
- MODEL_TENSOR.POSNET_ATTN_Q: (
- "backbone.posnet.{bid}.q", # wavtokenizer
- ),
- MODEL_TENSOR.POSNET_ATTN_K: (
- "backbone.posnet.{bid}.k", # wavtokenizer
- ),
- MODEL_TENSOR.POSNET_ATTN_V: (
- "backbone.posnet.{bid}.v", # wavtokenizer
- ),
- MODEL_TENSOR.POSNET_ATTN_OUT: (
- "backbone.posnet.{bid}.proj_out", # wavtokenizer
- ),
- MODEL_TENSOR.SHORTCONV_CONV: (
- "model.layers.{bid}.conv.conv",
- ),
- MODEL_TENSOR.SHORTCONV_INPROJ: (
- "model.layers.{bid}.conv.in_proj",
- ),
- MODEL_TENSOR.SHORTCONV_OUTPROJ: (
- "model.layers.{bid}.conv.out_proj",
- ),
- #############################################################################
- ## Vision encoder
- MODEL_TENSOR.V_MMPROJ: (
- "multi_modal_projector.linear_{bid}",
- "visual.merger.mlp.{bid}", # qwen2vl
- ),
- MODEL_TENSOR.V_MMPROJ_FC: (
- "model.connector.modality_projection.proj", # SmolVLM
- ),
- MODEL_TENSOR.V_MMPROJ_MLP: (
- "model.mm_projector.mlp.mlp.{bid}",
- "vision_model.vision_adapter.mlp.fc{bid}", # llama 4
- "mlp1.{bid}", # InternVL
- ),
- MODEL_TENSOR.V_MMPROJ_PEG: (
- "model.mm_projector.peg.peg.{bid}",
- ),
- MODEL_TENSOR.V_ENC_EMBD_CLS: (
- "vision_tower.vision_model.embeddings.class_embedding",
- "model.vision_tower.embeddings.cls_token", # Intern-S1
- "vision_model.class_embedding", # llama 4
- ),
- MODEL_TENSOR.V_ENC_EMBD_PATCH: (
- "vision_tower.vision_model.embeddings.patch_embedding",
- "model.vision_tower.embeddings.patch_embeddings.projection", # Intern-S1
- "vpm.embeddings.patch_embedding",
- "model.vision_model.embeddings.patch_embedding", # SmolVLM
- "vision_tower.patch_conv", # pixtral-hf
- "vision_encoder.patch_conv", # pixtral
- "vision_model.patch_embedding.linear", # llama 4
- "visual.patch_embed.proj", # qwen2vl
- "vision_tower.patch_embed.proj", # kimi-vl
- ),
- MODEL_TENSOR.V_ENC_EMBD_POS: (
- "vision_tower.vision_model.embeddings.position_embedding",
- "model.vision_tower.embeddings.position_embeddings", # Intern-S1
- "vpm.embeddings.position_embedding",
- "model.vision_model.embeddings.position_embedding", # SmolVLM
- "vision_model.positional_embedding_vlm", # llama 4
- "vision_tower.patch_embed.pos_emb", # kimi-vl
- ),
- MODEL_TENSOR.V_ENC_ATTN_Q: (
- "vision_tower.vision_model.encoder.layers.{bid}.self_attn.q_proj",
- "model.vision_tower.encoder.layer.{bid}.attention.q_proj", # Intern-S1
- "vpm.encoder.layers.{bid}.self_attn.q_proj",
- "model.vision_model.encoder.layers.{bid}.self_attn.q_proj", # SmolVLM
- "vision_model.model.layers.{bid}.self_attn.q_proj", # llama4
- "vision_tower.transformer.layers.{bid}.attention.q_proj", # pixtral-hf
- "vision_encoder.transformer.layers.{bid}.attention.wq", # pixtral
- "visual.blocks.{bid}.attn.q", # qwen2vl, generated
- "vision_tower.encoder.blocks.{bid}.wq", # kimi-vl, generated
- ),
- MODEL_TENSOR.V_ENC_ATTN_Q_NORM: (
- "vision_tower.vision_model.encoder.layers.{bid}.attn.q_norm", # InternVL
- "model.vision_tower.encoder.layer.{bid}.attention.q_norm", # Intern-S1
- ),
- MODEL_TENSOR.V_ENC_ATTN_K: (
- "vision_tower.vision_model.encoder.layers.{bid}.self_attn.k_proj",
- "model.vision_tower.encoder.layer.{bid}.attention.k_proj", # Intern-S1
- "vpm.encoder.layers.{bid}.self_attn.k_proj",
- "model.vision_model.encoder.layers.{bid}.self_attn.k_proj", # SmolVLM
- "vision_model.model.layers.{bid}.self_attn.k_proj", # llama4
- "vision_tower.transformer.layers.{bid}.attention.k_proj", # pixtral-hf
- "vision_encoder.transformer.layers.{bid}.attention.wk", # pixtral
- "visual.blocks.{bid}.attn.k", # qwen2vl, generated
- "vision_tower.encoder.blocks.{bid}.wk", # kimi-vl, generated
- ),
- MODEL_TENSOR.V_ENC_ATTN_K_NORM: (
- "vision_tower.vision_model.encoder.layers.{bid}.attn.k_norm", # InternVL
- "model.vision_tower.encoder.layer.{bid}.attention.k_norm", # Intern-S1
- ),
- MODEL_TENSOR.V_ENC_ATTN_V: (
- "vision_tower.vision_model.encoder.layers.{bid}.self_attn.v_proj",
- "model.vision_tower.encoder.layer.{bid}.attention.v_proj", # Intern-S1
- "vpm.encoder.layers.{bid}.self_attn.v_proj",
- "model.vision_model.encoder.layers.{bid}.self_attn.v_proj", # SmolVLM
- "vision_model.model.layers.{bid}.self_attn.v_proj", # llama4
- "vision_tower.transformer.layers.{bid}.attention.v_proj", # pixtral-hf
- "vision_encoder.transformer.layers.{bid}.attention.wv", # pixtral
- "visual.blocks.{bid}.attn.v", # qwen2vl, generated
- "vision_tower.encoder.blocks.{bid}.wv", # kimi-vl, generated
- ),
- MODEL_TENSOR.V_ENC_INPUT_NORM: (
- "vision_tower.vision_model.encoder.layers.{bid}.layer_norm1",
- "vision_tower.vision_model.encoder.layers.{bid}.norm1", # InternVL
- "model.vision_tower.encoder.layer.{bid}.layernorm_before", # Intern-S1
- "vpm.encoder.layers.{bid}.layer_norm1",
- "model.vision_model.encoder.layers.{bid}.layer_norm1", # SmolVLM
- "vision_tower.transformer.layers.{bid}.attention_norm", # pixtral-hf
- "vision_encoder.transformer.layers.{bid}.attention_norm", # pixtral
- "vision_model.model.layers.{bid}.input_layernorm", # llama4
- "visual.blocks.{bid}.norm1", # qwen2vl
- "vision_tower.encoder.blocks.{bid}.norm0", # kimi-vl (norm0/norm1)
- ),
- MODEL_TENSOR.V_ENC_ATTN_O: (
- "vision_tower.vision_model.encoder.layers.{bid}.self_attn.out_proj",
- "vision_tower.vision_model.encoder.layers.{bid}.attn.proj", # InternVL
- "model.vision_tower.encoder.layer.{bid}.attention.projection_layer", # Intern-S1
- "vpm.encoder.layers.{bid}.self_attn.out_proj",
- "model.vision_model.encoder.layers.{bid}.self_attn.out_proj", # SmolVLM
- "vision_model.model.layers.{bid}.self_attn.o_proj", # llama4
- "vision_tower.transformer.layers.{bid}.attention.o_proj", # pixtral-hf
- "vision_encoder.transformer.layers.{bid}.attention.wo", # pixtral
- "visual.blocks.{bid}.attn.proj", # qwen2vl
- "vision_tower.encoder.blocks.{bid}.wo", # kimi-vl
- ),
- MODEL_TENSOR.V_ENC_POST_ATTN_NORM: (
- "vision_tower.vision_model.encoder.layers.{bid}.layer_norm2",
- "vision_tower.vision_model.encoder.layers.{bid}.norm2", # InternVL
- "model.vision_tower.encoder.layer.{bid}.layernorm_after", # Intern-S1
- "vpm.encoder.layers.{bid}.layer_norm2",
- "model.vision_model.encoder.layers.{bid}.layer_norm2", # SmolVLM
- "vision_model.model.layers.{bid}.post_attention_layernorm", # llama4
- "vision_tower.transformer.layers.{bid}.ffn_norm", # pixtral-hf
- "vision_encoder.transformer.layers.{bid}.ffn_norm", # pixtral
- "visual.blocks.{bid}.norm2", # qwen2vl
- "vision_tower.encoder.blocks.{bid}.norm1", # kimi-vl (norm0/norm1)
- ),
- MODEL_TENSOR.V_ENC_FFN_UP: (
- "vision_tower.vision_model.encoder.layers.{bid}.mlp.fc1",
- "model.vision_tower.encoder.layer.{bid}.mlp.fc1", # Intern-S1
- "vpm.encoder.layers.{bid}.mlp.fc1",
- "model.vision_model.encoder.layers.{bid}.mlp.fc1", # SmolVLM, gemma3
- "vision_tower.transformer.layers.{bid}.feed_forward.up_proj", # pixtral-hf
- "vision_encoder.transformer.layers.{bid}.feed_forward.w3", # pixtral
- "vision_model.model.layers.{bid}.mlp.fc1", # llama4
- "visual.blocks.{bid}.mlp.fc1", # qwen2vl
- "visual.blocks.{bid}.mlp.up_proj", # qwen2.5vl
- "vision_tower.encoder.blocks.{bid}.mlp.fc0", # kimi-vl (fc0/fc1)
- ),
- MODEL_TENSOR.V_ENC_FFN_GATE: (
- "vision_tower.transformer.layers.{bid}.feed_forward.gate_proj", # pixtral-hf
- "vision_encoder.transformer.layers.{bid}.feed_forward.w1", # pixtral
- "visual.blocks.{bid}.mlp.gate_proj", # qwen2.5vl
- ),
- MODEL_TENSOR.V_ENC_FFN_DOWN: (
- "vision_tower.vision_model.encoder.layers.{bid}.mlp.fc2",
- "model.vision_tower.encoder.layer.{bid}.mlp.fc2", # Intern-S1
- "vpm.encoder.layers.{bid}.mlp.fc2",
- "model.vision_model.encoder.layers.{bid}.mlp.fc2", # SmolVLM, gemma3
- "vision_tower.transformer.layers.{bid}.feed_forward.down_proj", # pixtral-hf
- "vision_encoder.transformer.layers.{bid}.feed_forward.w2", # pixtral
- "vision_model.model.layers.{bid}.mlp.fc2", # llama4
- "visual.blocks.{bid}.mlp.fc2", # qwen2vl
- "visual.blocks.{bid}.mlp.down_proj", # qwen2.5vl
- "vision_tower.encoder.blocks.{bid}.mlp.fc1", # kimi-vl (fc0/fc1)
- ),
- MODEL_TENSOR.V_LAYER_SCALE_1: (
- "vision_tower.vision_model.encoder.layers.{bid}.ls1", # InternVL
- "model.vision_tower.encoder.layer.{bid}.lambda_1", # Intern-S1
- ),
- MODEL_TENSOR.V_LAYER_SCALE_2: (
- "vision_tower.vision_model.encoder.layers.{bid}.ls2", # InternVL
- "model.vision_tower.encoder.layer.{bid}.lambda_2", # Intern-S1
- ),
- MODEL_TENSOR.V_PRE_NORM: (
- "vision_tower.vision_model.pre_layrnorm",
- "vision_tower.ln_pre", # pixtral-hf
- "vision_encoder.ln_pre", # pixtral
- "vision_model.layernorm_pre", # llama4
- ),
- MODEL_TENSOR.V_POST_NORM: (
- "vision_tower.vision_model.post_layernorm",
- "model.vision_model.post_layernorm", # SmolVLM
- "vision_model.layernorm_post", # llama4
- "visual.merger.ln_q", # qwen2vl
- "vision_tower.encoder.final_layernorm", # kimi-vl
- ),
- MODEL_TENSOR.V_MM_INP_PROJ: (
- "multi_modal_projector.mm_input_projection",
- ),
- MODEL_TENSOR.V_MM_INP_NORM: (
- "multi_modal_projector.norm",
- "multi_modal_projector.layer_norm",
- "multi_modal_projector.pre_norm",
- "pre_mm_projector_norm",
- ),
- MODEL_TENSOR.V_MM_SOFT_EMB_NORM: (
- "multi_modal_projector.mm_soft_emb_norm",
- ),
- MODEL_TENSOR.V_RESMPL_POS_EMBD_K: (
- "resampler.pos_embed_k",
- ),
- MODEL_TENSOR.V_RESMPL_ATTN_Q: (
- "resampler.attn.in_proj_q", # tensor generated from resampler.attn.in_proj
- ),
- MODEL_TENSOR.V_RESMPL_ATTN_K: (
- "resampler.attn.in_proj_k", # tensor generated from resampler.attn.in_proj
- ),
- MODEL_TENSOR.V_RESMPL_ATTN_V: (
- "resampler.attn.in_proj_v", # tensor generated from resampler.attn.in_proj
- ),
- MODEL_TENSOR.V_RESMPL_ATTN_OUT: (
- "resampler.attn.out_proj",
- ),
- MODEL_TENSOR.V_RESMPL_KV: (
- "resampler.kv_proj",
- ),
- MODEL_TENSOR.V_RESMPL_POST_NORM: (
- "resampler.ln_post",
- ),
- MODEL_TENSOR.V_RESMPL_KV_NORM: (
- "resampler.ln_kv",
- ),
- MODEL_TENSOR.V_RESMPL_Q_NORM: (
- "resampler.ln_q",
- ),
- MODEL_TENSOR.V_RESMPL_PROJ: (
- "resampler.proj",
- ),
- MODEL_TENSOR.V_RESMPL_QUERY: (
- "resampler.query",
- ),
- MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK: (
- "v.token_embd.img_break", # for pixtral, this is a generated vector
- ),
- MODEL_TENSOR.V_MM_PATCH_MERGER: (
- "multi_modal_projector.patch_merger.merging_layer", # mistral small 3.1 - hf
- "patch_merger.merging_layer", # mistral
- ),
- # audio (mtmd)
- MODEL_TENSOR.A_ENC_EMBD_POS: (
- "audio_tower.embed_positions", # ultravox
- ),
- MODEL_TENSOR.A_ENC_CONV1D: (
- "audio_tower.conv{bid}", # ultravox
- ),
- MODEL_TENSOR.A_PRE_NORM: (),
- MODEL_TENSOR.A_POST_NORM: (
- "audio_tower.layer_norm", # ultravox
- "audio_tower.ln_post", # qwen2omni
- ),
- MODEL_TENSOR.A_ENC_ATTN_Q: (
- "audio_tower.layers.{bid}.self_attn.q_proj", # ultravox
- ),
- MODEL_TENSOR.A_ENC_ATTN_K: (
- "audio_tower.layers.{bid}.self_attn.k_proj", # ultravox
- ),
- MODEL_TENSOR.A_ENC_ATTN_V: (
- "audio_tower.layers.{bid}.self_attn.v_proj", # ultravox
- ),
- MODEL_TENSOR.A_ENC_INPUT_NORM: (
- "audio_tower.layers.{bid}.self_attn_layer_norm", # ultravox
- ),
- MODEL_TENSOR.A_ENC_OUTPUT: (
- "audio_tower.layers.{bid}.self_attn.out_proj", # ultravox
- ),
- MODEL_TENSOR.A_ENC_OUTPUT_NORM: (
- "audio_tower.layers.{bid}.final_layer_norm", # ultravox
- ),
- MODEL_TENSOR.A_ENC_FFN_UP: (
- "audio_tower.layers.{bid}.fc1", # ultravox
- ),
- MODEL_TENSOR.A_ENC_FFN_GATE: (),
- MODEL_TENSOR.A_ENC_FFN_DOWN: (
- "audio_tower.layers.{bid}.fc2", # ultravox
- ),
- # note: some tensors below has "audio." pseudo-prefix, to prevent conflicts with vision tensors
- # this prefix is added in the conversion code in modify_tensors()
- MODEL_TENSOR.A_MMPROJ: (
- "audio.multi_modal_projector.linear_{bid}", # ultravox
- ),
- MODEL_TENSOR.A_MMPROJ_FC: (
- "audio.multi_modal_projector.linear", # qwen2audio
- "audio_tower.proj", # qwen2omni
- ),
- MODEL_TENSOR.A_MM_NORM_PRE: (
- "audio.multi_modal_projector.ln_pre", # ultravox
- ),
- MODEL_TENSOR.A_MM_NORM_MID: (
- "audio.multi_modal_projector.ln_mid", # ultravox
- ),
- # NextN/MTP tensors for GLM4_MOE
- MODEL_TENSOR.NEXTN_EH_PROJ: (
- "model.layers.{bid}.eh_proj",
- ),
- MODEL_TENSOR.NEXTN_EMBED_TOKENS: (
- "model.layers.{bid}.embed_tokens",
- ),
- MODEL_TENSOR.NEXTN_ENORM: (
- "model.layers.{bid}.enorm",
- ),
- MODEL_TENSOR.NEXTN_HNORM: (
- "model.layers.{bid}.hnorm",
- ),
- MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD: (
- "model.layers.{bid}.shared_head.head",
- ),
- MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: (
- "model.layers.{bid}.shared_head.norm",
- ),
- }
- # architecture-specific block mappings
- arch_block_mappings_cfg: dict[MODEL_ARCH, dict[MODEL_TENSOR, tuple[str, ...]]] = {
- MODEL_ARCH.ARCTIC: {
- MODEL_TENSOR.FFN_NORM: (
- "model.layers.{bid}.residual_layernorm",
- ),
- MODEL_TENSOR.FFN_NORM_EXP: (
- "model.layers.{bid}.post_attention_layernorm",
- ),
- },
- }
- mapping: dict[str, tuple[MODEL_TENSOR, str]]
- def __init__(self, arch: MODEL_ARCH, n_blocks: int):
- self.mapping = {}
- for tensor, keys in self.mappings_cfg.items():
- if tensor not in MODEL_TENSORS[arch]:
- continue
- tensor_name = TENSOR_NAMES[tensor]
- self.mapping[tensor_name] = (tensor, tensor_name)
- for key in keys:
- self.mapping[key] = (tensor, tensor_name)
- if arch in self.arch_block_mappings_cfg:
- self.block_mappings_cfg.update(self.arch_block_mappings_cfg[arch])
- for bid in range(n_blocks):
- for tensor, keys in self.block_mappings_cfg.items():
- if tensor not in MODEL_TENSORS[arch]:
- continue
- tensor_name = TENSOR_NAMES[tensor].format(bid = bid)
- self.mapping[tensor_name] = (tensor, tensor_name)
- for key in keys:
- key = key.format(bid = bid)
- self.mapping[key] = (tensor, tensor_name)
- def get_type_and_name(self, key: str, try_suffixes: Sequence[str] = ()) -> tuple[MODEL_TENSOR, str] | None:
- result = self.mapping.get(key)
- if result is not None:
- return result
- for suffix in try_suffixes:
- if key.endswith(suffix):
- result = self.mapping.get(key[:-len(suffix)])
- if result is not None:
- return result[0], result[1] + suffix
- return None
- def get_name(self, key: str, try_suffixes: Sequence[str] = ()) -> str | None:
- result = self.get_type_and_name(key, try_suffixes = try_suffixes)
- if result is None:
- return None
- return result[1]
- def get_type(self, key: str, try_suffixes: Sequence[str] = ()) -> MODEL_TENSOR | None:
- result = self.get_type_and_name(key, try_suffixes = try_suffixes)
- if result is None:
- return None
- return result[0]
- def __getitem__(self, key: str) -> str:
- try:
- return self.mapping[key][1]
- except KeyError:
- raise KeyError(key)
- def __contains__(self, key: str) -> bool:
- return key in self.mapping
- def __repr__(self) -> str:
- return repr(self.mapping)
- def get_tensor_name_map(arch: MODEL_ARCH, n_blocks: int) -> TensorNameMap:
- return TensorNameMap(arch, n_blocks)
|