# 优化后的训练命令
export NCCL_IB_DISABLE=1
torchrun --nproc_per_node=4 train_mini_bert.py \
  --shards_index /root/autodl-tmp/pretraining_data/shards/pretraining_dataset.index.json \
  --tokenizer_dir /root/autodl-tmp/pretraining_data/tokenizer \
  --output_dir /root/autodl-tmp/mini-bert-output \
  --max_steps 50000 \
  --batch_size 32 \
  --grad_accum 2 \
  --num_workers 4 \
  --lr 1e-4 \
  --warmup_ratio 0.1 \
  --weight_decay 0.01 \
  --save_every 5000 \
  --log_every 50