SIGSEGV when training on TPU

[FabianBell]

❓ Questions and Help

Before asking:

1. Try to find answers to your questions in the Lightning Forum!
2. Search for similar issues.
3. Search the docs.

What is your question?

I tried to apply the reformer model on a sentiment analysis task and train it on a tpu. I get a

ProcessExitedException: process X terminated with signal SIGSEGV

What did I do wrong?

Code

You can find my code in a colab notebook here.

What have you tried?

I tried to stick to notebook 1 for the general setup and notebook 2 for the tpu setup. I saw #1956 and #2124 however it does not work with the latest version (1.0.4).

What's your environment?

• OS: [e.g. iOS, Linux, Win] Linux
• Packaging [e.g. pip, conda] pip
• Version [e.g. 0.5.2.1] 1.0.4

[github-actions]

Hi! thanks for your contribution!, great first issue!

[rohitgr7]

cc @lezwon

[lezwon]

This seems to be an XLA issue and is tracked here pytorch/xla#1775

@FabianBell Mind adding the following at the beginning of your notebook and trying?

import os
os.environ['XLA_USE_32BIT_LONG'] = '1'
os.environ['TRIM_GRAPH_SIZE'] = '1000000'

[FabianBell]

@lezwon thank you for your help.

I changed the notebook but I still get the same error.

[lezwon]

I get a different error: Notebook

RuntimeError: Expected tensor for argument #1 'indices' to have scalar type Long; but got XLAIntType instead (while checking arguments for embedding)

Will look into this.

[lezwon]

@FabianBell I'm not able to figure out the root cause of the error mentioned above. I think it might be similar to this one: huggingface/transformers#2952

[FabianBell]

@lezwon thank you for your help. I followed notebook and I ended up with the same error. I do not think that it is a pytorch lightning problem. I will therefore close this issue here.