code: def validate_args(args, defaults={}):
      # 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 * args.context_parallel_size)
          == 0
      ), (
          "world size ({}) is not divisible by tensor parallel size ({}) times "
          "pipeline parallel size ({}) times context parallel size ({})".format(
              args.world_size,
              args.tensor_model_parallel_size,
              args.pipeline_model_parallel_size,
              args.context_parallel_size,
          )
      )
      args.data_parallel_size = args.world_size // (
          model_parallel_size * args.context_parallel_size
      )
      if args.rank == 0:
          print(
              "using world size: {}, data-parallel size: {}, "
              "context-parallel size: {} "
              "tensor-model-parallel size: {}, "
              "pipeline-model-parallel size: {} ".format(
                  args.world_size,
                  args.data_parallel_size,
                  args.context_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
                  )
              )
  
      if args.tp_comm_overlap:
          assert (
              args.sequence_parallel == True
          ), "Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled"
  
      # 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:
          if args.rank == 0:
              print(
                  "--checkpoint-activations is no longer valid, use --recompute-activations, "
                  "or, for more control, --recompute-granularity and --recompute-method."
              )
          exit()
      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, None) 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.transformer_pipeline_model_parallel_size
              == 0
          ), "number of layers should be divisible by the pipeline parallel size"
          num_layers_per_pipeline_stage = (
              args.num_layers
              // args.transformer_pipeline_model_parallel_size
          )
          assert (
              num_layers_per_pipeline_stage
              % args.num_layers_per_virtual_pipeline_stage
              == 0
          ), "number of layers per pipeline stage must be divisible number of layers per virtual pipeline stage"
          args.virtual_pipeline_model_parallel_size = (
              num_layers_per_pipeline_stage
              // args.num_layers_per_virtual_pipeline_stage
          )
      else:
          args.virtual_pipeline_model_parallel_size = None
          # Overlap P2P communication is disabled if not using the interleaved schedule.
          args.overlap_p2p_comm = False
          if args.rank == 0:
              print(
                  "WARNING: Setting args.overlap_p2p_comm to False since non-interleaved "
                  "schedule does not support overlapping p2p communication"
              )
  
      if args.overlap_param_gather:
          assert (
              args.use_distributed_optimizer
          ), "--overlap-param-gather only supported with distributed optimizer"
          assert (
              args.overlap_grad_reduce
          ), "--overlap-grad-reduce should be turned on when using --overlap-param-gather"
  
      # 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 args.dataloader_type is None:
          args.dataloader_type = "single"
  
      # Consumed tokens.
      args.consumed_train_samples = 0
      args.consumed_valid_samples = 0
  
      # Support for variable sequence lengths across batches/microbatches.
      # set it if the dataloader supports generation of variable sequence lengths
      # across batches/microbatches. Due to additional communication overhead
      # during pipeline parallelism, it should not be set if sequence length
      # is constant during training.
      args.variable_seq_lengths = False
  
      # 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"
              )
  
      if args.num_layers is not None:
          assert (
              args.encoder_num_layers is None
          ), "cannot have both num-layers and encoder-num-layers specified"
          args.encoder_num_layers = args.num_layers
      else:
          assert (
              args.encoder_num_layers is not None
          ), "either num-layers or encoder-num-layers should be specified"
          args.num_layers = args.encoder_num_layers
  
      # 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:
          if args.swiglu:
              # reduce the dimnesion for MLP since projections happens on
              # two linear layers. this keeps the number of paramters in
              # the same ballpark as the counterpart with 4*h size
              # we keep it a multiple of 64, which means the actual tensor size
              # will be a multiple of 64 / tp_size
              args.ffn_hidden_size = (
                  int((4 * args.hidden_size * 2 / 3) / 64) * 64
              )
          else:
              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.moe_grouped_gemm:
          assert (
              args.bf16
          ), "Currently GroupedGEMM for MoE only supports bf16 dtype."
          dc = torch.cuda.get_device_capability()
          assert (
              dc[0] >= 8
          ), "Unsupported compute capability for GroupedGEMM kernels."
  
      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, TORCH_MINOR) >= (1, 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 sequence parallelism when tp=1
      # to avoid change in numerics when
      # sequence_parallelism is enabled.
      if args.tensor_model_parallel_size == 1:
          args.sequence_parallel = False
  
      # disable async_tensor_model_parallel_allreduce when
      # model parallel memory optimization is enabled
      if args.sequence_parallel:
          args.async_tensor_model_parallel_allreduce = False
  
      if os.environ.get("CUDA_DEVICE_MAX_CONNECTIONS") != "1":
          if args.sequence_parallel:
              raise RuntimeError(
                  "Using sequence parallelism requires setting the environment variable "
                  "CUDA_DEVICE_MAX_CONNECTIONS to 1"
              )
          if args.async_tensor_model_parallel_allreduce:
              raise RuntimeError(
                  "Using async gradient all reduce requires setting the environment "
                  "variable CUDA_DEVICE_MAX_CONNECTIONS to 1"
              )
  
      # Disable bias gelu fusion if we are disabling bias altogether
      if not args.add_bias_linear:
          args.bias_gelu_fusion = False
  
      # Retro checks.
      if args.retro_add_retriever:
          # Sequence parallelism unsupported.
          assert (
              not args.sequence_parallel
          ), "retro currently does not support sequence parallelism."
  
          # Pipeline parallelism unsupported.
          assert (
              args.pipeline_model_parallel_size == 1
          ), "retro currently does not support pipeline parallelism."
  
      # Load retro args (used by both Retro & GPT).
      if args.retro_workdir:
          retro_args_path = get_retro_args_path(args.retro_workdir)
          assert os.path.exists(
              retro_args_path
          ), "retro workdir missing args.json"
          with open(retro_args_path) as f:
              retro_args = types.SimpleNamespace(**json.load(f))
              retro_args.retro_return_doc_ids = (
                  args.retro_return_doc_ids
              )
              retro_args.retro_gpt_retrieved_length = (
                  args.retro_num_retrieved_chunks
                  * retro_args.retro_gpt_chunk_length
              )
              set_retro_args(retro_args)
  
      # Legacy RoPE arguments
      if args.use_rotary_position_embeddings:
          args.position_embedding_type = "rope"
  
      # Would just need to add 'NoPE' as a position_embedding_type to support this, but for now
      # don't allow it to keep things simple
      if (
          not args.add_position_embedding
          and args.position_embedding_type != "rope"
      ):
          raise RuntimeError(
              "--no-position-embedding is deprecated, use --position-embedding-type"
          )
  
      # MoE Spec check
      if args.num_experts is not None:
          assert (
              args.spec is None
          ), "Model Spec must be None when using MoEs"
  
      # Expert parallelism check
      if args.expert_model_parallel_size > 1:
          assert (
              args.num_experts is not None
          ), "num_experts must be non None to use expert model parallelism"
          assert (
              args.num_experts % args.expert_model_parallel_size == 0
          ), "Number of experts should be a multiple of expert model parallel_size."
          assert (
              not args.use_distributed_optimizer
          ), "Expert parallelism is not suppored with distributed optimizer."
          assert (
              not args.fp16
          ), "Expert parallelism is not supported with fp16 training."
          if args.tensor_model_parallel_size > 1:
              assert (
                  args.sequence_parallel
              ), "When using expert parallelism and tensor parallelism, sequence parallelism must be used."
  
      # Print arguments.
      _print_args("arguments", args)
      retro_args = get_retro_args()
      if retro_args and args != retro_args:
          _print_args(
              "retro arguments",
              types.SimpleNamespace(
                  **{
                      k: v
                      for k, v in vars(retro_args).items()
                      if k.startswith("retro")
                  },
                  rank=args.rank,
              ),
          )
  
      return args
  
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 2 for calling a global function with example usage: validate_args(args, defaults) and returns: args