pretrain_argument.json

{
    "model_name_or_path": "gpt2-medium-en",
    "tokenizer_name_or_path": "gpt2-medium-en",
    "input_dir": "./data",
    "output_dir": "./checkpoints/gpt2_pretrain_ckpts",
    "per_device_train_batch_size": 1,
    "gradient_accumulation_steps": 2,
    "per_device_eval_batch_size": 1,
    "tensor_parallel_degree": 1,
    "pipeline_parallel_degree": 1,
    "sharding": "stage2",
    "virtual_pp_degree": 1,
    "sequence_parallel": 0,   
    "use_flash_attention": true,
    "max_seq_length": 1024,
    "learning_rate": 3e-05,
    "min_learning_rate": 3e-06,
    "warmup_steps": 30,
    "logging_steps": 1,
    "max_steps": 10000,
    "save_steps": 5000,
    "eval_steps": 1000,
    "weight_decay": 0.01,
    "fp16": true,
    "fp16_opt_level": "O2",
    "warmup_ratio": 0.01,
    "max_grad_norm": 1.0,
    "dataloader_num_workers": 1,
    "continue_training": 1,
    "do_train": true,
    "do_eval": true,
    "do_predict": true,
    "disable_tqdm": true,
    "recompute": false,
    "distributed_dataloader": 1,
    "unified_checkpoint": true,
    "recompute_granularity": "full",
    "save_total_limit": 2
}