code: def main():
      parser = ArgumentParser()
      parser.add_argument(
          "--model_file",
          type=str,
          default=None,
          required=False,
          help="Path to source .nemo file",
      )
      parser.add_argument(
          "--target_file",
          type=str,
          required=True,
          help="Path to write target .nemo file",
      )
      parser.add_argument(
          "--tensor_model_parallel_size",
          type=int,
          default=-1,
          required=False,
          help="TP size of source model",
      )
      parser.add_argument(
          "--target_tensor_model_parallel_size",
          type=int,
          required=True,
          help="TP size of target model",
      )
      parser.add_argument(
          "--pipeline_model_parallel_size",
          type=int,
          default=-1,
          required=False,
          help="PP size of source model",
      )
      parser.add_argument(
          "--target_pipeline_model_parallel_size",
          type=int,
          required=True,
          help="PP size of target model",
      )
      parser.add_argument(
          "--target_pipeline_model_parallel_split_rank",
          type=int,
          default=0,
          help="PP rank to split for Enc-Dec models",
      )
      parser.add_argument(
          "--virtual_pipeline_model_parallel_size",
          type=int,
          default=None,
          help="Virtual Pipeline parallelism size",
      )
      parser.add_argument(
          "--ckpt_name",
          type=str,
          default=None,
          help="Checkpoint name to load from for Virtual Parallel",
      )
      parser.add_argument(
          "--model_class",
          type=str,
          default="nemo.collections.nlp.models.language_modeling.megatron_gpt_model.MegatronGPTModel",
          help="NeMo model class. This script should support all NeMo megatron models that use Tensor Parallel",
      )
      parser.add_argument(
          "--precision",
          default=16,
          help="PyTorch Lightning Trainer precision flag",
      )
      parser.add_argument(
          "--num_gpu_per_node",
          default=8,
          type=int,
          help="Number of GPUs per node",
      )
      parser.add_argument(
          "--megatron_legacy",
          action="store_true",
          help="Converter for legacy megatron modles that have different q,k,v weight splits",
      )
      parser.add_argument(
          "--tokenizer_model_path",
          type=str,
          required=False,
          default=None,
          help="Path to the tokenizer model path if your model uses a tokenizer model as an artifact. This is needed if your model uses a sentencepiece tokenizer.",
      )
      parser.add_argument(
          "--tokenizer_vocab_file",
          type=str,
          required=False,
          default=None,
          help="Path to the tokenizer model path if your model uses a tokenizer model as an artifact. This is needed if your model uses a sentencepiece tokenizer.",
      )
      parser.add_argument(
          "--hparams_file",
          type=str,
          default=None,
          help="Path to hparams file from PTL training",
      )
      parser.add_argument(
          "--tp_conversion_only",
          action="store_true",
          help="Only convert TP model to TP model",
      )
      parser.add_argument(
          "--model_extracted_dir",
          type=str,
          default=None,
          help="Path to pre-extracted model directory",
      )
  
      args = parser.parse_args()
  
      precision = args.precision
      num_gpu_per_node = int(args.num_gpu_per_node)
      if args.precision in ["32", "16"]:
          precision = int(float(args.precision))
  
      if precision in ["bf16", "bf16-mixed"]:
          if (
              torch.cuda.is_available()
              and torch.cuda.is_bf16_supported()
          ):
              pass
          else:
              logging.warning(
                  "BF16 is not supported on this device. Using FP16 instead."
              )
              precision = precision[2:]
  
      if precision == 32:
          dtype = torch.float32
      elif precision in [16, "16", "16-mixed"]:
          dtype = torch.float16
      elif precision in ["bf16", "bf16-mixed"]:
          dtype = torch.bfloat16
      else:
          dtype = torch.float32  # fallback
  
      # Built target directory if it does not exist
      target_dir = os.path.split(args.target_file)[0]
      if not os.path.exists(target_dir):
          os.makedirs(target_dir, exist_ok=True)
  
      tp_size = args.tensor_model_parallel_size
      tgt_tp_size = args.target_tensor_model_parallel_size
      pp_size = args.pipeline_model_parallel_size
      tgt_pp_size = args.target_pipeline_model_parallel_size
      pipeline_model_parallel_split_rank = (
          args.target_pipeline_model_parallel_split_rank
      )
      vp_size = args.virtual_pipeline_model_parallel_size
      if vp_size is None:
          vp_size = 1
  
      convert_vp = vp_size > 1
      if convert_vp:
          from megatron.core import parallel_state
  
          parallel_state.set_virtual_pipeline_model_parallel_world_size(
              vp_size
          )
  
          hparams_filepath = args.hparams_file
          if hparams_filepath is None:
              logging.warning(
                  "\n\n\n!!!!!!!!!\n"
                  "You are converting a model with virtual pipeline parallelism enabled, \n"
                  "but have not passed `hparams_file` argument. \n"
                  "This will cause each ckpt file to be temporarily laoded onto GPU memory!\n\n"
                  "It is highly recommended to pass `hparams_file` argument to avoid this.\n"
              )
      else:
          hparams_filepath = None
  
      # Import the class of the model
      cls = model_utils.import_class_by_path(args.model_class)
  
      if args.model_file is None and args.model_extracted_dir is None:
          raise ValueError(
              "Cannot pass model_file and model_extracted_dir as None at the same time."
          )
  
      tmp_cfg = cls.restore_from(
          restore_path=args.model_file,
          trainer=Trainer(
              devices=1,
              strategy=NLPDDPStrategy(),
              accelerator="cpu",
              precision=precision,
          ),
          map_location=torch.device("cpu"),
          return_config=True,
      )
      plugins = []
      if precision in [16, "16", "bf16", "16-mixed", "bf16-mixed"]:
          scaler = None
          if precision in [16, "16", "16-mixed"]:
              scaler = GradScaler(
                  init_scale=tmp_cfg.get(
                      "native_amp_init_scale", 2**32
                  ),
                  growth_interval=tmp_cfg.get(
                      "native_amp_growth_interval", 1000
                  ),
                  hysteresis=tmp_cfg.get("hysteresis", 2),
              )
              # MixedPrecisionPlugin in PTL >= 2.0 requires precision to be 16-mixed or bf16-mixed
              plugin_precision = "16-mixed"
          else:
              plugin_precision = "bf16-mixed"
  
          if tmp_cfg.get("megatron_amp_O2", False):
              plugins.append(
                  MegatronHalfPrecisionPlugin(
                      precision=plugin_precision,
                      device="cuda",
                      scaler=scaler,
                  )
              )
          else:
              plugins.append(
                  PipelineMixedPrecisionPlugin(
                      precision=plugin_precision,
                      device="cuda",
                      scaler=scaler,
                  )
              )
      trainer = Trainer(
          plugins=plugins,
          devices=1,
          strategy=NLPDDPStrategy(),
          accelerator="cpu",
          precision=precision,
      )
  
      if tp_size < 0 or pp_size < 0:
          logging.info(
              f"Loading model config from {args.model_file} to get TP and PP size"
          )
          model_config_internal = cls.restore_from(
              restore_path=args.model_file,
              trainer=trainer,
              map_location=torch.device("cpu"),
              return_config=True,
          )
  
          tp_size = model_config_internal.get(
              "tensor_model_parallel_size", 1
          )
          pp_size = model_config_internal.get(
              "pipeline_model_parallel_size", 1
          )
  
      # Check if TP conversion only
      tp_conversion_only = args.tp_conversion_only
      if tp_conversion_only:
          logging.info("Converting TP model to TP model only")
  
          if pp_size > 1:
              raise ValueError(
                  "Provided `--tp_conversion_only` but `--pipeline_model_parallel_size` > 1"
              )
  
          if tgt_pp_size > 1:
              raise ValueError(
                  "Provided `--tp_conversion_only` but `--target_pipeline_model_parallel_size` > 1"
              )
  
          if pipeline_model_parallel_split_rank > 0:
              raise ValueError(
                  "Provided `--tp_conversion_only` but `--target_pipeline_model_parallel_split_rank` > 0"
              )
  
          # Force PP size to 1
          pp_size = 1
          tgt_pp_size = 1
          pipeline_model_parallel_split_rank = 0
  
      if vp_size is None or vp_size < 0:
          vp_size = 1
  
      app_state = AppState()
      app_state.data_parallel_rank = 0
      app_state.pipeline_model_parallel_size = pp_size
      app_state.tensor_model_parallel_size = tp_size
  
      if vp_size > 1:
          app_state.virtual_pipeline_model_parallel_size = vp_size
      app_state.model_parallel_size = (
          app_state.pipeline_model_parallel_size
          * app_state.tensor_model_parallel_size
      )
  
      world_size = (
          pp_size * tp_size
      )  # pseudo world size for simulating load of a specific rank on a single gpu
  
      app_state.tensor_model_parallel_rank = 0
      app_state.pipeline_model_parallel_rank = 0
  
      if vp_size > 1:
          set_virtual_parallel_rank_safely(0)
  
      # Extract tokenizer artifact from the model to temp directory
      logging.info("Extracting tokenizer artifact from NeMo file...")
      temp_dir = tempfile.mkdtemp()
      tokenizer_model_path = None
      with tarfile.open(args.model_file, "r") as tar:
          for member in tar.getmembers():
              if ".model" in member.name:
                  extracted_file = tar.extractfile(member)
                  extracted_file_path = os.path.join(
                      temp_dir, member.name
                  )
  
                  if tokenizer_model_path is None:
                      logging.info(
                          f"Found tokenizer. Extracting {member.name} to {extracted_file_path}"
                      )
  
                      tokenizer_model_path = extracted_file_path
                      with open(extracted_file_path, "wb") as f:
                          f.write(extracted_file.read())
                  else:
                      if args.tokenizer_model_path is None:
                          logging.warning(
                              f"\n\nFound multiple tokenizer artifacts in the model file.\n"
                              f"Using only {tokenizer_model_path}.\n"
                              f"If this is incorrect, manually pass the correct tokenizer using "
                              f"`--tokenizer_model_path`.\n\n"
                          )
  
      # If input model has TP > 1 or PP > 1
      # Reconstruct the model to have TP = 1 and PP = 1
      # Note that this is a forward loop that will process PP [0..N] TP [0..M] in sequential order.
      if tp_size > 1 or pp_size > 1:
          partitions = {}  # 3d list of VP x PP x TP
          model = None
  
          # Build partitions structure
          for vp_idx in range(vp_size):
              partitions[vp_idx] = []  # Build first layer - VP
  
              for pp_idx in range(pp_size):
                  # For each VP, build PP x TP holder
                  partitions[vp_idx].append({})
                  partitions[vp_idx][pp_idx] = []
  
          for vp_rank in range(vp_size):
              if vp_size > 1:
                  set_virtual_parallel_rank_safely(vp_rank)
  
              for pp_rank in range(pp_size):
                  app_state.pipeline_model_parallel_rank = pp_rank
  
                  for tp_rank in range(tp_size):
                      app_state.tensor_model_parallel_rank = tp_rank
  
                      logging.info(
                          f"Loading ------------ PP Rank: {pp_rank} TP Rank: {tp_rank}"
                      )
  
                      # Override flag that forces Model to use AppState instead of Trainer
                      # to determine the world size, global and local rank
                      # Used for simulating load of a specific rank on a single gpu
                      os.environ[
                          NEMO_MEGATRON_MODEL_PARALLEL_APPSTATE_OVERRIDE
                      ] = "true"
  
                      # Compute the global rank to load the correct subset of parameters
                      global_rank = pp_rank * tp_size + tp_rank
  
                      # Update AppState
                      app_state.world_size = world_size
                      app_state.global_rank = global_rank
                      app_state.local_rank = (
                          global_rank % num_gpu_per_node
                      )
                      app_state.pipeline_model_parallel_size = pp_size
                      app_state.tensor_model_parallel_size = tp_size
                      app_state.pipeline_model_parallel_split_rank = (
                          pipeline_model_parallel_split_rank
                      )
                      app_state.model_parallel_size = (
                          app_state.pipeline_model_parallel_size
                          * app_state.tensor_model_parallel_size
                      )
  
                      if vp_size > 1:
                          set_virtual_parallel_rank_safely(vp_rank)
  
                      if vp_rank == 0:
                          save_restore_connector = (
                              NLPSaveRestoreConnector()
                          )
  
                          if args.model_extracted_dir is not None:
                              logging.info(
                                  f"Using extracted model directory: {args.model_extracted_dir}"
                              )
                              save_restore_connector.model_extracted_dir = (
                                  args.model_extracted_dir
                              )
  
                          if args.model_file is not None:
                              model_filepath = args.model_file
                          else:
                              model_filepath = args.model_extracted_dir
  
                          if vp_size == 1:
                              # Get model config
                              tmp_cfg = cls.restore_from(
                                  restore_path=model_filepath,
                                  trainer=trainer,
                                  map_location=torch.device("cpu"),
                                  save_restore_connector=save_restore_connector,
                                  return_config=True,
                              )
  
                              # Force model onto CPU
                              tmp_cfg, restore_dict = force_cpu_model(
                                  tmp_cfg
                              )
  
                              # Restore model
                              model = cls.restore_from(
                                  restore_path=model_filepath,
                                  trainer=trainer,
                                  map_location=torch.device("cpu"),
                                  save_restore_connector=save_restore_connector,
                                  override_config_path=tmp_cfg,
                              )
                              model.freeze()
  
                              # Restore model config
                              restore_model_config(
                                  model.cfg, restore_dict
                              )
  
                          else:
                              if args.ckpt_name is None:
                                  raise ValueError(
                                      "For Virtual Parallel, ckpt name is required.\n"
                                      "Please provide `--ckpt_name` argument."
                                  )
  
                              # inject model parallel rank
                              checkpoint_path = model_utils.inject_model_parallel_rank(
                                  os.path.join(
                                      model_filepath, args.ckpt_name
                                  )
                              )
  
                              vp_state_dict = torch.load(
                                  checkpoint_path, map_location="cpu"
                              )
  
                              if hparams_filepath is not None:
                                  # Force the model onto CPU
                                  tmp_cfg = OmegaConf.load(
                                      hparams_filepath
                                  )
                                  (
                                      tmp_cfg,
                                      restore_dict,
                                  ) = force_cpu_model(tmp_cfg)
  
                                  with tempfile.NamedTemporaryFile(
                                      mode="w",
                                      encoding="utf-8",
                                      suffix=".yml",
                                  ) as tmp:
                                      OmegaConf.save(
                                          tmp_cfg, tmp, resolve=True
                                      )
                                      tmp.seek(0)
  
                                      model = cls.load_from_checkpoint(
                                          checkpoint_path=checkpoint_path,
                                          trainer=trainer,
                                          map_location=torch.device(
                                              "cpu"
                                          ),
                                          hparams_file=tmp.name,
                                      )
                                      model.freeze()
  
                                      restore_model_config(
                                          model.cfg, restore_dict
                                      )
  
                              else:
                                  model = cls.load_from_checkpoint(
                                      checkpoint_path=checkpoint_path,
                                      trainer=trainer,
                                      map_location=torch.device("cpu"),
                                  )
                                  model.freeze()
  
                          model.to(dtype=dtype)
  
                          # Reset env flag
                          os.environ.pop(
                              NEMO_MEGATRON_MODEL_PARALLEL_APPSTATE_OVERRIDE,
                              None,
                          )
  
                          logging.info(
                              f"<<<<<<<< LOADED MODEL PP={pp_rank + 1} TP={tp_rank + 1} | "
                              f"GLOBAL RANK = {global_rank} >>>>>>>>>"
                          )
  
                          # Save the parameters
                          if vp_size == 1:
                              params = [p for p in model.parameters()]
                              partitions[vp_rank][pp_rank].append(
                                  params
                              )  # vp_rank = 0
  
                          else:
                              vp_params_tmp = []
                              for vp_idx in range(vp_size):
                                  set_virtual_parallel_rank_safely(
                                      vp_idx
                                  )
                                  vp_params = vp_state_dict[
                                      f"model{vp_idx}"
                                  ]
                                  model.model[
                                      vp_idx
                                  ].module.load_state_dict(
                                      vp_params, strict=True
                                  )
                                  model.model[vp_idx].module.to("cpu")
                                  params = [
                                      p
                                      for p in model.model[
                                          vp_idx
                                      ].module.parameters()
                                  ]
                                  vp_params_tmp.append(params)
                                  # partitions[pp_rank][vp_idx].append(params)
  
                              for vp_idx in range(vp_size):
                                  partitions[vp_idx][pp_rank].append(
                                      vp_params_tmp[vp_idx]
                                  )
  
                              del vp_params_tmp
                              set_virtual_parallel_rank_safely(0)
  
                          # app_state is being updated incorrectly during restore
                          app_state.data_parallel_rank = 0
                          app_state.pipeline_model_parallel_rank = (
                              pp_rank
                          )
                          app_state.tensor_model_parallel_rank = tp_rank
                          app_state.pipeline_model_parallel_size = (
                              pp_size
                          )
                          app_state.tensor_model_parallel_size = tp_size
                          app_state.model_parallel_size = (
                              app_state.pipeline_model_parallel_size
                              * app_state.tensor_model_parallel_size
                          )
  
                          if vp_size > 1:
                              app_state.virtual_pipeline_model_parallel_size = (
                                  vp_size
                              )
                              set_virtual_parallel_rank_safely(vp_rank)
  
          # Build a unified model with PP 1 TP 1
          with open_dict(model.cfg):
              model.cfg.tensor_model_parallel_size = 1
              model.cfg.pipeline_model_parallel_size = 1
              model.cfg.virtual_pipeline_model_parallel_size = None
  
          app_state.global_rank = 0
          app_state.local_rank = 0
          app_state.data_parallel_rank = 0
          app_state.pipeline_model_parallel_rank = 0
          app_state.tensor_model_parallel_rank = 0
          app_state.pipeline_model_parallel_size = 1
          app_state.tensor_model_parallel_size = 1
          app_state.model_parallel_size = 1
  
          if vp_size > 1:
              set_virtual_parallel_rank_safely(None)
  
          trainer = Trainer(
              plugins=plugins,
              devices=1,
              strategy=NLPDDPStrategy(),
              accelerator="cpu",
              precision=precision,
          )
  
          with open_dict(model.cfg):
              if args.tokenizer_model_path is not None:
                  model.cfg.tokenizer.model = args.tokenizer_model_path
              if args.tokenizer_vocab_file is not None:
                  model.cfg.tokenizer.vocab_file = (
                      args.tokenizer_vocab_file
                  )
  
              model.cfg, restore_dict = force_cpu_model(model.cfg)
  
              # Remove Virtual Parallelism
              model.cfg.virtual_pipeline_model_parallel_size = None
  
          logging.info(
              f"<<<<<<<< Building TP 1 PP 1 base model >>>>>>>>>"
          )
          model = cls(model.cfg, trainer)  # type: nn.Module
          model.freeze()
          model = model.to("cpu")
          model._save_restore_connector = NLPSaveRestoreConnector()
  
          restore_model_config(model.cfg, restore_dict)
  
          vp_param_count = 0
          for vp in range(vp_size):
              for pp in range(pp_size):
                  for tp in range(tp_size):
                      vp_param_count += len(partitions[vp][pp][tp])
  
          if vp_size > 1:
              logging.debug(
                  f"Total params in TP PP VP = 1 : {len(list(model.parameters()))}"
              )
              logging.debug(
                  f"Total params in VP PP TP (og): {vp_param_count}"
              )
  
          # Flatten Virtual Pipeline
          if vp_size == 1:
              # unpack vp container, pack pp tp container
              partitions = partitions[0]
              partitions = {
                  idx: val for idx, val in enumerate(partitions)
              }
          else:
              flat_partitions = {idx: [] for idx in range(pp_size)}
  
              """
              Under VP convention
              Notation :
              Stage  = PP rank
              Number = GPT model / layer index
              Ignore TP - every PP has all TP corresponding to that PP
              chunk_index = the physical index of any [] in the list. Ex idx = 2 in below map corresponds to [2: PP 0 VP 1]]
  
  
              For a PP 2 VP 4 model with 8 GPT layers-
  
              Indices
              # Stage 0: [0:PP 0 VP 0]  [2:PP 0 VP 1]  [4:PP 0 VP 2]  [6:PP 0 VP 3]
              # Stage 1: [1:PP 1 VP 0]  [3:PP 1 VP 1]  [5:PP 1 VP 2]  [7:PP 1 VP 3]
  
              after conversion will become
  
              # Stage 0: [0,1,2,3:PP 0]
              # Stage 1: [4,5,6,7:PP 1]
  
              """
              pp_index = 0
              chunk_counter = 0
              tp_cache = [[] for _ in range(tp_size)]
  
              for vp in range(vp_size):
                  for pp in range(pp_size):
                      # Gather all TP under this VP PP combination.
                      # We will accumulate TP parameters from multiple layers in this cache.
                      for tp in range(tp_size):
                          tp_cache[tp].extend(partitions[vp][pp][tp])
  
                      # This counter indexes the global selection of a VP PP combination in the above map
                      chunk_counter += 1
  
                      # Log the mapping from old VP x PP to new PP index
                      logging.info(
                          f"VP Conversion - vp: {vp} pp: {pp} -> pp_idx: {pp_index}"
                      )
  
                      # Every vp_size chunks, we can fill a new PP index in the flat_partitions
                      if chunk_counter % vp_size == 0:
                          flat_partitions[pp_index].extend(tp_cache)
                          tp_cache = [[] for _ in range(tp_size)]
                          pp_index += 1
  
                          logging.debug(
                              f"VP merge step: \n"
                              f"vp: {vp} pp: {pp} pp_idx: {pp_index - 1} "
                              f"len(flat_partitions): {len(flat_partitions[pp_index - 1])}"
                          )
  
              logging.debug(
                  f"PP Size len(flat partitions) : {len(flat_partitions)}"
              )
              logging.debug(
                  f"TP Size len(flat partitions[0]): {len(flat_partitions[0])}"
              )
              logging.debug(
                  f"Layers  len(flat partitions[0][0]) : {len(flat_partitions[0][0])}"
              )
  
              partitions = flat_partitions
              del tp_cache
  
          if tgt_tp_size > 1 or tgt_pp_size > 1:
              merge_partition(model, partitions)
          else:
              # Write out the PP 1 TP 1 model to disk
              merge_partition(model, partitions, args.target_file)
  
          # Empty cache memory of all parameters from all PP TP partitions
          partitions.clear()
  
      else:
          # If input model has TP = 1 and PP = 1
          app_state.model_parallel_size = 1
  
          save_restore_connector = NLPSaveRestoreConnector()
  
          if args.model_extracted_dir is not None:
              logging.info(
                  f"Using extracted model directory: {args.model_extracted_dir}"
              )
              save_restore_connector.model_extracted_dir = (
                  args.model_extracted_dir
              )
  
          if args.model_file is not None:
              model_filepath = args.model_file
          else:
              model_filepath = args.model_extracted_dir
  
          tmp_cfg = cls.restore_from(
              restore_path=model_filepath,
              trainer=trainer,
              map_location=torch.device("cpu"),
              save_restore_connector=save_restore_connector,
              return_config=True,
          )
  
          tmp_cfg, restore_dict = force_cpu_model(tmp_cfg)
  
          model = cls.restore_from(
              restore_path=model_filepath,
              trainer=trainer,
              map_location=torch.device("cpu"),
              save_restore_connector=save_restore_connector,
              override_config_path=tmp_cfg,
          )
          model.to(dtype=dtype)
  
          restore_model_config(model.cfg, restore_dict)
  
      # If target model has TP > 1 or PP > 1
      if tgt_pp_size > 1 or tgt_tp_size > 1:
          # Preserve the TP 1 PP 1 model parameters and names
          global_params = []
          global_params.append(
              [p for n, p in model.named_parameters()]
          )  # params
          global_params.append(
              [n for n, p in model.named_parameters()]
          )  # names
  
          logging.debug("Global parameters:")
          for idx, (name, p) in enumerate(
              zip(global_params[1], global_params[0])
          ):
              logging.debug(f"{name} - {p.shape}")
  
          logging.info(
              f"TP 1 PP 1 Number of Parameters : {len(global_params[0])}"
          )
  
          world_size = (
              tgt_pp_size * tgt_tp_size
          )  # pseudo world size for simulating load of a specific rank on a single gpu
          new_global_batch_size = (
              model.cfg.micro_batch_size * world_size
          )
          old_global_batch_size = model.cfg.get(
              "global_batch_size", model.cfg.micro_batch_size
          )
  
          global_offset = (
              len(global_params[0]) - 1
          )  # -1 cause this indexes the array, range [0, L-1]
          logging.info(
              f"Final layer offset for parameters: {global_offset}"
          )
  
          for pp_rank in range(
              tgt_pp_size - 1, -1, -1
          ):  # reverse order
              with open_dict(model.cfg):
                  model.cfg.pipeline_model_parallel_size = tgt_pp_size
                  model.cfg.tensor_model_parallel_size = tgt_tp_size
  
                  if "pipeline_model_parallel_split_rank" in model.cfg:
                      if pipeline_model_parallel_split_rank > 0:
                          model.cfg.pipeline_model_parallel_split_rank = (
                              pipeline_model_parallel_split_rank
                          )
                      elif pp_size > 1:
                          logging.warning(
                              f"Model config has `pipeline_model_parallel_split_rank` set to "
                              f"{model.cfg.pipeline_model_parallel_split_rank} and target PP "
                              f"size is {tgt_pp_size}. "
                              f"Provided `pipeline_model_parallel_split_rank` is "
                              f"{pipeline_model_parallel_split_rank}. "
                              f"Be careful that the model config is correct "
                              f"if encoder-decoder models are being converted."
                          )
  
                  model.cfg.global_batch_size = (
                      old_global_batch_size  # Used for restoration
                  )
  
              # Override flag that forces Model to use AppState instead of Trainer
              # to determine the world size, global and local rank
              # Used for simulating load of a specific rank on a single gpu
              os.environ[
                  NEMO_MEGATRON_MODEL_PARALLEL_APPSTATE_OVERRIDE
              ] = "true"
  
              # Compute the global rank
              global_rank = (
                  pp_rank * tgt_tp_size + 0
              )  # tp_rank = 0 needed just for modules, all TP will be merged to this PP rank
  
              # Update AppState
              app_state.world_size = world_size
              app_state.global_rank = global_rank
              app_state.local_rank = global_rank % num_gpu_per_node
              app_state.pipeline_model_parallel_size = tgt_pp_size
              app_state.tensor_model_parallel_size = tgt_tp_size
              app_state.model_parallel_size = (
                  app_state.pipeline_model_parallel_size
                  * app_state.tensor_model_parallel_size
              )
  
              trainer = Trainer(
                  plugins=plugins,
                  devices=1,
                  strategy=NLPDDPStrategy(),
                  accelerator="cpu",
                  precision=precision,
              )
              if args.tokenizer_model_path is not None:
                  with open_dict(model.cfg):
                      model.cfg.tokenizer.model = (
                          args.tokenizer_model_path
                      )
  
              else:
                  if tokenizer_model_path is None:
                      logging.warning(
                          "Could not extract tokenizer model file from checkpoint."
                      )
  
                  else:
                      # Extract tokenizer info
                      with open_dict(model.cfg):
                          model.cfg.tokenizer.model = (
                              tokenizer_model_path
                          )
  
              model.cfg, restore_dict = force_cpu_model(model.cfg)
  
              model = cls(model.cfg, trainer)
              model = model.to("cpu")
              model._save_restore_connector = NLPSaveRestoreConnector()
              model.freeze()
              model.to(dtype=dtype)
  
              restore_model_config(model.cfg, restore_dict)
  
              # Update global batch size
              if (
                  old_global_batch_size % new_global_batch_size != 0
                  or old_global_batch_size < new_global_batch_size
              ):
                  logging.info(
                      f"Global batch size {old_global_batch_size} is not divisible by new global batch size {new_global_batch_size}."
                      f" The model config will be updated with new global batch size {new_global_batch_size}."
                  )
                  with open_dict(model.cfg):
                      model.cfg.global_batch_size = (
                          new_global_batch_size
                      )
  
              logging.info(
                  f"Global rank: {global_rank} Local rank: {app_state.local_rank} World size: {world_size}"
              )
              logging.info(f"PP rank: {pp_rank} TP rank: {0}")
              logging.info(
                  f"TP 1 PP 1 Number of Layers : {len(global_params[0])}"
              )
              logging.info(
                  f"Remaining layer offset for parameters: {global_offset}"
              )
              logging.info("\n")
  
              # Special case for TP conversion only mode
              if tp_conversion_only:
                  logging.info(
                      f"Skipping PP split due to flag `--tp_conversion_only`"
                  )
  
                  split_tp_partition_only(
                      model,
                      global_params,
                      tgt_tp_size,
                      args.target_file,
                      args.megatron_legacy,
                  )
                  break
  
              global_offset = split_partition(
                  model,
                  global_params,
                  tgt_pp_size,
                  tgt_tp_size,
                  pp_rank,
                  global_offset,
                  pipeline_model_parallel_split_rank,
                  args.target_file,
                  args.megatron_legacy,
              )
  
              # Reset env flag
              os.environ.pop(
                  NEMO_MEGATRON_MODEL_PARALLEL_APPSTATE_OVERRIDE, None
              )
  
          # Check if invalid global offset - after all PP splits, global offset should be -1
          if global_offset < -1 and not tp_conversion_only:
              raise ValueError(
                  f"Invalid global offset {global_offset} found for global rank {app_state.global_rank} "
                  f"and local rank {app_state.local_rank}. Should be -1 if all parameters have been assigned. "
                  f"Currently, seems some parameters were duplicated."
              )
          elif global_offset > -1 and not tp_conversion_only:
              logging.error("\n")
              logging.error("!" * 80)
              logging.error(
                  "Error: Some parameters were not correctly added to model partitions."
              )
              logging.error(
                  "Below is list of parameters skipped in reverse order: "
              )
  
              for param_id in range(global_offset, -1, -1):
                  logging.error(
                      f"Param ID: {param_id} : {global_params[1][param_id]} {global_params[0][param_id].shape}"
                  )
              logging.error("!" * 80)
  
              raise ValueError(
                  f"Invalid global offset {global_offset} found for global rank {app_state.global_rank} "
                  f"and local rank {app_state.local_rank}. Should be -1 if all parameters have been assigned. "
                  f"Currently, seems some parameters were not assigned."
              )
  
      logging.info("Successfully finished changing partitions!")
  
      if temp_dir is not None:
          shutil.rmtree(temp_dir, ignore_errors=True)
  
imports: import os
  import shutil
  import tarfile
  import tempfile
  from argparse import ArgumentParser
  from typing import Dict, List
  import torch
  import torch.nn
  from omegaconf import OmegaConf, open_dict
  from pytorch_lightning import Trainer
  from nemo.collections.nlp.parts.nlp_overrides import NEMO_MEGATRON_MODEL_PARALLEL_APPSTATE_OVERRIDE, GradScaler, MegatronHalfPrecisionPlugin, NLPDDPStrategy, NLPSaveRestoreConnector, PipelineMixedPrecisionPlugin
  from nemo.utils import logging, model_utils
  from nemo.utils.app_state import AppState
label_desc: Use image node_id 11 for calling a global function with example usage: main() without return types