code: def add_optimizer_state(
      lm_checkpoint, new_checkpoint, megatron_amp_O2=True
  ):
      # this method is to convert lm_checkpoint optimizer states for nemo checkpoint
      OPTIMIZER_KEY = "optimizer"
      FP32_FP16_KEY = "fp32_from_fp16_params"
      NEW_OPTIMIZER_KEY = "optimizer_states"
      STEP_KEY = "iteration"
      NEW_STEP_KEY = "global_step"
      LR_SCHEDULER = "lr_scheduler"
      NEW_LR_SCHEDULER = "lr_schedulers"
      if (
          OPTIMIZER_KEY in lm_checkpoint
          and OPTIMIZER_KEY in lm_checkpoint[OPTIMIZER_KEY]
      ):
          opt_state = lm_checkpoint[OPTIMIZER_KEY][OPTIMIZER_KEY]
          if megatron_amp_O2:
              opt_dict = dict()
              if LR_SCHEDULER in lm_checkpoint:
                  sched = lm_checkpoint[LR_SCHEDULER]
                  for param_group in opt_state["param_groups"]:
                      param_group["initial_lr"] = sched["max_lr"]
              if FP32_FP16_KEY in lm_checkpoint[OPTIMIZER_KEY]:
                  fp32_state = lm_checkpoint[OPTIMIZER_KEY][
                      FP32_FP16_KEY
                  ]
                  opt_dict[FP32_FP16_KEY] = fp32_state
              opt_dict[OPTIMIZER_KEY] = opt_state
              new_checkpoint[NEW_OPTIMIZER_KEY] = [opt_dict]
          else:
              new_checkpoint[NEW_OPTIMIZER_KEY] = [opt_state]
  
      if STEP_KEY in lm_checkpoint:
          new_checkpoint[NEW_STEP_KEY] = lm_checkpoint[STEP_KEY]
          new_checkpoint["epoch"] = 1  # always one epoch
      if LR_SCHEDULER in lm_checkpoint:
          gbs = lm_checkpoint["args"].global_batch_size
          sched = lm_checkpoint[LR_SCHEDULER]
          content = OrderedDict()
          content["max_steps"] = (
              int(sched["decay_steps"]) // gbs
              + sched["warmup_steps"] // gbs
          )
          content["warmup_steps"] = int(sched["warmup_steps"]) // gbs
          content["constant_steps"] = 0  # no such conf in lm checkpoint
          content["decay_steps"] = int(sched["decay_steps"]) // gbs
          content["min_lr"] = sched["min_lr"]
          if OPTIMIZER_KEY in lm_checkpoint:
              content["base_lrs"] = [
                  i["initial_lr"]
                  for i in new_checkpoint["optimizer_states"][0][
                      "optimizer"
                  ]["param_groups"]
              ]
              content["last_epoch"] = int(sched["num_steps"]) // gbs
              content["_last_lr"] = [
                  i["lr"]
                  for i in new_checkpoint["optimizer_states"][0][
                      "optimizer"
                  ]["param_groups"]
              ]
          else:
              content["base_lrs"] = [sched["max_lr"]]
              content["last_epoch"] = int(sched["num_steps"]) // gbs
          content["_step_count"] = int(sched["num_steps"]) // gbs
          content["verbose"] = False
          content["_get_lr_called_within_step"] = False
          new_checkpoint[NEW_LR_SCHEDULER] = [content]
  
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 4 for calling a global function with example usage: add_optimizer_state(lm_checkpoint, new_checkpoint, megatron_amp_O2) without return types