Suggestion to Out-Of-Memory Error (with 4 GPU)

[KaiserWhoLearns]

I am trying to only use T5-small for fine-tuning model on a custom dataset (in tsv format) with 4 GPUs but keep getting Out-Of-Memory error.

Allocator (GPU_0_bfc) ran out of memory trying to allocate 16.00MiB (rounded to 16777216)requested by op while_loop/while/decoder/block_000/layer_002/einsum_1/parallel_0/Mul
If the cause is memory fragmentation maybe the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' will improve the situation. 

I believe this is very unlikely as my GPU should have more than enough memory to fine-tune the small model, so there should be something wrong with the command I used:

t5_mesh_transformer  \
  --model_dir="${MODEL_DIR}" \
  --t5_tfds_data_dir="${DATA_DIR}" \
  --gin_file="dataset.gin" \
  --gin_param="utils.run.mesh_shape = 'model:4,batch:1'" \
  --gin_param="utils.run.mesh_devices = ['gpu:0','gpu:1','gpu:2','gpu:3']" \
  --gin_param="utils.run.train_dataset_fn = @t5.models.mesh_transformer.tsv_dataset_fn" \
  --gin_param="tsv_dataset_fn.filename = '${base_dir}/data/spider/random_split/train.tsv'" \
  --gin_file="gs://t5-data/pretrained_models/small/operative_config.gin" \
  --gin_param="utils.run.init_checkpoint='gs://t5-data/pretrained_models/small/model.ckpt-1000000'" \
  --gin_file="learning_rate_schedules/constant_0_001.gin" \
  --gin_param="batch_size = 1"

Any suggestion on what I can modify to resolve the issue?
(It looks like it is out of memory on the first GPU, so maybe I happened to used a invalid command for multi-GPU?)