Training Configuration

This document provides a comprehensive reference for all training and evaluation parameters available in ERNIEKit. It covers:

• Basic model configuration and training setup
• Evaluation metrics and strategies
• Performance optimization techniques
• Distributed training configurations
• Memory optimization options
• Checkpoint saving strategies
• Acceleration methods
• Mixed precision training settings
• Specialized configurations for SFT, LoRA, DPO and FP8 training

Each parameter is documented with its type, default value and detailed description to help developers properly configure their training jobs.

1. General Configuration

1.1 Basic Configuration

Parameter	Type	Default	Description
model_name_or_path	str	Required	Model name or local model path for the model and tokenizer
hidden_dropout_prob	float	0.0	Dropout probability for hidden layers
attention_probs_dropout_prob	float	0.0	Dropout probability for attention layers
dropout_warmup_steps	int	0	Warmup steps for dropout. Dropout probability increases linearly during warmup and disables afterward. Set to 0 to disable dropout.
weight_quantize_algo	str	Required	Model quantization algorithm. Options: weight_only_mix (expert weights as int4, other linear layers as int8) or weight_only_int8 (all linear layers as int8) or fp8_linear
output_dir	str	Required	Directory to save model files, checkpoints, tokenizers, and evaluation results
logging_steps	int	Required	Logging interval. Decrease for more frequent updates.
logging_dir	str	Required	Log directory (defaults to output_dir if unspecified)
do_eval	bool	False	Enable model evaluation
do_train	bool	False	Enable training
disable_tqdm	bool	False	Disable tqdm progress bar for estimating total training time
continue_training	bool	True	Load pretrained weights to continue training
from_hf_hub	bool	False	Downloading model from HuggingFace Hub
from_aistudio	bool	False	Downloading model from Aistudio
from_modelscope	bool	False	Downloading model from ModelScope

1.2 Evaluation

