| # Pre-Training | |
| ## Preparation | |
| You should prepare training data following [PREPARE_DATA.md](prepare_data.md), and make sure that the environment variable `DATA_PATH` is indeed the location path that stores pre-training data. | |
| ``` | |
| echo $DATA_PATH | |
| ``` | |
| ## Training on singe node | |
| ``` | |
| python -m torch.distributed.launch --nproc_per_node=8 --master_port=$PORT \ | |
| main.py --num-gpus $GPUS --config-file ${CONFIG} OUTPUT_DIR $WORK_DIR | |
| ``` | |
| where `$GPUS` is GPU number, `${CONFIG}` is the configuration file, `$PORT` is the specified available port used for distributed training, and `$WORK_DIR` is the directory used to store checkpoint and training log. | |
| For exmaple, the command for pre-training a Uni-Perceiver-Tiny model with `configs/BERT_L12_H192_experiments/7tasks_berttiny_training.yaml` is as folllowing: | |
| ``` | |
| python -m torch.distributed.launch --nproc_per_node=8 --master_port=26511 \ | |
| main.py --num-gpus 8 --config-file configs/BERT_L12_H192_experiments/7tasks_berttiny_training.yaml OUTPUT_DIR work_dirs/exp_demo_log | |
| ``` | |
| Another training example with gradient accumulation : | |
| ``` | |
| python -m torch.distributed.launch --nproc_per_node=4 --master_port=26511 \ | |
| main.py --num-gpus 4 --config-file configs/BERT_L12_H384_experiments/in1k_training.yaml SOLVER.ACCUM_ITER 2 OUTPUT_DIR work_dirs/deepspeed_moe/BERT_L12_H384_experiments/debug | |
| ``` | |
| ## Evaluation without any tuning | |
| You can evaluate the pre-training tasks by adding the `--eval-only` argument. | |
| ``` | |
| python -m torch.distributed.launch --nproc_per_node=8 --master_port=$PORT \ | |
| main.py --num-gpus $GPUS --config-file ${CONFIG} --eval-only OUTPUT_DIR $WORK_DIR | |
| ``` | |
| ## Training on multiple nodes | |
| For example, the command for training Uni-Perceiver on 2 nodes of each with 8 GPUs is as following: | |
| On node 1: | |
| ``` | |
| MASTER_ADDR=<IP address of node 1> NODE_RANK=0 GPUS_PER_NODE=8 | |
| python -m torch.distributed.launch --nproc_per_node=8 --master_port=$PORT \ | |
| main.py --num-gpus $GPUS --config-file ${CONFIG} OUTPUT_DIR $WORK_DIR | |
| ``` | |
| On node 2: | |
| ``` | |
| MASTER_ADDR=<IP address of node 1> NODE_RANK=1 GPUS_PER_NODE=8 | |
| python -m torch.distributed.launch --nproc_per_node=8 --master_port=$PORT \ | |
| main.py --num-gpus $GPUS --config-file ${CONFIG} OUTPUT_DIR $WORK_DIR | |
| ``` | |
| ## Training on slurm cluster | |
| If you are using slurm cluster, you can simply run the following command to train Uni-Perceiver on `GPUS/8` nodes with `GPUS` GPUs: | |
| ``` | |
| sh run.sh ${CONFIG} ${JOBNAME} ${GPUS} ${PARTITION} | |
| ``` | |
| * Note: you should change the `DATA_PATH` in the script `./run.sh` before your training. | |
| ## Pre-Training of Uni-Perceiver models | |
| To save the computation cost, Uni-Perceiver and Uni-Perceiver-MoE are both pre-trained in a two-stage way: | |
| they are pre-trained with the image resolution of 160x160 firstly, and then are pre-trained for another 10% of total iterations on a higher resolution of 224x224. | |
| The two-stage training strategy makes our training more effective. | |
| ### Uni-Perceiver | |
| Take __Uni-Perceiver-Base__ as an example, the 1-st pre-training stage can be conducted as | |
| ``` | |
| sh run.sh configs/BERT_L12_H768_experiments/16tasks_training_basedense_stage1_64gpu.yaml base_pretrain_stage1 64 partitionname | |
| ``` | |
| After the 1-stage, you can run the 2-nd stage pre-training as | |
| ``` | |
| sh run.sh configs/BERT_L12_H768_experiments/16tasks_training_basedense_stage2_64gpu.yaml base_pretrain_stage2 64 partitionname MODEL.WEIGHTS work_dirs/BERT_L12_H768_experiments/16tasks_training_basedense_stage1_64gpu/base_pretrain_stage1/model_Epoch_200000_Iter_0199999.pth | |
| ``` | |
| ### Uni-Perceiver-MoE | |
| The __Uni-Perceiver-MoE__ model can also be pre-trained in a similar way, which also follows two-stage pre-training. | |
| ``` | |
| sh run.sh configs/BERT_L12_H768_experiments/16tasks_training_basemoe_stage1_56gpu.yaml base_moe_pretrain_stage1 56 partitionname | |
| ``` | |
| ``` | |
| sh run.sh configs/BERT_L12_H768_experiments/16tasks_training_basemoe_stage2_56gpu.yaml base_moe_pretrain_stage2 56 partitionname MODEL.WEIGHTS work_dirs/BERT_L12_H768_experiments/16tasks_training_basemoe_stage1_56gpu/base_moe_pretrain_stage1/model_Epoch_200000_Iter_0199999.pth | |
| ``` | |
| By the way, you should adjust the training iteration and learning scheduler accordingly as you use a different number of GPUs. |