configs

Configuration and parameters

GPT-NeoX parameters are defined in a YAML configuration file which is passed to the deepy.py launcher - for examples see the files contained in this folder. Parameters originate from either the DeepSpeed runner CLI (DSL), DeepSpeed configuration file (DSC), Megatron-LM CLI (Meg) or are GPT-NeoX (NeoX) modifications.

Example Configuration (GPT3 Small):

Below is an example configuration .yaml to train a ~160M parameter GPT model. This readme will go through each section in the configuration and the options available.

# GPT-3 pretraining setup
{
   # parallelism settings ( you will want to change these based on your cluster setup, ideally scheduling pipeline stages
   # across the node boundaries )
   "pipe-parallel-size": 1,
   "model-parallel-size": 1,

   # model settings
   "num-layers": 12,
   "hidden-size": 768,
   "num-attention-heads": 12,
   "seq-length": 2048,
   "max-position-embeddings": 2048,
   "norm": "rmsnorm",
   "pos-emb": "none",
   "no-weight-tying": true,
    # this should provide some speedup but takes a while to build, set to true if desired
   "scaled-upper-triang-masked-softmax-fusion": false,
   "train-iters": 320000,

   # optimizer settings
   "optimizer": {
     "type": "Adam",
     "params": {
       "lr": 0.0006,
       "max_grad_norm": 1.0,
       "betas": [0.9, 0.95]
     }
   },
   "zero_optimization": {
    "stage": 0,
    "allgather_partitions": True,
    "allgather_bucket_size": 500000000,
    "overlap_comm": True,
    "reduce_scatter": True,
    "reduce_bucket_size": 500000000,
    "contiguous_gradients": True,
    "cpu_offload": False
  },

   # batch / data settings
   "train_micro_batch_size_per_gpu": 4,
   "gradient_accumulation_steps": 1,
   "data-impl": "mmap",
   "split": "949,50,1",

   # activation checkpointing
   "checkpoint-activations": true,
   "checkpoint-num-layers": 1,
   "partition-activations": true,
   "synchronize-each-layer": true,

   # regularization
   "gradient_clipping": 1.0,
   "weight-decay": 0,
   "hidden-dropout": 0,
   "attention-dropout": 0,

   # precision settings
   "fp16": { 
     "enabled": true,
     "loss_scale": 0,
     "loss_scale_window": 1000,
     "hysteresis": 2,
     "min_loss_scale": 1
   },

   # lr decay settings
   "lr-decay-iters": 320000,
   "lr-decay-style": "cosine",
   "warmup": 0.01,
  
   # misc. training settings
   "distributed-backend": "nccl",
   "save-interval": 10000,
   "eval-interval": 1000,
   "eval-iters": 10,

   # logging
   "log-interval": 100,
   "steps_per_print": 10,
   "keep-last-n-checkpoints": 4,
   "wall_clock_breakdown": true,
}

Parallelism Settings:

The parallelism settings are left at 1 in all configs, as the settings you want will be highly dependent on your compute setup and network topology. We have found it best to do model parallelism within a node, and schedule pipeline stages across node boundaries.

   "pipe-parallel-size": 1,
   "model-parallel-size": 1,

These can be set to any integer between 0 and num_gpus, and num_gpus must be divisible by pipe_parallel_size * model_parallel_size.

All Parallelism Settings are below:

Origin	Parameter name	Default value	Description
Meg	pipe-parallel-size	0	Number of pipeline parallel stages. Disable with 0.
Meg	model-parallel-size	1	Size of the model parallelism.
Meg	pipe-partition-method	"type:transformer"	method used to distribute model layers across pipeline stages. Choose from "parameters", which balances the number of parameters on each pipeline stage, "uniform", which naively balances the number of layers per stage, or "type:[regex]" (in our case this will basically only be "type:transformer"), which balances layers whose class names match [regex]

Model Settings:

   # model settings
   "num-layers": 12,
   "hidden-size": 768,
   "num-attention-heads": 12,
   "seq-length": 2048,
   "max-position-embeddings": 2048,
   "norm": "rmsnorm",
   "pos-emb": "none",
   "no-weight-tying": true,
    # this should provide some speedup but takes a while to build, set to true if desired
   "scaled-upper-triang-masked-softmax-fusion": false,
   "train-iters": 320000,

