python train_mini_bert.py \
  --shards_index /root/autodl-tmp/pretraining_data/shards/pretraining_dataset.index.json \
  --tokenizer_dir /data/tokenizer \
  --output_dir /root/autodl-tmp/mini-bert-output \
  --max_steps 100000 --batch_size 32 --grad_accum 2 --num_workers 4 \
  --lr 2e-4 --warmup_ratio 0.06 --save_every 5000 --log_every 100