code: def merge(
      base_model_state_dict: Dict[str, Any],
      lora_state_dict: Dict[int, Any],
      tp: int,
      num_layers: int,
      curr_rank: int,
      mcore: bool,
  ):
      """
      Iterate through all the self_attention.query_key_value projection feedforward weights in all the layers.
      Collect the corresponding lora weights for each layer and across tp ranks.
      Computes the "full rank" weight from the two low-rank weights and add it to the self_attention.query_key_value weight.
      Args:
          base_model_state_dict: A state_dict for the base model for the current rank.
          lora_state_dict: A complete set of weights for the lora model across all tp ranks. They key for this dict is an int tp rank.
          tp: the tensor_model_parallel_size for the base_model (and the lora model)
          num_layers: the number of layers in the base_model to iterate over.
          curr_rank: current tp rank of the base model which is being merged with Lora.
          mcore: whether the model uses megatron core.
      """
  
      for nl in range(num_layers):
          if mcore:
              key_self_attn_kqv = f"model.decoder.layers.{nl}.self_attention.linear_qkv.weight"
              key_lora_in = f"model.decoder.layers.{nl}.self_attention.adapter_layer.lora_kqv_adapter.linear_in.weight"
              key_lora_out = f"model.decoder.layers.{nl}.self_attention.adapter_layer.lora_kqv_adapter.linear_out.weight"
          else:
              key_self_attn_kqv = f"model.language_model.encoder.layers.{nl}.self_attention.query_key_value.weight"
              key_lora_in = f"model.language_model.encoder.layers.{nl}.self_attention.adapter_layer.lora_kqv_adapter.linear_in.weight"
              key_lora_out = f"model.language_model.encoder.layers.{nl}.self_attention.adapter_layer.lora_kqv_adapter.linear_out.weight"
          wt_lora_in = torch.cat(
              [lora_state_dict[_tp][key_lora_in] for _tp in range(tp)],
              dim=0,
          )
          wt_lora_out = lora_state_dict[curr_rank][key_lora_out]
          wt_self_attn = base_model_state_dict[key_self_attn_kqv]
          wt_lora = wt_lora_out @ wt_lora_in
          base_model_state_dict[
              key_self_attn_kqv
          ] = wt_self_attn + wt_lora.type_as(wt_self_attn)
          print("mergeing for weight", key_self_attn_kqv)
      return base_model_state_dict
  
imports: import os
  import tempfile
  from typing import Any, Dict
  import torch
  from omegaconf import OmegaConf, open_dict
  from pytorch_lightning.trainer.trainer import Trainer
  from torch.utils.data import DataLoader, Dataset
  from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
  from nemo.collections.nlp.models.language_modeling.megatron_gpt_sft_model import MegatronGPTSFTModel
  from nemo.collections.nlp.modules.common.megatron.megatron_init import fake_initialize_model_parallel
  from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy, NLPSaveRestoreConnector
  from nemo.core.config import hydra_runner
  from nemo.utils import logging
  from nemo.utils.app_state import AppState
  from nemo.utils.model_utils import inject_model_parallel_rank
label_desc: Use image node_id 3 for calling a global function with example usage: merge(base_model_state_dict, lora_state_dict, tp, num_layers, curr_rank, mcore) and returns: base_model_state_dict