Trying the BERT fine tunning on GPU

steffel

I have successfully fine-tuned the BERT model using my CPU. I would like to try fine-tuning using my GPU an RTX4090 with 24GB RAM. I am getting an out-of-memory error. I am a little surprised but I would like to make sure the memory is truly the issue and it’s not a misconfiguration.
The out-of-memory error raised on loading the model

{:ok, model} = Bumblebee.load_model({:hf, "bert-base-cased"}, spec: spec)

Thanks

19 comments

/nx #bumblebee

1 2633 19

2023-09-14 16:00:33 UTC

Showing Posts 1 to 10

Show Best Posts
Show All (oldest first)
Show All (newest first)

outlog

have you configured it to run on GPU and not CPU?

eg the appropriate XLA_TARGET ?

maybe post more code..

Post #1

steffel

The code is the one from the example here with the difference I have updated the libraries in order to be able to use cuda 12.

XLA_TARGET is set to cuda120 with my version of cuda being cuda 12.2. Should I downgrade to cuda 12.0? From my understanding cuda120 is CUDA 12.0+ so 12.0 and above.
Running nvidia-smi I can see the GPU memory is filled up.

Post #2

steffel

I have managed to get it to run using XLA_TARGET=cuda so it will compile it’s own version.

To get it to work, the batch size needs to be decreased, on my GPU it’s a batch size of 4.

Post #3

jonatanklosko

Creator of Livebook

@steffel this is very interesting, the model is like 0.5gb, so loading it should be far from running out of memory. Note that XLA preallocates memory upfront, so your GPU memory usage will bump to a high value, but an OOM is definitely unexpected.

What is the CUDA and cuDNN version you built with locally to make it work?

Post #4

jonatanklosko

Creator of Livebook

@steffel can you please run with XLA_ARCHIVE_URL=https://static.jonatanklosko.com/builds/xla_extension-x86_64-linux-gnu-cuda120.tar.gz and see if it makes a difference?

Post #5

steffel

Regarding the versions:
Cuda compilation tools, release 12.2, V12.2.140
Build cuda_12.2.r12.2/compiler.33191640_0
cudnn: v8

Post #6

steffel

The model loads without any memory issues.
Now when it comes to training the model I am using mrm8488/codebert-base-finetuned-detect-insecure-code that is based on roberta-base when compiling myself I am not able to push the sequence_length up to 1024 (batch_size=1) with your version it works. I would like to know what is different.

Post #7

jonatanklosko

Creator of Livebook

Oh, so just to make sure we are on the same page, the original OOM was during loading or training?

The archive I sent is precompiled using CUDA 12.1, rather than 12.0. According to Jax README it should still be compatible with CUDA 12.0, in which case we would be fine precompiling with 12.1. I will need to do some further tests.

Post #8

steffel

First I had an OOM during loading.
I managed to solve it using my own built version of xla extension.
I was able to train using a batch_size=1 and sequence_length=64 or 128.
With sequence_length above 128, I am getting an OOM error during training.
Something like that:

** (RuntimeError) Out of memory while trying to allocate 4290931080 bytes.
BufferAssignment OOM Debugging.
BufferAssignment stats:
             parameter allocation:    1.86GiB
              constant allocation:     2.2KiB
        maybe_live_out allocation:    1.86GiB
     preallocated temp allocation:    4.00GiB
  preallocated temp fragmentation:     9.1KiB (0.00%)
                 total allocation:    7.71GiB
              total fragmentation:  656.76MiB (8.32%)

Post #9