unit_test / data /pretraining.md
herrius's picture
Upload 259 files
32b542e
# 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.