OOM ERROR when using local batch size=128 on TPUv3-8

[GhaliaRehawi]

Hi, I am trying to train XLNet on protein sequences. I am running into OOM error when running the script train.py using a TPUv3-8 with train_batch_size=128. (I also get OOM error with train batch size 64, 48, but not with 32, 16). In the paper it is mentioned: "Specifically, we train on 512 TPU v3 chips for 500K steps with an Adam weight decay optimizer, linear learning rate decay, and a batch size of 8192, which takes about 5.5 days." If I understand this correctly then the local batch size used is also 128= (8192/(512/8)) and I shouldn't get an OOM error. for context, am using TPUv3-8 (version 1.14.1.dev20190518) and a cloud instance vm both in us-central1-a and Tensorflow version 1.13.1 For the data preprocessing I am using the script data_utils and it runs with no problem. Here are the command am using for both preprocessing and training :

python xlnet/data_utils.py
--use_tpu=True
--save_dir=proc_data_bsz128/example
--bsz_per_host=128
--num_core_per_host=8
--seq_len=512
--reuse_len=256
--input_glob=testdata_xlnet.txt
--num_passes=20
--bi_data=True
--sp_path=sp.model
--mask_alpha=6
--mask_beta=1
--uncased=False
--num_predict=85

python xlnet/train.py
--use_tpu=True
--tpu=name
--record_info_dir=$DATA_DIR
--save_steps=1000
--model_dir=$MODEL_DIR
--train_batch_size=128
--seq_len=512
--reuse_len=256
--mem_len=384
--perm_size=256
--n_layer=24
--d_model=1024
--d_embed=1024
--n_head=16
--d_head=64
--d_inner=4096
--untie_r=True
--mask_alpha=6
--mask_beta=1
--num_predict=85

$DATA_DIR and $MODEL_DIR are google bucket directories. Is there something am missing here? Thanks for your help in advance.

[fangli80]

I think they mean 512 TPUs and each TPU has 8 cores. So each TPU core has two sequences