code: def nemo_export(cfg):
      """Convert a nemo model into .onnx ONNX format."""
      nemo_in = None
      if cfg.gpt_model_file:
          nemo_in = cfg.gpt_model_file
      elif cfg.checkpoint_dir:
          nemo_in = os.path.join(
              cfg.checkpoint_dir, cfg.checkpoint_name
          )
      assert (
          nemo_in is not None
      ), "NeMo model not provided. Please provide the path to the .nemo or .ckpt file"
  
      onnx_out = cfg.onnx_model_file
  
      trainer = Trainer(strategy=NLPDDPStrategy(), **cfg.trainer)
      assert (
          cfg.trainer.devices * cfg.trainer.num_nodes
          == cfg.tensor_model_parallel_size
          * cfg.pipeline_model_parallel_size
      ), "devices * num_nodes should equal tensor_model_parallel_size * pipeline_model_parallel_size"
  
      logging.info("Restoring NeMo model from '{}'".format(nemo_in))
      try:
          if cfg.gpt_model_file:
              save_restore_connector = NLPSaveRestoreConnector()
              if os.path.isdir(cfg.gpt_model_file):
                  save_restore_connector.model_extracted_dir = (
                      cfg.gpt_model_file
                  )
  
              pretrained_cfg = ModelPT.restore_from(
                  restore_path=cfg.gpt_model_file,
                  trainer=trainer,
                  return_config=True,
                  save_restore_connector=save_restore_connector,
              )
              OmegaConf.set_struct(pretrained_cfg, True)
              with open_dict(pretrained_cfg):
                  pretrained_cfg.sequence_parallel = False
                  pretrained_cfg.activations_checkpoint_granularity = (
                      None
                  )
                  pretrained_cfg.activations_checkpoint_method = None
                  pretrained_cfg.precision = trainer.precision
                  if trainer.precision == "16":
                      pretrained_cfg.megatron_amp_O2 = False
              model = ModelPT.restore_from(
                  restore_path=cfg.gpt_model_file,
                  trainer=trainer,
                  override_config_path=pretrained_cfg,
                  save_restore_connector=save_restore_connector,
              )
          elif cfg.checkpoint_dir:
              app_state = AppState()
              if (
                  cfg.tensor_model_parallel_size > 1
                  or cfg.pipeline_model_parallel_size > 1
              ):
                  app_state.model_parallel_size = (
                      cfg.tensor_model_parallel_size
                      * cfg.pipeline_model_parallel_size
                  )
                  app_state.tensor_model_parallel_size = (
                      cfg.tensor_model_parallel_size
                  )
                  app_state.pipeline_model_parallel_size = (
                      cfg.pipeline_model_parallel_size
                  )
                  (
                      app_state.tensor_model_parallel_rank,
                      app_state.pipeline_model_parallel_rank,
                      app_state.model_parallel_size,
                      app_state.data_parallel_size,
                      app_state.pipeline_model_parallel_split_rank,
                      app_state.virtual_pipeline_model_parallel_rank,
                  ) = fake_initialize_model_parallel(
                      world_size=app_state.model_parallel_size,
                      rank=trainer.global_rank,
                      tensor_model_parallel_size_=cfg.tensor_model_parallel_size,
                      pipeline_model_parallel_size_=cfg.pipeline_model_parallel_size,
                      pipeline_model_parallel_split_rank_=cfg.pipeline_model_parallel_split_rank,
                  )
              checkpoint_path = inject_model_parallel_rank(
                  os.path.join(cfg.checkpoint_dir, cfg.checkpoint_name)
              )
              model_cls = get_model_class(cfg)
              model = model_cls.load_from_checkpoint(
                  checkpoint_path,
                  hparams_file=cfg.hparams_file,
                  trainer=trainer,
              )
          else:
              raise ValueError(
                  "need at least a nemo file or checkpoint dir"
              )
      except Exception as e:
          logging.error(
              "Failed to restore model from NeMo file : {}. Please make sure you have the latest NeMo package installed with [all] dependencies.".format(
                  nemo_in
              )
          )
          raise e
  
      logging.info(
          "Model {} restored from '{}'".format(
              model.__class__.__name__, nemo_in
          )
      )
  
      # Export
      check_trace = cfg.export_options.runtime_check
  
      try:
          model.to(device=cfg.export_options.device).freeze()
          model.eval()
          model.export(
              onnx_out,
              onnx_opset_version=cfg.export_options.onnx_opset,
              do_constant_folding=cfg.export_options.do_constant_folding,
              dynamic_axes={
                  "input_ids": {0: "sequence", 1: "batch"},
                  "position_ids": {0: "sequence", 1: "batch"},
                  "logits": {0: "sequence", 1: "batch"},
              },
              check_trace=check_trace,
              check_tolerance=cfg.export_options.check_tolerance,
              verbose=cfg.export_options.verbose,
          )
      except Exception as e:
          logging.error(
              "Export failed. Please make sure your NeMo model class ({}) has working export() and that you have the latest NeMo package installed with [all] dependencies.".format(
                  model.__class__
              )
          )
          raise e
  
file_path: NeMo/examples/nlp/language_modeling/megatron_export.py
imports: import os
  from omegaconf import OmegaConf, open_dict
  from pytorch_lightning import Trainer
  from nemo.collections.nlp.models.language_modeling.megatron_bart_model import MegatronBARTModel
  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.models.language_modeling.megatron_retrieval_model import MegatronRetrievalModel
  from nemo.collections.nlp.models.language_modeling.megatron_t5_model import MegatronT5Model
  from nemo.collections.nlp.models.machine_translation.megatron_nmt_model import MegatronNMTModel
  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 import ModelPT
  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 2 for calling a global function with example usage: nemo_export(cfg) without return types