o }oi @sddlZddlmZmZddlmZddlmZmZm Z ddl Z ddl Z ddl m Z ddl mZddlmZddlmZdd lmZmZdd lmZdd lmZdd lmZmZmZmZdd lm Z ddl!m"Z"m#Z#ddl$m%Z%ddl&m'Z'ddl(m)Z)eGdddeZ*eGdddeZ+eGdddeZ,eGddde,Z-e".edGddde"j/defZ0e"1edGddde"j/ed fZ2d!efd"d#Z3d$d%Z4d&d'Z5d(d)Z6d*e"j7fd+d,Z8d*e"j7fd-d.Z9d*e"j7fd/d0Z:d1d2Z;d3ed4ed5ed>ed5e/zMLlamaConfig11B.default_factorylanguage_model_configcC tddS)Nivision_chunk_sizerrrrr r!1 vision_model_configN __name__ __module__ __qualname__rr%rr__annotations__r+rrrr r- rc@r)MLlamaConfig11BInstructcCrrrrrrr r!7r"z MLlamaConfig11BInstruct.r#r%cCr&)N0r'r)rrrr r!9r*r+Nr,rrrr r25r1r2c@r)MLlamaConfig90BcCstddddddS)N ip@P) hidden_sizeffn_hidden_sizenum_attention_heads num_layersnum_cross_attention_layersrrrrr r!@szMLlamaConfig90B.r#r%cCs tdddS)Nr3r5)r(text_hidden_sizer)rrrr r!Is r+Nr,rrrr r4=s  r4c@s eZdZdS)MLlamaConfig90BInstructN)r-r.r/rrrr r?Msr?hfcseZdZdefddZddeedeffdd Zdedefd d Zd d Z e dddZ e de fddZ deefddZdeefddZZS)HFMLlamaImporterreturncCst|j|jdS)N) tokenizer)rconfigrCselfrrr initTszHFMLlamaImporter.initN base_pathcst|}|Sr)super local_path)rFrH output_path __class__rr rJWs zHFMLlamaImporter.local_pathrKcCsptjt|dd}t|}t|}|}||}|||| ||t d|t ||~~|S)Nauto torch_dtypez/Converted Mllama model to Nemo, model saved to ) r from_pretrainedstr_rename_xattn_layer_nums_hf state_dictrrG nemo_setup convert_state nemo_saveprintr)rFrKsourcerTtargettrainerrrr apply\s     zHFMLlamaImporter.applycCsi}g}|ddddddddd d d d d d |tjddtdtjddtdtjddtdtjddtdtjddtdtjddtdtjddtdgd}|id|d d!|d"d#|d$d%|d&d'|d(d)|d*d+|d,d-|d.d/|d0d1|d2d3|d4d5|d6d7|d8d9|d:d;|d<d=|d>d?|d@idA|dBdC|dDdE|dFdG|dHdI|dJdK|dLdM|dNdO|dPdQ|dRdS|dTdU|dVdW|dXdY|dZd[|d\d]|d^d_d`dadb|tjdc|ddt dtjde|dft dtjdg|dht dgtj ||||diS)jNAlanguage_model.decoder.layers.*.self_attention.linear_proj.weightHlanguage_model.decoder.xattn_layers.*.cross_attention.linear_proj.weightElanguage_model.decoder.xattn_layers.*.cross_attention.linear_q.weight-language_model.decoder.final_layernorm.weight"language_model.output_layer.weight@language_model.decoder.layers.*.mlp.linear_fc1.layer_norm_weight5language_model.decoder.layers.*.mlp.linear_fc2.weightKlanguage_model.decoder.layers.*.self_attention.linear_qkv.layer_norm_weightHlanguage_model.decoder.xattn_layers.*.cross_attention.k_layernorm.weightzPlanguage_model.decoder.xattn_layers.*.cross_attention.linear_q.layer_norm_weightHlanguage_model.decoder.xattn_layers.*.cross_attention.q_layernorm.weightzFlanguage_model.decoder.xattn_layers.*.mlp.linear_fc1.layer_norm_weightz;language_model.decoder.xattn_layers.*.mlp.linear_fc2.weight) 5language_model.model.layers.*.self_attn.o_proj.weightz.global_transformer.layers.*.self_attention.linear_proj.weight2vision_model.global_transformer.layers.*.gate_attn&.global_transformer.layers.*.gate_attn1vision_model.global_transformer.layers.*.gate_ffn%.global_transformer.layers.*.gate_ffn=vision_model.global_transformer.layers.*.input_layernorm.bias1.global_transformer.layers.*.input_layernorm.bias?vision_model.global_transformer.layers.*.input_layernorm.weight3.global_transformer.layers.*.input_layernorm.weightFvision_model.global_transformer.layers.*.post_attention_layernorm.bias3.global_transformer.layers.*.pre_mlp_layernorm.biasHvision_model.global_transformer.layers.*.post_attention_layernorm.weight5.global_transformer.layers.*.pre_mlp_layernorm.weight5vision_model.global_transformer.layers.*.mlp.fc1.bias0.global_transformer.layers.*.mlp.linear_fc1.bias7vision_model.global_transformer.layers.*.mlp.fc1.weight2.global_transformer.layers.*.mlp.linear_fc1.weight5vision_model.global_transformer.layers.*.mlp.fc2.bias0.global_transformer.layers.*.mlp.linear_fc2.bias7vision_model.global_transformer.layers.*.mlp.fc2.weight2.global_transformer.layers.*.mlp.linear_fc2.weight9vision_model.transformer.layers.*.self_attn.o_proj.weight7.transformer.layers.*.self_attention.linear_proj.weight6vision_model.transformer.layers.*.input_layernorm.bias*.transformer.layers.*.input_layernorm.bias8vision_model.transformer.layers.*.input_layernorm.weight,.transformer.layers.*.input_layernorm.weight?vision_model.transformer.layers.*.post_attention_layernorm.bias,.transformer.layers.*.pre_mlp_layernorm.biasAvision_model.transformer.layers.*.post_attention_layernorm.weight..transformer.layers.*.pre_mlp_layernorm.weight.vision_model.transformer.layers.*.mlp.fc1.bias).transformer.layers.*.mlp.linear_fc1.bias0vision_model.transformer.layers.*.mlp.fc1.weight+.transformer.layers.*.mlp.linear_fc1.weight.vision_model.transformer.layers.*.mlp.fc2.bias).transformer.layers.*.mlp.linear_fc2.bias0vision_model.transformer.layers.*.mlp.fc2.weight+.transformer.layers.*.mlp.linear_fc2.weightvision_model.class_embedding.class_embedding1vision_model.gated_positional_embedding.embedding.positional_embedding=vision_model.gated_positional_embedding.tile_embedding.weight'.gated_tile_positional_embedding.weight,vision_model.gated_positional_embedding.gate .gated_positional_embedding_gate vision_model.layernorm_post.bias .ln_post.bias"vision_model.layernorm_post.weight.ln_post.weightvision_model.layernorm_pre.bias .ln_pre.bias!vision_model.layernorm_pre.weight.ln_pre.weight._calculate_num_layersr5) rotary_base seq_lengthr<r=r9r:r;num_query_groups vocab_sizefp16bf16 params_dtype)r text_config rope_thetarrrr9intermediate_sizer;num_key_value_headsrrtorchfloat16bfloat16)rFrYrrrr rs$   z'HFMLlamaImporter._language_model_configc CsJt|jj|jj|jj|jj|jj|jjt|t j kt|t j kt|d S)N) r<r9r;r(vision_max_num_chunksr>rrr) r vision_configrr9attention_heads image_size max_num_tilesrrrrr)rFrYrrr r$s  z%HFMLlamaImporter._vision_model_configr)rBr)r-r.r/rrGrrrJr\rVpropertyrCrrDrrrr __classcell__rrrLr rARs  rArc@sxeZdZdZejfdddZdedefddZd d Z e dd d Z dede e e ffddZddZe dddZdS)HFMLlamaExportera Exporter class for converting NeMo MLlama model to HuggingFace format. Inherits: io.ModelConnector: Connector interface to handle setup, save, and load using the Lightning framework. Methods: init: Initializes a new HuggingFace MLlama model instance. apply: Converts the NeMo model to HuggingFace format and saves it. convert_state: Maps and transforms the state dictionary from NeMo to HuggingFace format. config: Generates and returns the HuggingFace MLlama config for the model. rBr cCsFddlm}|tj|j|dWdS1swYdS)a- Initializes a HuggingFace MllamaForConditionalGeneration model. Args: dtype: The data type to use for the model (default: torch.bfloat16) Returns: MllamaForConditionalGeneration: A HuggingFace MLlama model initialized with the configuration. r)no_init_weightsrON)transformers.modeling_utilsrr _from_configrD)rFdtyperrrr rGAs $zHFMLlamaExporter.initrKcCstd||\}}tdtd|}td||||}|}||z |jj|WntyFt dYnwt d||S)a2 Converts the NeMo MLlama model to HuggingFace format and saves it to the specified path. Args: output_path (Path): The path where the converted HuggingFace model will be saved. Returns: Path: The output path where the HuggingFace model was saved. z8Loading MLlama NeMo checkpoint. This may take a while...zMLlama NeMo checkpoint loaded.zInitializing the HF model..zStart Converting the model..zFailed to save tokenizerz Converted MLlama model saved to ) rinfo ckpt_loadrGrVcpusave_pretrainedrC ExceptionwarningrX)rFrKrY source_configrZrrr r\Ps      zHFMLlamaExporter.applyc Cs|||}i}g}|ddddddddd d d |tjd d tdtjddtdtjddtdtjddtdtjddtdtjddt dgd}|i|dd|dd|dd|d d!|d"d#|d$d%|d&d'|d(d)|d*d+|d,d-|d.d/|d0d1|d2d3|d4d5|d6d7|d8d9|d:d;i|d<d=|d>d?|d@dA|dBdC|dDdE|dFdG|dHdI|dJdK|dLdM|dNdO|dPdQ|dRdS|dTdU|dVdW|dXdYdZd[d\d]|tj|d^d_t dtj|d`dat dtj|dbdct dgtj ||||ddS)ea Maps and transforms the state dictionary from NeMo to HuggingFace format. Args: source: The source NeMo model. target: The target HuggingFace model. Returns: The target HuggingFace model with the converted state. rgrlrjrkz6language_model.model.layers.*.cross_attn.q_norm.weightz6language_model.model.layers.*.cross_attn.q_proj.weightz6language_model.model.layers.*.cross_attn.k_norm.weightz6language_model.model.layers.*.cross_attn.o_proj.weightrhri) r]rdrbrcrfr_rer^r`rarmz2language_model.model.layers.*.cross_attn_attn_gaternrrz1language_model.model.layers.*.cross_attn_mlp_gatertrsrvrurw)z6language_model.model.layers.*.cross_attn.k_proj.weightz6language_model.model.layers.*.cross_attn.v_proj.weightryrxrzr|r{r~r}rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr) _modify_mllama_source_staterrrr _export_gate_export_text_qkv_export_simple_split_export_text_kv_export_embedding_hf_export_vision_qkv_export_patch_embedding_hfr)rFrYrZrrrrrrr rVms  /                                !"&  zHFMLlamaExporter.convert_statercCstjt|ddjS)z~ Gets the tokenizer from the loaded model context. Returns: The tokenizer specification. modelsubpath)r load_contextrRrCrErrr rCszHFMLlamaExporter.tokenizerpathc Cstjt|dd}|d}i}|jj}|jj}t|d}|D]\\}} d|vr)q |dd} d| vr[| d |ksA| d |kr[t | d D]} | | || dd t| <qHnd | vrxt | d D]} | | || dd t| <qf| || <q ||fS) a This function loads the state dict directly from a distributed checkpoint, and modify the state dict so that it is consistent with the key names you would get from loading the checkpoint into a model. This is a more memory-efficient method to obtain a state dict without initializing the nemo model. Args: path (Path): The path from which the model will be loaded. Returns ------- Tuple[Dict, Dict]: The loaded state dict and the yaml config dict. model.configr weightsT _extra_statezmodule.layersrlayers.zglobal_transformer.layers) rr rRr%r<r+ritemsreplacesizerange) rFr rDdist_ckpt_folderrTlangauge_layers vision_layersdistributed_model_weightskrnew_kirrr rs&   $   zHFMLlamaExporter.ckpt_loadcsfdd}|j}|j|j|j|j}d}i}t|D]}|dd} |dddkr| d} ||d| d||d |d<||d| d ||d |d <||d| d ||d |d <||d| d ||d |d <q|| d} ||d | d||d |d<||d | d ||d |d <||d | d ||d |d <||d | d ||d |d <q|D]\} } | || <qi}|D]\} } d| vr| |td|| <q| || <qt|}|S)a# - Modify state dict to integrate cross-attention layers into self-attention layer. e.g. 11B: 32 self-attn + 8 cross-attn -> 40 layers, 90B: 80 self-attn + 20 cross-attn -> 100 layers - Change the layer index to match the cross_attention_layers in the model config. e.g. 11B: [3, 7, 11, 15, 19, 23, 27, 31] -> [3, 8, 13, 18, 23, 28, 33, 38] Args: state_dict: Source model state dict source_config: Model config dict Returns: _ModelState: Modified state csTt|d}|d}|ddkr||}d|dStd|dd)Nr.zUnexpected layer_num: z0 (does not align with cross_attention_frequency=))intgroup ValueError)match layer_numx_num new_layer_numcross_attention_frequencyrr convert_layer_num9s  zGHFMLlamaExporter._modify_mllama_source_state..convert_layer_numzlanguage_model.decoderrrrz.xattn_layers.z!.mlp.linear_fc1.layer_norm_weightz.layers.z.mlp.linear_fc2.weightz+.cross_attention.linear_q.layer_norm_weightz,.self_attention.linear_qkv.layer_norm_weightz.mlp.linear_fc1.weight xattn_layersz \.(\d+)\.) r%r<r=rpoprresubr)rFrTrr+rtotal_num_layerprefixnew_state_dictr cross_num xattn_index attn_indexrrrYrr)r r*sX        z,HFMLlamaExporter._modify_mllama_source_stateHFMllamaConfigcCstjt|dd}|j}|j}t|j|j|j|j |j dd}ddt | |j D}t|j|j|j |j||j|j|j|j|jddd d d d gd dd }t||ddS)z Generates the configuration for the HuggingFace MLlama model based on the NeMo model. Returns: HFMllamaConfig: A configuration object for the HuggingFace MLlama model. r r r)rr9rrrrPcSsg|]\}}||qSrr).0rxrrr sz+HFMLlamaExporter.config..g @g@g?r5llama3)factorhigh_freq_factorlow_freq_factor original_max_position_embeddings rope_type)iii ) rrtie_word_embeddingsrr9rr;rr rope_scaling eos_token_idrP)rrrP)rr rRr+r%r r<r9r;r(r enumerate_init_fusion_scheduler=r r#share_embeddings_and_output_weightsr:rrr6)rFrYr+language_configrrrrrr rDwsD zHFMLlamaExporter.configN)rBr )rBr)rBr6)r-r.r/__doc__rrrGrr\rVrrCrrrrrDrrrr r2s   "Mrr rYcsNfdd}di}|D]\}}d|vr ||td||<q|||<q|S)Ncsbt|d}|dd}|dddkr%|d}d|dS||d}d|dS)Nrrrz xattn_layers.r r)r"r#)r%r&r3r(r)rr r+s    z6_rename_xattn_layer_nums_hf..convert_layer_numlanguage_modelzlayers\.(\d+)\.)rr.r/)rYr+ output_dictrrrr)r rSs   rScCstj||jddddS)Nrdim)rsplitshapearrr rsrcCs||jddS)NrreshaperOrPrrr rsrcC |ddSNrrrgaterrr rrrctxc C6|jjj}|j}|j}|j}|j}t|||||||Sr)rZrDr+r;r kv_channelsr9 _merge_qkv) rYqrrrhead_numr head_sizer9rrr r rc CrZr)rZrDr%r;rr[r9r\) rYr]rrrr^rr_r9rrr rr`rcCs4|jjj}|j}|j}|j}|j}t||||||Sr)rZrDr%r;rr[r9 _merge_kv)rYrrrr^rr_r9rrr rs rcCtj||fddS)NrrLrcat)rQbrrr rrrrrr^rr_r9c Cs|}||f|dd}|j|}|j|}tj||fdd}|jdg|ddR}|jdks8J|j|jdd|ksFJ|j|jd|ksRJ|j|jd|dks`J|j||d||g}|S)NrrLrRrr)rviewrstackrTndimrO) rrr^rr_r9old_tensor_shapenew_kv_tensor_shape kv_weightsrrr ras  rar]cCsl||}|}||f|dd} ||f|dd} |j| }|j| }|j| }g} t|D]<} | || || d|ddddf| || | dddddf| || | dddddfq1t| } | jdks}J| j| jd|d|ksJ| j| jd|ksJ| j| jd|dksJ| j| ||d||g} | S)Nrrrrf) rrgrappendrrdrirOrT)r]rrr^rr_r9heads_per_grouprjnew_q_tensor_shaperk qkv_weights_lr qkv_weightsrrr r\s&     ,$&  r\kvc Csdd|}||||g}td|d}td|d}||d|} ||d|} | | fS)NrfrrrR)rTraranger) rrr^rr_r9 kv_total_dim linear_kvk_slicev_slicek_projv_projrrr _split_kv(srzc s|||d|}||||g}tfddt|D}t|d}td|d} ||d|} ||d|} || d|} | | | fS)Nrfcs,g|]}td|d|qS)rf)rrs)r7rrnrr r9Bsz_split_qkv..rrR)rTrrdrrsr) qkvr^rr_r9 qkv_total_dim linear_qkvq_slicervrwq_projrxryrr{r _split_qkv<s   rcCrUrVrrWrrr rQrrcCs||jddddS)NrrrSrPrrr rUsrcCs4|jjj}|j}|j}|j}||}t|||||Sr)rZrDrrr9r)rYr|rr^rr9r_rrr rY rcC4|jjj}|j}|j}|j}||}t|||||Sr)rZrDrr;rr9rz)rYrrrr^rr9r_rrr rcrrcCrr)rZrDrr;rr9r)rYr|rr^rr9r_rrr rmrrcCstj|ddd\}}||fS)z\Splits NeMo's fused MLP linear_fc1 weight into gate_proj and up_proj for HuggingFace format.rfrrL)rchunk) linear_fc1 gate_projup_projrrr rwsrcCrb)z?Transforms the word embeddings from NeMo to HuggingFace format.rrLrc)word_embeddingslearnable_embeddingrrr r}rr)Hr. dataclassesrrpathlibrtypingrrrrtorch.distributedmegatron.core.transformerrr rr r6r /transformers.models.mllama.configuration_mllamar r =nemo.collections.common.tokenizers.huggingface.auto_tokenizerr1nemo.collections.common.tokenizers.tokenizer_specr&nemo.collections.vlm.mllama.model.baserrrr.nemo.export.trt_llm.nemo_ckpt_loader.nemo_filernemo.lightningrrnemo.lightning.io.staternemo.lightning.pytorch.utilsr nemo.utilsrrr2r4r?model_importerModelConnectorrAmodel_exporterrrSrrr TransformCTXrrrrr"rar\rzrrrrrrrrrrrr s             `w     "