code: def load_from_checkpoint(
      cls,
      checkpoint_path: str,
      map_location: Any = None,
      hparams_file: Optional[str] = None,
      strict: bool = True,
      **kwargs,
  ):
      """
      Loads Megatron_LM checkpoints, convert it, with some maintenance of restoration.
      For documentation, please refer to LightningModule.load_from_checkpoin() documentation.
      """
      checkpoint = None
      try:
          cls._set_model_restore_state(is_being_restored=True)
          # TODO: replace with proper PTL API
  
          with pl_legacy_patch():
              if map_location is not None:
                  old_checkpoint = pl_load(
                      checkpoint_path, map_location=map_location
                  )
              else:
                  old_checkpoint = pl_load(
                      checkpoint_path,
                      map_location=lambda storage, loc: storage,
                  )
  
          total_params = [0]
          checkpoint = OrderedDict()
          checkpoint["state_dict"] = OrderedDict()
          parse_weights(
              old_checkpoint["model"],
              "",
              total_params,
              checkpoint["state_dict"],
              translator=kwargs["translator"],
          )
          print(
              "converted {:.2f}M parameters".format(
                  total_params[0] / 1e6
              )
          )
  
          if hparams_file is not None:
              extension = hparams_file.split(".")[-1]
              if extension.lower() == "csv":
                  hparams = load_hparams_from_tags_csv(hparams_file)
              elif extension.lower() in ("yml", "yaml"):
                  hparams = load_hparams_from_yaml(hparams_file)
              else:
                  raise ValueError(
                      ".csv, .yml or .yaml is required for `hparams_file`"
                  )
  
              hparams["on_gpu"] = False
  
              # overwrite hparams by the given file
              checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY] = hparams
  
          check_point_version = old_checkpoint.get(
              "checkpoint_version", 0
          )
          if check_point_version < 3:
              # need to do the transpose of query_key_value variables
              if hparams_file is not None:
                  np = hparams["cfg"]["num_attention_heads"]
              elif (
                  "config" in old_checkpoint
                  and "num-attention-heads" in old_checkpoint["config"]
              ):
                  np = old_checkpoint["config"]["num-attention-heads"]
  
              else:
                  logging.warning(
                      "cannot determine the number attention heads"
                  )
                  raise ValueError(
                      "need to know number of attention heads"
                  )
  
              if check_point_version == 0:
                  # 3, np, hn -> np, 3, hn
                  for key in checkpoint["state_dict"]:
                      if key.find("query_key_value") >= 0:
                          weight = checkpoint["state_dict"][key]
                          if len(weight.size()) == 2:
                              # weight
                              weight = weight.view(
                                  3, np, -1, weight.size()[-1]
                              )
                              weight = weight.transpose(
                                  0, 1
                              ).contiguous()
                              checkpoint["state_dict"][
                                  key
                              ] = weight.view(-1, weight.size()[-1])
                          else:
                              # biase
                              weight = weight.view(3, np, -1)
                              weight = weight.transpose(
                                  0, 1
                              ).contiguous()
                              checkpoint["state_dict"][
                                  key
                              ] = weight.view(-1)
              elif check_point_version == 1:
                  # np, hn, 3 -> np, 3, hn
                  for key in checkpoint["state_dict"]:
                      if key.find("query_key_value") >= 0:
                          weight = checkpoint["state_dict"][key]
                          if len(weight.size()) == 2:
                              # weight
                              weight = weight.view(
                                  np, -1, 3, weight.size()[-1]
                              )
                              weight = weight.transpose(
                                  1, 2
                              ).contiguous()
                              checkpoint["state_dict"][key] = weight
                          else:
                              # biase
                              weight = weight.view(np, -1, 3)
                              weight = weight.transpose(
                                  1, 2
                              ).contiguous()
                              checkpoint["state_dict"][key] = weight
  
          # for past checkpoint need to add the new key
          if cls.CHECKPOINT_HYPER_PARAMS_KEY not in checkpoint:
              checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY] = {}
          # override the hparams with values that were passed in
          # TODO: can we do this without overriding?
          config_kwargs = kwargs.copy()
          if "trainer" in config_kwargs:
              config_kwargs.pop("trainer")
          checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY].update(
              config_kwargs
          )
          add_optimizer_state(old_checkpoint, checkpoint)
          consumed = None
          if "args" in old_checkpoint and hasattr(
              old_checkpoint["args"], "consumed_train_samples"
          ):
              consumed = getattr(
                  old_checkpoint["args"], "consumed_train_samples"
              )
          steps = None
          if "iteration" in old_checkpoint:
              steps = old_checkpoint["iteration"]
      finally:
          cls._set_model_restore_state(is_being_restored=False)
      logging.warning(
          f"the checkpoint version is {check_point_version}"
      )
      return checkpoint, consumed, steps, check_point_version
  
imports: import importlib
  import os
  import pathlib
  import sys
  from argparse import ArgumentParser
  from collections import OrderedDict
  from typing import Any, Optional
  import torch
  from lightning_fabric.utilities.cloud_io import _load
  from megatron.core import parallel_state
  from pytorch_lightning.core.saving import _load_state
  from pytorch_lightning.core.saving import load_hparams_from_tags_csv, load_hparams_from_yaml
  from pytorch_lightning.trainer.trainer import Trainer
  from pytorch_lightning.utilities.migration import pl_legacy_patch
  from nemo.collections.nlp.models.language_modeling.megatron_bert_model import MegatronBertModel
  from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
  from nemo.collections.nlp.modules.common.megatron.megatron_init import initialize_model_parallel_for_nemo
  from nemo.collections.nlp.parts.nlp_overrides import NLPSaveRestoreConnector
  from nemo.utils import AppState, logging
  from nemo.utils.distributed import initialize_distributed
  from nemo.utils.model_utils import inject_model_parallel_rank, uninject_model_parallel_rank
label_desc: Use image node_id 6 for calling a global function with example usage: load_from_checkpoint(cls, checkpoint_path, map_location, hparams_file, strict) and returns: checkpoint,  consumed,  steps,  check_point_version