Your code is loading the full model into a single GPU so having multiple GPUs does not prevent out of memory errors. By default, transformer models only have DDP (distributed data parallel) so each GPU has a copy of your model for speeding up training. Thus the maximum VRAM that you are allowed to use is the max of a single GPU or 16 GB. The moment OOM happens for 1 GPU, it will happen for all others.

To split your model so that you are training a single model with 4 GPU, you need to set up a different type of model parallelism that splits the model into multiple shards and having each GPU train one shard and then having the GPUs communicate with each other to combine the result into a single training loop.

ZeRO DDP or Fully Sharded Data Parallel is what you are looking for.