Distributed Machine Learning with Python

Distributed Machine Learning with Python

Section 1 – Data Parallelism
Section 2 – Model Parallelism
Section 3 – Advanced Parallelism Paradigms

Single-node training error – out of memory

Giant NLP models, such as BERT, are often hard to train using a single GPU (that is, single-node). The main reason is due to the limited on-device memory size.

Here, we will first fine-tune the BERT model using a single GPU. The dataset we will use is SQuAD 2.0. It often throws an Out-of-Memory (OOM) error due to the giant model size and huge intermediate results size.

Second, we will use a state-of-the-art GPU and try our best to pack the relatively small BERT-base model inside a single GPU. Then, we will carefully adjust the batch size to avoid an OOM error.

Fine-tuning BERT on a single GPU

The first thing we need to do is to install the transformers library on our machine. Here, we use the transformers library provided by Hugging Face. The following command is how we install it on an Ubuntu machine using PyTorch:

$ pip install transformers

Please make sure you are installing the correct transformers version (...

Unlock full access

Continue reading with a subscription

Packt gives you instant online access to a library of over 7,500 practical eBooks and videos, constantly updated with the latest in tech

End of Section 3

Your notes and bookmarks