The codebase is mainly built with following libraries:
-
Python 3.6 or higher
-
PyTorch and torchvision.
We can successfully reproduce the main results under two settings below:
Tesla A100 (40G): CUDA 11.1 + PyTorch 1.8.0 + torchvision 0.9.0
Tesla V100 (32G): CUDA 10.1 + PyTorch 1.6.0 + torchvision 0.7.0 -
DS_BUILD_OPS=1 pip install deepspeed
- We recommend you to use
PyTorch >= 1.8.0
. - We observed accidental interrupt in the last epoch when conducted the pre-training experiments on V100 GPUs (PyTorch 1.6.0). This interrupt is caused by the scheduler of learning rate. We naively set
--epochs 801
to walk away from issue :)