Trainer with TPUs
Hello, I’ve been trying to train a BERT-based model for SQuAD with the Trainer API, and I would like to use TPUs from COLAB to speed-up the training.I have followed the official notebook for training a SQuAD model, and then, I have followed the suggestions on this older topic:Tpu TrainerWhich says the only thing to edit is to set a max_length parameter for the padding, which I already did.So what I did is basically:Install the xla librarySet max_length for paddingSet TPU environment in COLABBut I don’t get any speedup for training, am I missing something? This is my code:My codeEdit: I made some progress! I have to restart the runtime as soon as I install the XLA library, otherwise there are compatibility issues with older versions.But still, I didn’t get the speed-up I hoped for, if I do 0.8 iterations/second with a GPU, I only get 3.2 iterations/second with TPU, is it normal?Also I get the following data:Num examples = 66107Num Epochs = 3Instantaneous batch size per device = 16Total train batch size (w. parallel, distributed & accumulation) = 16Is it normal? To my understanding, each TPU should have 8 cores, so if I have a batch size of 16 per each devices, the total train batch size should be 128? How do I use all the cores? Thanks
2022-04-12T16:21:02Z
Looking at the source code, it might be that you need to specify thetpu_num_cores, assuming that the notebook/script in your case does not set this automatically.github.comhuggingface/transformers/blob/cc034f72eb6137f4c550e911fba67f8a0e1e98fa/src/transformers/training_args.py#L258-L259tpu_num_cores (`int`, *optional*):When training on TPU, the number of TPU cores (automatically passed by launcher script).
2022-04-12T20:42:00Z
Thank you for the answer Bram! Unfortunately that’s not the case, I think by default the parameter is set to 8, since even tuning it manually doesn’t change anything.
2022-04-13T08:34:36Z
In that case, I am not sure if there are any other parameters that may improve your speed more. Perhaps someone else can help!