code: def _add_training_args(parser):
      group = parser.add_argument_group(title="training")
  
      group.add_argument(
          "--micro-batch-size",
          type=int,
          default=None,
          help="Batch size per model instance (local batch size). "
          "Global batch size is local batch size times data "
          "parallel size times number of micro batches.",
      )
      group.add_argument(
          "--batch-size",
          type=int,
          default=None,
          help="Old batch size parameter, do not use. "
          "Use --micro-batch-size instead",
      )
      group.add_argument(
          "--global-batch-size",
          type=int,
          default=None,
          help="Training batch size. If set, it should be a "
          "multiple of micro-batch-size times data-parallel-size. "
          "If this value is None, then "
          "use micro-batch-size * data-parallel-size as the "
          "global batch size. This choice will result in 1 for "
          "number of micro-batches.",
      )
      group.add_argument(
          "--rampup-batch-size",
          nargs="*",
          default=None,
          help="Batch size ramp up with the following values:"
          "  --rampup-batch-size <start batch size> "
          "                      <batch size incerement> "
          "                      <ramp-up samples> "
          "For example:"
          "   --rampup-batch-size 16 8 300000 \ "
          "   --global-batch-size 1024"
          "will start with global batch size 16 and over "
          " (1024 - 16) / 8 = 126 intervals will increase"
          "the batch size linearly to 1024. In each interval"
          "we will use approximately 300000 / 126 = 2380 samples.",
      )
      group.add_argument(
          "--recompute-activations",
          action="store_true",
          help="recompute activation to allow for training "
          "with larger models, sequences, and batch sizes.",
      )
      group.add_argument(
          "--recompute-granularity",
          type=str,
          default=None,
          choices=["full", "selective"],
          help="Checkpoint activations to allow for training "
          "with larger models, sequences, and batch sizes. "
          "It is supported at two granularities 1) full: "
          "whole transformer layer is recomputed, "
          "2) selective: core attention part of the transformer "
          "layer is recomputed.",
      )
      group.add_argument(
          "--no-check-for-nan-in-loss-and-grad",
          action="store_false",
          help="Check for NaNs in loss and grad",
          dest="check_for_nan_in_loss_and_grad",
      )
      group.add_argument(
          "--distribute-saved-activations",
          action="store_true",
          help="If set, distribute recomputed activations "
          "across model parallel group.",
      )
      group.add_argument(
          "--recompute-method",
          type=str,
          default=None,
          choices=["uniform", "block"],
          help="1) uniform: uniformly divide the total number of "
          "Transformer layers and recompute the input activation of "
          "each divided chunk at specified granularity, "
          "2) recompute the input activations of only a set number of "
          "individual Transformer layers per pipeline stage and do the "
          "rest without any recomputing at specified granularity"
          "default) do not apply activations recompute to any layers",
      )
      group.add_argument(
          "--recompute-num-layers",
          type=int,
          default=None,
          help="1) uniform: the number of Transformer layers in each "
          "uniformly divided recompute unit, "
          "2) block: the number of individual Transformer layers "
          "to recompute within each pipeline stage.",
      )
      group.add_argument(
          "--no-clone-scatter-output-in-embedding",
          action="store_false",
          help="If not set, clone the output of the scatter in embedding layer to GC original tensor.",
          dest="clone_scatter_output_in_embedding",
      )
      group.add_argument(
          "--profile",
          action="store_true",
          help="Enable nsys profiling. When using this option, nsys "
          "options should be specified in commandline. An example "
          "nsys commandline is `nsys profile -s none -t nvtx,cuda "
          "-o <path/to/output_file> --force-overwrite true "
          "--capture-range=cudaProfilerApi "
          "--capture-range-end=stop`.",
      )
      group.add_argument(
          "--profile-step-start",
          type=int,
          default=10,
          help="Global step to start profiling.",
      )
      group.add_argument(
          "--profile-step-end",
          type=int,
          default=12,
          help="Global step to stop profiling.",
      )
      group.add_argument(
          "--profile-ranks",
          nargs="+",
          type=int,
          default=[0],
          help="Global ranks to profile.",
      )
      group.add_argument(
          "--tp-comm-overlap",
          action="store_true",
          help="Enables the "
          " overlap of Tensor parallel communication and GEMM kernels.",
      )
      group.add_argument(
          "--tp-comm-overlap-cfg",
          type=str,
          default=None,
          help="Config file when tp_comm_overlap is enabled.",
      )
      group.add_argument(
          "--disable-tp-comm-split-ag",
          action="store_false",
          help="Disables the All-Gather overlap with fprop GEMM.",
          dest="tp_comm_split_ag",
      )
      group.add_argument(
          "--disable-tp-comm-split-rs",
          action="store_false",
          help="Disables the Reduce-Scatter overlap with fprop GEMM.",
          dest="tp_comm_split_rs",
      )
      group.add_argument(
          "--disable-tp-comm-bulk-dgrad",
          action="store_false",
          help="Disables the All-Gather overlap with bprop activation gradient GEMM.",
          dest="tp_comm_bulk_dgrad",
      )
      group.add_argument(
          "--disable-tp-comm-bulk-wgrad",
          action="store_false",
          help="Disables the Reduce-Scatter overlap with bprop weight gradient GEMM.",
          dest="tp_comm_bulk_wgrad",
      )
  
      # deprecated
      group.add_argument(
          "--checkpoint-activations",
          action="store_true",
          help="Checkpoint activation to allow for training "
          "with larger models, sequences, and batch sizes.",
      )
      group.add_argument(
          "--train-iters",
          type=int,
          default=None,
          help="Total number of iterations to train over all "
          "training runs. Note that either train-iters or "
          "train-samples should be provided.",
      )
      group.add_argument(
          "--train-samples",
          type=int,
          default=None,
          help="Total number of samples to train over all "
          "training runs. Note that either train-iters or "
          "train-samples should be provided.",
      )
      group.add_argument(
          "--log-interval",
          type=int,
          default=100,
          help="Report loss and timing interval.",
      )
      group.add_argument(
          "--exit-interval",
          type=int,
          default=None,
          help="Exit the program after the iteration is divisible "
          "by this value.",
      )
      group.add_argument(
          "--exit-duration-in-mins",
          type=int,
          default=None,
          help="Exit the program after this many minutes.",
      )
      group.add_argument(
          "--exit-signal-handler",
          action="store_true",
          help="Dynamically save the checkpoint and shutdown the "
          "training if SIGTERM is received",
      )
      group.add_argument(
          "--tensorboard-dir",
          type=str,
          default=None,
          help="Write TensorBoard logs to this directory.",
      )
      group.add_argument(
          "--no-masked-softmax-fusion",
          action="store_false",
          help="Disable fusion of query_key_value scaling, "
          "masking, and softmax.",
          dest="masked_softmax_fusion",
      )
      group.add_argument(
          "--no-bias-gelu-fusion",
          action="store_false",
          help="Disable bias and gelu fusion.",
          dest="bias_gelu_fusion",
      )
      group.add_argument(
          "--no-bias-dropout-fusion",
          action="store_false",
          help="Disable bias and dropout fusion.",
          dest="bias_dropout_fusion",
      )
      group.add_argument(
          "--use-flash-attn",
          action="store_true",
          help="use FlashAttention implementation of attention. "
          "https://arxiv.org/abs/2205.14135",
      )
      group.add_argument(
          "--disable-bias-linear",
          action="store_false",
          help="Disable bias in the linear layers",
          dest="add_bias_linear",
      )
      group.add_argument(
          "--optimizer",
          type=str,
          default="adam",
          choices=["adam", "sgd"],
          help="Optimizer function",
      )
      group.add_argument(
          "--dataloader-type",
          type=str,
          default=None,
          choices=["single", "cyclic"],
          help="Single pass vs multiple pass data loader",
      )
      group.add_argument(
          "--no-async-tensor-model-parallel-allreduce",
          action="store_false",
          help="Disable asynchronous execution of "
          "tensor-model-parallel all-reduce with weight "
          "gradient compuation of a column-linear layer.",
          dest="async_tensor_model_parallel_allreduce",
      )
      group.add_argument(
          "--no-persist-layer-norm",
          action="store_true",
          help="Disable using persistent fused layer norm kernel. "
          "This kernel supports only a set of hidden sizes. Please "
          "check persist_ln_hidden_sizes if your hidden "
          "size is supported.",
      )
      group.add_argument(
          "--sequence-parallel",
          action="store_true",
          help="Enable sequence parallel optimization.",
      )
      group.add_argument(
          "--no-gradient-accumulation-fusion",
          action="store_false",
          help="Disable fusing gradient accumulation to weight "
          "gradient computation of linear layers",
          dest="gradient_accumulation_fusion",
      )
      group.add_argument(
          "--use-mcore-models",
          action="store_true",
          help="Use the implementation from megatron core",
      )
      group.add_argument(
          "--manual-gc",
          action="store_true",
          help="Disable the threshold-based default garbage "
          "collector and trigger the garbage collection manually. "
          "Manual garbage collection helps to align the timing of "
          "the collection across ranks which mitigates the impact "
          "of CPU-associated jitters. When the manual gc is enabled, "
          "garbage collection is performed only at the start and the "
          "end of the validation routine by default.",
      )
      group.add_argument(
          "--manual-gc-interval",
          type=int,
          default=0,
          help="Training step interval to trigger manual garbage "
          "collection. When the value is set to 0, garbage "
          "collection is not triggered between training steps.",
      )
      group.add_argument(
          "--no-manual-gc-eval",
          action="store_false",
          help="When using manual garbage collection, disable "
          "garbage collection at the start and the end of each "
          "evaluation run.",
          dest="manual_gc_eval",
      )
  
      return parser
  
imports: import argparse
  import dataclasses
  import json
  import os
  import torch
  import types
  import torch.nn.functional
  from megatron.global_vars import set_retro_args, get_retro_args
  from tools.retro.utils import get_args_path
  from megatron.core.models.retro import RetroConfig
  from megatron.core.transformer import TransformerConfig
label_desc: Use image node_id 12 for calling a global function with example usage: _add_training_args(parser) and returns: parser