"""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)

parser = _add_inference_args(parser)

parser = _add_transformer_engine_args(parser)

parser = _add_retro_args(parser)

# 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()

# Args from environment

args.rank = int(os.getenv('RANK', '0'))

args.world_size = int(os.getenv("WORLD_SIZE", '1'))

# 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)

# --data-path and --train-weighted-splits-paths

message = "Data loading Mode 1: --data-path and --split "\

"and Mode 2: --(train|valid|test)-weighted-split-paths"\

"are mutually exclusive i.e. cannot be set together."

if args.data_path:

assert args.train_weighted_split_paths is None, message

setattr(args, "valid_weighted_split_names", None)

setattr(args, "valid_weighted_split_weights", None)

setattr(args, "valid_weighted_split_splits", None)

setattr(args, "test_weighted_split_names", None)

setattr(args, "test_weighted_split_weights", None)

setattr(args, "test_weighted_split_splits", None)

# args.split default value in the args is None it is set here in order

# to check that it does not to overlap with the 2nd mode of data loading

if args.split is None:

args.split = "969, 30, 1"

if args.train_weighted_split_paths or args.valid_weighted_split_paths or \

args.test_weighted_split_paths:

assert args.data_path is None and args.split is None, message

# 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) 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.transformer_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

# If we use the distributed optimizer, we need to have local DDP

# and we should make sure use-contiguous-buffers-in-local-ddp is on.

if args.use_distributed_optimizer:

assert args.DDP_impl == 'local'

assert args.use_contiguous_buffers_in_local_ddp

# 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

# 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:

args.ffn_hidden_size = 4 * args.hidden_size

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

assert not args.glu_activation, \

"--glu-activation is not compatible with --swiglu. Use only one of the two."

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

# NOTE: this was before integrating alibi

# 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.position_embedding_type == PositionEmbeddingType.absolute or args.position_embedding_type == PositionEmbeddingType.alibi:

assert args.max_position_embeddings is not None

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

else:

assert args.max_position_embeddings is None

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.'

# Activation function

if args.glu_activation is not None and args.bias_gelu_fusion:

raise ValueError("if glu-activation is used, please set --no-bias-gelu-fusion")

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)

# Weights and Biases

if args.wandb_entity_name or args.wandb_project_name:

assert args.wandb_entity_name and args.wandb_project_name, \

"Both entity and project name must be set in order to report to wandb"

# Local-rank from environment variable (if using torchrun)

if args.local_rank is None and "LOCAL_RANK" in os.environ:

args.local_rank = int(os.environ["LOCAL_RANK"])

# Tranformer-Engine/FP8 related checking

if args.fp8_e4m3 or args.fp8_hybrid:

assert args.transformer_impl == 'transformer_engine', \

'transformer-engine required for fp8 training and inference'

assert not (args.fp8_e4m3 and args.fp8_hybrid), \

'cannot train with both fp8 e4m3 and hybrid formatting'

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 args.use_flash_attn:

assert not args.reset_attention_mask, \

"Flash Attention doesn't support arbitrary attention masks. Please turn off reset-attention-mask"

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

# Load retro args.

if args.retro_workdir:

retro_args_path = get_retro_args_path(args.retro_workdir)

if os.path.exists(retro_args_path):

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)

# 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))

"""Print arguments."""

if args.rank == 0:

print(f'------------------------ {title} ------------------------',

flush=True)

str_list = []

for arg in vars(args):

dots = '.' * (48 - len(arg))

str_list.append(' {} {} {}'.format(arg, dots, getattr(args, arg)))

for arg in sorted(str_list, key=lambda x: x.lower()):

print(arg, flush=True)

print(f'-------------------- end of {title} ---------------------',

group = parser.add_argument_group(title='Transformer-Engine')

group.add_argument('--fp8-e4m3', action='store_true',

help='E4M3 TransformerLayer', dest='fp8_e4m3')

group.add_argument('--fp8-hybrid', action='store_true',

help='Hybrid FP8 TransformerLayer', dest='fp8_hybrid')

group.add_argument('--no-fp8-wgrad', action='store_false',

help='Execute wgrad in higher precision even for FP8 runs', dest='fp8_wgrad')

group.add_argument('--fp8-margin', type=int, default=0,

help='Scaling margin for fp8', dest='fp8_margin')

group.add_argument('--fp8-interval', type=int, default=1,

help='Scaling update interval for fp8', dest='fp8_interval')

group.add_argument('--transformer-impl', default='local',

choices=['local', 'transformer_engine'],

help='Which Transformer implementation to use.',

dest='transformer_impl')

group.add_argument('--fp8-amax-history-len', type=int, default=1,

help='Number of steps for which amax history is recorded per tensor',

dest='fp8_amax_history_len')

group.add_argument('--fp8-amax-compute-algo', default='most_recent',

choices=['most_recent', 'max'],

help='Algorithm for computing amax from history',

dest='fp8_amax_compute_algo')

group = parser.add_argument_group(title='inference')

group.add_argument('--inference-batch-times-seqlen-threshold',

type=int, default=512,

help='During inference, if batch-size times '

'sequence-length is smaller than this threshold '

'then we will not use pipelining, otherwise we will.')

group.add_argument('--max-tokens-to-oom',

type=int, default=12000,

help='Maximum number of tokens during inference'

'tokens here is # in prompt + # to generate'

'Allows us to throw an error before OOM crashes server')

group.add_argument('--output-bert-embeddings', action='store_true',

help='Output Bert embeddings (via mean pooling) from '

'model, rather than its binary head output or entire '

'hidden batch.')

group.add_argument('--bert-embedder-type', default="megatron",

choices=["megatron", "huggingface"],

help='Select either Megatron or Huggingface as the '

'Bert embedder.')

group = parser.add_argument_group(title='retro')

group.add_argument('--retro-workdir', default=None,

help='Retro working directory, which contains the '

'preprocessed data for for pretraining. This directory '

'is built during preprocessing (see '

'tools/retro/README.md), and contains subdirectories '

'for the chunk database and pretraining neighbors.')

group.add_argument('--retro-add-retriever',

action='store_true', default=False,

help='Add a retriever to the transformer, for use in '

'pretraining a Retro model.')

group.add_argument('--retro-cyclic-train-iters', type=int, default=None,

help='Set number of training iterations for cyclic '

'Retro training.')

group.add_argument('--retro-encoder-layers', type=int, default=2,

help='Number of layers to use for the retrieval '

'encoder.')

group.add_argument('--retro-encoder-hidden-dropout',

type=float, default=0.1, help='Hidden dropout for '

'retrieval encoder.')

group.add_argument('--retro-encoder-attention-dropout',

type=float, default=0.1, help='Attention dropout for '

'retrieval encoder.')

group.add_argument("--retro-num-neighbors", type=int, default=2,

help='Number of neighbors to retrieve during '

'pretraining.')

group.add_argument("--retro-num-retrieved-chunks", type=int, default=2,

help='Number of chunks to retrieve from the retrieval '

'database.')

group.add_argument("--retro-return-doc-ids", action="store_true",

help="Turn this on when preprocessing retro data.")

# Enforce argument naming convention.

for action in group._group_actions:

prefix = action.dest.split("_")[0]

assert prefix == "retro", \

"Retro args must be prefixed with '--retro-*', for consistent " \

"styling. Please fix '%s'." % ", ".join(action.option_strings)

group = parser.add_argument_group(title='network size')

group.add_argument('--num-layers', type=int, default=None,

help='Number of transformer layers.')

group.add_argument('--encoder-num-layers', type=int, default=None,

help='Number of encoder transformer layers.')

group.add_argument('--decoder-num-layers', type=int, default=None,

help='Number of decoder transformer layers.')

group.add_argument('--hidden-size', type=int, default=None,

help='Tansformer hidden size.')

group.add_argument('--ffn-hidden-size', type=int, default=None,

help='Transformer Feed-Forward Network hidden size. '

'This is set to 4*hidden-size if not provided')

group.add_argument('--num-attention-heads', type=int, default=None,

help='Number of transformer attention heads.')

group.add_argument('--kv-channels', type=int, default=None,

help='Projection weights dimension in multi-head '

'attention. This is set to '

' args.hidden_size // args.num_attention_heads '

'if not provided.')

group.add_argument('--attention-head-type', type=str, default=None,

choices=['multihead', 'multiquery'],

help='Type of attention heads. `multihead` is the standard multi-head attention.'

'`multiquery` shares the values and keys across attention heads')

group.add_argument('--max-position-embeddings', type=int, default=None,

help='Maximum number of position embeddings to use. '

'This is the size of position embedding.')

group.add_argument('--use-rotary-position-embeddings', action='store_true',

help='Use rotary positional embeddings or not')

group.add_argument('--rotary-percent', type=float, default=1.0,

help='Percent of rotary dimension to use, default 100%')

group.add_argument('--no-position-embedding',

action='store_false',

help='Disable position embedding.',

dest='add_position_embedding')

group.add_argument('--make-vocab-size-divisible-by', type=int, default=128,

help='Pad the vocab size to be divisible by this value.'

'This is added for computational efficieny reasons.')

group.add_argument('--layernorm-epsilon', type=float, default=1e-5,

help='Layer norm epsilon.')

group.add_argument('--apply-layernorm-1p', action='store_true',

help='Adjust LayerNorm weights such that they are centered '

'around zero. This improves numerical stability.')

group.add_argument('--apply-residual-connection-post-layernorm',

action='store_true',

help='If set, use original BERT residula connection '

'ordering.')

group.add_argument('--openai-gelu', action='store_true',

help='Use OpenAIs GeLU implementation. This option'

'should not be used unless for backward compatibility'

'reasons.')

group.add_argument('--squared-relu', action='store_true',

help='Use squared relu activation instead of default gelu')

group.add_argument('--swiglu', action='store_true',

help='Use gated linear units and SiLU activation instead of default gelu')

group.add_argument('--onnx-safe', type=bool, required=False,

help='Use workarounds for known problems with '

'Torch ONNX exporter')

group.add_argument('--bert-no-binary-head', action='store_false',

help='Disable BERT binary head.',

dest='bert_binary_head')

group.add_argument('--position-embedding-type', type=lambda x: PositionEmbeddingType[x],

choices=list(PositionEmbeddingType),

default=PositionEmbeddingType.absolute,

help='Define position embedding type ("absolute" | "rotary" | "alibi"). "absolute" by default.'

)

group.add_argument('--glu-activation', type=str,

choices=megatron.model.glu_activations.GLU_ACTIVATIONS.keys(),

help='GLU activations to use.'

)

group.add_argument('--num-experts', type=int, default=None,

help='Number of Experts in Switch Transformer (None means no Switch)')

group.add_argument('--untie-embeddings-and-output-weights', action='store_true',

help='Untie embeddings and output weights.'),

group = parser.add_argument_group(title='logging')

group.add_argument('--log-params-norm', action='store_true',

help='If set, calculate and log parameters norm.')

group.add_argument('--log-num-zeros-in-grad', action='store_true',

help='If set, calculate and log the number of zeros in gradient.')

group.add_argument('--timing-log-level', type=int,

default=0, choices=range(0,3),

help='Granularity level to measure and report timing. '

' 0: report only iteration time and make sure timing '

' does not introduce extra overhead.'

' 1: report timing for operations that are executed '

' very limited times (basically once) during '

' each iteration (such as gradient all-reduce) '

' 2: report timing for operations that migh be '

' executed numerous times during each iteration. '

'Note that setting the level to 1 or 2 might '

'cause increase in iteration time.')

group.add_argument('--no-barrier-with-level-1-timing', action='store_false',

help='If not set, use barrier with level 1 time '

'measurements. Note that this is up to the user '

'to make sure calling barrier with their timers '

'will not result in hangs. This can happen if for '

'example the user adds a level 1 timer that is not '

'called by all ranks.',

dest='barrier_with_L1_time')

group.add_argument('--timing-log-option', type=str, default='minmax',

choices=['max', 'minmax', 'all'],

help='Options for logging timing:'

' max: report the max timing across all ranks'

' minmax: report min and max timings across all ranks'

' all: report timings of all ranks.')

group.add_argument('--tensorboard-log-interval', type=int, default=1,

help='Report to tensorboard interval.')

group.add_argument('--tensorboard-queue-size', type=int, default=1000,

help='Size of the tensorboard queue for pending events '

'and summaries before one of the ‘add’ calls forces a '

'flush to disk.')

group.add_argument('--log-timers-to-tensorboard', action='store_true',

help='If set, write timers to tensorboard.')

group.add_argument('--log-batch-size-to-tensorboard', action='store_true',

help='If set, write batch-size to tensorboard.')

group.add_argument('--no-log-learnig-rate-to-tensorboard',

action='store_false',

help='Disable learning rate logging to tensorboard.',

dest='log_learning_rate_to_tensorboard')

group.add_argument('--no-log-loss-scale-to-tensorboard',

action='store_false',

help='Disable loss-scale logging to tensorboard.',

dest='log_loss_scale_to_tensorboard')

group.add_argument('--log-validation-ppl-to-tensorboard',

action='store_true',

help='If set, write validation perplexity to '

'tensorboard.')

group.add_argument('--log-memory-to-tensorboard',

action='store_true',

help='Enable memory logging to tensorboard.')

group.add_argument('--log-world-size-to-tensorboard',

action='store_true',

help='Enable world size logging to tensorboard.')

group.add_argument('--wandb-entity-name', type=str, default=None,

help="Name of wandb entity for reporting")

group.add_argument('--wandb-project-name', type=str, default=None,

help="Name of wandb project")

group.add_argument('--transformer-timers', action='store_true',

help="If set, activate the timers within the transformer layers."

"Only for debugging, as this slows down the model.")

group.add_argument('--structured-logs', action="store_true",

help='Add timestamp and worker name to stdout and stderr.')

group.add_argument('--structured-logs-dir', type=str, default=None,

help='Directory to save the logs.')

group = parser.add_argument_group(title='regularization')

group.add_argument('--attention-dropout', type=float, default=0.1,

help='Post attention dropout probability.')

group.add_argument('--hidden-dropout', type=float, default=0.1,

help='Dropout probability for hidden state transformer.')

group.add_argument('--weight-decay', type=float, default=0.01,

help='Weight decay coefficient for L2 regularization.')

group.add_argument('--start-weight-decay', type=float,

help='Initial weight decay coefficient for L2 regularization.')

group.add_argument('--end-weight-decay', type=float,

help='End of run weight decay coefficient for L2 regularization.')

group.add_argument('--weight-decay-incr-style', type=str, default='constant',

choices=['constant', 'linear', 'cosine'],

help='Weight decay increment function.')

group.add_argument('--clip-grad', type=float, default=1.0,

help='Gradient clipping based on global L2 norm.')

group.add_argument('--adam-beta1', type=float, default=0.9,

help='First coefficient for computing running averages '

'of gradient and its square')

group.add_argument('--adam-beta2', type=float, default=0.999,

help='Second coefficient for computing running averages '

'of gradient and its square')

group.add_argument('--adam-eps', type=float, default=1e-08,

help='Term added to the denominator to improve'

'numerical stability')

group.add_argument('--sgd-momentum', type=float, default=0.9,

help='Momentum factor for sgd')

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('--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=1,

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.')

# 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 = parser.add_argument_group(title='initialization')

group.add_argument('--seed', type=int, default=1234,

help='Random seed used for python, numpy, '

'pytorch, and cuda.')

group.add_argument('--data-parallel-random-init', action='store_true',

help='Enable random initialization of params '

'across data parallel ranks')

group.add_argument('--init-method-std', type=float, default=0.02,

help='Standard deviation of the zero mean normal '

'distribution used for weight initialization.')

group.add_argument('--init-method-xavier-uniform', action='store_true',

help='Enable Xavier uniform parameter initialization')

group = parser.add_argument_group(title='learning rate')

group.add_argument('--lr', type=float, default=None,

help='Initial learning rate. Depending on decay style '

'and initial warmup, the learing rate at each '

'iteration would be different.')

group.add_argument('--lr-decay-style', type=str, default='linear',

choices=['constant', 'linear', 'cosine', 'inverse-square-root'],

help='Learning rate decay function.')

group.add_argument('--lr-decay-iters', type=int, default=None,

help='number of iterations to decay learning rate over,'

' If None defaults to `--train-iters`')

group.add_argument('--lr-decay-samples', type=int, default=None,

help='number of samples to decay learning rate over,'

' If None defaults to `--train-samples`')

group.add_argument('--lr-warmup-fraction', type=float, default=None,

help='fraction of lr-warmup-(iters/samples) to use '

'for warmup (as a float)')

group.add_argument('--lr-warmup-iters', type=int, default=0,

help='number of iterations to linearly warmup '

'learning rate over.')

group.add_argument('--lr-warmup-samples', type=int, default=0,

help='number of samples to linearly warmup '

'learning rate over.')

group.add_argument('--warmup', type=int, default=None,

help='Old lr warmup argument, do not use. Use one of the'

'--lr-warmup-* arguments above')

group.add_argument('--min-lr', type=float, default=0.0,

help='Minumum value for learning rate. The scheduler'

'clip values below this threshold.')

group.add_argument('--override-opt_param-scheduler', action='store_true',

help='Reset the values of the scheduler (learning rate,'

'warmup iterations, minimum learning rate, maximum '

'number of iterations, and decay style from input '

'arguments and ignore values from checkpoints. Note'

'that all the above values will be reset.')

group.add_argument('--use-checkpoint-opt_param-scheduler', action='store_true',

help='Use checkpoint to set the values of the scheduler '

'(learning rate, warmup iterations, minimum learning '

'rate, maximum number of iterations, and decay style '

'from checkpoint and ignore input arguments.')

group = parser.add_argument_group(title='checkpointing')

group.add_argument('--save', type=str, default=None,

help='Output directory to save checkpoints to.')

group.add_argument('--save-interval', type=int, default=None,

help='Number of iterations between checkpoint saves.')

group.add_argument('--no-save-optim', action='store_true', default=None,

help='Do not save current optimizer.')

group.add_argument('--no-save-rng', action='store_true', default=None,

help='Do not save current rng state.')

group.add_argument('--load', type=str, default=None,

help='Directory containing a model checkpoint.')

group.add_argument('--no-load-optim', action='store_true', default=None,

help='Do not load optimizer when loading checkpoint.')

group.add_argument('--no-load-rng', action='store_true', default=None,

help='Do not load rng state when loading checkpoint.')

group.add_argument('--finetune', action='store_true',

help='Load model for finetuning. Do not load optimizer '

'or rng state from checkpoint and set iteration to 0. '

'Assumed when loading a release checkpoint.')

group.add_argument('--no-initialization', action='store_false',

help='Do not perform initialization when building model, '

'can reduce startup time when definitely loading from a '

'checkpoint',

dest='perform_initialization')

group.add_argument('--use-checkpoint-args', action='store_true',

help='Override any command line arguments with arguments '

'from the checkpoint')

group.add_argument('--exit-on-missing-checkpoint', action='store_true',

help="If '--load' is set, but checkpoint is not found "

"(e.g., path typo), then exit instead of random "

"initialization.")

group = parser.add_argument_group(title='mixed precision')

group.add_argument('--fp16', action='store_true',

help='Run model in fp16 mode.')

group.add_argument('--bf16', action='store_true',

help='Run model in bfloat16 mode.')

group.add_argument('--loss-scale', type=float, default=None,

help='Static loss scaling, positive power of 2 '

'values can improve fp16 convergence. If None, dynamic'

'loss scaling is used.')

group.add_argument('--initial-loss-scale', type=float, default=2**32,

help='Initial loss-scale for dynamic loss scaling.')

group.add_argument('--min-loss-scale', type=float, default=1.0,

help='Minimum loss scale for dynamic loss scale.')

group.add_argument('--loss-scale-window', type=float, default=1000,

help='Window over which to raise/lower dynamic scale.')

group.add_argument('--hysteresis', type=int, default=2,

help='hysteresis for dynamic loss scaling')

group.add_argument('--fp32-residual-connection', action='store_true',

help='Move residual connections to fp32.')

group.add_argument('--no-query-key-layer-scaling', action='store_false',

help='Do not scale Q * K^T by 1 / layer-number.',

dest='apply_query_key_layer_scaling')

group.add_argument('--attention-softmax-in-fp32', action='store_true',

help='Run attention masking and softmax in fp32. '

'This flag is ignored unless '

'--no-query-key-layer-scaling is specified.')

group.add_argument('--accumulate-allreduce-grads-in-fp32',

action='store_true',

help='Gradient accumulation and all-reduce in fp32.')