o TÃizã@s8ddlTddlmZGdd„deƒZGdd„deƒZdS)é)Ú*)ÚLayerContainerc @s`eZdZUeed<eed<eed<eed<eed<e ed<e ed<dd d d d d dddddœ Z dS)ÚMixtralTransformerContainerÚqkv_wÚ attn_out_wÚmoe_gateÚ moe_mlp_1Ú moe_mlp_2Úattn_norm_gammaÚmlp_norm_gammazattn_norm_gamma.paramszmlp_norm_gamma.paramszqkv_w.q_paramszqkv_w.k_paramszqkv_w.v_paramszattn_out_w.paramszmoe_gate.paramszmoe_mlp_1.gating_expertszmoe_mlp_1.up_expertszmoe_mlp_2.experts) zinput_layernorm.weightzpost_attention_layernorm.weightzself_attn.q_proj.weightzself_attn.k_proj.weightzself_attn.v_proj.weightzself_attn.o_proj.weightzblock_sparse_moe.gate.weightz$block_sparse_moe.experts.*.w1.weightz$block_sparse_moe.experts.*.w3.weightz$block_sparse_moe.experts.*.w2.weightN) Ú__name__Ú __module__Ú __qualname__ÚUnfusedQKVParameterÚ__annotations__ÚAttentionOutputParameterÚMoEGatingWeightParameterÚUnfusedMoEGatedMLPParameterÚUnfusedMoEMLP2ParameterÚ NormParameterÚ PARAM_MAPPING©rrúr/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/inference/v2/model_implementations/mixtral/container.pyr s&  örc@s2eZdZUeed<eed<eed<ddddœZdS) ÚMixtralNonTransformerContainerÚword_embÚ word_unembedÚ final_normzword_emb.paramszword_unembed.paramszfinal_norm.params)zmodel.embed_tokens.weightzlm_head.weightzmodel.norm.weightN)r r rÚEmbeddingParameterrÚUnembedParameterrrrrrrr$s  ýrN)Ú>deepspeed.inference.v2.model_implementations.common_parametersÚAdeepspeed.inference.v2.model_implementations.layer_container_baserrrrrrrÚs