An example of some basic settings used to configure your model's architecture and number of training steps.

All available options for model settings are below:

Origin	Parameter name	Default value	Description
Meg	train-iters		Total number of iterations to train over all training runs.
Meg	num-layers		Number of transformer layers.
Meg	hidden-size		Transformer hidden size.
Meg	num-attention-heads		Number of transformer attention heads.
Meg	seq-length		Maximum sequence length to process.
Meg	max-position-embeddings		Maximum number of position embeddings to use. This is the size of position embedding.
NeoX	norm	layernorm	Normalization layer to use. Choose from "layernorm", "rmsnorm" and "scalenorm".
Meg	layernorm-epsilon	1e-05	Layer norm epsilon.
NeoX	rms-norm-epsilon	1e-8	Root mean squared norm epsilon
NeoX	scalenorm-epsilon	1e-8	Scalenorm epsilon
NeoX	pos-emb	learned	Type of positional embedding to use - choose from 'learned', 'sinusoidal', 'rpe', 'rotary', 'none'
NeoX	rpe-num-buckets	32	T5 relative positional encoding number of buckets, default 32.
NeoX	rpe-max-distance	128	T5 relative positional encoding max distance, default 128.
NeoX	no-weight-tying	false	Disables weight tying between embedding weights and final Linear layer
NeoX	geglu	false	Enable geglu activation function (WARNING: will increase memory usage, adjust embd dims accordingly)
NeoX	sparsity	none	Sparse attention layer configuration: none = all regular attn, all = all sparse attn, interspersed = sparse on odd layers, dense on even.
Meg	num-unique-layers		Number of unique transformer layers. num-layers should be divisible by this value. Currently only has an effect when pipe_parallel_size=0.
Meg	param-sharing-style	grouped	Ordering of the shared parameters. For example, for a num-layers=4 and --num-unique-layers=2, we will have the following ordering for two unique layers 1 and 2-: grouped: [1, 2, 1, 2] and spaced: [1, 1, 2, 2].
Meg	make-vocab-size-divisible-by	128	Pad the vocab size to be divisible by this value. This is added for computational efficiency reasons.
Meg	apply-residual-connection-post-layernorm	false	If set, use original BERT residual connection ordering.
Meg	openai-gelu	false	Use OpenAIs GeLU implementation. This option should not be used unless for backward compatibility reasons.
Meg	scaled-upper-triang-masked-softmax-fusion	false	Enable fusion of query_key_value_scaling time (upper diagonal) masking and softmax.
Meg	scaled-masked-softmax-fusion	false	Enable fusion of query_key_value_scaling general masking and softmax.
Meg	bias-gelu-fusion	false	Enable bias and gelu fusion.
Meg	bias-dropout-fusion	false	Enable bias and dropout fusion.
Meg	fp16-lm-cross-entropy	false	Move the cross entropy unreduced loss calculation for lm head to fp16.
Meg	init-method-std	0.02	Standard deviation of the zero mean normal distribution used for weight initialization.
Meg	apply-query-key-layer-scaling	false	Scale Q * K^T by 1 / layer-number. If this flag is set, then it will automatically set attention-softmax-in-fp32 to true
Meg	use-cpu-initialization	false	If set, affine parallel weights initialization uses CPU
Meg	attention-softmax-in-fp32	false	Run attention masking and softmax in fp32.
Meg	fp32-allreduce	false	All-reduce in fp32

Optimizer Settings:

Our optimizer configuration has a similar syntax to deepspeed's. Different optimizers will have different arguments for "params". Learning rate should be configured from here using the "lr" field of optimizer["params"].

  # optimizer settings
   "optimizer": {
     "type": "Adam",
     "params": {
       "lr": 0.0006,
       "max_grad_norm": 1.0,
       "betas": [0.9, 0.95]
     }
   }

Available optimizer types are:

