in archived/smp-gpt-sharded-data-parallel/train.py [0:0]
def parse_args(): # pylint: disable=too-many-statements
"""Parse args."""
parser = argparse.ArgumentParser()
# hyperparameters sent by the client are passed as command-line arguments to the script.
opt_grp = parser.add_argument_group(
title="optimization", description="arguments for optimization"
)
opt_grp.add_argument(
"--train_batch_size",
type=int,
default=4,
help="batch size per dp rank, for tensor parallelism degree 8 with pipeline parallel degree 1 this means 8*this batch size per node", # pylint: disable=line-too-long
)
opt_grp.add_argument("--val_batch_size", type=int, default=4)
opt_grp.add_argument("--max_steps", "--max_training_steps", type=int, default=5000)
opt_grp.add_argument("--seed", type=int, default=12345)
opt_grp.add_argument("--same_seed", type=int, default=0)
opt_grp.add_argument("--n_gpus", type=str, default=os.environ["SM_NUM_GPUS"])
opt_grp.add_argument("--fp16", default=0, type=int, help="automatic mixed precision training")
opt_grp.add_argument("--bf16", default=0, type=int, help="automatic mixed precision training")
opt_grp.add_argument("--sharded_data_parallel_degree", default=1, type=int)
opt_grp.add_argument("--ddp_dist_backend", type=str, default="auto")
opt_grp.add_argument("--grad_clip", default=1.0, type=float, help="gradient clipping")
opt_grp.add_argument("--weight_decay", default=0.01, type=float, help="weight decay")
opt_grp.add_argument(
"--beta1", default=0.9, type=float, help="beta1 parameter for Adam optimizer"
)
opt_grp.add_argument(
"--beta2", default=0.95, type=float, help="beta2 parameter for Adam optimizer"
)
opt_grp.add_argument(
"--activation_checkpointing",
type=int,
default=1,
help="enable gradient checkpointing to reduce memory consumption",
)
parser.add_argument(
"--logging_freq", type=int, default=1, help="number of iterations between logging"
)
parser.add_argument(
"--log_param_norms",
type=int,
default=0,
help="to log param norms with logging_freq frequency, currently works only for sharded data parallel jobs", # pylint: disable=line-too-long
)
parser.add_argument(
"--log_reduced_training_loss",
type=int,
default=0,
help="to log training loss after reducing across all data parallel ranks with logging_freq frequency", # pylint: disable=line-too-long
)
# I/O
io_grp = parser.add_argument_group(title="io", description="location for input and output")
io_grp.add_argument("--use_bert_data", type=int, default=0, help="use bert data for training")
io_grp.add_argument("--zipped_data", type=int, default=1, help="input data is zipped files")
io_grp.add_argument(
"--epochs", type=int, default=3, help="times of iterating over the training dataset"
)
io_grp.add_argument("--output-data-dir", type=str, default=os.environ["SM_OUTPUT_DATA_DIR"])
io_grp.add_argument(
"--checkpoint-dir",
type=str,
default="/opt/ml/checkpoints",
help="Saves partial checkpoints (model, optimizer) to this dir, and loads latest checkpoint from this if load_partial is specified.", # pylint: disable=line-too-long
)
io_grp.add_argument(
"--model-dir",
type=str,
default=os.environ["SM_MODEL_DIR"],
help="Saves full model for inference to this dir. Also used if load_full is given to load the model. Note the lack of optimizer state here.", # pylint: disable=line-too-long
)
io_grp.add_argument("--training-dir", type=str, default=os.environ["SM_CHANNEL_TRAIN"])
io_grp.add_argument("--test-dir", type=str, default=os.environ["SM_CHANNEL_TEST"])
io_grp.add_argument(
"--parallel_proc_data_processing",
type=int,
default=0,
help="Load data in parallel with a different process. At any point a process can have two files in memory. With tensor parallelism, each of the 8 processes on an instance will then have 2 files in memory. Depending on file sizes this may or may not be feasible. With pipeline parallelism this was not a problem as only 1 rank on an instance loaded data.", # pylint: disable=line-too-long
)
io_grp.add_argument(
"--save_final_full_model",
type=int,
default=0,
help="Enabling this will save a combined model only at the end",
)
io_grp.add_argument("--load_partial", type=int, default=0, help="Load from partial checkpoints")
io_grp.add_argument("--load_full", type=int, default=0, help="Load from full checkpoints")
io_grp.add_argument(
"--logits_output", type=str, default="", help="Path to save logits and loss"
)
io_grp.add_argument("--prescaled_batch", type=int, default=1, help="use prescaled batch")
# configure model size
model_grp = parser.add_argument_group(
title="model", description="arguments to describe model configuration"
)
model_grp.add_argument(
"--fine_tune",
type=int,
default=0,
help="Fine-tune model from checkpoint or pretrained model",
)
model_grp.add_argument("--model_name", type=str, default="", help="HF model name")
model_grp.add_argument("--max_context_width", type=int, default=1024)
model_grp.add_argument("--vocab_size", type=int, default=50264)
model_grp.add_argument("--hidden_width", type=int, default=768)
model_grp.add_argument("--intermediate_size", type=int, default=2048)
model_grp.add_argument("--num_layers", type=int, default=12)
model_grp.add_argument("--num_heads", type=int, default=12)
model_grp.add_argument("--resid_pdrop", type=float, default=0.1)
model_grp.add_argument("--embd_pdrop", type=float, default=0.1)
model_grp.add_argument("--attn_pdrop", type=float, default=0.1)
model_grp.add_argument("--alibi", type=float, default=0)
model_grp.add_argument("--summary_first_pdrop", type=float, default=0.1)
model_grp.add_argument("--use_adamw", type=int, default=0, help="Use adamw optimizer")
model_grp.add_argument(
"--use_distributed_transformer", type=int, default=1, help="Use distributed transformer"
)
model_grp.add_argument(
"--checkpoint_sublayers",
type=int,
default=0,
help="Apply activation checkpointing to submodules of each transformer layer",
)
model_grp.add_argument("--initializer_range", type=float, default=0.02)
smp_grp = parser.add_argument_group(title="smp", description="smp")
smp_grp.add_argument("--tensor_parallel_degree", type=int, default=1)
smp_grp.add_argument("--pipeline_parallel_degree", type=int, default=1)
smp_grp.add_argument("--microbatches", type=int, default=1)
smp_grp.add_argument("--active_microbatches", type=int, default=None)
smp_grp.add_argument("--optimize", type=str, default="speed")
smp_grp.add_argument("--activation_strategy", type=str, default="each")
smp_grp.add_argument("--shard_optimizer_state", type=int, default=0)
smp_grp.add_argument("--offload_activations", type=int, default=0)
smp_grp.add_argument("--fast_mode", type=int, default=0)
smp_grp.add_argument("--static_mode", type=int, default=0)
smp_grp.add_argument("--delayed_param", type=int, default=0)
smp_grp.add_argument("--same_partition_load", type=int, default=0)
smp_grp.add_argument(
"--attention_in_fp32",
type=int,
default=0,
help="When using FP16 and if the activations overflow, doing the attention computation in fp32 may help. But note that this can substantially increase memory usage and reduce performance. We recommend using bf16 instead which is more numerically stable and would not need this.", # pylint: disable=line-too-long
)
smp_grp.add_argument(
"--residual_addition_in_fp32",
type=int,
default=0,
help="When using FP16 and if the activations overflow, adding residuals in fp32 may help. But note that this can substantially increase memory usage and reduce performance. We recommend using bf16 instead which is more numerically stable and would not need this.", # pylint: disable=line-too-long
)
smp_grp.add_argument("--placement_strategy", type=str, default="cluster")
smp_grp.add_argument("--activation_loading_horizon", type=int, default=4)
smp_grp.add_argument("--skip_tracing", type=int, default=0)
smp_grp.add_argument("--query_key_layer_scaling", type=int, default=0)
smp_grp.add_argument("--fused_softmax", type=int, default=1)
smp_grp.add_argument("--flash_attention", type=int, default=1)
smp_grp.add_argument("--fused_dropout", type=int, default=0)
smp_grp.add_argument("--fused_bias_gelu", type=int, default=1)
smp_grp.add_argument("--gradient_accumulation", type=int, default=1)
smp_grp.add_argument("--model_type", type=str, default="gpt2")
smp_grp.add_argument("--rotary_pct", type=float, default=0.25)
smp_grp.add_argument("--rotary_emb_base", type=int, default=10000)
parser.add_argument(
"--num_kept_checkpoints",
type=int,
default=5,
help="how many checkpoints to keep before deleting",
)
parser.add_argument(
"--checkpoint_freq",
type=int,
default=10000,
help="number of iterations between checkpointing",
)
parser.add_argument(
"--validation_freq",
type=int,
default=None,
help="number of iterations to print validation loss",
)
parser.add_argument(
"--validation_batches",
type=int,
default=10,
help="number of batches to estimate validation loss",
)
parser.add_argument(
"--manual_partition",
type=int,
default=0,
help="evenly distribute layers across the partitions",
)
parser.add_argument(
"--partition_assignment",
type=str,
default="",
help="number of transformer layers assigned to each partition",
)
parser.add_argument(
"--preserve_np_state",
type=int,
default=0,
help="Perserve the numpy random state between validation",
)
parser.add_argument(
"--fast_validation",
type=int,
default=1,
help="Running validation only with the last data file for faster speed",
)
parser.add_argument(
"--gather_if_shard",
type=int,
default=1,
help="When sharding opt states is enabled, gather the opt checkpoint to rdp rank 0 during saving", # pylint: disable=line-too-long
)
parser.add_argument(
"--clean_cache",
type=int,
default=0,
help="Clean torch reserved memory at he end of every step",
)
parser.add_argument("--use_fsx", type=int, default=0, help="Using FSx for checkpointing")
parser.add_argument(
"--enable_memory_profiling", type=int, default=0, help="Enable memory profile"
)
# learning rate
lr_grp = parser.add_argument_group(
title="lr", description="arguments for learning rate schedule"
)
lr_grp.add_argument("--lr", type=float, default=None, help="Initial learning rate.")
lr_grp.add_argument(
"--lr_decay_style",
type=str,
default="linear",
choices=["constant", "linear", "cosine", "exponential", "plateau"],
help="Learning rate decay function.",
)
lr_grp.add_argument(
"--lr_decay_iters",
type=int,
default=None,
help="number of iterations to decay learning rate over," " If None defaults to train iters",
)
lr_grp.add_argument(
"--min_lr",
type=float,
default=0.0,
help="Minumum value for learning rate. The scheduler" "clip values below this threshold.",
)
lr_grp.add_argument(
"--warmup",
type=float,
default=0.01,
help="Percentage of total iterations to warmup on "
"(.01 = 1 percent of all training iters).",
)
lr_grp.add_argument(
"--plateau",
type=float,
default=0.4,
help="Percentage of total iterations to keep at max if using plateau lr",
)
ci_grp = parser.add_argument_group(title="ci", description="ci related settings")
ci_grp.add_argument("--ci", default=False, action="store_true", help="Whether enable ci")
ci_grp.add_argument("--time_to_train", type=int, help="time to train threshold")
ci_grp.add_argument("--throughput", type=float, help="throughput threshold")
ci_grp.add_argument("--loss", type=float, help="loss threshold")
args, _ = parser.parse_known_args()
return args