Views
No views yet
MixtralForCausalLM(
(model): MixtralModel(
(embed_tokens): Embedding(32000, 1024)
(layers): ModuleList(
(0-1): 2 x MixtralDecoderLayer(
(self_attn): MixtralAttention(
(q_proj): Linear(in_features=1024, out_features=1024, bias=False)
(k_proj): Linear(in_features=1024, out_features=256, bias=False)
(v_proj): Linear(in_features=1024, out_features=256, bias=False)
(o_proj): Linear(in_features=1024, out_features=1024, bias=False)
)
(block_sparse_moe): MixtralSparseMoeBlock(
(gate): Linear(in_features=1024, out_features=8, bias=False)
(experts): ModuleList(
(0-7): 8 x MixtralBlockSparseTop2MLP(
(w1): Linear(in_features=1024, out_features=3584, bias=False)
(w2): Linear(in_features=3584, out_features=1024, bias=False)
(w3): Linear(in_features=1024, out_features=3584, bias=False)
(act_fn): SiLUActivation()
)
)
)
(input_layernorm): MixtralRMSNorm((1024,), eps=1e-05)
(post_attention_layernorm): MixtralRMSNorm((1024,), eps=1e-05)
)
)
(norm): MixtralRMSNorm((1024,), eps=1e-05)
(rotary_emb): MixtralRotaryEmbedding()
)
(lm_head): Linear(in_features=1024, out_features=32000, bias=False)
)MixtralForCausalLM(
(model): MixtralModel(
(embed_tokens): Embedding(32000, 1024)
(layers): ModuleList(
(0-1): 2 x MixtralDecoderLayer(
(self_attn): MixtralAttention(
(q_proj): Linear(in_features=1024, out_features=1024, bias=False)
(k_proj): Linear(in_features=1024, out_features=256, bias=False)
(v_proj): Linear(in_features=1024, out_features=256, bias=False)
(o_proj): Linear(in_features=1024, out_features=1024, bias=False)
)
(mlp): MixtralSparseMoeBlock(
(gate): MixtralTopKRouter()
(experts): MixtralExperts(
(act_fn): SiLUActivation()
)
)
(input_layernorm): MixtralRMSNorm((1024,), eps=1e-05)
(post_attention_layernorm): MixtralRMSNorm((1024,), eps=1e-05)
)
)
(norm): MixtralRMSNorm((1024,), eps=1e-05)
(rotary_emb): MixtralRotaryEmbedding()
)
(lm_head): Linear(in_features=1024, out_features=32000, bias=False)
)1 "mixtral": [
2 WeightRenaming(".block_sparse_moe.gate", ".mlp.gate"),
3 WeightConverter(
4 source_patterns=[
5 "block_sparse_moe.experts.*.w1.weight",
6 "block_sparse_moe.experts.*.w3.weight",
7 ], # you give me a list of 2 keys, I collect a list of a list of tensors
8 target_patterns="mlp.experts.gate_up_proj", # target key gets the list of two tensors
9 operations=[
10 MergeModulelist(
11 dim=0
12 ), # each process has two lists of tensors, we cat each list. -> we end up with 2 tensors
13 Concatenate(dim=1), # each process has 2 tensors, gate and up, we concat them into gate_up
14 ], # we want the loading to add this shard operation here. Though we can't shard after concats and merge, needs to be first
15 ),
16 WeightConverter(
17 source_patterns=[
18 "block_sparse_moe.experts.*.w2.weight",
19 ],
20 target_patterns="mlp.experts.down_proj", # target key gets the list of two tensors
21 operations=[
22 MergeModulelist(
23 dim=0
24 ), # each process has two lists of tensors, we cat each list. -> we end up with 2 tensors
25 ], # we want the loading to add this shard operation here. Though we can't shard after concats and merge, needs to be first
26 ),
27 ],