Parameter	Type	Default	Description
per_device_eval_batch_size	int	Required	Evaluation batch size (micro batch size)
eval_dataset_path	str	Required	Path to evaluation dataset (see sft-eval.jsonl
eval_dataset_prob	str	1.0	Evaluation dataset sampling probability.
eval_dataset_type	str	erniekit	Evaluation dataset type.
eval_steps	int	Required	Evaluation interval steps
evaluation_strategy	str	"steps"	Evaluation strategy. "steps" enables periodic evaluation
max_evaluate_steps	int	1	Maximum steps per evaluation (if positive)

1.3 Training Performance

Parameter	Type	Default	Description
train_dataset_path	str	Required	Training dataset path (see sft-train.jsonl)
train_dataset_prob	str	1.0	Training dataset sampling probability.
train_dataset_type	str	erniekit	Training dataset type.
max_steps	int	Required	Maximum training steps (overrides num_train_epochs if set)
num_train_epochs	int	Required	Training epochs
per_device_train_batch_size	int	Required	Training batch size (micro batch size). Global batch size = DP sharding micro_batch_size * gradient_accumulation_steps
gradient_accumulation_steps	int	Required	Gradient accumulation steps
weight_decay	float	0.0	AdamW optimizer weight decay
seed	int	42	Random seed
max_seq_len	int	Required	Maximum token length. Reduce if OOM occurs when increasing GBS.
learning_rate	float	Required	Learning rate (SFT: 3e-5, DPO: 1e-6, SFT-LoRA: 3e-4, DPO-LoRA: 1e-5)
warmup_steps	int	Required	Warmup steps (typically 1%-10% of max_steps)
lr_scheduler_type	str	linear	Learning rate scheduler (linear/cosine/polynomial/constant/constant_with_warmup)
min_lr	float	0.0	Minimum learning rate (cosine scheduler only)
layerwise_lr_decay_bound	float	1.0	Layerwise LR decay factor (0,1]. 1 means no decay.
random_shuffle	bool	True	Enable dataset shuffling
num_cycles	float	0.5	Cosine scheduler: number of waves
lr_end	float	1e-7	Polynomial scheduler: final LR
power	float	1.0	Polynomial scheduler: power
adam_beta1	float	0.9	AdamW beta1
adam_beta2	float	0.999	AdamW beta2
adam_epsilon	float	1e-8	AdamW epsilon

1.4 Distributed Training

Parameter	Type	Default	Description
tensor_parallel_degree	int	Required	Tensor parallelism degree
tensor_parallel_config	str	Required	Recommended: "sync_param sync_grad sync_moment"
tensor_parallel_output	bool	True	Enable parallel output for last Transformer layer to save memory
pipeline_parallel_degree	int	Required	Pipeline parallelism degree
pipeline_parallel_config	str	Required	Recommended: "disable_partial_send_recv enable_clear_every_step_cache enable_delay_scale_loss enable_overlap_p2p_comm best_unbalanced_scheduler"
pp_seg_method	str	Required	Pipeline layer segmentation method
virtual_pp_degree	int	1	Virtual pipeline degree (effective when pipeline_parallel_degree > 1)
add_tail_layers	int	0	Add EmptyLayers after DecodeLayer for virtual pipeline requirements
sharding_parallel_degree	int	Required	Sharding parallelism degree
sharding_parallel_config	str	Required	Recommended: "enable_stage1_overlap enable_release_grads"
sharding	str	Required	Sharding stage (stage1: optimizer, stage2: gradients, stage3: parameters)
sequence_parallel	bool	True	Enable sequence parallelism
moe_group	str	"dummy"	MoE communication group ("mp" for training, "dummy" for inference)

1.5 Memory Optimization

Parameter	Type	Default	Description
release_grads	bool	False	Release gradients after each iteration to reduce peak memory
use_sparse_head_and_loss_fn	bool	False	Use sparse LM Head and loss function
use_fused_head_and_loss_fn	bool	False	Fuse LM head and CrossEntropyLoss to save memory
use_attn_mask_startend_row_indices	bool	True	Use sparse mask representation with start row indices
recompute_use_reentrant	bool	False	Recompute implementation (PyLayer if True, hooks if False)
recompute	bool	False	Enable gradient checkpointing
recompute_granularity	str	"full"	Recompute granularity ("full"/"full_attn"/"core_attn")
offload_optim	bool	False	Offload optimizer to CPU

1.6 Checkpoint

Parameter	Type	Default	Description
save_steps	int	Required	Checkpoint save interval (when save_strategy=="steps")
save_strategy	str	"no"	Checkpoint save strategy
unified_checkpoint	bool	True	Use unified checkpoint format
unified_checkpoint_config	str	""	See Unified Checkpoint
disable_ckpt_quant	bool	False	See Unified Checkpoint
ignore_save_lr_and_optim	bool	False	Skip saving optimizer states
ignore_load_lr_and_optim	bool	False	Skip loading optimizer states
save_total_limit	int	None	Maximum number of checkpoints to keep

1.7 Acceleration

Parameter	Type	Default	Description
use_flash_attention	bool	True	Enable FlashAttention
use_sparse_flash_attn	bool	True	Enable FlashMask (requires use_attn_mask_startend_row_indices)
fuse_rope	bool	False	Fuse rotary position embedding
fuse_linear	bool	False	F fuse linear operations
greedy_intokens	bool	True	Enable greedy token-based packing. Instead of sequential sampling, a global buffer of samples is maintained and greedily packed into sequences to maximize token utilization and minimize padding.
dataloader_num_workers	int	1	Dataloader subprocess count (0 to disable)
distributed_dataloader	int	0	Use distributed dataloader for large datasets
moe_multimodal_dispatch_use_allgather	str	v2-alltoall-unpad	Optimize MoE layer with allgather+unpad

1.8 Mixed Precision (Recommended Defaults)

Parameter	Type	Default	Description
bf16	bool	False	Enable BF16 training
fp16_opt_level	str	O1	AMP level (O2 converts params to float16/bfloat16)
scale_loss	int	2 ** 15	Loss scaling factor for float16
amp_custom_white_list	str	Required	AMP O2 whitelist (e.g., "lookup_table flash_attn matmul")
amp_custom_black_list	str	Required	AMP O2 blacklist (e.g., "reduce_sum elementwise_div")
amp_master_grad	bool	False	Maintain float32 gradients for AMP O2

2. Specialized Configurations

2.1 SFT

Parameter	Type	Default	Description
num_samples_each_epoch	int	6000000	Virtual epoch size (recommend keeping default)

2.2 LoRA

Parameter	Type	Default	Description
lora_rank	int	8	LoRA rank (typical: 8/16/32. Higher improves quality but increases memory)
lora_alpha	float	-1	LoRA scaling factor (scaling = alpha/rank or alpha/sqrt(rank) for rslora)
rslora	bool	False	Enable rslora scaling (recommended for rank ≥64)
lora_plus_scale	float	1	LoRA+ learning rate multiplier (recommended: 4-16)
rslora_plus	bool	False	Enhanced LoRA (improves performance but may cause forgetting)
lora	bool	False	Enable LoRA training

2.3 DPO

Parameter	Type	Default	Description
beta	float	0.1	DPO loss temperature
simpo_gamma	float	0.5	SimPO loss gamma
offset_alpha	float	0.0	Score-based DPO loss offset
max_prompt_len	int	2048	Maximum prompt length (truncated beyond max_seq_len-10)
loss_type	str	sigmoid	Preference loss type (sigmoid/ipo/kto_pair)
pref_loss_ratio	float	1.0	Preference loss weight
sft_loss_ratio	float	0.0	Chosen data SFT loss weight
label_smoothing	float	0.0	Label smoothing for sigmoid loss
reference_free	bool	False	Disable reference model
ref_model_update_steps	int	-1	Reference model update interval (-1 to disable)

2.4 FP8 Training

Parameter	Type	Default	Description
apply_hadamard	bool	True	Use Hadamard transform for FP8 precision
use_lowprecision_moment	bool	False	Use BF16 optimizer momentum (recommended for FP8)
tensorwise_offload_optimizer	bool	False	Offload optimizer to reduce memory
apply_online_actscale_step	bool	200	Dynamic quantization scale steps
optim_shard_num	int	1	Split optimizer state files during saving to avoid memory OOM. Works only when unified_checkpoint_config: ignore_merge_optimizer.