code: def parse_args(
      extra_args_provider=None,
      defaults={},
      override_args={},
      ignore_unknown_args=False,
  ):
      """Parse all arguments."""
      parser = argparse.ArgumentParser(
          description="Megatron-LM Arguments", allow_abbrev=False
      )
  
      # Standard arguments.
      parser = _add_network_size_args(parser)
      parser = _add_regularization_args(parser)
      parser = _add_training_args(parser)
      parser = _add_initialization_args(parser)
      parser = _add_learning_rate_args(parser)
      parser = _add_checkpointing_args(parser)
      parser = _add_mixed_precision_args(parser)
      parser = _add_distributed_args(parser)
      parser = _add_validation_args(parser)
      parser = _add_data_args(parser)
      parser = _add_autoresume_args(parser)
      parser = _add_biencoder_args(parser)
      parser = _add_vision_args(parser)
      parser = _add_logging_args(parser)
  
      # NOTE(mkozuki): This option is added to investigate the potential of `torch.autograd.graph.save_on_cpu()`.
      # ref: https://pytorch.org/docs/stable/autograd.html#torch.autograd.graph.save_on_cpu.
      parser.add_argument(
          "--cpu-offload",
          action="store_true",
          default=False,
          help="Turns on CPU offloading",
      )
  
      # Custom arguments.
      if extra_args_provider is not None:
          parser = extra_args_provider(parser)
  
      # Parse.
      if ignore_unknown_args:
          args, _ = parser.parse_known_args()
      else:
          args = parser.parse_args()
  
      # Distributed args.
      args.rank = int(os.getenv("RANK", "0"))
      args.world_size = int(os.getenv("WORLD_SIZE", "1"))
  
      for key in override_args:
          setattr(args, key, override_args[key])
  
      # Tensor model parallel size.
      args.tensor_model_parallel_size = min(
          args.tensor_model_parallel_size, args.world_size
      )
      assert args.world_size % args.tensor_model_parallel_size == 0, (
          "world size"
          " ({}) is not divisible by tensor model parallel size ({})".format(
              args.world_size, args.tensor_model_parallel_size
          )
      )
  
      # Pipeline model parallel size.
      args.pipeline_model_parallel_size = min(
          args.pipeline_model_parallel_size,
          (args.world_size // args.tensor_model_parallel_size),
      )
  
      args.transformer_pipeline_model_parallel_size = (
          args.pipeline_model_parallel_size - 1
          if args.standalone_embedding_stage
          else args.pipeline_model_parallel_size
      )
      # Checks.
      model_parallel_size = (
          args.pipeline_model_parallel_size
          * args.tensor_model_parallel_size
      )
      assert args.world_size % model_parallel_size == 0, (
          "world size is not"
          " divisible by tensor parallel size ({}) times pipeline parallel "
          "size ({})".format(
              args.world_size,
              args.tensor_model_parallel_size,
              args.pipeline_model_parallel_size,
          )
      )
      args.data_parallel_size = args.world_size // model_parallel_size
      if args.rank == 0:
          print(
              "using world size: {}, data-parallel-size: {}, "
              "tensor-model-parallel size: {}, "
              "pipeline-model-parallel size: {} ".format(
                  args.world_size,
                  args.data_parallel_size,
                  args.tensor_model_parallel_size,
                  args.pipeline_model_parallel_size,
              ),
              flush=True,
          )
      if args.pipeline_model_parallel_size > 1:
          if args.pipeline_model_parallel_split_rank is not None:
              assert (
                  args.pipeline_model_parallel_split_rank
                  < args.pipeline_model_parallel_size
              ), (
                  "split rank needs"
                  " to be less than pipeline model parallel size ({})".format(
                      args.pipeline_model_parallel_size
                  )
              )
  
      # Deprecated arguments
      assert args.batch_size is None, (
          "--batch-size argument is no longer "
          "valid, use --micro-batch-size instead"
      )
      del args.batch_size
      assert args.warmup is None, (
          "--warmup argument is no longer valid, use "
          "--lr-warmup-fraction instead"
      )
      del args.warmup
      assert args.model_parallel_size is None, (
          "--model-parallel-size is no "
          "longer valid, use --tensor-model-parallel-size instead"
      )
      del args.model_parallel_size
      if args.checkpoint_activations:
          args.recompute_granularity = "full"
          args.recompute_method = "uniform"
          if args.rank == 0:
              print(
                  "--checkpoint-activations is no longer valid, "
                  "use --recompute-granularity and --recompute-method  instead. "
                  "Defaulting to recompute-granularity=full and recompute-method=uniform."
              )
      del args.checkpoint_activations
  
      if args.recompute_activations:
          args.recompute_granularity = "selective"
      del args.recompute_activations
  
      # Set input defaults.
      for key in defaults:
          # For default to be valid, it should not be provided in the
          # arguments that are passed to the program. We check this by
          # ensuring the arg is set to None.
          if getattr(args, key) is not None:
              if args.rank == 0:
                  print(
                      "WARNING: overriding default arguments for {key}:{v} \
                         with {key}:{v2}".format(
                          key=key,
                          v=defaults[key],
                          v2=getattr(args, key),
                      ),
                      flush=True,
                  )
          else:
              setattr(args, key, defaults[key])
  
      # Batch size.
      assert args.micro_batch_size is not None
      assert args.micro_batch_size > 0
      if args.global_batch_size is None:
          args.global_batch_size = (
              args.micro_batch_size * args.data_parallel_size
          )
          if args.rank == 0:
              print(
                  "setting global batch size to {}".format(
                      args.global_batch_size
                  ),
                  flush=True,
              )
      assert args.global_batch_size > 0
      if args.num_layers_per_virtual_pipeline_stage is not None:
          assert args.pipeline_model_parallel_size > 2, (
              "pipeline-model-parallel size should be greater than 2 with "
              "interleaved schedule"
          )
          assert (
              args.num_layers
              % args.num_layers_per_virtual_pipeline_stage
              == 0
          ), (
              "number of layers is not divisible by number of layers per virtual "
              "pipeline stage"
          )
          args.virtual_pipeline_model_parallel_size = (
              args.num_layers // args.pipeline_model_parallel_size
          ) // args.num_layers_per_virtual_pipeline_stage
      else:
          args.virtual_pipeline_model_parallel_size = None
  
      # Parameters dtype.
      args.params_dtype = torch.float
      if args.fp16:
          assert not args.bf16
          args.params_dtype = torch.half
      if args.bf16:
          assert not args.fp16
          args.params_dtype = torch.bfloat16
          # bfloat16 requires gradient accumulation and all-reduce to
          # be done in fp32.
          if not args.accumulate_allreduce_grads_in_fp32:
              args.accumulate_allreduce_grads_in_fp32 = True
              if args.rank == 0:
                  print(
                      "accumulate and all-reduce gradients in fp32 for "
                      "bfloat16 data type.",
                      flush=True,
                  )
  
      if args.rank == 0:
          print(
              "using {} for parameters ...".format(args.params_dtype),
              flush=True,
          )
  
      # If we do accumulation and all-reduces in fp32, we need to have local DDP
      # and we should make sure use-contiguous-buffers-in-local-ddp is not off.
      if args.accumulate_allreduce_grads_in_fp32:
          assert args.DDP_impl == "local"
          assert args.use_contiguous_buffers_in_local_ddp
      else:
          if args.gradient_accumulation_fusion:
              args.gradient_accumulation_fusion = False
              if args.rank == 0:
                  print(
                      "Gradient accumulation fusion to linear layer weight "
                      "gradient computation is supported only with fp32 "
                      "gradient accumulation. Setting gradient_accumulation_fusion "
                      "to False",
                      flush=True,
                  )
  
      # For torch DDP, we do not use contiguous buffer
      if args.DDP_impl == "torch":
          args.use_contiguous_buffers_in_local_ddp = False
  
      if args.dataloader_type is None:
          args.dataloader_type = "single"
  
      # Consumed tokens.
      args.consumed_train_samples = 0
      args.consumed_valid_samples = 0
  
      # Iteration-based training.
      if args.train_iters:
          # If we use iteration-based training, make sure the
          # sample-based options are off.
          assert (
              args.train_samples is None
          ), "expected iteration-based training"
          assert (
              args.lr_decay_samples is None
          ), "expected iteration-based learning rate decay"
          assert (
              args.lr_warmup_samples == 0
          ), "expected iteration-based learning rate warmup"
          assert (
              args.rampup_batch_size is None
          ), "expected no batch-size rampup for iteration-based training"
          if args.lr_warmup_fraction is not None:
              assert (
                  args.lr_warmup_iters == 0
              ), "can only specify one of lr-warmup-fraction and lr-warmup-iters"
  
      # Sample-based training.
      if args.train_samples:
          # If we use sample-based training, make sure the
          # iteration-based options are off.
          assert (
              args.train_iters is None
          ), "expected sample-based training"
          assert (
              args.lr_decay_iters is None
          ), "expected sample-based learning rate decay"
          assert (
              args.lr_warmup_iters == 0
          ), "expected sample-based learnig rate warmup"
          if args.lr_warmup_fraction is not None:
              assert args.lr_warmup_samples == 0, (
                  "can only specify one of lr-warmup-fraction "
                  "and lr-warmup-samples"
              )
  
      # Check required arguments.
      required_args = [
          "num_layers",
          "hidden_size",
          "num_attention_heads",
          "max_position_embeddings",
      ]
      for req_arg in required_args:
          _check_arg_is_not_none(args, req_arg)
  
      # Checks.
      if args.ffn_hidden_size is None:
          args.ffn_hidden_size = 4 * args.hidden_size
  
      if args.kv_channels is None:
          assert args.hidden_size % args.num_attention_heads == 0
          args.kv_channels = (
              args.hidden_size // args.num_attention_heads
          )
  
      if args.seq_length is not None:
          assert args.encoder_seq_length is None
          args.encoder_seq_length = args.seq_length
      else:
          assert args.encoder_seq_length is not None
          args.seq_length = args.encoder_seq_length
  
      if args.seq_length is not None:
          assert args.max_position_embeddings >= args.seq_length
      if args.decoder_seq_length is not None:
          assert args.max_position_embeddings >= args.decoder_seq_length
      if args.lr is not None:
          assert args.min_lr <= args.lr
      if args.save is not None:
          assert args.save_interval is not None
      # Mixed precision checks.
      if args.fp16_lm_cross_entropy:
          assert (
              args.fp16
          ), "lm cross entropy in fp16 only support in fp16 mode."
      if args.fp32_residual_connection:
          assert (
              args.fp16 or args.bf16
          ), "residual connection in fp32 only supported when using fp16 or bf16."
  
      if args.weight_decay_incr_style == "constant":
          assert args.start_weight_decay is None
          assert args.end_weight_decay is None
          args.start_weight_decay = args.weight_decay
          args.end_weight_decay = args.weight_decay
      else:
          assert args.start_weight_decay is not None
          assert args.end_weight_decay is not None
  
      TORCH_MAJOR = int(torch.__version__.split(".")[0])
      TORCH_MINOR = int(torch.__version__.split(".")[1])
      # Persistent fused layer norm.
      if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 11):
          args.no_persist_layer_norm = True
          if args.rank == 0:
              print(
                  "Persistent fused layer norm kernel is supported from "
                  "pytorch v1.11 (nvidia pytorch container paired with v1.11). "
                  "Defaulting to no_persist_layer_norm=True"
              )
  
      # Activation recomputing.
      if args.distribute_saved_activations:
          assert args.tensor_model_parallel_size > 1, (
              "can distribute "
              "recomputed activations only across tensor model "
              "parallel groups"
          )
          assert args.recompute_granularity == "full", (
              "distributed recompute activations is only "
              "application to full recompute granularity"
          )
          assert args.recompute_method is not None, (
              "for distributed recompute activations to work you "
              "need to use a recompute method "
          )
          assert TORCH_MAJOR >= 1 and TORCH_MINOR >= 10, (
              "distributed recompute activations are supported for pytorch "
              "v1.10 and above (Nvidia Pytorch container >= 21.07). Current "
              "pytorch version is v%s.%s." % (TORCH_MAJOR, TORCH_MINOR)
          )
  
      if args.recompute_granularity == "selective":
          assert args.recompute_method is None, (
              "recompute method is not yet supported for "
              "selective recomputing granularity"
          )
  
      # disable async_tensor_model_parallel_allreduce when
      # model parallel memory optimization is enabled
      if args.sequence_parallel:
          args.async_tensor_model_parallel_allreduce = False
  
      _print_args(args)
      return args
  
imports: import argparse
  import os
  import torch
label_desc: Use image node_id 1 for calling a global function with example usage: parse_args(extra_args_provider, defaults, override_args, ignore_unknown_args) and returns: args