Trainer on colab TPU error: process X terminated with signal SIGSEGV

[VictorCallejas]

I am trying to train an image encoder with pytorch-lightning on colab using TPU(8 cores).

I am following this demo notebook: https://colab.research.google.com/drive/1-_LKx4HwAxl5M6xPJmqAAu444LTDQoa3#scrollTo=dEeUzX_5aLrX

LIbraries version:
torch: 1.5.0
torchvision_ 0.6.0
pytorch-lightning: 0.7.5
pytorch-xla:1.6

I have also tried with nightly versions and older, but same error.

When running:

trainer = pl.Trainer(num_tpu_cores=8, progress_bar_refresh_rate=10, max_epochs=10)
# Run lr finder
lr_finder = trainer.lr_find(model)

Error:

---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-25-8e434d4d50a5> in <module>()
      2 
      3 # Run lr finder
----> 4 lr_finder = trainer.lr_find(model)
      5 
      6 fig = lr_finder.plot(suggest=True)

4 frames
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/lr_finder.py in lr_find(self, model, train_dataloader, min_lr, max_lr, num_training, mode, num_accumulation_steps)
    151 
    152         # Fit, lr & loss logged in callback
--> 153         self.fit(model, train_dataloader=train_dataloader)
    154 
    155         # Prompt if we stopped early

/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders)
    775 
    776             # train
--> 777             xmp.spawn(self.tpu_train, args=(model,), nprocs=self.num_tpu_cores, start_method=start_method)
    778 
    779             # load weights if not interrupted

/usr/local/lib/python3.6/dist-packages/torch_xla/distributed/xla_multiprocessing.py in spawn(fn, args, nprocs, join, daemon, start_method)
    180         join=join,
    181         daemon=daemon,
--> 182         start_method=start_method)

/usr/local/lib/python3.6/dist-packages/torch/multiprocessing/spawn.py in start_processes(fn, args, nprocs, join, daemon, start_method)
    156 
    157     # Loop on join until it returns True or raises an exception.
--> 158     while not context.join():
    159         pass
    160 

/usr/local/lib/python3.6/dist-packages/torch/multiprocessing/spawn.py in join(self, timeout)
    106                 raise Exception(
    107                     "process %d terminated with signal %s" %
--> 108                     (error_index, name)
    109                 )
    110             else:

Exception: process 2 terminated with signal SIGSEGV

My notebook on Gist here: https://colab.research.google.com/gist/VictorCallejas/10e4c39fc25051012ae28a2a7261f814/untitled.ipynb

It seems it is raising the exception because the other processes are not joining, but I have not a clue why. I have everything as in the demo notebook.

Thanks.

[github-actions]

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

[williamFalcon]

The problem here is the learning rate finder doesn't work with anything multi-process... @SkafteNicki .

[SkafteNicki]

This is weird. It is correct that learning rate finder does not have multi process support, but that is because the state of the search is destroyed when self.fit() finishes. However, here the problem seems to be earlier during the fit

@VictorCallejas do trainer.fit work normally for you?

[VictorCallejas]

The error is when running trainer.fit(model), when using trainer.lr_find(model) or trainer.fit(model) alone.

Yes, it works on cpu or gpu.

[SkafteNicki]

Then I guess that the problem is unrelated to the learning rate finder, if standard trainer.fit(model) fails.

[Borda]

that is spawn issue, shall be fixed in #2632