# How to use TPU for model training using example script run\_mlm.py

**URL:** <https://discuss.huggingface.co/t/how-to-use-tpu-for-model-training-using-example-script-run-mlm-py/15510>\
**Category:** Beginners\
**Created:** [March 8, 2022, 2:58pm UTC](https://discuss.huggingface.co/t/how-to-use-tpu-for-model-training-using-example-script-run-mlm-py/15510 "2022-03-08T14:58:18Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![silkski](https://avatars.discourse-cdn.com/v4/letter/s/e9bcb4/32.png) [@silkski](https://discuss.huggingface.co/u/silkski)\
**Post date:** [March 8, 2022, 2:58pm UTC](https://discuss.huggingface.co/t/how-to-use-tpu-for-model-training-using-example-script-run-mlm-py/15510/1 "2022-03-08T14:58:18Z")

</div>

Hi there,

I’ve been trying to use colab TPU for training BERT and other language models using the example script provided ‘run\_mlm.py’.

I’ve used the code for loading a TPU in colab from the original T5 repo here ([Google Colab](https://colab.research.google.com/github/google-research/text-to-text-transfer-transformer/blob/master/notebooks/t5-trivia.ipynb#scrollTo=Set_Up))

Code provides `TPU_ADDRESS` from `TPU_ADDRESS = tpu.get_master()`, which looks something like this: `grpc://10.121.143.146:8470`

As stated in the docs ([https://github.com/huggingface/transformers/tree/master/examples/tensorflow/language-modeling#multi-gpu-and-tpu-usage](https://github.com/huggingface/transformers/tree/master/examples/tensorflow/language-modeling#multi-gpu-and-tpu-usage)) TPU can be used by passing the `--tpu` argument, however doing this I’m getting the following error:

```auto
!python /content/transformers/examples/tensorflow/language-modeling/run_mlm.py \
--model_name_or_path bert-base-cased \
--validation_split_percentage 20 \
--line_by_line \
--learning_rate 2e-5 \
--per_device_train_batch_size 16 \
--per_device_eval_batch_size 32 \
--num_train_epochs 4 \
--output_dir /content/output \
--train_file /content/text.csv \
--tpu TPU ADDRESS

```

> run\_mlm.py: error: ambiguous option: --tpu could match --tpu\_num\_cores, --tpu\_metrics\_debug, --tpu\_name, --tpu\_zone

Apparently, there is no option for passing the tpu\_address.

I’m not sure how to get `tpu_name`, so that I can pass that either?

Any help is appreciated.

---

<div class="post-metadata">

**Author:** ![GenV](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/genv/32/12132_2.png) [@GenV](https://discuss.huggingface.co/u/GenV)\
**Post date:** [March 8, 2022, 3:53pm UTC](https://discuss.huggingface.co/t/how-to-use-tpu-for-model-training-using-example-script-run-mlm-py/15510/2 "2022-03-08T15:53:11Z")

</div>

Hi,  
I think you need to run this script (on Colab, with TPU session active):

`!python xla_spawn.py --num_cores 8 run_mlm.py --model_name_or_path bert-base-cased \ --tpu_num_cores 8 \ --all other params`

---

<div class="post-metadata">

**Author:** ![timodim](https://avatars.discourse-cdn.com/v4/letter/t/4af34b/32.png) [@timodim](https://discuss.huggingface.co/u/timodim)\
**Post date:** [March 11, 2022, 11:17am UTC](https://discuss.huggingface.co/t/how-to-use-tpu-for-model-training-using-example-script-run-mlm-py/15510/3 "2022-03-11T11:17:54Z")

</div>

Hi Gennaro,

Appreciate your response. Where exactly can I get this xla\_spawn.py script from? Thanks!

EDIT: I actually managed to get the script to run just by passing the `--tpu_num_cores 8` parameter to the script. Thank you so much for your help!

---

<div class="post-metadata">

**Author:** ![GenV](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/genv/32/12132_2.png) [@GenV](https://discuss.huggingface.co/u/GenV)\
**Post date:** [March 11, 2022, 11:36am UTC](https://discuss.huggingface.co/t/how-to-use-tpu-for-model-training-using-example-script-run-mlm-py/15510/4 "2022-03-11T11:36:33Z")

</div>

Hi @timodim,  
[here](https://github.com/huggingface/transformers/blob/master/examples/legacy/seq2seq/xla_spawn.py) the xla\_spawn script. Glad it is working!

---

<div class="post-metadata">

**Author:** ![kechan](https://avatars.discourse-cdn.com/v4/letter/k/aca169/32.png) [@kechan](https://discuss.huggingface.co/u/kechan)\
**Post date:** [June 18, 2022, 9:02pm UTC](https://discuss.huggingface.co/t/how-to-use-tpu-for-model-training-using-example-script-run-mlm-py/15510/5 "2022-06-18T21:02:44Z")

</div>

I am trying to train a TFDistilBertForTokenClassification with TPU + Colab. I work around the problem of readying a TF dataset on GCS, but the model still fails to train. Prominently:

> (0) INVALID\_ARGUMENT: {{function\_node \_\_inference\_train\_function\_26618}} Detected unsupported operations when trying to compile graph tf\_distil\_bert\_for\_token\_classification\_cond\_true\_22902 on XLA\_TPU\_JIT: PrintV2 (No registered ‘PrintV2’ OpKernel for XLA\_TPU\_JIT devices compatible with node {{node tf\_distil\_bert\_for\_token\_classification/cond/PrintV2}}){{node tf\_distil\_bert\_for\_token\_classification/cond/PrintV2}}  
> One approach is to outside compile the unsupported ops to run on CPUs by enabling soft placement `tf.config.set_soft_device_placement(True)`. This has a potential performance penalty.

It seems to be using some unsupported op TPU doesn’t like. Any idea how to hack this?