pretrain.sh

#!/bin/bash
echo "Export params ..."
export HIP_VISIBLE_DEVICES=0,1,2,3 # 自行修改为训练的卡号和数量

torchrun --nproc_per_node {number of gpus} \
    -m FlagEmbedding.baai_general_embedding.retromae_pretrain.run \
    --output_dir {path to save model} \
    --model_name_or_path  {path to finetune model} \
    --train_data {path of train datasets}  \
    --learning_rate 2e-5 \
    --num_train_epochs 2 \
    --per_device_train_batch_size {large batch size; set 1 for toy data} \
    --dataloader_drop_last True \
    --max_seq_length 512 \
    --logging_steps 10 \
    --dataloader_num_workers 12