DeepSpeed Apex Automatic Mixed Precision

Install

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

Modify Config

deepspeed_config = {
    'train_batch_size': BATCH_SIZE,
    'gradient_clipping': GRAD_CLIP_NORM,
    'amp': {
        'enabled': True,
        'opt_level': 'O1',
    }
}