Use cases and training parameter details for Pai-Megatron-Patch - Platform For AI

Use cases

Black-box accelerated fine-tuning
Black-box accelerated pre-training
Whitebox implementation of accelerated fine-tuning based on the Finetuner code template
Whitebox implementation of accelerated pre-training based on the Pretrainer code template

Parameter name	Value type	Required	Default value	Enumeration	Description
--micro-batch-size	int	Yes	None	N/A	The batch size for each card.
--global-batch-size	int	Yes	None	N/A	The total batch size for distributed training.
--tokenizer-type	str	No	None	BertWordPieceLowerCase BertWordPieceCase GPT2BPETokenizer	The type of the tokenizer.
--split	str	No	969, 30, 1	N/A	The split ratio for the pre-training, validation, and test datasets.
--data-impl	str	No	mmap	lazy cached mmap infer	The implementation method for the indexed dataset during pre-training.
--data-path	str	Yes	None	N/A	The file path of the pre-training dataset.
--data-dir	str	No	None	N/A	The folder path of the fine-tuning dataset.
--data-name	str	Yes	None	N/A	The file name of the fine-tuning dataset.

Parameter name	Value type	Required	Default value	Enumeration	Description
--pretrained-model-name-or-path	str	Yes	None	N/A	The name or path of the pre-trained model.
--num-layers	int	Yes	None	N/A	The number of layers.
--hidden-size	int	Yes	None	N/A	The dimension of the hidden layer.
--num-attention-heads	int	Yes	None	N/A	The number of heads in the self-attention layer.
--max-position-embeddings	int	Yes	None	N/A	The sequence length for position embeddings.
--seq-length	int	Yes	None	N/A	The sequence length.

Parameter name	Value type	Required	Default value	Enumeration	Description
--task	str	Yes	None	sequence_classification token_classification question_answering masked_lm casual_lm seq2seq_lm pretraining	The name of the training task to execute.
--save	str	Yes	None	N/A	The folder path to save the model.
--lr	float	Yes	None	N/A	The learning rate.
--lr-decay-style	str	Yes	linear	constant linear cosine	The learning rate decay style.
--weight-decay	float	Yes	0.01	N/A	The L2 regularization weight decay coefficient.
--clip-grad	float	No	1	N/A	Gradient clipping.
--lr-warmup-fraction	float	No	None	N/A	The learning rate warmup fraction.
--train-iters	int	Yes	None	N/A	The number of training iterations.
--epochs	int	No	None	N/A	The number of training epochs.
--log-interval	int	No	100	N/A	The interval for log printing.

Parameter name	Value type	Required	Default value	Enumeration	Description
--num-experts	int	No	None	N/A	The number of experts in the MLP layer.

Parameter name	Value type	Required	Default value	Enumeration	Description
--mixed-precision	bool	No	None	N/A	Enables FP16 mixed-precision training.

Note: This must be used with the Trainer, PreTrainer, or Finetuner provided by Pai-Megatron-Patch.

Parameter name	Value type	Required	Default value	Enumeration	Description
--oss-memory-optimization	bool	No	N/A	N/A	Enables partitioning of the optimizer state.
--oss-sdp-memory-optimization	bool	No	N/A	N/A	Enables partitioning of the optimizer state and gradients.
--fsdp-memory-optimization	bool	No	N/A	N/A	Enables partitioning of the optimizer state, gradients, and model parameters.
--zero-1-memory-optimization	bool	No	N/A	N/A	Enables partitioning of the optimizer state using ZeRO.
--zero-2-memory-optimization	bool	No	N/A	N/A	Enables partitioning of the optimizer state and gradients using ZeRO.
--zero-3-memory-optimization	bool	No	N/A	N/A	Enables partitioning of the optimizer state, gradients, and model parameters using ZeRO.

Note: ZeRO memory optimization can only be used with the Trainer. OSS, SDP, and FSDP memory optimizations can be used in code without the Trainer.

Parameter name	Value type	Required	Default value	Enumeration	Description
--tensor-model-parallel-size	int	No	1	N/A	tensor parallelism
--pipeline-model-parallel-size	int	No	1	N/A	Pipeline Parallelism

Notes:

Parameter name	Value type	Required	Default value	Enumeration	Description
--onnx-runtime-training	bool	No	None	N/A	Enables computational graph optimization using ONNX Runtime.

Parameter name	Value type	Required	Default value	Enumeration	Description
--cpu-offload	bool	No	None	N/A	Enabling CPU load training

Note: CPU offload training for memory optimization can only be used with ZeRO. It cannot be used with OSS, SDP, or FSDP model state partitioning.

Parameter name	Value type	Required	Default value	Enumeration	Description
--checkpoint-activations	bool	No	None	N/A	Enables activation checkpointing.

Parameter name	Value type	Required	Default value	Enumeration	Description
--micro-batch-size	int	Yes	1	N/A	Denominator for gradient accumulation
--global-batch-size	int	Yes	1	N/A	Gradient accumulation numerator

Notes:

This applies only to iteration-based pre-training with PreTrainer and Finetuner. Gradient accumulation is not supported for epoch-based fine-tuning.
The number of accumulation steps is automatically calculated based on the batch size and the number of ranks.

Parameter name	Value type	Required	Default value	Enumeration	Description
--data-iterator dynamic-shape	str	No	None	dynamic-shape fixed-shape	Data iterator for dynamic shapes

Note: This applies only to pre-training with PreTrainer and is not supported for fine-tuning.

Parameter name	Value type	Required	Default value	Enumeration	Description
--optimizers	str	Yes	apex_adam	apex_adam apex_lamb	Apex-fused Adam/LAMB