1
$\begingroup$

I'm working on an NLP task, using BERT, and I have a little doubt about GPU memory.

I already made a model (using DistilBERT) since I had out-of-memory problems with tensorflow on a RTX3090 (24gb gpu's ram, but ~20.5gb usable) with BERT base model.

To make it working, I limited my data to 1.1 milion of sentences in training set (truncating sentences at 128 words), and like 300k in validation, but using an high batch size (256).

Now I have the possibility to retrain the model on a Nvidia A100 (with 40gb gpu's ram), so it's time to use BERT base, and not the distilled version.

My question is, if I reduce the batch size (e.g. from 256 to 64), will I have some possibilities to increase the size of my training data (e.g. from 1.1 to 2-3 milions), the lenght of sentences (e.g. from 128 to 256, or 198) and use the bert base (which has a lot of trainable params more than distilled version) on the 40gb of the A100, or it's probably that I will get an OOM error?

I ask this because I haven't unlimited tries on this cluster, since I'm not alone using it (plus I have to prepare data differently in each case, and it has a quite high size), so I would have an estimation on what could happen.

Improve this question
$\endgroup$
7
  • 1
    $\begingroup$ Only the current batch should be loaded in GPU RAM, so you should not need to reduce your training data size (assuming your data loading and training routines are implemented properly). $\endgroup$
    – noe
    Commented May 11, 2022 at 16:04
  • $\begingroup$ @noe For now, I pre-tokenized my data and saved ids, attention mask and label as 3 differents npy files, then I load them in my training script and covert them in tf tensors with train_ids = tf.convert_to_tensor(np.load(ids_file.npy)) . Are tf tensors kept in gpu ram? $\endgroup$
    – Gius
    Commented May 11, 2022 at 16:13
  • $\begingroup$ It depends on the device you configure (e.g. with with tf.device(...):). You can check whether your tensor is in the GPU RAM with tensor.device. $\endgroup$
    – noe
    Commented May 11, 2022 at 16:40
  • $\begingroup$ @noe I checked with tensor.device and it says it's in GPU, so I tried to set it on CPU with with tf.device('CPU:0'): ..., and it says now it's in CPU. So in this way I keep my whole tensor in system ram and GPU ram will be occupied only by the current batch and model params? $\endgroup$
    – Gius
    Commented May 12, 2022 at 6:20
  • $\begingroup$ About "this way I keep my whole tensor in system ram", you already know it is the case. About "GPU ram will be occupied only by the current batch and model params", it depends on how the rest of your code handles devices. $\endgroup$
    – noe
    Commented May 12, 2022 at 7:56

1 Answer 1

Reset to default
3
$\begingroup$

As you pointed out in your comments, you pre-tokenized the data and kept in in tensors in GPU memory.

Only the current batch should be loaded in GPU RAM, so you should not need to reduce your training data size (assuming your data loading and training routines are implemented properly). To keep you training data tensor in CPU, you can use with tf.device(...):.

However, take into account that the size of the training data can also be huge for the size of the CPU memory. A typical approach for this is to save the token IDs on disk and then load them from there.

Improve this answer
$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.