llama3_sft.yaml

### model
model_name_or_path: meta-llama/Llama-3.2-1B-Instruct

### method
stage: sft
do_train: true
finetuning_type: full  # full or lora
dataset: alpaca_en_demo
template: llama3
cutoff_len: 2048
preprocessing_num_workers: 8

### output
output_dir: saves/mbridge/llama3_sft
logging_steps: 1
overwrite_output_dir: true

### train
per_device_train_batch_size: 1
gradient_accumulation_steps: 1
num_train_epochs: 3
max_steps: 5 # max_steps is used to limit the number of steps to train, if set, num_train_epochs will be ignored
save_steps: 3000
learning_rate: 5.0e-6
lr_scheduler_type: cosine
warmup_steps: 10
bf16: true

### megatron bridge
tensor_model_parallel_size: 2
pipeline_model_parallel_size: 1
context_parallel_size: 1
sequence_parallel: true
use_distributed_optimizer: true
overlap_param_gather: true
overlap_grad_reduce: true
mixed_precision: bf16_mixed
export_hf_on_finish: true