• "Adam": regular Adam optimizer
• "OneBitAdam": Deepspeed's OneBitAdam optimizer. To use 1-bit adam, you'll also need to add the freeze_step, cuda_aware, and comm_backend_name fields, like so:

   "optimizer": {
     "type": "OneBitAdam",
     "params": {
       "lr": 0.0001,
       "freeze_step": 23000,
       "betas": [0.9, 0.95],
       "cuda_aware": false,
       "comm_backend_name": "nccl"
     }

• "CPU_Adam"/"CPU_torch_adam": Adam optimizer on CPU. Either megatron's version ("CPU_Adam") or torch's ("CPU_torch_adam")
• "SM3": SM3 or Memory adaptive efficient optimization optimizer. We have found this doesn't work well with fp16 training.
• "madgrad_wd": MADGRAD or [A Momentumized, Adaptive, Dual Averaged Gradient Method for Stochastic Optimizer] weight decay has been implemented AdamW style instead of the original madgrad Adam style. https://arxiv.org/abs/2101.11075

ZeRO Optimization:

  "zero_optimization": {
        "stage": 0,
        "allgather_partitions": True,
        "allgather_bucket_size": 500000000,
        "overlap_comm": True,
        "reduce_scatter": True,
        "reduce_bucket_size": 500000000,
        "contiguous_gradients": True,
        "cpu_offload": False
  },
  "zero_allow_untested_optimizer": false,

ZeRO optimization in NeoX is currently configured identically to how deepspeed configures it, please see the deepspeed docs for more information.

If you want to combine an optimizer untested by DeepSpeed with ZeRO (i.e, not ADAM or LAMB), you must pass "zero_allow_untested_optimizer": true outside of the "zero_optimization" dictionary (see above).

N.B - ZeRO stages 2+ are incompatible with pipeline parallelism. Please set "pipe-parallel-size" to 0 if you want to use ZeRO stage 2 or more.

Batch Size Settings:

   # batch / data settings
   "train_micro_batch_size_per_gpu": 4,
   "gradient_accumulation_steps": 1,

Our global batch size configuration follows deepspeed's and can be configured in a number of ways. At least any one of "train_batch_size" and "train_micro_batch_size_per_gpu".

• "train_batch_size": The effective training batch size. This is the amount of data samples that leads to one step of model update. train_batch_size is aggregated by the batch size that a single GPU processes in one forward/backward pass (a.k.a., train_step_batch_size), the gradient accumulation steps (a.k.a., gradient_accumulation_steps), and the number of GPUs.
• "train_micro_batch_size_per_gpu"": Batch size to be processed by one GPU in one step (without gradient accumulation). When specified, gradient_accumulation_steps is automatically calculated using train_batch_size and number of GPUs.
• "gradient_accumulation_steps": Number of training steps to accumulate gradients before averaging and applying them. This feature is sometimes useful to improve scalability since it results in less frequent communication of gradients between steps. Another impact of this feature is the ability to train with larger batch sizes per GPU. When specified, train_step_batch_size is automatically calculated using train_batch_size and number of GPUs.

Dataset / Tokenizer / Checkpoint / Logging Settings:

   "data-impl": "mmap",
   "split": "949,50,1",
   # Suggested data paths when using GPT-NeoX locally
   "data-path": "data/enron/enron_text_document",
   #"train-data-path": "data/train/train_text_document",
   #"test-data-path": "data/test/test_text_document",
   #"valid-data-path": "data/valid/valid_text_document",
   "vocab-file": "data/gpt2-vocab.json",
   "merge-file": "data/gpt2-merges.txt",
   "save": "checkpoints",
   "load": "checkpoints",
   "tensorboard-dir": "tensorboard",
   "log-dir": "logs",
   "save-interval": 10000,
   "eval-interval": 1000,
   "eval-iters": 10,

These are the settings used to control the dataloading, train/test/val splits, tokenization and checkpointing. All available options for these settings are below:

Origin	Parameter name	Default value	Description
Meg	data-path		Path to combined dataset to split.
Meg	data-impl	infer	Implementation of indexed datasets.
Meg	mmap-warmup	false	Warm up mmap files.
Meg	save		Output directory to save checkpoints to.
Meg	load		Directory containing a model checkpoint.
Meg	save-interval		Number of iterations between checkpoint saves.
Meg	seed	1234	Random seed used for python, numpy, pytorch, and cuda.
Meg	no-save-optim	false	Do not save current optimizer.
Meg	no-save-rng	false	Do not save current rng state.
Meg	no-load-optim	false	Do not load optimizer when loading checkpoint.
Meg	no-load-rng	false	Do not load rng state when loading checkpoint.
Meg	finetune	false	Load model for finetuning. Do not load optimizer or rng state from checkpoint and set iteration to 0. Assumed when loading a release checkpoint.
Meg	eval-iters	100	Number of iterations to run for evaluationvalidation/test for.
Meg	eval-interval	1000	Interval between running evaluation on validation set.
Meg	split	969, 30, 1	Comma-separated list of proportions for training, validation, and test split. For example the split 90,5,5 will use 90% of data for training, 5% for validation and 5% for test.
Meg	vocab-file		Path to the vocab file.
Meg	merge-file		Path to the BPE merge file.
NeoX	log-dir		Directory to save logs to.
NeoX	tensorboard-dir		Write TensorBoard logs to this directory.
Meg	num-workers	2	Dataloader number of workers.
DSC	steps_per_print	10	Print train loss every N steps
DSC	wall_clock_breakdown	false	Enable timing of the latency of forward/backward/update training phases
DSC	dump_state	false	Print out state information of DeepSpeed object after initialization
Meg	tokenizer-type		What type of tokenizer to use. DEPRECATED - currently only GPT2Tokenizer is available.
Meg	exit-interval		Exit the program after the iteration is divisible by this value.

LR Scheduler settings

   "lr-decay-iters": 320000,
   "lr-decay-style": "cosine",
   "warmup": 0.01,

Settings used to modify the learning rate over time. All settings available:

Origin	Parameter name	Default value	Description
Meg	lr-decay-style	linear	Learning rate decay function. Choose from 'constant', 'linear', 'cosine', 'exponential'.
Meg	lr-decay-iters		Number of iterations to decay learning rate over, If None defaults to --train-iters
Meg	min-lr	0.0	Minumum value for learning rate. The scheduler clips values below this threshold.
Meg	warmup	0.01	Percentage of total iterations to warmup on (.01 = 1 percent of all training iters).
Meg	override-lr-scheduler	false	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.
Meg	use-checkpoint-lr-scheduler	false	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.

N.B - OneBitAdam requires you to use deepspeed's internal lr scheduler because reasons. Currently the lr decay style defaults to deepspeed's WarmupDecayLR. min lr, lr and warmup should still be configured as above. We're working on making this more flexible.

Regularization Settings:

   "gradient_clipping": 1.0,
   "weight-decay": 0,
   "hidden-dropout": 0,
   "attention-dropout": 0,

Various settings to regularize the model. At larger scales, we find that regularization only slows down training and has little to negative effect. At smaller scales and on smaller datasets, however, it can improve performance. All regularization settings below:

Origin	Parameter name	Default value	Description
DSC	gradient_clipping	0	Enable gradient clipping with provided value
Meg	attention-dropout	0.1	Post attention dropout probability.
Meg	hidden-dropout	0.1	Dropout probability for hidden state transformer.
Meg	weight-decay	0.01	Weight decay coefficient for L2 regularization.

Activation Checkpointing Settings:

   "checkpoint-activations": true,
   "checkpoint-num-layers": 1,
   "partition-activations": true,
   "synchronize-each-layer": true,

Checkpointing works by trading compute for memory. Rather than storing all intermediate activations of the entire computation graph for computing backward, the checkpointed part does not save intermediate activations, and instead recomputes them in backward pass. All options for configuring activation checkpointing are below:

Origin	Parameter name	Default value	Description
Meg	checkpoint-activations	false	Checkpoint activation to allow for training with larger models, sequences, and batch sizes.
Meg	checkpoint-num-layers	1	Chunk size (number of layers) for checkpointing.
Meg	distribute-checkpointed-activations	false	If set, distribute checkpointed activations across model parallel group.
Meg	deepspeed-activation-checkpointing	false	Uses activation checkpointing from deepspeed
Meg	contiguous-checkpointing	false	Contiguous memory checkpointing for activations.
Meg	checkpoint-in-cpu	false	Move the activation checkpoints to CPU.
Meg	synchronize-each-layer	false	does a synchronize at the beginning and end of each checkpointed layer.
Meg	profile-backward	false	Enables backward pass profiling for checkpointed layers.
Meg	partition-activations	false	Partition Activations across GPUs before checkpointing.

Mixed Precision Training Settings:

gpt-neox's mixed precision training is configured identically to DeepSpeed's, please see their documentation for more information. An example config for fp16 training:

   "fp16": { 
     "enabled": true,
     "loss_scale": 0,
     "loss_scale_window": 1000,
     "hysteresis": 2,
     "min_loss_scale": 1
   },

To train in fp32, simply set fp16["enabled"] to false.

Deepspeed Launcher Options

Origin	Parameter name	Default value	Description
DSL	hostfile	/job/hostfile	Hostfile path (in MPI style) that defines the resource pool available to the job (e.g., worker-0 slots=4)
DSL	include		Specify hardware resources to use during execution. String format is NODE_SPEC[@NODE_SPEC ...] where NODE_SPEC=NAME[:SLOT[,SLOT ...]]. If :SLOT is omitted, include all slots on that host. Example: "worker-0@worker-1:0,2" will use all slots. on worker-0 and slots [0, 2] on worker-1.
DSL	exclude		Specify hardware resources to NOT use during execution. Same format as include.
DSL	num_nodes	-1	Total number of worker nodes to run on, this will use the top N hosts from the given hostfile. -1 will use all.
DSL	num_gpus	-1	Max number of GPUs to use on each node, will use [0:N) GPU ids on each node. -1 will use all.
DSL	master_port	29500	Port used by PyTorch distributed for communication during training.
DSL	master_addr		IP address of node 0, will be inferred via 'hostname -I' if not specified.
DSL	launcher	pdsh	Launcher backend for multi-node training. Options currently include PDSH, OpenMPI, MVAPICH.
DSL	detect_nvlink_pairs	false	If true, autodetects nvlink pairs and remaps cuda visible devices to place them next to each other. This is an Eleuther addition to deepspeed, and should speed up model parallel training on setups with nvlink pairs when mp=2.

Other Configuration Options

And finally, a few leftover options that don't fit into any particular category.

Origin	Parameter name	Default value	Description
Meg	distributed-backend	nccl	Which backend to use for distributed training.
Meg	local_rank		local rank passed from distributed launcher.
Meg	lazy-mpu-init		If set to True, initialize_megatron() skips DDP initialization and returns function to complete it instead. Also turns on use-cpu-initialization flag. This is for external DDP manager.
Meg	short-seq-prob	0.1	Probability of producing a short sequence.
Meg	reset-position-ids	false	Reset posistion ids after end-of-document token.
Meg	reset-attention-mask	false	Reset self attention maske after end-of-document token.
Meg	eod-mask-loss	false	Mask loss for the end of document tokens.
Meg	adlr-autoresume	false	Enable auto-resume on adlr cluster.
Meg	adlr-autoresume-interval	1000	Intervals over which check for auto-resume termination signal
Meg	seed	1234	Random seed used for python, numpy, pytorch, and cuda.
NeoX	deepspeed_mpi	false	Run via MPI, this will attempt to discover the necessary variables to initialize torch distributed from the MPI environment
DSC	prescale_gradients	false	Scale gradients before doing allreduce.
DSC	gradient_predivide_factor	1.0	Before gradient averaging predivide gradients by a specified factor, can sometimes help with fp16 stability when scaling to large numbers of GPUs
DSC	sparse_gradients	false	Enable sparse compression of torch.nn.Embedding gradients.
DSC	amp		Dictionary as described in Deepspeed documentation.
DSC	flops_profiler		Dictionary as described in Deepspeed documentation.